# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: E501, F841 import tvm import tvm.testing from tvm import relax, topi from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T def _check(mod_actual, mod_expected): mod_actual = relax.transform.AnnotateTIROpPattern()(mod_actual) mod_actual = relax.transform.FuseOps()(mod_actual) mod_expected = relax.transform.AnnotateTIROpPattern()(mod_expected) tvm.ir.assert_structural_equal(mod_actual, mod_expected) def test_fuse_simple(): """Simple testcase.""" def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor([10, 20], "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv1 = bb.emit_te(topi.exp, lv0) gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor([10, 20], "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, p0) lv1 = bb.emit_te(topi.exp, lv0) gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) bb.emit_func_output(gv) fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze") x = relax.Var("x", R.Tensor([10, 20], "float32")) with bb.function("main", [x]): with bb.dataflow(): gv = bb.emit_output( relax.Call(fused_add_exp_squeeze, [x, relax.const(1, "float32")]) ) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_conv2d_fuse(): """Test fusion case of conv2d""" def before(dtype): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype)) w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype)) w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype)) with bb.function("main", [x, w1, w2, w3]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype)) lv1 = bb.emit_te(topi.nn.conv2d, lv0, w1, strides=1, padding=1, dilation=1) # this is the next dominator. lv2 = bb.emit_te(topi.add, relax.const(1, dtype), lv1) lv3 = bb.emit_te(topi.add, lv1, lv2) # second path lv4 = bb.emit_te(topi.nn.conv2d, lv3, w2, strides=1, padding=0, dilation=1) lv5 = bb.emit_te(topi.nn.conv2d, lv3, w3, strides=1, padding=1, dilation=1) gv = bb.emit_output(bb.call_te(topi.add, lv4, lv5)) bb.emit_func_output(gv) return bb.get() def expected(dtype): bb = relax.BlockBuilder() # Grouped function 1 x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) w = relax.Var("w", R.Tensor((16, 16, 3, 3), dtype)) p0 = relax.Var("p0", R.Tensor((), dtype)) with bb.function( "fused_conv2d_add1_add2", [x, w, p0], attrs={"Primitive": True}, private=True ): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.conv2d, x, w, strides=1, padding=1, dilation=1, primfunc_name_hint="conv2d", ) lv1 = bb.emit_te(topi.add, p0, lv0, primfunc_name_hint="add1") gv = bb.emit_output(bb.call_te(topi.add, lv0, lv1, primfunc_name_hint="add2")) bb.emit_func_output(gv) # Grouped function 2 x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) w = relax.Var("w", R.Tensor((16, 16, 1, 1), dtype)) y = relax.Var("y", R.Tensor((1, 16, 64, 64), dtype)) with bb.function("fused_conv2d1_add2", [x, w, y], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.conv2d, x, w, strides=1, padding=0, dilation=1, primfunc_name_hint="conv2d1", ) gv = bb.emit_output(bb.call_te(topi.add, lv0, y, primfunc_name_hint="add2")) bb.emit_func_output(gv) # Get the global variables of the grouped functions mod = bb.get() fused_conv2d_add1_add2 = mod.get_global_var("fused_conv2d_add1_add2") fused_conv2d1_add2 = mod.get_global_var("fused_conv2d1_add2") # Main function x = relax.Var("x", R.Tensor((1, 16, 64, 64), dtype)) w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), dtype)) w2 = relax.Var("w2", R.Tensor((16, 16, 1, 1), dtype)) w3 = relax.Var("w3", R.Tensor((16, 16, 3, 3), dtype)) with bb.function("main", [x, w1, w2, w3]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, dtype)) lv1 = bb.emit(relax.Call(fused_conv2d_add1_add2, [lv0, w1, relax.const(1, dtype)])) lv2 = bb.emit_te( topi.nn.conv2d, lv1, w3, strides=1, padding=1, dilation=1, ) gv = bb.emit_output(relax.Call(fused_conv2d1_add2, [lv1, w2, lv2])) bb.emit_func_output(gv) return bb.get() _check(before("float32"), expected("float32")) _check(before("float16"), expected("float16")) _check(before("int8"), expected("int8")) def test_concatenate(): """Test fusion case involving concat op and Tuple node""" def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.pool2d, x, kernel=(2, 2), stride=(2, 2), dilation=(1, 1), padding=(0, 0, 0, 0), pool_type="max", ) lv1 = bb.emit_te(topi.nn.upsampling, lv0, scale_h=2.0, scale_w=2.0) lv2 = bb.emit_te(topi.concatenate, (lv1, x), axis=1) gv = bb.emit_output(bb.call_te(topi.add, lv2, relax.const(1, "float32"))) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() # Grouped function x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) w = relax.Var("w", R.Tensor((1, 16, 32, 32), "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) with bb.function( "fused_upsampling_concatenate_add", [w, x, p0], attrs={"Primitive": True}, private=True ): with bb.dataflow(): lv0 = bb.emit_te(topi.nn.upsampling, w, scale_h=2.0, scale_w=2.0) lv1 = bb.emit_te(topi.concatenate, (lv0, x), axis=1) gv = bb.emit_output(bb.call_te(topi.add, lv1, p0)) bb.emit_func_output(gv) # Get the global variables of the grouped functions fused_upsampling_concatenate_add = bb.get().get_global_var( "fused_upsampling_concatenate_add" ) # Main function x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.pool2d, x, kernel=(2, 2), stride=(2, 2), dilation=(1, 1), padding=(0, 0, 0, 0), pool_type="max", ) gv = bb.emit_output( relax.Call( fused_upsampling_concatenate_add, (lv0, x, relax.const(1, "float32")) ) ) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_tuple_root(): """Test fusion case where Tuple node is the root in its group""" def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.pool2d, x, kernel=(2, 2), stride=(2, 2), dilation=(1, 1), padding=(0, 0, 0, 0), pool_type="max", ) lv1 = bb.emit_te(topi.nn.upsampling, lv0, scale_h=2.0, scale_w=2.0) gv = bb.emit_output((lv1, x)) bb.emit_func_output(gv) return bb.get() # The fusion is supposed to make no change. _check(before(), before()) def test_fuse_tuple_get_elemwise(): def before(dim: int): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, dim), "float32")) w = relax.Var("w", R.Tensor((3 * dim, dim), "float32")) with bb.function("main", [x, w]): with bb.dataflow(): lv0 = bb.emit_te(topi.nn.dense, x, w) lv1 = bb.emit_te(topi.split, lv0, indices_or_sections=3, axis=1) lv2 = bb.emit(relax.TupleGetItem(lv1, 0)) lv3 = bb.emit_te(topi.sigmoid, lv2) lv4 = bb.emit(relax.TupleGetItem(lv1, 1)) lv5 = bb.emit_te(topi.tanh, lv4) lv6 = bb.emit(relax.TupleGetItem(lv1, 2)) lv7 = bb.emit_te(topi.exp, lv6) lv8 = bb.emit_te(topi.multiply, lv5, lv7) gv = bb.emit_output(bb.call_te(topi.add, lv3, lv8)) bb.emit_func_output(gv) return bb.get() def expected(dim: int): bb = relax.BlockBuilder() # Grouped function dense = relax.Var("dense", R.Tensor((1, 3 * dim), "float32")) with bb.function( "fused_split_sigmoid_tanh_exp_multiply_add", [dense], attrs={"Primitive": True}, private=True, ): with bb.dataflow(): lv0 = bb.emit_te(topi.split, dense, indices_or_sections=3, axis=1) lv1 = bb.emit(relax.TupleGetItem(lv0, 0)) lv2 = bb.emit_te(topi.sigmoid, lv1) lv3 = bb.emit(relax.TupleGetItem(lv0, 1)) lv4 = bb.emit_te(topi.tanh, lv3) lv5 = bb.emit(relax.TupleGetItem(lv0, 2)) lv6 = bb.emit_te(topi.exp, lv5) lv7 = bb.emit_te(topi.multiply, lv4, lv6) gv = bb.emit_output(bb.call_te(topi.add, lv2, lv7)) bb.emit_func_output(gv) # Get the global variables of the grouped functions fused_split_sigmoid_tanh_exp_multiply_add = bb.get().get_global_var( "fused_split_sigmoid_tanh_exp_multiply_add" ) # Main function x = relax.Var("x", R.Tensor((1, dim), "float32")) w = relax.Var("w", R.Tensor((3 * dim, dim), "float32")) with bb.function("main", [x, w]): with bb.dataflow(): lv0 = bb.emit_te(topi.nn.dense, x, w) gv = bb.emit_output(relax.Call(fused_split_sigmoid_tanh_exp_multiply_add, (lv0,))) bb.emit_func_output(gv) return bb.get() dim = 10 _check(before(dim), expected(dim)) def test_tuple_get_root(): def before(dim: int): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 3 * dim), "float32")) w = relax.Var("w", R.Tensor((dim, dim), "float32")) with bb.function("main", [x, w]): with bb.dataflow(): lv0 = bb.emit_te(topi.split, x, indices_or_sections=3, axis=1) lv1 = bb.emit(relax.TupleGetItem(lv0, 0)) gv = bb.emit_output(bb.call_te(topi.nn.dense, lv1, w)) bb.emit_func_output(gv) return bb.get() def expected(dim: int): bb = relax.BlockBuilder() # Grouped function x = relax.Var("x", R.Tensor((1, 3 * dim), "float32")) with bb.function("fused_split", [x], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te(topi.split, x, indices_or_sections=3, axis=1) gv = bb.emit_output(relax.TupleGetItem(lv0, 0)) bb.emit_func_output(gv) # Get the global variables of the grouped functions fused_split = bb.get().get_global_var("fused_split") # Main function x = relax.Var("x", R.Tensor((1, 3 * dim), "float32")) w = relax.Var("w", R.Tensor((dim, dim), "float32")) with bb.function("main", [x, w]): with bb.dataflow(): lv0 = bb.emit(relax.Call(fused_split, (x,))) gv = bb.emit_output(bb.call_te(topi.nn.dense, lv0, w)) bb.emit_func_output(gv) return bb.get() dim = 10 _check(before(dim), expected(dim)) def test_tuple_intermediate(): def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.squeeze, x) lv1 = bb.emit_te(topi.add, lv0, relax.const(1, "float32")) lv2 = bb.emit_te(topi.squeeze, lv0) lv3 = bb.emit_te(topi.add, lv2, relax.const(1, "float32")) lv4 = bb.emit_te(topi.add, lv3, relax.const(1, "float32")) lv5 = bb.emit_te(topi.add, lv0, relax.const(1, "float32")) lv6 = bb.emit_te(topi.concatenate, (lv1, lv4, lv5), axis=1) lv7 = bb.emit_te(topi.squeeze, lv6) gv = bb.emit_output(bb.call_te(topi.add, lv7, relax.const(1, "float32"))) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() # Grouped function x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) p1 = relax.Var("p1", R.Tensor((), "float32")) p2 = relax.Var("p2", R.Tensor((), "float32")) p3 = relax.Var("p3", R.Tensor((), "float32")) p4 = relax.Var("p4", R.Tensor((), "float32")) with bb.function( "fused_squeeze_add_squeeze1_add_add_add_concatenate_squeeze2_add1", [x, p0, p1, p2, p3, p4], attrs={"Primitive": True}, private=True, ): with bb.dataflow(): lv0 = bb.emit_te(topi.squeeze, x) lv1 = bb.emit_te(topi.add, lv0, p0) lv2 = bb.emit_te(topi.squeeze, lv0) lv3 = bb.emit_te(topi.add, lv2, p1) lv4 = bb.emit_te(topi.add, lv3, p2) lv5 = bb.emit_te(topi.add, lv0, p3) lv6 = bb.emit_te(topi.concatenate, (lv1, lv4, lv5), axis=1) lv7 = bb.emit_te(topi.squeeze, lv6) gv = bb.emit_output(bb.call_te(topi.add, lv7, p4)) bb.emit_func_output(gv) # Get the global variables of the grouped functions fused_func = bb.get().get_global_var( "fused_squeeze_add_squeeze1_add_add_add_concatenate_squeeze2_add1" ) # Main func x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): gv = bb.emit_output( relax.Call( fused_func, ( x, relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), ), ) ) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_tuple_consecutive(): def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv1 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv2 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv3 = bb.emit_te(topi.concatenate, (lv0, lv1, lv2), axis=1) lv4 = bb.emit_te(topi.add, lv3, relax.const(1, "float32")) lv5 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv6 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv7 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv8 = bb.emit_te(topi.concatenate, (lv5, lv6, lv7), axis=1) lv9 = bb.emit_te(topi.add, lv8, relax.const(1, "float32")) lv10 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv11 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv12 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv13 = bb.emit_te(topi.concatenate, (lv10, lv11, lv12), axis=1) lv14 = bb.emit_te(topi.add, lv13, relax.const(1, "float32")) lv15 = bb.emit_te(topi.concatenate, (lv4, lv9, lv14), axis=1) lv16 = bb.emit_te( topi.nn.pool2d, lv15, kernel=(2, 2), stride=(2, 2), dilation=(1, 1), padding=(0, 0, 0, 0), pool_type="max", ) lv17 = bb.emit_te(topi.add, lv16, relax.const(1, "float32")) lv18 = bb.emit_te(topi.add, lv17, relax.const(1, "float32")) gv = bb.emit_output((lv17, lv18)) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() # Grouped function 1 x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) p1 = relax.Var("p1", R.Tensor((), "float32")) p2 = relax.Var("p2", R.Tensor((), "float32")) p3 = relax.Var("p3", R.Tensor((), "float32")) p4 = relax.Var("p4", R.Tensor((), "float32")) p5 = relax.Var("p5", R.Tensor((), "float32")) p6 = relax.Var("p6", R.Tensor((), "float32")) p7 = relax.Var("p7", R.Tensor((), "float32")) p8 = relax.Var("p8", R.Tensor((), "float32")) p9 = relax.Var("p9", R.Tensor((), "float32")) p10 = relax.Var("p10", R.Tensor((), "float32")) p11 = relax.Var("p11", R.Tensor((), "float32")) with bb.function( "fused_add_add_add_concatenate_add1_add_add_add_concatenate_add1_add_add_add_concatenate_add1_concatenate1", [x, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9, p10, p11], attrs={"Primitive": True}, private=True, ): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, p0) lv1 = bb.emit_te(topi.add, x, p1) lv2 = bb.emit_te(topi.add, x, p2) lv3 = bb.emit_te(topi.concatenate, (lv0, lv1, lv2), axis=1) lv4 = bb.emit_te(topi.add, lv3, p3) lv5 = bb.emit_te(topi.add, x, p4) lv6 = bb.emit_te(topi.add, x, p5) lv7 = bb.emit_te(topi.add, x, p6) lv8 = bb.emit_te(topi.concatenate, (lv5, lv6, lv7), axis=1) lv9 = bb.emit_te(topi.add, lv8, p7) lv10 = bb.emit_te(topi.add, x, p8) lv11 = bb.emit_te(topi.add, x, p9) lv12 = bb.emit_te(topi.add, x, p10) lv13 = bb.emit_te(topi.concatenate, (lv10, lv11, lv12), axis=1) lv14 = bb.emit_te(topi.add, lv13, p11) gv = bb.emit_output(bb.call_te(topi.concatenate, (lv4, lv9, lv14), axis=1)) bb.emit_func_output(gv) # Grouped function 2 concat = relax.Var("concat", R.Tensor((1, 144, 64, 64), "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) with bb.function( "fused_pool2d_add2", [concat, p0], attrs={"Primitive": True}, private=True ): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.pool2d, concat, kernel=(2, 2), stride=(2, 2), dilation=(1, 1), padding=(0, 0, 0, 0), pool_type="max", ) gv = bb.emit_output(bb.call_te(topi.add, lv0, p0)) bb.emit_func_output(gv) # Get the global variables of the grouped functions mod = bb.get() fused_func1 = mod.get_global_var( "fused_add_add_add_concatenate_add1_add_add_add_concatenate_add1_add_add_add_concatenate_add1_concatenate1" ) fused_func2 = mod.get_global_var("fused_pool2d_add2") # Main function x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit( relax.Call( fused_func1, ( x, relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), relax.const(1, "float32"), ), ) ) lv1 = bb.emit(relax.Call(fused_func2, (lv0, relax.const(1, "float32")))) lv2 = bb.emit_te(topi.add, lv1, relax.const(1, "float32")) gv = bb.emit_output((lv1, lv2)) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_inception_like(): def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) w0 = relax.Var("w0", R.Tensor((16, 16, 3, 3), "float32")) w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), "float32")) w2 = relax.Var("w2", R.Tensor((16, 32, 3, 3), "float32")) w3 = relax.Var("w3", R.Tensor((16, 32, 3, 3), "float32")) with bb.function("main", [x, w0, w1, w2, w3]): with bb.dataflow(): lv0 = bb.emit_te(topi.nn.conv2d, x, w0, strides=1, padding=1, dilation=1) lv1 = bb.emit_te(topi.nn.relu, lv0) lv2 = bb.emit_te(topi.nn.conv2d, x, w1, strides=1, padding=1, dilation=1) lv3 = bb.emit_te(topi.nn.relu, lv2) lv4 = bb.emit_te(topi.concatenate, (lv1, lv3), axis=1) lv5 = bb.emit_te(topi.nn.conv2d, lv4, w2, strides=1, padding=1, dilation=1) lv6 = bb.emit_te(topi.nn.relu, lv5) lv7 = bb.emit_te(topi.nn.conv2d, lv4, w3, strides=1, padding=1, dilation=1) lv8 = bb.emit_te(topi.nn.relu, lv7) gv = bb.emit_output(bb.call_te(topi.concatenate, (lv6, lv8), axis=1)) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() # Grouped function 1 x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) w = relax.Var("w", R.Tensor((16, 16, 3, 3), "float32")) with bb.function("fused_conv2d_relu", [x, w], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.conv2d, x, w, strides=1, padding=1, dilation=1, primfunc_name_hint="conv2d", ) gv = bb.emit_output(bb.call_te(topi.nn.relu, lv0)) bb.emit_func_output(gv) # Grouped function 2 x = relax.Var("x", R.Tensor((1, 32, 64, 64), "float32")) w = relax.Var("w", R.Tensor((16, 32, 3, 3), "float32")) with bb.function("fused_conv2d1_relu", [x, w], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te( topi.nn.conv2d, x, w, strides=1, padding=1, dilation=1, primfunc_name_hint="conv2d1", ) gv = bb.emit_output(bb.call_te(topi.nn.relu, lv0)) bb.emit_func_output(gv) # Get the global variables of the grouped functions mod = bb.get() fused_conv2d_relu1 = mod.get_global_var("fused_conv2d_relu") fused_conv2d_relu2 = mod.get_global_var("fused_conv2d1_relu") # Main function x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) w0 = relax.Var("w0", R.Tensor((16, 16, 3, 3), "float32")) w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), "float32")) w2 = relax.Var("w2", R.Tensor((16, 32, 3, 3), "float32")) w3 = relax.Var("w3", R.Tensor((16, 32, 3, 3), "float32")) with bb.function("main", [x, w0, w1, w2, w3]): with bb.dataflow(): lv0 = bb.emit(relax.Call(fused_conv2d_relu1, (x, w0))) lv1 = bb.emit(relax.Call(fused_conv2d_relu1, (x, w1))) lv2 = bb.emit_te(topi.concatenate, (lv0, lv1), axis=1) lv3 = bb.emit(relax.Call(fused_conv2d_relu2, (lv2, w2))) lv4 = bb.emit(relax.Call(fused_conv2d_relu2, (lv2, w3))) gv = bb.emit_output(bb.call_te(topi.concatenate, (lv3, lv4), axis=1)) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_fuse_parallel_injective(): def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((10, 20), "int32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, "int32")) lv1 = bb.emit_te(topi.squeeze, lv0) lv2 = bb.emit_te(topi.transpose, lv0, axes=[1, 0]) lv3 = bb.emit_te(topi.transpose, lv2, axes=[1, 0]) gv = bb.emit_output(bb.call_te(topi.left_shift, lv1, lv3)) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() # Grouped function x = relax.Var("x", R.Tensor((10, 20), "int32")) p0 = relax.Var("p0", R.Tensor((), "int32")) with bb.function( "fused_add_squeeze_transpose_transpose1_left_shift", [x, p0], attrs={"Primitive": True}, private=True, ): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, p0) lv1 = bb.emit_te(topi.squeeze, lv0) lv2 = bb.emit_te(topi.transpose, lv0, axes=[1, 0]) lv3 = bb.emit_te(topi.transpose, lv2, axes=[1, 0], primfunc_name_hint="transpose1") gv = bb.emit_output(bb.call_te(topi.left_shift, lv1, lv3)) bb.emit_func_output(gv) # Get the global variables of the grouped functions fused_func = bb.get().get_global_var("fused_add_squeeze_transpose_transpose1_left_shift") # Main function x = relax.Var("x", R.Tensor((10, 20), "int32")) with bb.function("main", [x]): with bb.dataflow(): gv = bb.emit_output(relax.Call(fused_func, (x, relax.const(1, "int32")))) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_softmax(): """Test if softmax can be fused with following ops.""" def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor((16, 16), "float32")) with bb.function("main", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.nn.softmax, x) gv = bb.emit_output(bb.call_te(topi.cast, lv0, dtype="float16")) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() # Grouped function x = relax.Var("x", R.Tensor((16, 16), "float32")) with bb.function("fused_softmax_cast", [x], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te(topi.nn.softmax, x) gv = bb.emit_output(bb.call_te(topi.cast, lv0, dtype="float16")) bb.emit_func_output(gv) # Get the global variables of the grouped functions fused_func = bb.get().get_global_var("fused_softmax_cast") # Main function x = relax.Var("x", R.Tensor((16, 16), "float32")) with bb.function("main", [x]): with bb.dataflow(): gv = bb.emit_output(relax.Call(fused_func, (x,))) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_multiple_relax_functions(): def before(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor([10, 20], "float32")) with bb.function("func1", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv1 = bb.emit_te(topi.exp, lv0) gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) bb.emit_func_output(gv) x = relax.Var("x", R.Tensor([20, 10], "float32")) with bb.function("func2", [x]): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32")) lv1 = bb.emit_te(topi.exp, lv0) gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) bb.emit_func_output(gv) return bb.get() def expected(): bb = relax.BlockBuilder() x = relax.Var("x", R.Tensor([10, 20], "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) with bb.function("fused_add_exp_squeeze", [x, p0], attrs={"Primitive": True}, private=True): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, p0) lv1 = bb.emit_te(topi.exp, lv0) gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) bb.emit_func_output(gv) fused_add_exp_squeeze = bb.get().get_global_var("fused_add_exp_squeeze") x = relax.Var("x", R.Tensor([20, 10], "float32")) p0 = relax.Var("p0", R.Tensor((), "float32")) with bb.function( "fused_add1_exp1_squeeze1", [x, p0], attrs={"Primitive": True}, private=True ): with bb.dataflow(): lv0 = bb.emit_te(topi.add, x, p0) lv1 = bb.emit_te(topi.exp, lv0) gv = bb.emit_output(bb.call_te(topi.squeeze, lv1)) bb.emit_func_output(gv) fused_add1_exp1_squeeze1 = bb.get().get_global_var("fused_add1_exp1_squeeze1") x = relax.Var("x", R.Tensor([10, 20], "float32")) with bb.function("func1", [x]): with bb.dataflow(): gv = bb.emit_output( relax.Call(fused_add_exp_squeeze, [x, relax.const(1, "float32")]) ) bb.emit_func_output(gv) x = relax.Var("x", R.Tensor([20, 10], "float32")) with bb.function("func2", [x]): with bb.dataflow(): gv = bb.emit_output( relax.Call(fused_add1_exp1_squeeze1, [x, relax.const(1, "float32")]) ) bb.emit_func_output(gv) return bb.get() _check(before(), expected()) def test_skip_call_dps_packed(): @I.ir_module(s_tir=True) class Module: @R.function def main(x: R.Tensor((2, 3), "float32")): with R.dataflow(): y = R.call_dps_packed("func_packed_dps", x, R.Tensor((2, 3), "float32")) R.output(y) return y # FuseOps should does no change to it. _check(Module, Module) def test_edge_with_call_dps_packed(): @I.ir_module(s_tir=True) class Module: @R.function def main(x: R.Tensor((2, 3), "float32")): cls = Module with R.dataflow(): a = R.call_tir(cls.exp, (x,), out_ty=R.Tensor((2, 3), "float32")) b = R.call_tir(cls.exp, (a,), out_ty=R.Tensor((2, 3), "float32")) c = R.call_dps_packed("packed_dps", (a,), out_ty=R.Tensor((2, 3), "float32")) R.output(b, c) return R.tuple(b, c) @T.prim_func(private=True, s_tir=True) def exp(A: T.Buffer((2, 3), "float32"), B: T.Buffer((2, 3), "float32")): T.evaluate(0) # FuseOps should does no change to it. _check(Module, Module) def test_layer_norm_silu(): # fmt: off @I.ir_module(s_tir=True) class Module: @R.function def main(x: R.Tensor((1, 512, 64, 64), "float32"), mean: R.Tensor((64, 64), "float32"), var: R.Tensor((64, 64), "float32")): cls = Module with R.dataflow(): gv0 = R.call_tir(cls.layer_norm, (x, mean, var), out_ty=R.Tensor((1, 512, 64, 64), 'float32')) gv1 = R.call_tir(cls.relu, gv0, out_ty=R.Tensor((1, 512, 64, 64), "float32")) R.output(gv1) return gv1 @T.prim_func(private=True, s_tir=True) def layer_norm(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), gamma: T.Buffer((T.int64(64), T.int64(64)), "float32"), beta: T.Buffer((T.int64(64), T.int64(64)), "float32"), T_layer_norm: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")): rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer([T.int64(64), T.int64(64)], dtype="float32") rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer([T.int64(64), T.int64(64)], dtype="float32") for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)): with T.sblock("rxplaceholder_red_temp"): ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3]) T.reads(A[ax0, ax1, k2, k3]) T.writes(rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1]) with T.init(): rxplaceholder_red_temp_v0[ax0, ax1] = T.float32(0) rxplaceholder_red_temp_v1[ax0, ax1] = T.float32(0) v_rxplaceholder_red_temp_v0: T.let[T.float32] = rxplaceholder_red_temp_v0[ax0, ax1] + A[ax0, ax1, k2, k3] v_rxplaceholder_red_temp_v1: T.let[T.float32] = rxplaceholder_red_temp_v1[ax0, ax1] + A[ax0, ax1, k2, k3] * A[ax0, ax1, k2, k3] rxplaceholder_red_temp_v0[ax0, ax1] = v_rxplaceholder_red_temp_v0 rxplaceholder_red_temp_v1[ax0, ax1] = v_rxplaceholder_red_temp_v1 for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)): with T.sblock("T_layer_norm"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(A[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], gamma[ax2, ax3], beta[ax2, ax3]) T.writes(T_layer_norm[ax0, ax1, ax2, ax3]) T_layer_norm[ax0, ax1, ax2, ax3] = (A[ax0, ax1, ax2, ax3] - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05)) * T.rsqrt(rxplaceholder_red_temp_v1[ax0, ax1] * T.float32(0.05) - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05) * (rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.05)) + T.float32(1e-05), dtype="float32") * gamma[ax2, ax3] + beta[ax2, ax3] @T.prim_func(private=True, s_tir=True) def relu(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), B: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")): for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)): with T.sblock("relu"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(A[v_i0, v_i1, v_i2, v_i3]) T.writes(B[v_i0, v_i1, v_i2, v_i3]) B[v_i0, v_i1, v_i2, v_i3] = T.max(A[v_i0, v_i1, v_i2, v_i3], T.float32(0)) @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def layer_norm(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), gamma: T.Buffer((T.int64(64), T.int64(64)), "float32"), beta: T.Buffer((T.int64(64), T.int64(64)), "float32"), T_layer_norm: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")): T.func_attr({"op_pattern": 4}) # with T.sblock("root"): rxplaceholder_red_temp_v0 = T.sblock_alloc_buffer((T.int64(64), T.int64(64))) rxplaceholder_red_temp_v1 = T.sblock_alloc_buffer((T.int64(64), T.int64(64))) for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)): with T.sblock("rxplaceholder_red_temp"): ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3]) T.reads(A[ax0, ax1, k2, k3]) T.writes(rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1]) with T.init(): rxplaceholder_red_temp_v0[ax0, ax1] = T.float32(0) rxplaceholder_red_temp_v1[ax0, ax1] = T.float32(0) v_rxplaceholder_red_temp_v0: T.let[T.float32] = rxplaceholder_red_temp_v0[ax0, ax1] + A[ax0, ax1, k2, k3] v_rxplaceholder_red_temp_v1: T.let[T.float32] = rxplaceholder_red_temp_v1[ax0, ax1] + A[ax0, ax1, k2, k3] * A[ax0, ax1, k2, k3] rxplaceholder_red_temp_v0[ax0, ax1] = v_rxplaceholder_red_temp_v0 rxplaceholder_red_temp_v1[ax0, ax1] = v_rxplaceholder_red_temp_v1 for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)): with T.sblock("T_layer_norm"): ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(A[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], gamma[ax2, ax3], beta[ax2, ax3]) T.writes(T_layer_norm[ax0, ax1, ax2, ax3]) T_layer_norm[ax0, ax1, ax2, ax3] = (A[ax0, ax1, ax2, ax3] - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003)) * T.rsqrt(rxplaceholder_red_temp_v1[ax0, ax1] * T.float32(0.050000000000000003) - rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003) * (rxplaceholder_red_temp_v0[ax0, ax1] * T.float32(0.050000000000000003)) + T.float32(1.0000000000000001e-05)) * gamma[ax2, ax3] + beta[ax2, ax3] @T.prim_func(private=True, s_tir=True) def relu(A: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32"), B: T.Buffer((T.int64(1), T.int64(512), T.int64(64), T.int64(64)), "float32")): T.func_attr({"op_pattern": 0}) # with T.sblock("root"): for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(512), T.int64(64), T.int64(64)): with T.sblock("relu"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(A[v_i0, v_i1, v_i2, v_i3]) T.writes(B[v_i0, v_i1, v_i2, v_i3]) B[v_i0, v_i1, v_i2, v_i3] = T.max(A[v_i0, v_i1, v_i2, v_i3], T.float32(0)) @R.function(private=True) def fused_layer_norm_relu(x: R.Tensor((1, 512, 64, 64), dtype="float32"), mean: R.Tensor((64, 64), dtype="float32"), var: R.Tensor((64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"): R.func_attr({"Primitive": True}) cls = Expected with R.dataflow(): gv0 = R.call_tir(cls.layer_norm, (x, mean, var), out_ty=R.Tensor((1, 512, 64, 64), 'float32')) gv = R.call_tir(cls.relu, (gv0,), out_ty=R.Tensor((1, 512, 64, 64), dtype="float32")) R.output(gv) return gv @R.function def main(x: R.Tensor((1, 512, 64, 64), dtype="float32"), mean: R.Tensor((64, 64), dtype="float32"), var: R.Tensor((64, 64), dtype="float32")) -> R.Tensor((1, 512, 64, 64), dtype="float32"): cls = Expected with R.dataflow(): gv: R.Tensor((1, 512, 64, 64), dtype="float32") = cls.fused_layer_norm_relu(x, mean, var) R.output(gv) return gv # fmt: on _check(Module, Expected) def test_multiple_paths(): # fmt: off @I.ir_module(s_tir=True) class Module: @R.function def main( inp_0: R.Tensor((2, 320, 64, 64), dtype="float32"), inp_1: R.Tensor((2, 1280), dtype="float32"), w1: R.Tensor((320, 320, 3, 3), dtype="float32"), b1: R.Tensor((320,), "float32"), w2: R.Tensor((320, 1280), "float32"), b2: R.Tensor((320,), "float32"), ): R.func_attr({"num_input": 2}) with R.dataflow(): lv27: R.Tensor((2, 320, 64, 64), dtype="float32") = R.nn.conv2d(inp_0, w1, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype="float32") lv28: R.Tensor((1, 320, 1, 1), dtype="float32") = R.reshape(b1, R.shape([1, 320, 1, 1])) ## lv29: R.Tensor((2, 320, 64, 64), dtype="float32") = R.add(lv27, lv28) lv31: R.Tensor((1280, 320), dtype="float32") = R.permute_dims(w2, axes=None) ## lv32: R.Tensor((2, 320), dtype="float32") = R.matmul(inp_1, lv31, out_dtype="float32") lv33: R.Tensor((2, 320), dtype="float32") = R.add(lv32, b2) lv35: R.Tensor((2, 320, 1, 1), dtype="float32") = R.reshape(lv33, R.shape([2, 320, 1, 1])) lv36: R.Tensor((2, 320, 64, 64), dtype="float32") = R.add(lv29, lv35) gv = lv36 R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def add(rxplaceholder: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1), T.int64(320), T.int64(1), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32")): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(320), T.int64(64), T.int64(64)): with T.sblock("T_add"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder_1[T.int64(0), v_ax1, T.int64(0), T.int64(0)]) T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3]) T_add[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder_1[T.int64(0), v_ax1, T.int64(0), T.int64(0)] @T.prim_func(private=True, s_tir=True) def add1(rxplaceholder: T.Buffer((T.int64(2), T.int64(320)), "float32"), rxplaceholder_1: T.Buffer((T.int64(320),), "float32"), T_add: T.Buffer((T.int64(2), T.int64(320)), "float32")): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(2), T.int64(320)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1]) T.writes(T_add[v_ax0, v_ax1]) T_add[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax1] @T.prim_func(private=True, s_tir=True) def add2(rxplaceholder: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32"), rxplaceholder_1: T.Buffer((T.int64(2), T.int64(320), T.int64(1), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32")): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(320), T.int64(64), T.int64(64)): with T.sblock("T_add"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3], rxplaceholder_1[v_ax0, v_ax1, T.int64(0), T.int64(0)]) T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3]) T_add[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0, v_ax1, v_ax2, v_ax3] + rxplaceholder_1[v_ax0, v_ax1, T.int64(0), T.int64(0)] @T.prim_func(private=True, s_tir=True) def conv2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32"), rxplaceholder_1: T.Buffer((T.int64(320), T.int64(320), T.int64(3), T.int64(3)), "float32"), conv2d_nchw: T.Buffer((T.int64(2), T.int64(320), T.int64(64), T.int64(64)), "float32")): T.func_attr({"op_pattern": 4, "tirx.noalias": True}) pad_temp = T.sblock_alloc_buffer((T.int64(2), T.int64(320), T.int64(66), T.int64(66))) for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(320), T.int64(66), T.int64(66)): with T.sblock("pad_temp"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(rxplaceholder[v_i0, v_i1, v_i2 - T.int64(1), v_i3 - T.int64(1)]) T.writes(pad_temp[v_i0, v_i1, v_i2, v_i3]) pad_temp[v_i0, v_i1, v_i2, v_i3] = T.if_then_else(T.int64(1) <= v_i2 and v_i2 < T.int64(65) and T.int64(1) <= v_i3 and v_i3 < T.int64(65), rxplaceholder[v_i0, v_i1, v_i2 - T.int64(1), v_i3 - T.int64(1)], T.float32(0)) for nn, ff, yy, xx, rc, ry, rx in T.grid(T.int64(2), T.int64(320), T.int64(64), T.int64(64), T.int64(320), T.int64(3), T.int64(3)): with T.sblock("conv2d_nchw"): v_nn, v_ff, v_yy, v_xx, v_rc, v_ry, v_rx = T.axis.remap("SSSSRRR", [nn, ff, yy, xx, rc, ry, rx]) T.reads(pad_temp[v_nn, v_rc, v_yy + v_ry, v_xx + v_rx], rxplaceholder_1[v_ff, v_rc, v_ry, v_rx]) T.writes(conv2d_nchw[v_nn, v_ff, v_yy, v_xx]) with T.init(): conv2d_nchw[v_nn, v_ff, v_yy, v_xx] = T.float32(0) conv2d_nchw[v_nn, v_ff, v_yy, v_xx] = conv2d_nchw[v_nn, v_ff, v_yy, v_xx] + pad_temp[v_nn, v_rc, v_yy + v_ry, v_xx + v_rx] * rxplaceholder_1[v_ff, v_rc, v_ry, v_rx] @T.prim_func(private=True, s_tir=True) def matmul(rxplaceholder: T.Buffer((T.int64(2), T.int64(1280)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1280), T.int64(320)), "float32"), matmul: T.Buffer((T.int64(2), T.int64(320)), "float32")): T.func_attr({"op_pattern": 4, "tirx.noalias": True}) for i0, i1, k in T.grid(T.int64(2), T.int64(320), T.int64(1280)): with T.sblock("matmul"): v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k]) T.reads(rxplaceholder[v_i0, v_k], rxplaceholder_1[v_k, v_i1]) T.writes(matmul[v_i0, v_i1]) with T.init(): matmul[v_i0, v_i1] = T.float32(0) matmul[v_i0, v_i1] = matmul[v_i0, v_i1] + rxplaceholder[v_i0, v_k] * rxplaceholder_1[v_k, v_i1] @T.prim_func(private=True, s_tir=True) def reshape(rxplaceholder: T.Buffer((T.int64(320),), "float32"), T_reshape: T.Buffer((T.int64(1), T.int64(320), T.int64(1), T.int64(1)), "float32")): T.func_attr({"op_pattern": 2, "tirx.noalias": True}) for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(320), T.int64(1), T.int64(1)): with T.sblock("T_reshape"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[(v_ax1 + v_ax2 + v_ax3) % T.int64(320)]) T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3]) T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[(v_ax1 + v_ax2 + v_ax3) % T.int64(320)] @T.prim_func(private=True, s_tir=True) def reshape1(rxplaceholder: T.Buffer((T.int64(2), T.int64(320)), "float32"), T_reshape: T.Buffer((T.int64(2), T.int64(320), T.int64(1), T.int64(1)), "float32")): T.func_attr({"op_pattern": 2, "tirx.noalias": True}) for ax0, ax1, ax2, ax3 in T.grid(T.int64(2), T.int64(320), T.int64(1), T.int64(1)): with T.sblock("T_reshape"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(rxplaceholder[((v_ax1 + v_ax2 + v_ax3) // T.int64(320) + v_ax0) % T.int64(2), (v_ax1 + v_ax2 + v_ax3) % T.int64(320)]) T.writes(T_reshape[v_ax0, v_ax1, v_ax2, v_ax3]) T_reshape[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[((v_ax1 + v_ax2 + v_ax3) // T.int64(320) + v_ax0) % T.int64(2), (v_ax1 + v_ax2 + v_ax3) % T.int64(320)] @T.prim_func(private=True, s_tir=True) def transpose(rxplaceholder: T.Buffer((T.int64(320), T.int64(1280)), "float32"), T_transpose: T.Buffer((T.int64(1280), T.int64(320)), "float32")): T.func_attr({"op_pattern": 2, "tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(1280), T.int64(320)): with T.sblock("T_transpose"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(rxplaceholder[v_ax1, v_ax0]) T.writes(T_transpose[v_ax0, v_ax1]) T_transpose[v_ax0, v_ax1] = rxplaceholder[v_ax1, v_ax0] @R.function(private=True) def fused_conv2d_add_add2(inp_0: R.Tensor((2, 320, 64, 64), dtype="float32"), w1: R.Tensor((320, 320, 3, 3), dtype="float32"), lv28: R.Tensor((1, 320, 1, 1), dtype="float32"), lv35: R.Tensor((2, 320, 1, 1), dtype="float32")) -> R.Tensor((2, 320, 64, 64), dtype="float32"): R.func_attr({"Primitive": True}) cls = Expected with R.dataflow(): lv27 = R.call_tir(cls.conv2d, (inp_0, w1), out_ty=R.Tensor((2, 320, 64, 64), dtype="float32")) lv29 = R.call_tir(cls.add, (lv27, lv28), out_ty=R.Tensor((2, 320, 64, 64), dtype="float32")) gv = R.call_tir(cls.add2, (lv29, lv35), out_ty=R.Tensor((2, 320, 64, 64), dtype="float32")) R.output(gv) return gv @R.function(private=True) def fused_matmul_add1(inp_1: R.Tensor((2, 1280), dtype="float32"), lv31: R.Tensor((1280, 320), dtype="float32"), b2: R.Tensor((320,), dtype="float32")) -> R.Tensor((2, 320), dtype="float32"): cls = Expected R.func_attr({"Primitive": True}) with R.dataflow(): lv32 = R.call_tir(cls.matmul, (inp_1, lv31), out_ty=R.Tensor((2, 320), dtype="float32")) gv = R.call_tir(cls.add1, (lv32, b2), out_ty=R.Tensor((2, 320), dtype="float32")) R.output(gv) return gv @R.function def main(inp_0: R.Tensor((2, 320, 64, 64), dtype="float32"), inp_1: R.Tensor((2, 1280), dtype="float32"), w1: R.Tensor((320, 320, 3, 3), dtype="float32"), b1: R.Tensor((320,), dtype="float32"), w2: R.Tensor((320, 1280), dtype="float32"), b2: R.Tensor((320,), dtype="float32")) -> R.Tensor((2, 320, 64, 64), dtype="float32"): R.func_attr({"num_input": 2}) cls = Expected with R.dataflow(): lv28 = R.call_tir(cls.reshape, (b1,), out_ty=R.Tensor((1, 320, 1, 1), dtype="float32")) lv31 = R.call_tir(cls.transpose, (w2,), out_ty=R.Tensor((1280, 320), dtype="float32")) lv: R.Tensor((2, 320), dtype="float32") = cls.fused_matmul_add1(inp_1, lv31, b2) lv35 = R.call_tir(cls.reshape1, (lv,), out_ty=R.Tensor((2, 320, 1, 1), dtype="float32")) lv1: R.Tensor((2, 320, 64, 64), dtype="float32") = cls.fused_conv2d_add_add2(inp_0, w1, lv28, lv35) gv: R.Tensor((2, 320, 64, 64), dtype="float32") = lv1 R.output(gv) return gv # fmt: on mod = relax.transform.LegalizeOps()(Module) mod = relax.transform.AnnotateTIROpPattern()(mod) mod = relax.transform.FuseOps()(mod) tvm.ir.assert_structural_equal(mod, Expected) def test_dead_group(): # fmt: off @I.ir_module(s_tir=True) class Module: @R.function def main(inp_0: R.Tensor((1, 784), dtype="float32"), inp_1: R.Tensor((1, 128), dtype="float32"), linear1_bias: R.Tensor((128,), dtype="float32"), linear1_weight: R.Tensor((128, 784), dtype="float32"), linear2_bias: R.Tensor((10,), dtype="float32"), linear2_weight: R.Tensor((10, 128), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"): R.func_attr({"num_input": 1}) with R.dataflow(): lv: R.Tensor((784, 128), dtype="float32") = R.permute_dims(linear1_weight, axes=None) lv1: R.Tensor((1, 128), dtype="float32") = R.matmul(inp_0, lv, out_dtype="float32") lv2: R.Tensor((1, 128), dtype="float32") = R.add(lv1, linear1_bias) lv3: R.Tensor((1, 128), dtype="float32") = R.nn.relu(lv2) lv4: R.Tensor((128, 10), dtype="float32") = R.permute_dims(linear2_weight, axes=None) lv5: R.Tensor((1, 10), dtype="float32") = R.matmul(inp_1, lv4, out_dtype="float32") lv6: R.Tensor((1, 10), dtype="float32") = R.add(lv5, linear2_bias) gv: R.Tensor((1, 10), dtype="float32") = lv6 R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(128)), "float32"), rxplaceholder_1: T.Buffer((T.int64(128),), "float32"), T_add: T.Buffer((T.int64(1), T.int64(128)), "float32")): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1 in T.grid(T.int64(1), T.int64(128)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1]) T.writes(T_add[v_ax0, v_ax1]) T_add[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax1] @T.prim_func(private=True, s_tir=True) def add1(rxplaceholder: T.Buffer((T.int64(1), T.int64(10)), "float32"), rxplaceholder_1: T.Buffer((T.int64(10),), "float32"), T_add: T.Buffer((T.int64(1), T.int64(10)), "float32")): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1 in T.grid(T.int64(1), T.int64(10)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(rxplaceholder[v_ax0, v_ax1], rxplaceholder_1[v_ax1]) T.writes(T_add[v_ax0, v_ax1]) T_add[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] + rxplaceholder_1[v_ax1] @T.prim_func(private=True, s_tir=True) def matmul(rxplaceholder: T.Buffer((T.int64(1), T.int64(784)), "float32"), rxplaceholder_1: T.Buffer((T.int64(784), T.int64(128)), "float32"), matmul_1: T.Buffer((T.int64(1), T.int64(128)), "float32")): T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # with T.sblock("root"): for i0, i1, k in T.grid(T.int64(1), T.int64(128), T.int64(784)): with T.sblock("matmul"): v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k]) T.reads(rxplaceholder[v_i0, v_k], rxplaceholder_1[v_k, v_i1]) T.writes(matmul_1[v_i0, v_i1]) with T.init(): matmul_1[v_i0, v_i1] = T.float32(0) matmul_1[v_i0, v_i1] = matmul_1[v_i0, v_i1] + rxplaceholder[v_i0, v_k] * rxplaceholder_1[v_k, v_i1] @T.prim_func(private=True, s_tir=True) def matmul1(rxplaceholder: T.Buffer((T.int64(1), T.int64(128)), "float32"), rxplaceholder_1: T.Buffer((T.int64(128), T.int64(10)), "float32"), matmul: T.Buffer((T.int64(1), T.int64(10)), "float32")): T.func_attr({"op_pattern": 4, "tirx.noalias": True}) # with T.sblock("root"): for i0, i1, k in T.grid(T.int64(1), T.int64(10), T.int64(128)): with T.sblock("matmul"): v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k]) T.reads(rxplaceholder[v_i0, v_k], rxplaceholder_1[v_k, v_i1]) T.writes(matmul[v_i0, v_i1]) with T.init(): matmul[v_i0, v_i1] = T.float32(0) matmul[v_i0, v_i1] = matmul[v_i0, v_i1] + rxplaceholder[v_i0, v_k] * rxplaceholder_1[v_k, v_i1] @T.prim_func(private=True, s_tir=True) def relu(rxplaceholder: T.Buffer((T.int64(1), T.int64(128)), "float32"), compute: T.Buffer((T.int64(1), T.int64(128)), "float32")): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) # with T.sblock("root"): for i0, i1 in T.grid(T.int64(1), T.int64(128)): with T.sblock("compute"): v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) T.reads(rxplaceholder[v_i0, v_i1]) T.writes(compute[v_i0, v_i1]) compute[v_i0, v_i1] = T.max(rxplaceholder[v_i0, v_i1], T.float32(0)) @T.prim_func(private=True, s_tir=True) def transpose(rxplaceholder: T.Buffer((T.int64(128), T.int64(784)), "float32"), T_transpose: T.Buffer((T.int64(784), T.int64(128)), "float32")): T.func_attr({"op_pattern": 2, "tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1 in T.grid(T.int64(784), T.int64(128)): with T.sblock("T_transpose"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(rxplaceholder[v_ax1, v_ax0]) T.writes(T_transpose[v_ax0, v_ax1]) T_transpose[v_ax0, v_ax1] = rxplaceholder[v_ax1, v_ax0] @T.prim_func(private=True, s_tir=True) def transpose1(rxplaceholder: T.Buffer((T.int64(10), T.int64(128)), "float32"), T_transpose: T.Buffer((T.int64(128), T.int64(10)), "float32")): T.func_attr({"op_pattern": 2, "tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1 in T.grid(T.int64(128), T.int64(10)): with T.sblock("T_transpose"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(rxplaceholder[v_ax1, v_ax0]) T.writes(T_transpose[v_ax0, v_ax1]) T_transpose[v_ax0, v_ax1] = rxplaceholder[v_ax1, v_ax0] @R.function(private=True) def fused_matmul1_add1(inp_1: R.Tensor((1, 128), dtype="float32"), lv4: R.Tensor((128, 10), dtype="float32"), linear2_bias: R.Tensor((10,), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"): R.func_attr({"Primitive": True}) cls = Expected with R.dataflow(): lv5 = R.call_tir(cls.matmul1, (inp_1, lv4), out_ty=R.Tensor((1, 10), dtype="float32")) gv = R.call_tir(cls.add1, (lv5, linear2_bias), out_ty=R.Tensor((1, 10), dtype="float32")) R.output(gv) return gv @R.function def main(inp_0: R.Tensor((1, 784), dtype="float32"), inp_1: R.Tensor((1, 128), dtype="float32"), linear1_bias: R.Tensor((128,), dtype="float32"), linear1_weight: R.Tensor((128, 784), dtype="float32"), linear2_bias: R.Tensor((10,), dtype="float32"), linear2_weight: R.Tensor((10, 128), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"): R.func_attr({"num_input": 1}) cls = Expected with R.dataflow(): lv = R.call_tir(cls.transpose, (linear1_weight,), out_ty=R.Tensor((784, 128), dtype="float32")) lv4 = R.call_tir(cls.transpose1, (linear2_weight,), out_ty=R.Tensor((128, 10), dtype="float32")) lv_1: R.Tensor((1, 10), dtype="float32") = cls.fused_matmul1_add1(inp_1, lv4, linear2_bias) gv: R.Tensor((1, 10), dtype="float32") = lv_1 R.output(gv) return gv # fmt: on mod = relax.transform.LegalizeOps()(Module) _check(mod, Expected) def test_symbolic_shape_aware_fuse(): @I.ir_module(s_tir=True) class Before: @R.function def main(x: R.Tensor(["n", "m"], "float32")): with R.dataflow(): lv0 = R.emit_te(topi.add, x, R.const(1, "float32")) lv1 = R.emit_te(topi.exp, lv0) gv = R.emit_te(topi.squeeze, lv1) R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @R.function(private=True) def fused_add_exp_squeeze( x: R.Tensor(["n", "m"], "float32"), p0: R.Tensor([], "float32") ) -> R.Tensor(["n", "m"], dtype="float32"): R.func_attr({"Primitive": True}) with R.dataflow(): lv0 = R.emit_te(topi.add, x, p0) lv1 = R.emit_te(topi.exp, lv0) gv = R.emit_te(topi.squeeze, lv1) R.output(gv) return gv @R.function def main(x: R.Tensor(["n", "m"], "float32")) -> R.Tensor(["n", "m"], dtype="float32"): cls = Expected with R.dataflow(): gv = cls.fused_add_exp_squeeze(x, R.const(1, "float32")) R.output(gv) return gv _check(Before, Expected) def test_symbolic_shape_aware_fuse_2(): @I.ir_module(s_tir=True) class Before: @R.function def main(s: R.Shape(["n"])): n = T.int64() with R.dataflow(): lv0 = R.emit_te(topi.full, [n, n], "float32", 0) lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True) gv = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n]) R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @R.function(private=True) def fused_full_trilu_broadcast_to( s: R.Shape(["n"]), ) -> R.Tensor([1, 1, "n", "n"], "float32"): R.func_attr({"Primitive": True}) n = T.int64() with R.dataflow(): lv0 = R.emit_te(topi.full, [n, n], "float32", 0) lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True) gv = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n]) R.output(gv) return gv @R.function def main(s: R.Shape(["n"])) -> R.Tensor((1, 1, "n", "n"), dtype="float32"): cls = Expected n = T.int64() with R.dataflow(): gv: R.Tensor([1, 1, n, n], "float32") = cls.fused_full_trilu_broadcast_to( R.shape([n]) ) R.output(gv) return gv _check(Before, Expected) def test_shape_expr_arg(): @I.ir_module(s_tir=True) class Before: @R.function def main(s: R.Shape(["n"]), kv_cache: R.Any): n = T.int64() with R.dataflow(): lv0 = R.emit_te(topi.full, [n, n], "float32", 0) lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True) lv2 = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n]) gv = R.call_pure_packed( "vm.builtin.attention_kv_cache_view", kv_cache, R.shape([1 + n, 32, 128]), ty_args=(R.Tensor((1 + n, 32, 128), dtype="float32"),), ) R.output(gv, lv2) return gv, lv2 @I.ir_module(s_tir=True) class Expected: @R.function(private=True) def fused_full_trilu_broadcast_to( s: R.Shape(["n"]), ) -> R.Tensor([1, 1, "n", "n"], "float32"): R.func_attr({"Primitive": True}) n = T.int64() with R.dataflow(): lv0 = R.emit_te(topi.full, [n, n], "float32", 0) lv1 = R.emit_te(topi.trilu, lv0, tvm.tirx.const(1, "int32"), upper=True) gv = R.emit_te(topi.broadcast_to, lv1, [1, 1, n, n]) R.output(gv) return gv @R.function def main(s: R.Shape(["n"]), kv_cache: R.Any): cls = Expected n = T.int64() with R.dataflow(): lv: R.Tensor([1, 1, n, n], "float32") = cls.fused_full_trilu_broadcast_to( R.shape([n]) ) gv = R.call_pure_packed( "vm.builtin.attention_kv_cache_view", kv_cache, R.shape([1 + n, 32, 128]), ty_args=(R.Tensor((1 + n, 32, 128), dtype="float32"),), ) R.output(gv, lv) return gv, lv _check(Before, Expected) def test_skipping_match_cast(): @I.ir_module(s_tir=True) class Module: @R.function def main(A: R.Tensor((10, 20), dtype="float32")) -> R.Tensor(dtype="float32", ndim=2): m = T.int64() n = T.int64() with R.dataflow(): lv: R.Tensor((m, n), dtype="float32") = R.match_cast( A, R.Tensor((m, n), dtype="float32") ) gv: R.Tensor((m, n), dtype="float32") = lv R.output(gv) return gv _check(Module, Module) def test_skipping_primvalue(): @I.ir_module(s_tir=True) class Module: @R.function def main(inp: R.Tensor((2, 2), dtype="float32")) -> R.Tensor((2, 2), dtype="float32"): with R.dataflow(): lv = R.call_pure_packed( "my_func1", inp, R.prim_value(0), ty_args=[R.Tensor((2, 2), dtype="float32")] ) lv1 = R.call_pure_packed( "my_func2", lv, R.str("str"), ty_args=[R.Tensor((2, 2), dtype="float32")] ) gv = R.call_pure_packed( "my_func3", lv1, R.dtype("float32"), ty_args=[R.Tensor((2, 2), dtype="float32")], ) R.output(gv) return gv _check(Module, Module) def test_partially_used_tuple_param(): @I.ir_module(s_tir=True) class Module: @R.function def main( x: R.Tuple( R.Tensor((2,), "float32"), R.Tensor((2,), "float32"), R.Tensor((2,), "float32"), R.Tensor((2,), "float32"), R.Tensor((2,), "float32"), R.Tensor((2,), "float32"), ), ): with R.dataflow(): x0 = x[0] y0 = R.emit_te(topi.add, x0, R.const(1, "float32")) y1 = R.emit_te(topi.divide, y0, R.const(1, "float32")) gv = y1 R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @R.function(private=True) def fused_add_divide( x_0: R.Tensor((2,), dtype="float32"), param_0: R.Tensor((), dtype="float32"), param_1: R.Tensor((), dtype="float32"), ) -> R.Tensor((2,), dtype="float32"): R.func_attr({"Primitive": True}) with R.dataflow(): y0 = R.emit_te(topi.add, x_0, param_0) gv = R.emit_te(topi.divide, y0, param_1) R.output(gv) return gv @R.function def main( x: R.Tuple( R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), R.Tensor((2,), dtype="float32"), ), ) -> R.Tensor((2,), dtype="float32"): cls = Expected with R.dataflow(): lv: R.Tensor((2,), dtype="float32") = x[0] lv1: R.Tensor((2,), dtype="float32") = cls.fused_add_divide( lv, R.const(1, "float32"), R.const(1, "float32") ) gv: R.Tensor((2,), dtype="float32") = lv1 R.output(gv) return gv _check(Module, Expected) def test_call_tir_inplace(): @I.ir_module(s_tir=True) class Module: @T.prim_func(private=True, s_tir=True) def add( A: T.Buffer((T.int64(10), T.int64(20)), "float32"), B: T.Buffer((), "float32"), Out: T.Buffer((T.int64(10), T.int64(20)), "float32"), ): T.func_attr({"tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(10), T.int64(20)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[()]) T.writes(Out[v_ax0, v_ax1]) Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()] @T.prim_func(private=True, s_tir=True) def exp_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")): T.func_attr({"tirx.noalias": True}) for i0, i1 in T.grid(T.int64(10), T.int64(20)): with T.sblock("compute"): v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) T.reads(A[v_i0, v_i1]) T.writes(A[v_i0, v_i1]) A[v_i0, v_i1] = T.exp(A[v_i0, v_i1]) @T.prim_func(private=True, s_tir=True) def squeeze_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")): T.func_attr({"tirx.noalias": True}) for ax0, ax1 in T.grid(T.int64(10), T.int64(20)): with T.sblock("T_squeeze"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] @R.function def main( x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32") ) -> R.Tensor((10, 20), dtype="float32"): cls = Module with R.dataflow(): lv = R.call_tir( cls.add, (x, p0), out_ty=R.Tensor((10, 20), dtype="float32"), ) lv1 = R.call_tir_inplace( cls.exp_inplace, (lv,), inplace_indices=[0], out_ty=R.Tensor((10, 20), dtype="float32"), ) gv = R.call_tir_inplace( cls.squeeze_inplace, (lv1,), inplace_indices=[0], out_ty=R.Tensor((10, 20), dtype="float32"), ) R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @T.prim_func(private=True, s_tir=True) def add( A: T.Buffer((T.int64(10), T.int64(20)), "float32"), B: T.Buffer((), "float32"), Out: T.Buffer((T.int64(10), T.int64(20)), "float32"), ): T.func_attr({"tirx.noalias": True, "op_pattern": 0}) for ax0, ax1 in T.grid(T.int64(10), T.int64(20)): with T.sblock("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1], B[()]) T.writes(Out[v_ax0, v_ax1]) Out[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[()] @T.prim_func(private=True, s_tir=True) def exp_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")): T.func_attr({"tirx.noalias": True, "op_pattern": 0}) for i0, i1 in T.grid(T.int64(10), T.int64(20)): with T.sblock("compute"): v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) T.reads(A[v_i0, v_i1]) T.writes(A[v_i0, v_i1]) A[v_i0, v_i1] = T.exp(A[v_i0, v_i1]) @T.prim_func(private=True, s_tir=True) def squeeze_inplace(A: T.Buffer((T.int64(10), T.int64(20)), "float32")): T.func_attr({"tirx.noalias": True, "op_pattern": 0}) for ax0, ax1 in T.grid(T.int64(10), T.int64(20)): with T.sblock("T_squeeze"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(A[v_ax0, v_ax1]) T.writes(A[v_ax0, v_ax1]) A[v_ax0, v_ax1] = A[v_ax0, v_ax1] @R.function(private=True) def fused_add_exp_inplace_squeeze_inplace( x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32") ) -> R.Tensor((10, 20), dtype="float32"): R.func_attr({"Primitive": True}) cls = Expected with R.dataflow(): lv = R.call_tir( cls.add, (x, p0), out_ty=R.Tensor((10, 20), dtype="float32"), ) lv1 = R.call_tir_inplace( cls.exp_inplace, (lv,), inplace_indices=[0], out_ty=R.Tensor((10, 20), dtype="float32"), ) gv = R.call_tir_inplace( cls.squeeze_inplace, (lv1,), inplace_indices=[0], out_ty=R.Tensor((10, 20), dtype="float32"), ) R.output(gv) return gv @R.function def main( x: R.Tensor((10, 20), dtype="float32"), p0: R.Tensor((), dtype="float32") ) -> R.Tensor((10, 20), dtype="float32"): cls = Expected with R.dataflow(): gv1: R.Tensor((10, 20), dtype="float32") = ( cls.fused_add_exp_inplace_squeeze_inplace(x, p0) ) R.output(gv1) return gv1 _check(Module, Expected) def test_packed_params(): # fmt: off @I.ir_module(s_tir=True) class Before: @T.prim_func(private=True, s_tir=True) def cast(lv: T.Buffer((T.int64(16), T.int64(16)), "float16"), compute: T.Buffer((T.int64(16), T.int64(16)), "float32")): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for i0, i1 in T.grid(T.int64(16), T.int64(16)): with T.sblock("compute"): v_i0, v_i1 = T.axis.remap("SS", [i0, i1]) T.reads(lv[v_i0, v_i1]) T.writes(compute[v_i0, v_i1]) compute[v_i0, v_i1] = T.Cast("float32", lv[v_i0, v_i1]) @T.prim_func(private=True, s_tir=True) def matmul(x: T.Buffer((T.int64(16), T.int64(16)), "float32"), lv2: T.Buffer((T.int64(16), T.int64(16)), "float32"), T_matmul: T.Buffer((T.int64(16), T.int64(16)), "float32")): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1, k in T.grid(T.int64(16), T.int64(16), T.int64(16)): with T.sblock("T_matmul"): v_ax0, v_ax1, v_k = T.axis.remap("SSR", [ax0, ax1, k]) T.reads(x[v_ax0, v_k], lv2[v_k, v_ax1]) T.writes(T_matmul[v_ax0, v_ax1]) with T.init(): T_matmul[v_ax0, v_ax1] = T.float32(0) T_matmul[v_ax0, v_ax1] = T_matmul[v_ax0, v_ax1] + x[v_ax0, v_k] * lv2[v_k, v_ax1] @R.function def main(x: R.Tensor((16, 16), dtype="float32"), packed_params: R.Tuple(R.Tensor((16, 16), dtype="float16"), R.Tensor((16, 16), dtype="float16"))) -> R.Tensor((16, 16), dtype="float32"): R.func_attr({"num_input": 1}) cls = Before with R.dataflow(): lv: R.Tensor((16, 16), dtype="float16") = packed_params[0] lv1: R.Tensor((16, 16), dtype="float16") = packed_params[1] lv2 = R.call_tir(cls.cast, (lv,), out_ty=R.Tensor((16, 16), dtype="float32")) lv3 = R.call_tir(cls.matmul, (x, lv2), out_ty=R.Tensor((16, 16), dtype="float32")) lv4 = R.call_tir(cls.cast, (lv1,), out_ty=R.Tensor((16, 16), dtype="float32")) lv5 = R.call_tir(cls.matmul, (lv3, lv4), out_ty=R.Tensor((16, 16), dtype="float32")) gv: R.Tensor((16, 16), dtype="float32") = lv5 R.output(gv) return gv # fmt: on Expected = Before _check(Before, Expected) if __name__ == "__main__": tvm.testing.main()