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