wehub-resource-sync
1226 lines
44 KiB
Python
1226 lines
44 KiB
Python
# 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__])
|