Simple testcase.
()
| 33 | |
| 34 | |
| 35 | def test_fuse_simple(): |
| 36 | """Simple testcase.""" |
| 37 | |
| 38 | def before(): |
| 39 | bb = relax.BlockBuilder() |
| 40 | x = relax.Var("x", R.Tensor([10, 20], "float32")) |
| 41 | with bb.function("main", [x]): |
| 42 | with bb.dataflow(): |
| 43 | lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32")) |
| 44 | lv1 = bb.emit_te(topi.exp, lv0) |
| 45 | gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) |
| 46 | bb.emit_func_output(gv) |
| 47 | |
| 48 | return bb.get() |
| 49 | |
| 50 | def expected(): |
| 51 | bb = relax.BlockBuilder() |
| 52 | x = relax.Var("x", R.Tensor([10, 20], "float32")) |
| 53 | p0 = relax.Var("p0", R.Tensor((), "float32")) |
| 54 | |
| 55 | with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True): |
| 56 | with bb.dataflow(): |
| 57 | lv0 = bb.emit_te(topi.add, x, p0) |
| 58 | lv1 = bb.emit_te(topi.exp, lv0) |
| 59 | gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) |
| 60 | bb.emit_func_output(gv) |
| 61 | fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze") |
| 62 | |
| 63 | x = relax.Var("x", R.Tensor([10, 20], "float32")) |
| 64 | with bb.function("main", [x]): |
| 65 | with bb.dataflow(): |
| 66 | gv = bb.emit_output( |
| 67 | relax.Call(fused_add_exp_squeeze, [x, relax.const(1, "float32")]) |
| 68 | ) |
| 69 | bb.emit_func_output(gv) |
| 70 | |
| 71 | return bb.get() |
| 72 | |
| 73 | _check(before(), expected()) |
| 74 | |
| 75 | |
| 76 | def test_conv2d_fuse(): |