Test fusion case of conv2d
()
| 74 | |
| 75 | |
| 76 | def test_conv2d_fuse(): |
| 77 | """Test fusion case of conv2d""" |
| 78 | |
| 79 | def before(dtype): |
| 80 | bb = relax.BlockBuilder() |
| 81 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) |
| 82 | w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype)) |
| 83 | w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype)) |
| 84 | w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype)) |
| 85 | with bb.function("main", [x, w1, w2, w3]): |
| 86 | with bb.dataflow(): |
| 87 | lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype)) |
| 88 | lv1 = bb.emit_te(topi.nn.conv2d, lv0, w1, strides=1, padding=1, dilation=1) |
| 89 | # this is the next dominator. |
| 90 | lv2 = bb.emit_te(topi.add, relax.const(1, dtype), lv1) |
| 91 | lv3 = bb.emit_te(topi.add, lv1, lv2) |
| 92 | # second path |
| 93 | lv4 = bb.emit_te(topi.nn.conv2d, lv3, w2, strides=1, padding=0, dilation=1) |
| 94 | lv5 = bb.emit_te(topi.nn.conv2d, lv3, w3, strides=1, padding=1, dilation=1) |
| 95 | gv = bb.emit_output(bb.call_te(topi.add, lv4, lv5)) |
| 96 | bb.emit_func_output(gv) |
| 97 | |
| 98 | return bb.get() |
| 99 | |
| 100 | def expected(dtype): |
| 101 | bb = relax.BlockBuilder() |
| 102 | |
| 103 | # Grouped function 1 |
| 104 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) |
| 105 | w = relax.Var("w", R.Tensor((16, 16, 3, 3), dtype)) |
| 106 | p0 = relax.Var("p0", R.Tensor((), dtype)) |
| 107 | with bb.function( |
| 108 | "fused_conv2d_add1_add2", [x, w, p0], attrs={"Primitive": True}, private=True |
| 109 | ): |
| 110 | with bb.dataflow(): |
| 111 | lv0 = bb.emit_te( |
| 112 | topi.nn.conv2d, |
| 113 | x, |
| 114 | w, |
| 115 | strides=1, |
| 116 | padding=1, |
| 117 | dilation=1, |
| 118 | primfunc_name_hint="conv2d", |
| 119 | ) |
| 120 | lv1 = bb.emit_te(topi.add, p0, lv0, primfunc_name_hint="add1") |
| 121 | gv = bb.emit_output(bb.call_te(topi.add, lv0, lv1, primfunc_name_hint="add2")) |
| 122 | bb.emit_func_output(gv) |
| 123 | |
| 124 | # Grouped function 2 |
| 125 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) |
| 126 | w = relax.Var("w", R.Tensor((16, 16, 1, 1), dtype)) |
| 127 | y = relax.Var("y", R.Tensor((1, 16, 64, 64), dtype)) |
| 128 | with bb.function("fused_conv2d1_add2", [x, w, y], attrs={"Primitive": True}, private=True): |
| 129 | with bb.dataflow(): |
| 130 | lv0 = bb.emit_te( |
| 131 | topi.nn.conv2d, |
| 132 | x, |
| 133 | w, |