# 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. import pytest import tvm from tvm import relax from tvm.script import ir as I from tvm.script import relax as R from tvm.script import tirx as T @tvm.script.ir_module class Conv2dReLUx2: @R.function def main( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), weight2: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): cls = Conv2dReLUx2 with R.dataflow(): lv: R.Tensor((1, 64, 56, 56), dtype="float32") = ( cls.fused_relax_nn_conv2d_relax_nn_relu(data, weight1) ) gv: R.Tensor((1, 64, 54, 54), dtype="float32") = ( cls.fused_relax_nn_conv2d_relax_nn_relu1(lv, weight2) ) R.output(gv) return gv @R.function(private=True) def fused_relax_nn_conv2d_relax_nn_relu( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight11: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "dnnl.conv2d_relu"}) with R.dataflow(): lv1: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d( data1, weight11, padding=[1, 1, 1, 1], ) gv1: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv1) R.output(gv1) return gv1 @R.function(private=True) def fused_relax_nn_conv2d_relax_nn_relu1( conv1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight21: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "dnnl.conv2d_relu"}) with R.dataflow(): lv2: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d( conv1, weight21, padding=[0, 0, 0, 0], ) gv2: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(lv2) R.output(gv2) return gv2 @tvm.script.ir_module class Conv2dReLUx2_merged: @R.function def main( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), weight2: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): cls = Conv2dReLUx2_merged with R.dataflow(): gv: R.Tensor((1, 64, 54, 54), dtype="float32") = ( cls.fused_relax_nn_conv2d_relax_nn_relu_relax_nn_conv2d_relax_nn_relu1_dnnl( data, weight1, weight2 ) ) R.output(gv) return gv @R.function def fused_relax_nn_conv2d_relax_nn_relu_relax_nn_conv2d_relax_nn_relu1_dnnl( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight11: R.Tensor((64, 64, 3, 3), dtype="float32"), weight21: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Codegen": "dnnl"}) @R.function def lv( data11: R.Tensor((1, 64, 56, 56), dtype="float32"), weight111: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): R.func_attr({"Composite": "dnnl.conv2d_relu"}) with R.dataflow(): lv1: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d( data11, weight111, padding=[1, 1, 1, 1], ) gv1: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv1) R.output(gv1) return gv1 lv2: R.Tensor((1, 64, 56, 56), dtype="float32") = lv(data1, weight11) @R.function def lv11( conv1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight211: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Composite": "dnnl.conv2d_relu"}) with R.dataflow(): lv21: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d( conv1, weight211, padding=[0, 0, 0, 0], ) gv2: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(lv21) R.output(gv2) return gv2 gv3: R.Tensor((1, 64, 54, 54), dtype="float32") = lv11(lv2, weight21) return gv3 @tvm.script.ir_module class Diamond: @R.function def main( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): cls = Diamond with R.dataflow(): lv2: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_conv2d( data, weight ) lv3: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_relu(lv2) lv4: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_gelu(lv2) gv2: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_add(lv3, lv4) R.output(gv2) return gv2 @R.function(private=True) def fused_relax_nn_gelu( lv: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.gelu"}) with R.dataflow(): gv: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.gelu(lv) R.output(gv) return gv @R.function(private=True) def fused_relax_nn_relu( lv1: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.relu"}) with R.dataflow(): gv1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(lv1) R.output(gv1) return gv1 @R.function(private=True) def fused_relax_add( lv5: R.Tensor((1, 64, 54, 54), dtype="float32"), gelu1: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.add"}) with R.dataflow(): gv3: R.Tensor((1, 64, 54, 54), dtype="float32") = R.add(lv5, gelu1) R.output(gv3) return gv3 @R.function(private=True) def fused_relax_nn_conv2d( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.conv2d"}) with R.dataflow(): gv4: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d( data1, weight1, padding=[0, 0, 0, 0], ) R.output(gv4) return gv4 @tvm.script.ir_module class Diamond_merged: @R.function def fused_relax_nn_conv2d_relax_nn_relu_relax_nn_gelu_relax_add_compiler_A( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): # function attr dict R.func_attr({"Codegen": "compiler_A"}) # block 0 @R.function def lv( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.conv2d"}) # block 0 with R.dataflow(): gv4: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d( data1, weight1, strides=[1, 1], padding=[0, 0, 0, 0], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype=None, ) R.output(gv4) return gv4 lv2: R.Tensor((1, 64, 54, 54), dtype="float32") = lv(data, weight) @R.function def lv1( lv11: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.relu"}) # block 0 with R.dataflow(): gv1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(lv11) R.output(gv1) return gv1 lv3: R.Tensor((1, 64, 54, 54), dtype="float32") = lv1(lv2) @R.function def lv21( lv4: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.gelu"}) # block 0 with R.dataflow(): gv: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.gelu(lv4) R.output(gv) return gv lv41: R.Tensor((1, 64, 54, 54), dtype="float32") = lv21(lv2) @R.function def lv31( lv5: R.Tensor((1, 64, 54, 54), dtype="float32"), gelu1: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.add"}) # block 0 with R.dataflow(): gv3: R.Tensor((1, 64, 54, 54), dtype="float32") = R.add(lv5, gelu1) R.output(gv3) return gv3 gv2: R.Tensor((1, 64, 54, 54), dtype="float32") = lv31(lv3, lv41) return gv2 @R.function def main( data2: R.Tensor((1, 64, 56, 56), dtype="float32"), weight2: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): cls = Diamond_merged with R.dataflow(): gv5: R.Tensor((1, 64, 54, 54), dtype="float32") = ( cls.fused_relax_nn_conv2d_relax_nn_relu_relax_nn_gelu_relax_add_compiler_A( data2, weight2 ) ) R.output(gv5) return gv5 @tvm.script.ir_module class Diamond_cyclic_dep: @R.function def main( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): cls = Diamond_cyclic_dep with R.dataflow(): lv2: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_conv2d( data, weight ) lv3: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_relu(lv2) lv4: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_gelu(lv2) gv2: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_add(lv3, lv4) R.output(gv2) return gv2 @R.function(private=True) def fused_relax_nn_gelu( lv: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_B.gelu"}) with R.dataflow(): gv: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.gelu(lv) R.output(gv) return gv @R.function(private=True) def fused_relax_nn_relu( lv1: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.relu"}) with R.dataflow(): gv1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(lv1) R.output(gv1) return gv1 @R.function(private=True) def fused_relax_add( lv5: R.Tensor((1, 64, 54, 54), dtype="float32"), gelu1: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.add"}) with R.dataflow(): gv3: R.Tensor((1, 64, 54, 54), dtype="float32") = R.add(lv5, gelu1) R.output(gv3) return gv3 @R.function(private=True) def fused_relax_nn_conv2d( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.conv2d"}) with R.dataflow(): gv4: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d( data1, weight1, padding=[0, 0, 0, 0], ) R.output(gv4) return gv4 @tvm.script.ir_module class Diamond_cyclic_dep_merged: @R.function def main( data2: R.Tensor((1, 64, 56, 56), dtype="float32"), weight2: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): cls = Diamond_cyclic_dep_merged with R.dataflow(): lv4: R.Tuple( R.Tensor((1, 64, 54, 54), dtype="float32"), R.Tensor((1, 64, 54, 54), dtype="float32"), ) = cls.fused_relax_nn_conv2d_relax_nn_relu_compiler_A(data2, weight2) lv12: R.Tensor((1, 64, 54, 54), dtype="float32") = lv4[0] lv22: R.Tensor((1, 64, 54, 54), dtype="float32") = lv4[1] lv31: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_nn_gelu1_compiler_B( lv12 ) gv5: R.Tensor((1, 64, 54, 54), dtype="float32") = cls.fused_relax_add1_compiler_A( lv22, lv31 ) R.output(gv5) return gv5 @R.function def fused_relax_nn_conv2d_relax_nn_relu_compiler_A( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tuple( R.Tensor((1, 64, 54, 54), dtype="float32"), R.Tensor((1, 64, 54, 54), dtype="float32") ): R.func_attr({"Codegen": "compiler_A"}) @R.function def lv( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Composite": "compiler_A.conv2d"}) with R.dataflow(): gv4: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.conv2d( data1, weight1, padding=[0, 0, 0, 0], ) R.output(gv4) return gv4 gv: R.Tensor((1, 64, 54, 54), dtype="float32") = lv(data, weight) @R.function def lv1( lv11: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Composite": "compiler_A.relu"}) with R.dataflow(): gv1: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.relu(lv11) R.output(gv1) return gv1 gv11: R.Tensor((1, 64, 54, 54), dtype="float32") = lv1(gv) return (gv, gv11) @R.function def fused_relax_nn_gelu1_compiler_B( lv2: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Codegen": "compiler_B"}) @R.function def lv21( lv3: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Composite": "compiler_B.gelu"}) with R.dataflow(): gv2: R.Tensor((1, 64, 54, 54), dtype="float32") = R.nn.gelu(lv3) R.output(gv2) return gv2 gv3: R.Tensor((1, 64, 54, 54), dtype="float32") = lv21(lv2) return gv3 @R.function def fused_relax_add1_compiler_A( lv32: R.Tensor((1, 64, 54, 54), dtype="float32"), lv41: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Codegen": "compiler_A"}) @R.function def lv33( lv5: R.Tensor((1, 64, 54, 54), dtype="float32"), gelu1: R.Tensor((1, 64, 54, 54), dtype="float32"), ) -> R.Tensor((1, 64, 54, 54), dtype="float32"): R.func_attr({"Composite": "compiler_A.add"}) with R.dataflow(): gv31: R.Tensor((1, 64, 54, 54), dtype="float32") = R.add(lv5, gelu1) R.output(gv31) return gv31 gv6: R.Tensor((1, 64, 54, 54), dtype="float32") = lv33(lv32, lv41) return gv6 @tvm.script.ir_module class MultipleProducers: @R.function def main( x1: R.Tensor((10,), dtype="float32"), x2: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): cls = MultipleProducers with R.dataflow(): lv1: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_relu(x1) lv2: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu(x2) lv3: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_relu(lv1) lv4: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu(lv2) gv1: R.Tensor((10,), dtype="float32") = cls.fused_relax_add(lv3, lv4) R.output(gv1) return gv1 @R.function(private=True) def fused_relax_nn_relu( x11: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.relu"}) with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.nn.relu(x11) R.output(gv2) return gv2 @R.function(private=True) def fused_relax_nn_gelu( x21: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.gelu"}) with R.dataflow(): gv3: R.Tensor((10,), dtype="float32") = R.nn.gelu(x21) R.output(gv3) return gv3 @R.function(private=True) def fused_relax_add( lv: R.Tensor((10,), dtype="float32"), gelu1: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.add"}) with R.dataflow(): gv: R.Tensor((10,), dtype="float32") = R.add(lv, gelu1) R.output(gv) return gv @tvm.script.ir_module class MultipleProducers_merged: @R.function def fused_relax_nn_relu_relax_nn_gelu_relax_nn_relu_relax_nn_gelu_relax_add_compiler_A( x1: R.Tensor((10,), dtype="float32"), x2: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Codegen": "compiler_A"}) # block 0 @R.function def lv(x11: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.relu"}) # block 0 with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.nn.relu(x11) R.output(gv2) return gv2 lv1: R.Tensor((10,), dtype="float32") = lv(x1) @R.function def lv11(x21: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.gelu"}) # block 0 with R.dataflow(): gv3: R.Tensor((10,), dtype="float32") = R.nn.gelu(x21) R.output(gv3) return gv3 lv2: R.Tensor((10,), dtype="float32") = lv11(x2) lv3: R.Tensor((10,), dtype="float32") = lv(lv1) lv4: R.Tensor((10,), dtype="float32") = lv11(lv2) @R.function def lv21( lv5: R.Tensor((10,), dtype="float32"), gelu1: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.add"}) # block 0 with R.dataflow(): gv: R.Tensor((10,), dtype="float32") = R.add(lv5, gelu1) R.output(gv) return gv gv1: R.Tensor((10,), dtype="float32") = lv21(lv3, lv4) return gv1 @R.function def main( x12: R.Tensor((10,), dtype="float32"), x22: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): cls = MultipleProducers_merged with R.dataflow(): gv4: R.Tensor((10,), dtype="float32") = ( cls.fused_relax_nn_relu_relax_nn_gelu_relax_nn_relu_relax_nn_gelu_relax_add_compiler_A( x12, x22 ) ) R.output(gv4) return gv4 @tvm.script.ir_module class MultipleProducersCyclic: @R.function def main(x1: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): cls = MultipleProducersCyclic with R.dataflow(): lv1: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_relu(x1) lv2: R.Tensor((10,), dtype="float32") = R.nn.relu(lv1) lv3: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu(lv2) gv1: R.Tensor((10,), dtype="float32") = cls.fused_relax_add(lv1, lv3) R.output(gv1) return gv1 @R.function(private=True) def fused_relax_nn_relu( x11: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.relu"}) with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.nn.relu(x11) R.output(gv2) return gv2 @R.function(private=True) def fused_relax_nn_gelu( x21: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.gelu"}) with R.dataflow(): gv3: R.Tensor((10,), dtype="float32") = R.nn.gelu(x21) R.output(gv3) return gv3 @R.function(private=True) def fused_relax_add( lv: R.Tensor((10,), dtype="float32"), gelu1: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.add"}) with R.dataflow(): gv: R.Tensor((10,), dtype="float32") = R.add(lv, gelu1) R.output(gv) return gv @tvm.script.ir_module class MultipleProducersCyclic_merged: @R.function def main(x1: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): cls = MultipleProducersCyclic_merged with R.dataflow(): lv: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_relu1_compiler_A(x1) lv2: R.Tensor((10,), dtype="float32") = R.nn.relu(lv) gv: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu_relax_add_compiler_A( lv2, lv ) R.output(gv) return gv @R.function def fused_relax_nn_relu1_compiler_A( x11: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Codegen": "compiler_A"}) # block 0 @R.function def lv1(x111: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.relu"}) # block 0 with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.nn.relu(x111) R.output(gv2) return gv2 gv1: R.Tensor((10,), dtype="float32") = lv1(x11) return gv1 @R.function def fused_relax_nn_gelu_relax_add_compiler_A( lv21: R.Tensor((10,), dtype="float32"), lv11: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Codegen": "compiler_A"}) # block 0 @R.function def lv12(x21: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.gelu"}) # block 0 with R.dataflow(): gv3: R.Tensor((10,), dtype="float32") = R.nn.gelu(x21) R.output(gv3) return gv3 lv3: R.Tensor((10,), dtype="float32") = lv12(lv21) @R.function def lv22( lv4: R.Tensor((10,), dtype="float32"), gelu1: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): # function attr dict R.func_attr({"Composite": "compiler_A.add"}) # block 0 with R.dataflow(): gv4: R.Tensor((10,), dtype="float32") = R.add(lv4, gelu1) R.output(gv4) return gv4 gv5: R.Tensor((10,), dtype="float32") = lv22(lv11, lv3) return gv5 @tvm.script.ir_module class MergeCompilerRegionsExample: @R.function def main( x1: R.Tensor((10,), dtype="float32"), x2: R.Tensor((10,), dtype="float32"), x3: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): cls = MergeCompilerRegionsExample with R.dataflow(): lv: R.Tensor((10,), dtype="float32") = cls.fused_relax_add(x1, x2) lv1: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu(x3) lv11: R.Tensor((10,), dtype="float32") = cls.fused_relax_add(lv, lv1) lv12: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu(lv11) lv2: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_relu(lv11) lv21: R.Tensor((10,), dtype="float32") = cls.fused_relax_add(lv12, lv2) gv1: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_relu(lv21) R.output(gv1) return gv1 @R.function(private=True) def fused_relax_nn_relu( add2: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.relu"}) with R.dataflow(): gv: R.Tensor((10,), dtype="float32") = R.nn.relu(add2) R.output(gv) return gv @R.function(private=True) def fused_relax_add( x11: R.Tensor((10,), dtype="float32"), x21: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_A.add"}) with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.add(x11, x21) R.output(gv2) return gv2 @R.function(private=True) def fused_relax_nn_gelu( x31: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Primitive": True, "Composite": "compiler_B.gelu"}) with R.dataflow(): gv3: R.Tensor((10,), dtype="float32") = R.nn.gelu(x31) R.output(gv3) return gv3 @tvm.script.ir_module class MergeCompilerRegionsExampleRef: @R.function def fused_relax_add_relax_add_relax_nn_relu_compiler_A( x1: R.Tensor((10,), dtype="float32"), x2: R.Tensor((10,), dtype="float32"), lv: R.Tensor((10,), dtype="float32"), ) -> R.Tuple(R.Tensor((10,), dtype="float32"), R.Tensor((10,), dtype="float32")): R.func_attr({"Codegen": "compiler_A"}) @R.function def lv1( x11: R.Tensor((10,), dtype="float32"), x21: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Composite": "compiler_A.add"}) with R.dataflow(): gv: R.Tensor((10,), dtype="float32") = R.add(x11, x21) R.output(gv) return gv lv2: R.Tensor((10,), dtype="float32") = lv1(x1, x2) gv1: R.Tensor((10,), dtype="float32") = lv1(lv2, lv) @R.function def lv11(add2: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Composite": "compiler_A.relu"}) with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.nn.relu(add2) R.output(gv2) return gv2 gv11: R.Tensor((10,), dtype="float32") = lv11(gv1) return (gv1, gv11) @R.function def fused_relax_add_relax_nn_relu_compiler_A( lv12: R.Tensor((10,), dtype="float32"), lv3: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Codegen": "compiler_A"}) @R.function def lv21( x11: R.Tensor((10,), dtype="float32"), x21: R.Tensor((10,), dtype="float32") ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Composite": "compiler_A.add"}) with R.dataflow(): gv: R.Tensor((10,), dtype="float32") = R.add(x11, x21) R.output(gv) return gv lv22: R.Tensor((10,), dtype="float32") = lv21(lv12, lv3) @R.function def lv31(add2: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Composite": "compiler_A.relu"}) with R.dataflow(): gv2: R.Tensor((10,), dtype="float32") = R.nn.relu(add2) R.output(gv2) return gv2 gv3: R.Tensor((10,), dtype="float32") = lv31(lv22) return gv3 @R.function def fused_relax_nn_gelu1_compiler_B( x3: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Codegen": "compiler_B"}) @R.function def lv4(x31: R.Tensor((10,), dtype="float32")) -> R.Tensor((10,), dtype="float32"): R.func_attr({"Composite": "compiler_B.gelu"}) with R.dataflow(): gv4: R.Tensor((10,), dtype="float32") = R.nn.gelu(x31) R.output(gv4) return gv4 gv5: R.Tensor((10,), dtype="float32") = lv4(x3) return gv5 @R.function def main( x12: R.Tensor((10,), dtype="float32"), x22: R.Tensor((10,), dtype="float32"), x32: R.Tensor((10,), dtype="float32"), ) -> R.Tensor((10,), dtype="float32"): cls = MergeCompilerRegionsExampleRef with R.dataflow(): lv5: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu1_compiler_B(x32) lv13: R.Tuple(R.Tensor((10,), dtype="float32"), R.Tensor((10,), dtype="float32")) = ( cls.fused_relax_add_relax_add_relax_nn_relu_compiler_A(x12, x22, lv5) ) lv23: R.Tensor((10,), dtype="float32") = lv13[0] lv32: R.Tensor((10,), dtype="float32") = lv13[1] lv41: R.Tensor((10,), dtype="float32") = cls.fused_relax_nn_gelu1_compiler_B(lv23) gv6: R.Tensor((10,), dtype="float32") = cls.fused_relax_add_relax_nn_relu_compiler_A( lv41, lv32 ) R.output(gv6) return gv6 @tvm.script.ir_module class ModuleWithNonComposite: @R.function def main( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): cls = ModuleWithNonComposite with R.dataflow(): lv: R.Tensor((1, 64, 56, 56), dtype="float32") = cls.fused_relax_nn_conv2d(data, weight) conv: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv) R.output(conv) return conv @R.function(private=True) def fused_relax_nn_conv2d( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): R.func_attr({"Composite": "tensorrt.conv2d", "Primitive": True}) with R.dataflow(): gv: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d( data1, weight1, padding=[1, 1, 1, 1], ) R.output(gv) return gv @tvm.script.ir_module class ModuleWithNonComposite_ref: @R.function def main( data: R.Tensor((1, 64, 56, 56), dtype="float32"), weight: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): cls = ModuleWithNonComposite_ref with R.dataflow(): lv: R.Tensor((1, 64, 56, 56), dtype="float32") = cls.fused_relax_nn_conv2d1_tensorrt( data, weight ) conv: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.relu(lv) R.output(conv) return conv @R.function def fused_relax_nn_conv2d1_tensorrt( data1: R.Tensor((1, 64, 56, 56), dtype="float32"), weight1: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): R.func_attr({"Codegen": "tensorrt"}) @R.function def lv1( data2: R.Tensor((1, 64, 56, 56), dtype="float32"), weight2: R.Tensor((64, 64, 3, 3), dtype="float32"), ) -> R.Tensor((1, 64, 56, 56), dtype="float32"): R.func_attr({"Composite": "tensorrt.conv2d"}) with R.dataflow(): gv: R.Tensor((1, 64, 56, 56), dtype="float32") = R.nn.conv2d( data2, weight2, padding=[1, 1, 1, 1], ) R.output(gv) return gv gv1: R.Tensor((1, 64, 56, 56), dtype="float32") = lv1(data1, weight1) return gv1 def check(mod, expected): partitioned = relax.transform.MergeCompositeFunctions()(mod) tvm.ir.assert_structural_equal(partitioned, expected) def test_conv2d_relu_x2(): check(Conv2dReLUx2, Conv2dReLUx2_merged) def test_diamond_cyclic_dep(): """ O = Offloaded to A X = Offloaded to B O O / \\ / \\ O X --> O + + X \\ / \\ / O O We cannot merge all 'O' since it would create a cyclic dependency between the group of `X`. """ check(Diamond_cyclic_dep, Diamond_cyclic_dep_merged) def test_diamond(): """ O = Offloaded to A O O / \\ / \\ O O --> O O \\ / \\ / O O """ check(Diamond, Diamond_merged) def test_merge_producers(): """ Test merging multiple producer groups into a single representative group. O O | | O O \\ / O """ check(MultipleProducers, MultipleProducers_merged) def test_merge_producers_cyclic_dep(): """ Test when multiple producer groups being blocked to merge due to circular dependency in the result. O |\\ | X | | | O |/ O """ check(MultipleProducersCyclic, MultipleProducersCyclic_merged) def test_merge_compiler_regions_example(): check( MergeCompilerRegionsExample, MergeCompilerRegionsExampleRef, ) def test_mixed_non_composite(): check(ModuleWithNonComposite, ModuleWithNonComposite_ref) def test_reshape(): # Verify that the non-CallNode input (shape in reshape) can be handled properly. @I.ir_module(s_tir=True) class Module: @R.function(private=True) def fused_relax_matmul( lv: R.Tensor((1, 784), dtype="float32"), lv1: R.Tensor((784, 512), dtype="float32") ) -> R.Tensor((1, 512), dtype="float32"): R.func_attr({"Composite": "tensorrt.matmul", "Primitive": True}) with R.dataflow(): gv: R.Tensor((1, 512), dtype="float32") = R.matmul(lv, lv1, out_dtype="float32") R.output(gv) return gv @R.function(private=True) def fused_relax_reshape( inp_0: R.Tensor((1, 1, 28, 28), dtype="float32"), param_0: R.Shape([1, 784]) ) -> R.Tensor((1, 784), dtype="float32"): R.func_attr({"Composite": "tensorrt.reshape", "Primitive": True}) with R.dataflow(): gv: R.Tensor((1, 784), dtype="float32") = R.reshape(inp_0, param_0) R.output(gv) return gv @R.function def main( inp_0: R.Tensor((1, 1, 28, 28), dtype="float32"), linear_relu_stack_0_weight: R.Tensor((512, 784), dtype="float32"), ) -> R.Tensor((1, 512), dtype="float32"): cls = Module with R.dataflow(): lv: R.Tensor((1, 784), dtype="float32") = cls.fused_relax_reshape( inp_0, R.shape([1, 784]) ) lv1: R.Tensor((784, 512), dtype="float32") = R.permute_dims( linear_relu_stack_0_weight, axes=None ) lv_1: R.Tensor((1, 512), dtype="float32") = cls.fused_relax_matmul(lv, lv1) gv: R.Tensor((1, 512), dtype="float32") = lv_1 R.output(gv) return gv @I.ir_module(s_tir=True) class Expected: @R.function def fused_relax_reshape_relax_matmul_tensorrt( inp_0: R.Tensor((1, 1, 28, 28), dtype="float32"), lv1: R.Tensor((784, 512), dtype="float32"), ) -> R.Tensor((1, 512), dtype="float32"): R.func_attr({"Codegen": "tensorrt"}) # from tvm.script import relax as R @R.function def lv_1( inp_0_1: R.Tensor((1, 1, 28, 28), dtype="float32"), param_0_1: R.Shape([1, 784]) ) -> R.Tensor((1, 784), dtype="float32"): R.func_attr({"Composite": "tensorrt.reshape"}) with R.dataflow(): gv: R.Tensor((1, 784), dtype="float32") = R.reshape(inp_0_1, param_0_1) R.output(gv) return gv lv_1: R.Tensor((1, 784), dtype="float32") = lv_1(inp_0, R.shape([1, 784])) @R.function def lv1_1_1( lv_2: R.Tensor((1, 784), dtype="float32"), lv1_2: R.Tensor((784, 512), dtype="float32"), ) -> R.Tensor((1, 512), dtype="float32"): R.func_attr({"Composite": "tensorrt.matmul"}) with R.dataflow(): gv: R.Tensor((1, 512), dtype="float32") = R.matmul( lv_2, lv1_2, out_dtype="float32" ) R.output(gv) return gv lv_2: R.Tensor((1, 512), dtype="float32") = lv1_1_1(lv_1, lv1) gv: R.Tensor((1, 512), dtype="float32") = lv_2 return gv @R.function def main( inp_0: R.Tensor((1, 1, 28, 28), dtype="float32"), linear_relu_stack_0_weight: R.Tensor((512, 784), dtype="float32"), ) -> R.Tensor((1, 512), dtype="float32"): cls = Expected with R.dataflow(): lv1: R.Tensor((784, 512), dtype="float32") = R.permute_dims( linear_relu_stack_0_weight, axes=None ) gv: R.Tensor((1, 512), dtype="float32") = ( cls.fused_relax_reshape_relax_matmul_tensorrt(inp_0, lv1) ) R.output(gv) return gv check(Module, Expected) def test_handle_existence_of_call_tir(): """MergeCompositeFunctions should accept R.call_tir as input No merging is required in this case, since the two composite functions have `R.call_tir` between them. This is a regression test, as previously the `Tuple` used to express of `R.call_tir` caused a segfault. """ @I.ir_module(s_tir=True) class Before: @R.function def main(A: R.Tensor([10], dtype="float32")) -> R.Tensor([10], dtype="float32"): cls = Before with R.dataflow(): B = cls.fused_relax_nn_relu(A) C = R.call_tir(cls.relu, (B,), out_ty=R.Tensor([10], dtype="float32")) D = cls.fused_relax_nn_gelu(C) R.output(D) return D @R.function(private=True) def fused_relax_nn_relu( Input: R.Tensor([10], dtype="float32"), ) -> R.Tensor([10], dtype="float32"): R.func_attr({"Composite": "compiler_A.relu", "Primitive": True}) with R.dataflow(): Output = R.nn.relu(Input) R.output(Output) return Output @T.prim_func(private=True, s_tir=True) def relu( Input: T.Buffer(T.int64(10), "float32"), Output: T.Buffer(T.int64(10), "float32"), ): T.func_attr({"tirx.noalias": True}) for i in range(T.int64(10)): with T.sblock("compute"): vi = T.axis.remap("S", [i]) Output[vi] = T.max(Input[vi], T.float32(0)) @R.function(private=True) def fused_relax_nn_gelu( Input: R.Tensor([10], dtype="float32"), ) -> R.Tensor([10], dtype="float32"): R.func_attr({"Composite": "compiler_A.gelu", "Primitive": True}) with R.dataflow(): Output = R.nn.gelu(Input) R.output(Output) return Output @I.ir_module(s_tir=True) class Expected: @R.function def main(A: R.Tensor([10], dtype="float32")) -> R.Tensor([10], dtype="float32"): cls = Expected with R.dataflow(): B = cls.fused_relax_nn_relu1_compiler_A(A) C = R.call_tir(cls.relu, (B,), out_ty=R.Tensor([10], dtype="float32")) D = cls.fused_relax_nn_gelu1_compiler_A(C) R.output(D) return D @R.function def fused_relax_nn_relu1_compiler_A( Input: R.Tensor([10], dtype="float32"), ) -> R.Tensor([10], dtype="float32"): R.func_attr({"Codegen": "compiler_A"}) @R.function def composite_lambda( Input: R.Tensor([10], dtype="float32"), ) -> R.Tensor([10], dtype="float32"): R.func_attr({"Composite": "compiler_A.relu"}) with R.dataflow(): Output = R.nn.relu(Input) R.output(Output) return Output Output = composite_lambda(Input) return Output @T.prim_func(private=True, s_tir=True) def relu( Input: T.Buffer(T.int64(10), "float32"), Output: T.Buffer(T.int64(10), "float32"), ): T.func_attr({"tirx.noalias": True}) for i in range(T.int64(10)): with T.sblock("compute"): vi = T.axis.remap("S", [i]) Output[vi] = T.max(Input[vi], T.float32(0)) @R.function def fused_relax_nn_gelu1_compiler_A( Input: R.Tensor([10], dtype="float32"), ) -> R.Tensor([10], dtype="float32"): R.func_attr({"Codegen": "compiler_A"}) @R.function def composite_lambda( Input: R.Tensor([10], dtype="float32"), ) -> R.Tensor([10], dtype="float32"): R.func_attr({"Composite": "compiler_A.gelu"}) with R.dataflow(): Output = R.nn.gelu(Input) R.output(Output) return Output Output = composite_lambda(Input) return Output After = relax.transform.MergeCompositeFunctions()(Before) tvm.ir.assert_structural_equal(Expected, After) if __name__ == "__main__": pytest.main([__file__])