| 682 | |
| 683 | |
| 684 | def test_fuse_parallel_injective(): |
| 685 | def before(): |
| 686 | bb = relax.BlockBuilder() |
| 687 | |
| 688 | x = relax.Var("x", R.Tensor((10, 20), "int32")) |
| 689 | with bb.function("main", [x]): |
| 690 | with bb.dataflow(): |
| 691 | lv0 = bb.emit_te(topi.add, x, relax.const(1, "int32")) |
| 692 | lv1 = bb.emit_te(topi.squeeze, lv0) |
| 693 | lv2 = bb.emit_te(topi.transpose, lv0, axes=[1, 0]) |
| 694 | lv3 = bb.emit_te(topi.transpose, lv2, axes=[1, 0]) |
| 695 | gv = bb.emit_output(bb.call_te(topi.left_shift, lv1, lv3)) |
| 696 | bb.emit_func_output(gv) |
| 697 | |
| 698 | return bb.get() |
| 699 | |
| 700 | def expected(): |
| 701 | bb = relax.BlockBuilder() |
| 702 | |
| 703 | # Grouped function |
| 704 | x = relax.Var("x", R.Tensor((10, 20), "int32")) |
| 705 | p0 = relax.Var("p0", R.Tensor((), "int32")) |
| 706 | with bb.function( |
| 707 | "fused_add_squeeze_transpose_transpose1_left_shift", |
| 708 | [x, p0], |
| 709 | attrs={"Primitive": True}, |
| 710 | private=True, |
| 711 | ): |
| 712 | with bb.dataflow(): |
| 713 | lv0 = bb.emit_te(topi.add, x, p0) |
| 714 | lv1 = bb.emit_te(topi.squeeze, lv0) |
| 715 | lv2 = bb.emit_te(topi.transpose, lv0, axes=[1, 0]) |
| 716 | lv3 = bb.emit_te(topi.transpose, lv2, axes=[1, 0], primfunc_name_hint="transpose1") |
| 717 | gv = bb.emit_output(bb.call_te(topi.left_shift, lv1, lv3)) |
| 718 | bb.emit_func_output(gv) |
| 719 | |
| 720 | # Get the global variables of the grouped functions |
| 721 | fused_func = bb.get().get_global_var("fused_add_squeeze_transpose_transpose1_left_shift") |
| 722 | |
| 723 | # Main function |
| 724 | x = relax.Var("x", R.Tensor((10, 20), "int32")) |
| 725 | with bb.function("main", [x]): |
| 726 | with bb.dataflow(): |
| 727 | gv = bb.emit_output(relax.Call(fused_func, (x, relax.const(1, "int32")))) |
| 728 | bb.emit_func_output(gv) |
| 729 | |
| 730 | return bb.get() |
| 731 | |
| 732 | _check(before(), expected()) |
| 733 | |
| 734 | |
| 735 | def test_softmax(): |