| 376 | |
| 377 | |
| 378 | def test_tuple_intermediate(): |
| 379 | def before(): |
| 380 | bb = relax.BlockBuilder() |
| 381 | |
| 382 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 383 | with bb.function("main", [x]): |
| 384 | with bb.dataflow(): |
| 385 | lv0 = bb.emit_te(topi.squeeze, x) |
| 386 | lv1 = bb.emit_te(topi.add, lv0, relax.const(1, "float32")) |
| 387 | lv2 = bb.emit_te(topi.squeeze, lv0) |
| 388 | lv3 = bb.emit_te(topi.add, lv2, relax.const(1, "float32")) |
| 389 | lv4 = bb.emit_te(topi.add, lv3, relax.const(1, "float32")) |
| 390 | lv5 = bb.emit_te(topi.add, lv0, relax.const(1, "float32")) |
| 391 | lv6 = bb.emit_te(topi.concatenate, (lv1, lv4, lv5), axis=1) |
| 392 | lv7 = bb.emit_te(topi.squeeze, lv6) |
| 393 | gv = bb.emit_output(bb.call_te(topi.add, lv7, relax.const(1, "float32"))) |
| 394 | bb.emit_func_output(gv) |
| 395 | |
| 396 | return bb.get() |
| 397 | |
| 398 | def expected(): |
| 399 | bb = relax.BlockBuilder() |
| 400 | |
| 401 | # Grouped function |
| 402 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 403 | p0 = relax.Var("p0", R.Tensor((), "float32")) |
| 404 | p1 = relax.Var("p1", R.Tensor((), "float32")) |
| 405 | p2 = relax.Var("p2", R.Tensor((), "float32")) |
| 406 | p3 = relax.Var("p3", R.Tensor((), "float32")) |
| 407 | p4 = relax.Var("p4", R.Tensor((), "float32")) |
| 408 | with bb.function( |
| 409 | "fused_squeeze_add_squeeze1_add_add_add_concatenate_squeeze2_add1", |
| 410 | [x, p0, p1, p2, p3, p4], |
| 411 | attrs={"Primitive": True}, |
| 412 | private=True, |
| 413 | ): |
| 414 | with bb.dataflow(): |
| 415 | lv0 = bb.emit_te(topi.squeeze, x) |
| 416 | lv1 = bb.emit_te(topi.add, lv0, p0) |
| 417 | lv2 = bb.emit_te(topi.squeeze, lv0) |
| 418 | lv3 = bb.emit_te(topi.add, lv2, p1) |
| 419 | lv4 = bb.emit_te(topi.add, lv3, p2) |
| 420 | lv5 = bb.emit_te(topi.add, lv0, p3) |
| 421 | lv6 = bb.emit_te(topi.concatenate, (lv1, lv4, lv5), axis=1) |
| 422 | lv7 = bb.emit_te(topi.squeeze, lv6) |
| 423 | gv = bb.emit_output(bb.call_te(topi.add, lv7, p4)) |
| 424 | bb.emit_func_output(gv) |
| 425 | |
| 426 | # Get the global variables of the grouped functions |
| 427 | fused_func = bb.get().get_global_var( |
| 428 | "fused_squeeze_add_squeeze1_add_add_add_concatenate_squeeze2_add1" |
| 429 | ) |
| 430 | |
| 431 | # Main func |
| 432 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 433 | with bb.function("main", [x]): |
| 434 | with bb.dataflow(): |
| 435 | gv = bb.emit_output( |