| 218 | |
| 219 | |
| 220 | def test_fuse_same_primfunc(): |
| 221 | def before(): |
| 222 | bb = relax.BlockBuilder() |
| 223 | x1 = relax.Var("x1", R.Tensor([10, 20], "float32")) |
| 224 | with bb.function("fused_exp_exp_squeeze", [x1], attrs={"Primitive": True}): |
| 225 | with bb.dataflow(): |
| 226 | lv1 = bb.emit_te(topi.exp, x1) |
| 227 | lv2 = bb.emit_te(topi.exp, lv1) |
| 228 | gv = bb.emit_output(bb.call_te(topi.squeeze, lv2)) |
| 229 | bb.emit_func_output(gv) |
| 230 | mod = bb.get() |
| 231 | |
| 232 | func_gv = mod.get_global_var("fused_exp_exp_squeeze") |
| 233 | x = relax.Var("x", R.Tensor([10, 20], "float32")) |
| 234 | with bb.function("main", [x]): |
| 235 | with bb.dataflow(): |
| 236 | lv = bb.emit(relax.Call(func_gv, [x])) |
| 237 | gv = bb.emit_output(lv) |
| 238 | bb.emit_func_output(gv) |
| 239 | return bb.get() |
| 240 | |
| 241 | def expected(): |
| 242 | def fused_exp_exp_squeeze(x): |
| 243 | exp = topi.exp(x) |
| 244 | exp = topi.exp(exp) |
| 245 | squeeze = topi.squeeze(exp) |
| 246 | return squeeze |
| 247 | |
| 248 | bb = relax.BlockBuilder() |
| 249 | x = relax.Var("x", R.Tensor([10, 20], "float32")) |
| 250 | with bb.function("main", [x]): |
| 251 | with bb.dataflow(): |
| 252 | lv = bb.call_te(fused_exp_exp_squeeze, x) |
| 253 | gv = bb.emit_output(lv) |
| 254 | bb.emit_func_output(gv) |
| 255 | return bb.get() |
| 256 | |
| 257 | _check(before(), expected()) |
| 258 | |
| 259 | |
| 260 | def test_fuse_with_tuple_as_param(): |