Test fusion case involving concat op and Tuple node
()
| 172 | |
| 173 | |
| 174 | def test_concatenate(): |
| 175 | """Test fusion case involving concat op and Tuple node""" |
| 176 | |
| 177 | def before(): |
| 178 | bb = relax.BlockBuilder() |
| 179 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 180 | with bb.function("main", [x]): |
| 181 | with bb.dataflow(): |
| 182 | lv0 = bb.emit_te( |
| 183 | topi.nn.pool2d, |
| 184 | x, |
| 185 | kernel=(2, 2), |
| 186 | stride=(2, 2), |
| 187 | dilation=(1, 1), |
| 188 | padding=(0, 0, 0, 0), |
| 189 | pool_type="max", |
| 190 | ) |
| 191 | lv1 = bb.emit_te(topi.nn.upsampling, lv0, scale_h=2.0, scale_w=2.0) |
| 192 | lv2 = bb.emit_te(topi.concatenate, (lv1, x), axis=1) |
| 193 | gv = bb.emit_output(bb.call_te(topi.add, lv2, relax.const(1, "float32"))) |
| 194 | bb.emit_func_output(gv) |
| 195 | |
| 196 | return bb.get() |
| 197 | |
| 198 | def expected(): |
| 199 | bb = relax.BlockBuilder() |
| 200 | |
| 201 | # Grouped function |
| 202 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 203 | w = relax.Var("w", R.Tensor((1, 16, 32, 32), "float32")) |
| 204 | p0 = relax.Var("p0", R.Tensor((), "float32")) |
| 205 | with bb.function( |
| 206 | "fused_upsampling_concatenate_add", [w, x, p0], attrs={"Primitive": True}, private=True |
| 207 | ): |
| 208 | with bb.dataflow(): |
| 209 | lv0 = bb.emit_te(topi.nn.upsampling, w, scale_h=2.0, scale_w=2.0) |
| 210 | lv1 = bb.emit_te(topi.concatenate, (lv0, x), axis=1) |
| 211 | gv = bb.emit_output(bb.call_te(topi.add, lv1, p0)) |
| 212 | bb.emit_func_output(gv) |
| 213 | |
| 214 | # Get the global variables of the grouped functions |
| 215 | fused_upsampling_concatenate_add = bb.get().get_global_var( |
| 216 | "fused_upsampling_concatenate_add" |
| 217 | ) |
| 218 | |
| 219 | # Main function |
| 220 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 221 | with bb.function("main", [x]): |
| 222 | with bb.dataflow(): |
| 223 | lv0 = bb.emit_te( |
| 224 | topi.nn.pool2d, |
| 225 | x, |
| 226 | kernel=(2, 2), |
| 227 | stride=(2, 2), |
| 228 | dilation=(1, 1), |
| 229 | padding=(0, 0, 0, 0), |
| 230 | pool_type="max", |
| 231 | ) |