Test if softmax can be fused with following ops.
()
| 733 | |
| 734 | |
| 735 | def test_softmax(): |
| 736 | """Test if softmax can be fused with following ops.""" |
| 737 | |
| 738 | def before(): |
| 739 | bb = relax.BlockBuilder() |
| 740 | |
| 741 | x = relax.Var("x", R.Tensor((16, 16), "float32")) |
| 742 | with bb.function("main", [x]): |
| 743 | with bb.dataflow(): |
| 744 | lv0 = bb.emit_te(topi.nn.softmax, x) |
| 745 | gv = bb.emit_output(bb.call_te(topi.cast, lv0, dtype="float16")) |
| 746 | bb.emit_func_output(gv) |
| 747 | |
| 748 | return bb.get() |
| 749 | |
| 750 | def expected(): |
| 751 | bb = relax.BlockBuilder() |
| 752 | |
| 753 | # Grouped function |
| 754 | x = relax.Var("x", R.Tensor((16, 16), "float32")) |
| 755 | with bb.function("fused_softmax_cast", [x], attrs={"Primitive": True}, private=True): |
| 756 | with bb.dataflow(): |
| 757 | lv0 = bb.emit_te(topi.nn.softmax, x) |
| 758 | gv = bb.emit_output(bb.call_te(topi.cast, lv0, dtype="float16")) |
| 759 | bb.emit_func_output(gv) |
| 760 | |
| 761 | # Get the global variables of the grouped functions |
| 762 | fused_func = bb.get().get_global_var("fused_softmax_cast") |
| 763 | |
| 764 | # Main function |
| 765 | x = relax.Var("x", R.Tensor((16, 16), "float32")) |
| 766 | with bb.function("main", [x]): |
| 767 | with bb.dataflow(): |
| 768 | gv = bb.emit_output(relax.Call(fused_func, (x,))) |
| 769 | bb.emit_func_output(gv) |
| 770 | |
| 771 | return bb.get() |
| 772 | |
| 773 | _check(before(), expected()) |
| 774 | |
| 775 | |
| 776 | def test_multiple_relax_functions(): |