You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/02/20 15:54:26 UTC
[tvm] 35/46: [Unity][Pass] Operator legalization (#14029)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 77604b464bfeda199fc66bcab1bde17bb0e5b111
Author: Ruihang Lai <ru...@cs.cmu.edu>
AuthorDate: Fri Feb 17 23:16:45 2023 -0500
[Unity][Pass] Operator legalization (#14029)
This PR is the operator legalization pass, which transforms high-level
operator calls to `call_tir`s of corresponding low-level TIR PrimFuncs.
- The legalization pass provides customizability, which enables people
to pass in a customized legalization map to override the default
legalization method.
- The legalization supports symbolic shape. (At this moment only pooling
does not support symbolic shape, as TOPI pooling does not support. This
needs to be fixed in followup PRs.)
Co-authored-by: Chaofan Lin <si...@sjtu.edu.cn>
Co-authored-by: Yixin Dong <ub...@gmail.com>
Co-authored-by: Siyuan Feng <Hz...@sjtu.edu.cn>
---
include/tvm/relax/op_attr_types.h | 9 +
include/tvm/relax/transform.h | 23 +
python/tvm/relax/__init__.py | 21 +-
python/tvm/relax/transform/__init__.py | 3 +
.../relax/transform/{ => legalize_ops}/__init__.py | 16 +-
python/tvm/relax/transform/legalize_ops/binary.py | 55 +
python/tvm/relax/transform/legalize_ops/common.py | 104 ++
.../tvm/relax/transform/legalize_ops/creation.py | 64 +
.../{__init__.py => legalize_ops/datatype.py} | 17 +-
python/tvm/relax/transform/legalize_ops/image.py | 39 +
python/tvm/relax/transform/legalize_ops/index.py | 63 +
.../relax/transform/legalize_ops/linear_algebra.py | 87 ++
.../tvm/relax/transform/legalize_ops/manipulate.py | 114 ++
python/tvm/relax/transform/legalize_ops/nn.py | 178 +++
.../{__init__.py => legalize_ops/search.py} | 7 +-
.../relax/transform/legalize_ops/statistical.py | 83 ++
python/tvm/relax/transform/legalize_ops/unary.py | 32 +
python/tvm/relax/transform/transform.py | 103 ++
src/relax/transform/legalize_ops.cc | 133 +++
tests/python/relax/test_transform_legalize_ops.py | 160 +++
.../relax/test_transform_legalize_ops_binary.py | 1251 ++++++++++++++++++++
.../test_transform_legalize_ops_create_datatype.py | 806 +++++++++++++
.../relax/test_transform_legalize_ops_image.py | 103 ++
..._transform_legalize_ops_index_linear_algebra.py | 401 +++++++
.../test_transform_legalize_ops_manipulate.py | 789 ++++++++++++
.../python/relax/test_transform_legalize_ops_nn.py | 1188 +++++++++++++++++++
...st_transform_legalize_ops_search_statistical.py | 793 +++++++++++++
.../relax/test_transform_legalize_ops_unary.py | 693 +++++++++++
28 files changed, 7315 insertions(+), 20 deletions(-)
diff --git a/include/tvm/relax/op_attr_types.h b/include/tvm/relax/op_attr_types.h
index e171a8d47b..a34cf251dc 100644
--- a/include/tvm/relax/op_attr_types.h
+++ b/include/tvm/relax/op_attr_types.h
@@ -49,6 +49,15 @@ using FInferStructInfo =
*/
using FCallPacked = String;
+/*!
+ * \brief The function type of a legalization function, which takes a
+ * BlockBuilder and the Call to be legalized, and outputs the legalization
+ * result Expr.
+ * \param bb The BlockBuilder context.
+ * \param call The call to be legalized.
+ */
+using FLegalize = runtime::TypedPackedFunc<Expr(const BlockBuilder& bb, const Call& call)>;
+
struct PrintAttrs : public tvm::AttrsNode<PrintAttrs> {
std::string format;
TVM_DECLARE_ATTRS(PrintAttrs, "relax.attrs.PrintAttrs") {
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index 7a4054d414..8b7c7880b9 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -142,6 +142,29 @@ TVM_DLL Pass FoldConstant();
*/
TVM_DLL Pass Normalize();
+/*!
+ * \brief Legalize high-level operator calls in Relax functions to call_tir
+ * with corresponding low-level TIR PrimFuncs.
+ *
+ * For each high-level operator, we register the way of legalizing it as a
+ * function, which takes a context BlockBuilder and the Call being legalized
+ * as input, and returns the legalized call. Here the input BlockBuilder is
+ * mainly used for adding the PrimFunc created by call_te into the context
+ * IRModule.
+ *
+ * The legalization function for each operator is registered as an attribute (with
+ * attribute key `FLegalize`) of the operator.
+ *
+ * For customizability, the user can pass their own legalization by an optional customized map,
+ * with the key to be the operator name and value to be the legalization function.
+ * The default legalization function will be overridden by the customized one.
+ *
+ * \param cmap The customized operator legalization function map. The customized function
+ * will override the default one.
+ * \return The Pass.
+ */
+TVM_DLL Pass LegalizeOps(Optional<Map<String, PackedFunc>> cmap);
+
} // namespace transform
} // namespace relax
} // namespace tvm
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index a6306b788e..cfcf7876dc 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -16,16 +16,6 @@
# under the License.
# pylint: disable=invalid-name, wrong-import-position
"""The Relax IR namespace containing the IR, type, operator, builder, vm, etc."""
-from . import exec_builder
-from . import expr
-from . import ty
-from . import analysis
-from . import transform
-from . import vm
-from . import block_builder
-from . import op
-from . import struct_info
-
# Expr
from .expr import (
Expr,
@@ -82,3 +72,14 @@ from .struct_info import (
TupleStructInfo,
FuncStructInfo,
)
+
+# Import submodules in the last to avoid dependency
+from . import exec_builder
+from . import expr
+from . import ty
+from . import analysis
+from . import transform
+from . import vm
+from . import block_builder
+from . import op
+from . import struct_info
diff --git a/python/tvm/relax/transform/__init__.py b/python/tvm/relax/transform/__init__.py
index eb4d5f710c..78f450b25c 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/__init__.py
@@ -18,3 +18,6 @@
"""Relax transformations. """
from .transform import *
+
+# Import to register the legalization functions.
+from . import legalize_ops
diff --git a/python/tvm/relax/transform/__init__.py b/python/tvm/relax/transform/legalize_ops/__init__.py
similarity index 71%
copy from python/tvm/relax/transform/__init__.py
copy to python/tvm/relax/transform/legalize_ops/__init__.py
index eb4d5f710c..3e57b815db 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/legalize_ops/__init__.py
@@ -14,7 +14,15 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax transformations. """
-
-from .transform import *
+"""Legalize high-level operator calls in Relax functions to call_tir."""
+from . import binary
+from . import creation
+from . import datatype
+from . import image
+from . import index
+from . import linear_algebra
+from . import manipulate
+from . import nn
+from . import search
+from . import statistical
+from . import unary
diff --git a/python/tvm/relax/transform/legalize_ops/binary.py b/python/tvm/relax/transform/legalize_ops/binary.py
new file mode 100644
index 0000000000..55b832021a
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/binary.py
@@ -0,0 +1,55 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for binary operators."""
+from tvm import topi
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import TEFunc, LegalizeFunc, _try_convert_to_scalar_const, register_legalize
+
+
+def _binary(te_func: TEFunc) -> LegalizeFunc:
+ """A common wrapper util for the legalization of binary operators.
+
+ It detects if one of the binary op arguments is a constant scalar. It so,
+ it extracts the scalar value to simplify the generated PrimFunc.
+ """
+
+ def binary_call_te(bb: BlockBuilder, call: Call) -> Expr:
+ # To simplify the created PrimFunc, we first check if arg1 is a constant scalar.
+ # If it is not, we then check if arg0 is a constant scalar.
+ arg0 = call.args[0]
+ arg1 = _try_convert_to_scalar_const(call.args[1])
+ if isinstance(arg1, Expr): # type: ignore
+ arg0 = _try_convert_to_scalar_const(arg0)
+ return bb.call_te(te_func, arg0, arg1)
+
+ return binary_call_te
+
+
+register_legalize("relax.add", _binary(topi.add))
+register_legalize("relax.divide", _binary(topi.divide))
+register_legalize("relax.floor_divide", _binary(topi.floor_divide))
+register_legalize("relax.multiply", _binary(topi.multiply))
+register_legalize("relax.subtract", _binary(topi.subtract))
+register_legalize("relax.equal", _binary(topi.equal))
+
+register_legalize("relax.greater", _binary(topi.greater))
+register_legalize("relax.greater_equal", _binary(topi.greater_equal))
+register_legalize("relax.less", _binary(topi.less))
+register_legalize("relax.less_equal", _binary(topi.less_equal))
+register_legalize("relax.not_equal", _binary(topi.not_equal))
diff --git a/python/tvm/relax/transform/legalize_ops/common.py b/python/tvm/relax/transform/legalize_ops/common.py
new file mode 100644
index 0000000000..85d7fba85c
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/common.py
@@ -0,0 +1,104 @@
+# 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.
+"""Common functionality for legalization."""
+from typing import Callable, Optional, Union
+
+import tvm
+from tvm import te
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr, Constant
+
+
+##################### Types #####################
+
+
+# The function type of a TE function, which accepts TE Tensors and
+# other attributes, and returns the output TE Tensor.
+TEFunc = Callable[..., te.Tensor]
+
+# The function type of a legalization function, which takes a
+# BlockBuilder and the Call to be legalized, and outputs the legalization
+# result Expr.
+LegalizeFunc = Callable[[BlockBuilder, Call], Expr]
+
+
+##################### Utilities #####################
+
+
+def _try_convert_to_scalar_const(expr: Expr) -> Union[Expr, bool, float, int]:
+ """Check if the input Expr is a scalar constant.
+ If it is, return its plain value.
+ If it is not, return the input expr.
+
+ Parameters
+ ----------
+ expr : Expr
+ The expr to be checked and converted.
+
+ Returns
+ --–----
+ ret : Union[Expr, bool, float, int]
+ Return a Python native value (int/float/bool) if the given
+ expr is a scalar constant. Or return the input itself
+ if it is not.
+ """
+ if isinstance(expr, Constant) and expr.struct_info.ndim == 0:
+ return expr.data.numpy()[()].item()
+ else:
+ return expr
+
+
+def _call_topi_without_attr(te_func: TEFunc, primfunc_name: Optional[str] = None) -> LegalizeFunc:
+ """A common wrapper util for the ops who has no attributes and whose
+ legalization is simply passing its arguments to some TE function.
+
+ Parameters
+ ----------
+ te_func : TEFunc
+ The input TE function which is to be converted to PrimFunc.
+
+ primfunc_name : Optional[str]
+ The name of the generated PrimFunc.
+ If it is not specified, the name of `te_func` will be used by default.
+
+ Returns
+ -------
+ func : LegalizeFunc
+ The legalization wrapper function, which wraps the input TE function.
+ """
+ if primfunc_name is None:
+ primfunc_name = te_func.__name__
+ return lambda bb, call: bb.call_te(te_func, *call.args, primfunc_name_hint=primfunc_name)
+
+
+##################### Decorators #####################
+
+_LEGALIZE_ATTR_NAME = "FLegalize"
+
+
+def register_legalize(op_name: str, legal_func: LegalizeFunc = None):
+ """Register legal transformation function for a Relax op.
+
+ Parameters
+ ----------
+ op_name : str
+ The name of the operator
+
+ legal_func: function (bb: BlockBuilder, call: Call) -> new_expr: Expr
+ The function for transforming an expr to another expr.
+ """
+ return tvm.ir.register_op_attr(op_name, _LEGALIZE_ATTR_NAME, legal_func)
diff --git a/python/tvm/relax/transform/legalize_ops/creation.py b/python/tvm/relax/transform/legalize_ops/creation.py
new file mode 100644
index 0000000000..38ce8427b7
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/creation.py
@@ -0,0 +1,64 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for creation operators."""
+from typing import Optional
+
+from tvm import topi, tir
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import LegalizeFunc, register_legalize, _try_convert_to_scalar_const
+
+
+def _full(is_like: bool, fill_value: Optional[float], primfunc_name: str) -> LegalizeFunc:
+ def full_call_te(bb: BlockBuilder, call: Call) -> Expr:
+ _fill_value = (
+ _try_convert_to_scalar_const(call.args[1]) if fill_value is None else fill_value
+ )
+
+ return bb.call_te(
+ topi.full,
+ call.args[0].struct_info.shape if is_like else call.args[0],
+ call.struct_info.dtype,
+ _fill_value,
+ primfunc_name_hint=primfunc_name,
+ )
+
+ return full_call_te
+
+
+def _tril_triu(is_upper: bool, primfunc_name: str) -> LegalizeFunc:
+ def tril_triu_call_te(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ topi.trilu,
+ call.args[0],
+ tir.const(call.attrs.k, "int32"),
+ upper=is_upper,
+ primfunc_name_hint=primfunc_name,
+ )
+
+ return tril_triu_call_te
+
+
+register_legalize("relax.full", _full(is_like=False, fill_value=None, primfunc_name="full"))
+register_legalize("relax.full_like", _full(is_like=True, fill_value=None, primfunc_name="full"))
+register_legalize("relax.ones", _full(is_like=False, fill_value=1.0, primfunc_name="ones"))
+register_legalize("relax.ones_like", _full(is_like=True, fill_value=1.0, primfunc_name="ones"))
+register_legalize("relax.zeros", _full(is_like=False, fill_value=0.0, primfunc_name="zeros"))
+register_legalize("relax.zeros_like", _full(is_like=True, fill_value=0.0, primfunc_name="zeros"))
+register_legalize("relax.tril", _tril_triu(is_upper=False, primfunc_name="tril"))
+register_legalize("relax.triu", _tril_triu(is_upper=True, primfunc_name="triu"))
diff --git a/python/tvm/relax/transform/__init__.py b/python/tvm/relax/transform/legalize_ops/datatype.py
similarity index 58%
copy from python/tvm/relax/transform/__init__.py
copy to python/tvm/relax/transform/legalize_ops/datatype.py
index eb4d5f710c..a71e8ca15e 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/legalize_ops/datatype.py
@@ -14,7 +14,18 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax transformations. """
+# pylint: disable=invalid-name
+"""Default legalization function for datatype operators."""
+from tvm import topi, relax
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import _try_convert_to_scalar_const, register_legalize
-from .transform import *
+
+@register_legalize("relax.astype")
+def _astype(bb: BlockBuilder, call: Call) -> Expr:
+ arg = _try_convert_to_scalar_const(call.args[0])
+ if isinstance(arg, Expr): # type: ignore
+ return bb.call_te(topi.cast, arg, call.attrs.dtype)
+ else:
+ return relax.const(arg, call.attrs.dtype)
diff --git a/python/tvm/relax/transform/legalize_ops/image.py b/python/tvm/relax/transform/legalize_ops/image.py
new file mode 100644
index 0000000000..1b2a342b0b
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/image.py
@@ -0,0 +1,39 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for image operators."""
+from tvm import topi
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import register_legalize
+
+
+@register_legalize("relax.image.resize2d")
+def _image_resize2d(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ topi.image.resize2d,
+ call.args[0],
+ roi=call.attrs.roi,
+ size=call.args[1],
+ layout=call.attrs.layout,
+ method=call.attrs.method,
+ coordinate_transformation_mode=call.attrs.coordinate_transformation_mode,
+ rounding_method=call.attrs.rounding_method,
+ bicubic_alpha=call.attrs.cubic_alpha,
+ bicubic_exclude=call.attrs.cubic_exclude,
+ extrapolation_value=call.attrs.extrapolation_value,
+ )
diff --git a/python/tvm/relax/transform/legalize_ops/index.py b/python/tvm/relax/transform/legalize_ops/index.py
new file mode 100644
index 0000000000..9ee6b28130
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/index.py
@@ -0,0 +1,63 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for index operators."""
+import logging
+
+from tvm import topi, tir
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import register_legalize
+
+
+@register_legalize("relax.take")
+def _take(bb: BlockBuilder, call: Call) -> Expr:
+ # Currently Relax `take` operator doesn't provide the mode choices and
+ # requires input indices to be in range.
+ # We use fast mode, which leads to runtime error whenever some index is
+ # out of bound.
+ return bb.call_te(topi.take, call.args[0], call.args[1], call.attrs.axis, mode="fast")
+
+
+@register_legalize("relax.strided_slice")
+def _strided_slice(bb: BlockBuilder, call: Call) -> Expr:
+ if not all(
+ [
+ isinstance(call.args[0].struct_info.shape.values[i.value], tir.IntImm)
+ for i in call.attrs.axes
+ ]
+ ):
+ logging.info(
+ "Cases where an axis with symbolic length is sliced are not able "
+ "to be legalized through TOPI"
+ )
+ return call
+
+ strides = (
+ [tir.IntImm("int64", 1)] * len(call.attrs.axes)
+ if call.attrs.strides is None
+ else call.attrs.strides
+ )
+ return bb.call_te(
+ topi.strided_slice,
+ call.args[0],
+ call.attrs.begin,
+ call.attrs.end,
+ strides,
+ call.attrs.axes,
+ slice_mode="end",
+ )
diff --git a/python/tvm/relax/transform/legalize_ops/linear_algebra.py b/python/tvm/relax/transform/legalize_ops/linear_algebra.py
new file mode 100644
index 0000000000..abe21d9fff
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/linear_algebra.py
@@ -0,0 +1,87 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for linear algebra operators."""
+from tvm import te, relax, tir
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import register_legalize
+
+
+@register_legalize("relax.matmul")
+def _matmul(bb: BlockBuilder, call: Call) -> Expr:
+ def te_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
+ a_shape = list(a.shape)
+ b_shape = list(b.shape)
+ a_prepended = False
+ b_appended = False
+ if len(a_shape) == 1:
+ a_prepended = True
+ a_shape.insert(0, 1)
+ if len(b_shape) == 1:
+ b_appended = True
+ b_shape.append(1)
+
+ is_a_larger = len(a_shape) > len(b_shape)
+ offset = len(a_shape) - len(b_shape) if is_a_larger else len(b_shape) - len(a_shape)
+
+ a_relax = relax.Var("a", relax.TensorStructInfo(a.shape))
+ b_relax = relax.Var("b", relax.TensorStructInfo(b.shape))
+ f_infer_sinfo = call.op.get_attr("FInferStructInfo")
+ output_shape = f_infer_sinfo(relax.op.matmul(a_relax, b_relax), bb).shape
+
+ def matmul_compute(*idx_spatial):
+ k = te.reduce_axis((0, a_shape[-1]), name="k")
+
+ def multiply_compute(idx_reduce):
+ a_indices = []
+ b_indices = []
+
+ for i in range(offset):
+ if is_a_larger:
+ a_indices.append(idx_spatial[i])
+ else:
+ b_indices.append(idx_spatial[i])
+ for i in range(offset, len(output_shape) - (2 - a_prepended - b_appended)):
+ a_dim = a_shape[i if is_a_larger else i - offset]
+ b_dim = b_shape[i if not is_a_larger else i - offset]
+ a_dim_is_one = isinstance(a_dim, tir.IntImm) and a_dim == 1
+ b_dim_is_one = isinstance(b_dim, tir.IntImm) and b_dim == 1
+ a_indices.append(0 if a_dim_is_one else idx_spatial[i])
+ b_indices.append(0 if b_dim_is_one else idx_spatial[i])
+ if not a_prepended:
+ a_indices.append(idx_spatial[-2 + b_appended])
+ a_indices.append(idx_reduce)
+ b_indices.append(idx_reduce)
+ if not b_appended:
+ b_indices.append(idx_spatial[-1])
+
+ dtype = call.attrs.out_dtype
+ if dtype != "":
+ return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
+ else:
+ return a(*a_indices) * b(*b_indices)
+
+ return te.sum(multiply_compute(k), axis=k)
+
+ return te.compute(
+ output_shape,
+ lambda *idx: matmul_compute(*idx), # pylint: disable=unnecessary-lambda
+ name="matmul",
+ )
+
+ return bb.call_te(te_matmul, call.args[0], call.args[1], primfunc_name_hint="matmul")
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py b/python/tvm/relax/transform/legalize_ops/manipulate.py
new file mode 100644
index 0000000000..76e3e74bab
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -0,0 +1,114 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for manipulate operators."""
+import logging
+
+import tvm
+from tvm import topi, tir, relax, te
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr, Var, Tuple, TupleGetItem
+from .common import TEFunc, LegalizeFunc, register_legalize
+
+
+def _reshape(
+ te_func: TEFunc, primfunc_name: str, is_collapse_sum_like: bool = False
+) -> LegalizeFunc:
+ def reshape_call_te(bb: BlockBuilder, call: Call):
+ tgt_shape = call.args[1].struct_info.shape if is_collapse_sum_like else call.args[1]
+ return bb.call_te(te_func, call.args[0], tgt_shape, primfunc_name_hint=primfunc_name)
+
+ return reshape_call_te
+
+
+register_legalize("relax.broadcast_to", _reshape(topi.broadcast_to, "broadcast_to"))
+register_legalize("relax.reshape", _reshape(topi.reshape, "reshape"))
+
+
+@register_legalize("relax.concat")
+def _concat(bb: BlockBuilder, call: Call) -> Expr:
+ t = call.args[0]
+ n_field = len(t.struct_info.fields)
+ while isinstance(t, Var):
+ binding = bb.lookup_binding(t)
+ if not isinstance(binding, (Tuple, Var)):
+ break
+ t = binding
+
+ assert isinstance(t, (Tuple, Var))
+ fields = (
+ t.fields if isinstance(t, Tuple) else [bb.emit(TupleGetItem(t, i)) for i in range(n_field)]
+ )
+ return bb.call_te(
+ topi.concatenate, fields, None if call.attrs.axis is None else call.attrs.axis.value
+ )
+
+
+@register_legalize("relax.expand_dims")
+def _expand_dims(bb: BlockBuilder, call: Call) -> Expr:
+ def te_expand_dims(data, axis):
+ data_relax = relax.Var("data", relax.TensorStructInfo(data.shape))
+ f_infer_sinfo = call.op.get_attr("FInferStructInfo")
+ output_shape = f_infer_sinfo(relax.op.expand_dims(data_relax, axis), bb).shape
+ output_ndim = len(output_shape)
+
+ data_dims = []
+ for i in range(output_ndim):
+ if i not in axis and (i - output_ndim) not in axis:
+ data_dims.append(i)
+ return te.compute(
+ output_shape,
+ lambda *idx: data(*[idx[dim] for dim in data_dims]),
+ name="expand_dims",
+ )
+
+ return bb.call_te(
+ te_expand_dims, call.args[0], call.attrs.axis, primfunc_name_hint="expand_dims"
+ )
+
+
+@register_legalize("relax.flatten")
+def _flatten(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(topi.reshape, call.args[0], call.struct_info.shape.values)
+
+
+@register_legalize("relax.permute_dims")
+def _permute_dims(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(topi.transpose, call.args[0], call.attrs.axes)
+
+
+@register_legalize("relax.split")
+def _split(bb: BlockBuilder, call: Call) -> Expr:
+ if isinstance(call.attrs.indices_or_sections, tir.IntImm):
+ indices_or_sections = call.attrs.indices_or_sections.value
+ modulo = tvm.arith.Analyzer().simplify(
+ call.args[0].struct_info.shape.values[call.attrs.axis] % indices_or_sections
+ )
+ if modulo != 0:
+ logging.info(
+ "Split cannot be legalized by TOPI when the axis being split has "
+ "length that not divisible by the input number of section."
+ )
+ return call
+ else:
+ indices_or_sections = call.attrs.indices_or_sections
+ return bb.call_te(topi.split, call.args[0], indices_or_sections, call.attrs.axis)
+
+
+@register_legalize("relax.squeeze")
+def _squeeze(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(topi.squeeze, call.args[0], call.attrs.axis)
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py b/python/tvm/relax/transform/legalize_ops/nn.py
new file mode 100644
index 0000000000..49f198306d
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -0,0 +1,178 @@
+# 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.
+# pylint: disable=invalid-name,unused-argument
+"""Default legalization function for neural network operators."""
+import logging
+
+from tvm import topi, tir, te
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import register_legalize, _call_topi_without_attr
+
+
+@register_legalize("relax.nn.conv2d")
+def _nn_conv2d(bb: BlockBuilder, call: Call) -> Expr:
+ if call.attrs.out_layout != call.attrs.data_layout:
+ logging.info(
+ "TOPI conv2d does not support different input-output "
+ "layouts, and thus cannot be legalized by TOPI"
+ )
+ return call
+ if len(call.attrs.data_layout) != 4 or len(call.attrs.kernel_layout) != 4:
+ logging.info(
+ "Conv2D where data layout or kernel layout have channel chunk "
+ "cannot be legalized by TOPI at this moment."
+ )
+ return call
+ if call.attrs.groups != 1:
+ data_layout = tir.layout(call.attrs.data_layout)
+ kernel_layout = tir.layout(call.attrs.kernel_layout)
+ ic = call.args[0].struct_info.shape.values[data_layout.index_of("C")]
+ oc = call.args[1].struct_info.shape.values[kernel_layout.index_of("O")]
+ if not isinstance(ic, tir.IntImm) or not isinstance(oc, tir.IntImm):
+ logging.info(
+ "Conv2D where number of groups is more than one and input or output "
+ "channel size is symbolic cannot be legalized by TOPI at this moment."
+ )
+ return call
+
+ return bb.call_te(
+ topi.nn.conv,
+ inp=call.args[0],
+ filt=call.args[1],
+ stride=call.attrs.strides,
+ padding=call.attrs.padding,
+ dilation=call.attrs.dilation,
+ groups=call.attrs.groups,
+ data_layout=call.attrs.data_layout,
+ kernel_layout=call.attrs.kernel_layout,
+ out_dtype=call.attrs.out_dtype if call.attrs.out_dtype != "" else None,
+ primfunc_name_hint="conv2d",
+ )
+
+
+@register_legalize("relax.nn.max_pool2d")
+def _nn_max_pool2d(bb: BlockBuilder, call: Call) -> Expr:
+ if call.attrs.out_layout != call.attrs.layout:
+ logging.info(
+ "TOPI max_pool2d does not support different input-output "
+ "layouts, and thus cannot be legalized by TOPI"
+ )
+ return call
+
+ return bb.call_te(
+ topi.nn.pool2d,
+ call.args[0],
+ kernel=call.attrs.pool_size,
+ stride=call.attrs.strides,
+ dilation=call.attrs.dilation,
+ padding=call.attrs.padding,
+ pool_type="max",
+ ceil_mode=call.attrs.ceil_mode,
+ layout=call.attrs.layout,
+ primfunc_name_hint="max_pool2d",
+ )
+
+
+@register_legalize("relax.nn.adaptive_avg_pool2d")
+def _nn_adaptive_avg_pool2d(bb: BlockBuilder, call: Call) -> Expr:
+ if call.attrs.out_layout != call.attrs.layout:
+ logging.info(
+ "TOPI adaptive_avg_pool2d does not support different input-output "
+ "layouts, and thus cannot be legalized by TOPI"
+ )
+ return call
+
+ def te_adaptive_avg_pool2d(data, output_size, layout_str):
+ if output_size is None:
+ layout = tir.layout(layout_str)
+ idx_H = layout.index_of("H")
+ idx_W = layout.index_of("W")
+ assert idx_H != -1 and idx_W != -1
+ output_size = (data.shape[idx_H], data.shape[idx_W])
+
+ return topi.nn.adaptive_pool(data, output_size, "avg", layout_str)
+
+ return bb.call_te(
+ te_adaptive_avg_pool2d,
+ call.args[0],
+ call.attrs.output_size,
+ call.attrs.layout,
+ primfunc_name_hint="adaptive_avg_pool2d",
+ )
+
+
+register_legalize("relax.nn.relu", _call_topi_without_attr(topi.nn.relu))
+
+
+@register_legalize("relax.nn.gelu")
+def _nn_gelu(bb: BlockBuilder, call: Call) -> Expr:
+ def te_gelu(x: te.Tensor):
+ dtype = x.dtype
+ return x * (
+ tir.const(0.5, dtype)
+ + topi.erf(x * tir.const(0.5**0.5, dtype)) * tir.const(0.5, dtype)
+ )
+
+ return bb.call_te(te_gelu, call.args[0], primfunc_name_hint="gelu")
+
+
+@register_legalize("relax.nn.silu")
+def _nn_silu(bb: BlockBuilder, call: Call) -> Expr:
+ def te_silu(x: te.Tensor):
+ return topi.multiply(x, topi.sigmoid(x))
+
+ return bb.call_te(te_silu, call.args[0], primfunc_name_hint="silu")
+
+
+@register_legalize("relax.nn.softmax")
+def _nn_softmax(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(topi.nn.softmax, call.args[0], call.attrs.axis)
+
+
+@register_legalize("relax.nn.batch_norm")
+def _nn_batch_norm(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ topi.nn.batch_norm,
+ data=call.args[0],
+ gamma=call.args[1],
+ beta=call.args[2],
+ moving_mean=call.args[3],
+ moving_var=call.args[4],
+ axis=call.attrs.axis,
+ epsilon=call.attrs.epsilon,
+ center=call.attrs.center,
+ scale=call.attrs.scale,
+ )
+
+
+@register_legalize("relax.nn.layer_norm")
+def _nn_layer_norm(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ topi.nn.layer_norm,
+ call.args[0],
+ call.args[1],
+ call.args[2],
+ axis=call.attrs.axes,
+ epsilon=call.attrs.epsilon,
+ )
+
+
+@register_legalize("relax.nn.dropout")
+def _nn_dropout(bb: BlockBuilder, call: Call) -> Expr:
+ logging.info("Dropout is handled by frontend translator at this moment and is not legalized.")
+ return call
diff --git a/python/tvm/relax/transform/__init__.py b/python/tvm/relax/transform/legalize_ops/search.py
similarity index 78%
copy from python/tvm/relax/transform/__init__.py
copy to python/tvm/relax/transform/legalize_ops/search.py
index eb4d5f710c..fb38d099c4 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/legalize_ops/search.py
@@ -14,7 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
-"""Relax transformations. """
+"""Default legalization function for search operators."""
+from tvm import topi
+from .common import _call_topi_without_attr, register_legalize
-from .transform import *
+register_legalize("relax.where", _call_topi_without_attr(topi.where))
diff --git a/python/tvm/relax/transform/legalize_ops/statistical.py b/python/tvm/relax/transform/legalize_ops/statistical.py
new file mode 100644
index 0000000000..3307d49f21
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/statistical.py
@@ -0,0 +1,83 @@
+# 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.
+# pylint: disable=invalid-name
+"""Default legalization function for statistical operators."""
+from typing import List
+from tvm import topi, tir, te
+from ...block_builder import BlockBuilder
+from ...expr import Call, Expr
+from .common import TEFunc, LegalizeFunc, register_legalize
+
+
+def _statistical(te_func: TEFunc) -> LegalizeFunc:
+ def statistical_call_te(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(te_func, call.args[0], call.attrs.axis, call.attrs.keepdims)
+
+ return statistical_call_te
+
+
+def _compute_shape_prod(x: te.Tensor, axis: List[tir.IntImm]) -> tir.PrimExpr:
+ shape_prod = tir.const(1, "int32")
+ axes = [_axis.value for _axis in axis] if axis is not None else range(0, len(x.shape))
+ for dim in axes:
+ shape_prod = shape_prod * x.shape[dim]
+ return shape_prod
+
+
+def _te_mean(x: te.Tensor, axis: List[tir.IntImm], keepdims: bool) -> te.Tensor:
+ shape_prod = _compute_shape_prod(x, axis)
+ res_sum = topi.sum(x, axis, keepdims)
+ return topi.divide(res_sum, shape_prod)
+
+
+def _te_variance(x: te.Tensor, axis: List[tir.IntImm], keepdims: bool) -> te.Tensor:
+ dev = x - _te_mean(x, axis, keepdims)
+ return _te_mean(dev * dev, axis, keepdims)
+
+
+@register_legalize("relax.mean")
+def _mean(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ _te_mean, call.args[0], call.attrs.axis, call.attrs.keepdims, primfunc_name_hint="mean"
+ )
+
+
+@register_legalize("relax.std")
+def _std(bb: BlockBuilder, call: Call) -> Expr:
+ def te_std(x: te.Tensor, axis: List[tir.IntImm], keepdims: bool) -> te.Tensor:
+ return topi.sqrt(_te_variance(x, axis, keepdims))
+
+ return bb.call_te(
+ te_std, call.args[0], call.attrs.axis, call.attrs.keepdims, primfunc_name_hint="std"
+ )
+
+
+@register_legalize("relax.variance")
+def _variance(bb: BlockBuilder, call: Call) -> Expr:
+ return bb.call_te(
+ _te_variance,
+ call.args[0],
+ call.attrs.axis,
+ call.attrs.keepdims,
+ primfunc_name_hint="variance",
+ )
+
+
+register_legalize("relax.max", _statistical(topi.max))
+register_legalize("relax.min", _statistical(topi.min))
+register_legalize("relax.prod", _statistical(topi.prod))
+register_legalize("relax.sum", _statistical(topi.sum))
diff --git a/python/tvm/relax/transform/legalize_ops/unary.py b/python/tvm/relax/transform/legalize_ops/unary.py
new file mode 100644
index 0000000000..cd29182c4d
--- /dev/null
+++ b/python/tvm/relax/transform/legalize_ops/unary.py
@@ -0,0 +1,32 @@
+# 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.
+"""Default legalization function for unary operators."""
+from tvm import topi
+from .common import _call_topi_without_attr, register_legalize
+
+# To avoid conflict of IRModule function name and libc function name, we add
+# "tir_" as the prefix of the generated PrimFunc name.
+register_legalize("relax.abs", _call_topi_without_attr(topi.abs, "tir_abs"))
+register_legalize("relax.cos", _call_topi_without_attr(topi.cos, "tir_cos"))
+register_legalize("relax.log", _call_topi_without_attr(topi.log, "tir_log"))
+register_legalize("relax.exp", _call_topi_without_attr(topi.exp, "tir_exp"))
+register_legalize("relax.negative", _call_topi_without_attr(topi.negative, "tir_negative"))
+register_legalize("relax.sigmoid", _call_topi_without_attr(topi.sigmoid, "tir_sigmoid"))
+register_legalize("relax.sin", _call_topi_without_attr(topi.sin, "tir_sin"))
+register_legalize("relax.sqrt", _call_topi_without_attr(topi.sqrt, "tir_sqrt"))
+register_legalize("relax.tanh", _call_topi_without_attr(topi.tanh, "tir_tanh"))
+register_legalize("relax.clip", _call_topi_without_attr(topi.clip, "tir_clip"))
diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
index 7fcf0b1121..4ba967935b 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -24,6 +24,7 @@ import numpy as np # type: ignore
import tvm.ir
from tvm.runtime import NDArray
from . import _ffi_api
+from .legalize_ops.common import LegalizeFunc
@tvm._ffi.register_object("relax.FunctionPass")
@@ -229,6 +230,108 @@ def FuseTIR() -> tvm.ir.transform.Pass:
return _ffi_api.FuseTIR() # type: ignore
+def LegalizeOps(customize_legalize_map: Optional[Dict[str, LegalizeFunc]] = None):
+ """Legalize high-level operator calls in Relax functions to call_tir
+ with corresponding low-level TIR PrimFuncs.
+
+ For each high-level operator, we register the way of legalizing it as a
+ function, which takes a context BlockBuilder and the Call being legalized
+ as input, and returns the legalized call. Here the input BlockBuilder is
+ mainly used for adding the PrimFunc created by call_te into the context
+ IRModule.
+
+ The legalization function for each operator is registered as an attribute (with
+ attribute key `FLegalize`) of the operator.
+
+ This pass provides customizability for users to use their own legalization
+ function for operators. The pass takes an optional customized map,
+ with the key to be the operator name (`str`) and value to be the function
+ (`LegalizeFunc`). The default legalization function will be overridden by the customized
+ one.
+
+ Parameters
+ ----------
+ customize_legalize_map : Optional[Dict[str, LegalizeFunc]]
+ The customized operator legalization function map. The customized function will override
+ the default one.
+
+ Returns
+ -------
+ ret : tvm.transform.Pass
+ The registered pass
+
+ Examples
+ --------
+ The following code shows how to use this pass:
+
+ .. code-block:: python
+
+ # Define the pass input IRModule
+ @tvm.script.ir_module
+ class Module:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
+ ) -> R.Tensor((2, 3), "float32"):
+ z: R.Tensor((2, 3), "float32") = R.add(x, y)
+ r: R.Tensor((2, 3), "float32") = R.multiply(y, z)
+ return r
+
+ # Define the customized legalization function for "relax.add"
+ def customize_legalize_add(bb: relax.BlockBuilder, call: relax.Call) -> relax.Expr:
+ from tvm import topi
+ return bb.call_te(topi.add, call.args[1], call.args[0])
+
+ # Apply the pass with the customized function to the module.
+ mod = LegalizeOps({"relax.add": customize_legalize_add})(Module)
+
+ Print out the result by `mod.show()`, we can see the IRModule after
+ legalization becomes
+
+ .. code-block:: python
+
+ @tvm.script.ir_module
+ class Module:
+ @R.function
+ def main(
+ x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 3), "float32")
+ ) -> R.Tensor((2, 3), "float32"):
+ z = R.call_tir(add, (y, x), (2, 3), dtype="float32")
+ r = R.call_tir(multiply, (y, z), (2, 3), dtype="float32")
+ return r
+
+ @T.prim_func
+ def add(
+ A: T.Buffer[(2, 3), "float32"],
+ B: T.Buffer[(2, 3), "float32"],
+ T_add: T.Buffer[(2, 3), "float32"],
+ ):
+ T.func_attr({"tir.noalias": True})
+ for ax0, ax1 in T.grid(2, 3):
+ with T.block("T_add"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
+ T.writes(T_add[v_ax0, v_ax1])
+ T_add[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1]
+
+ @T.prim_func
+ def multiply(
+ A: T.Buffer[(2, 3), "float32"],
+ B: T.Buffer[(2, 3), "float32"],
+ T_multiply: T.Buffer[(2, 3), "float32"],
+ ):
+ T.func_attr({"tir.noalias": True})
+ for ax0, ax1 in T.grid(2, 3):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(A[v_ax0, v_ax1], B[v_ax0, v_ax1])
+ T.writes(T_multiply[v_ax0, v_ax1])
+ T_multiply[v_ax0, v_ax1] = A[v_ax0, v_ax1] * B[v_ax0, v_ax1]
+ """
+
+ return _ffi_api.LegalizeOps(customize_legalize_map) # type: ignore
+
+
def MetaScheduleApplyDatabase(
work_dir: Optional[str] = None,
) -> tvm.ir.transform.Pass:
diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc
new file mode 100644
index 0000000000..f9a84c5361
--- /dev/null
+++ b/src/relax/transform/legalize_ops.cc
@@ -0,0 +1,133 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/legalize_ops.cc
+ * \brief Legalize high-level operator calls in Relax functions to call_tir
+ * with corresponding low-level TIR PrimFuncs.
+ */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/op_attr_types.h>
+#include <tvm/relax/transform.h>
+
+namespace tvm {
+namespace relax {
+
+/*!
+ * \brief Check if a given Tensor/Shape/TupleStructInfo contains shapes whose
+ * values are all known.
+ * \param sinfo The StructInfo to be checked.
+ * \return A boolean indicating the given struct info contains shape values that are all known.
+ */
+bool KnowAllShapeValues(const StructInfo& sinfo) {
+ if (const auto* tensor_sinfo = sinfo.as<TensorStructInfoNode>()) {
+ return tensor_sinfo->shape.defined() &&
+ tensor_sinfo->shape.value()->IsInstance<ShapeExprNode>();
+ } else if (const auto* shape_sinfo = sinfo.as<ShapeStructInfoNode>()) {
+ return shape_sinfo->values.defined();
+ } else if (const auto* tuple_sinfo = sinfo.as<TupleStructInfoNode>()) {
+ return std::all_of(tuple_sinfo->fields.begin(), tuple_sinfo->fields.end(),
+ [](StructInfo field_sinfo) { return KnowAllShapeValues(field_sinfo); });
+ } else if (sinfo.as<PrimStructInfoNode>()) {
+ return true;
+ } else {
+ return false;
+ }
+}
+
+class LegalizeMutator : public ExprMutator {
+ public:
+ explicit LegalizeMutator(const IRModule& mod, const Optional<Map<String, PackedFunc>>& cmap)
+ : ExprMutator(mod), mod_(std::move(mod)), cmap_(std::move(cmap)) {}
+
+ IRModule Transform() {
+ for (const auto& [gv, func] : mod_->functions) {
+ if (func->IsInstance<FunctionNode>()) {
+ auto updated_func = Downcast<Function>(this->VisitExpr(func));
+ builder_->UpdateFunction(gv, Downcast<BaseFunc>(updated_func));
+ }
+ }
+ return builder_->GetContextIRModule();
+ }
+
+ private:
+ using ExprMutator::VisitExpr_;
+
+ Expr VisitExpr_(const CallNode* call) final {
+ Call visited_call = Downcast<Call>(this->VisitExprPostOrder_(call));
+ static const auto& legalize_map = Op::GetAttrMap<FLegalize>("FLegalize");
+ static const Op& call_tir_op = Op::Get("relax.call_tir");
+ auto* op_node = visited_call->op.as<OpNode>();
+
+ // Not an OpNode
+ if (op_node == nullptr) {
+ return visited_call;
+ }
+
+ // Not all shape values are known
+ if (!std::all_of(visited_call->args.begin(), visited_call->args.end(),
+ [](Expr arg) { return KnowAllShapeValues(GetStructInfo(arg)); }) ||
+ !KnowAllShapeValues(GetStructInfo(visited_call))) {
+ return visited_call;
+ }
+
+ auto op = GetRef<Op>(op_node);
+
+ // Priority: customize > default.
+ // Check if it has customize legalization registered.
+ if (cmap_.defined() && cmap_.value().count(op->name)) {
+ return cmap_.value()[op->name](this->builder_, visited_call);
+ }
+ // Check if it has default legalization registered.
+ if (legalize_map.count(op)) {
+ return legalize_map[op](this->builder_, visited_call);
+ }
+
+ // No legalization.
+ if (op != call_tir_op) {
+ LOG(WARNING) << "No legalization func for " << op->name << " is found.";
+ }
+ return visited_call;
+ }
+
+ /*! \brief The context IRModule. */
+ IRModule mod_;
+ /*! \brief The customized legalization function map. */
+ Optional<Map<String, PackedFunc>> cmap_;
+};
+
+namespace transform {
+
+Pass LegalizeOps(Optional<Map<String, PackedFunc>> cmap) {
+ runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
+ [=](IRModule mod, PassContext pc) { return LegalizeMutator(mod, cmap).Transform(); };
+ return CreateModulePass(/*pass_function=*/pass_func,
+ /*opt_level=*/0,
+ /*pass_name=*/"LegalizeOps",
+ /*required=*/{});
+}
+
+TVM_REGISTER_GLOBAL("relax.transform.LegalizeOps").set_body_typed(LegalizeOps);
+
+} // namespace transform
+
+} // namespace relax
+} // namespace tvm
diff --git a/tests/python/relax/test_transform_legalize_ops.py b/tests/python/relax/test_transform_legalize_ops.py
new file mode 100644
index 0000000000..91f8cb4259
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops.py
@@ -0,0 +1,160 @@
+# 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 tvm
+from tvm import relax
+from tvm.relax.transform import LegalizeOps
+from tvm.relax.transform.legalize_ops.common import register_legalize
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+def test_customize_legalize():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv: R.Tensor((4, 3, 2, 3), "float32") = R.add(x, y)
+ return gv
+
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv = R.call_tir(add, (y, x), R.Tensor((4, 3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def add(rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_add"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3])
+ T.writes(T_add[ax0, ax1, ax2, ax3])
+ T_add[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] + rxplaceholder[T.int64(0), ax2, ax3]
+ # fmt: on
+
+ def customize_legalize_add(bb: relax.BlockBuilder, call: relax.Call):
+ from tvm import topi # pylint: disable=import-outside-toplevel
+
+ return bb.call_te(topi.add, call.args[1], call.args[0])
+
+ mod = LegalizeOps({"relax.add": customize_legalize_add})(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_legalize_multiple_types_of_call():
+ # fmt: off
+ @tvm.script.ir_module
+ class Before:
+ @R.function
+ def mul2(x: R.Tensor((3, 3), "float32")):
+ gv = R.multiply(x, R.const(2.0, "float32"))
+ return gv
+
+ @T.prim_func
+ def identity(rxplaceholder: T.Buffer((T.int64(3), T.int64(3)), "float32"), T_id: T.Buffer((T.int64(3), T.int64(3)), "float32")):
+ for ax0, ax1 in T.grid(T.int64(3), T.int64(3)):
+ with T.block("T_add"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(rxplaceholder[v_ax0, v_ax1])
+ T.writes(T_id[v_ax0, v_ax1])
+ T_id[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1]
+
+ @R.function
+ def main(x: R.Tensor((3, 3), "float32")):
+ gv: R.Tensor((3, 3), "float32") = mul2(x)
+ gv1 = R.call_tir(identity, gv, R.Tensor((3, 3), dtype="float32"))
+ gv2 = R.multiply(gv1, R.const(2.0, "float32"))
+ return gv2
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def mul2(x: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((3, 3), dtype="float32"):
+ gv = R.call_tir(multiply, (x,), R.Tensor((3, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def identity(rxplaceholder: T.Buffer((T.int64(3), T.int64(3)), "float32"), T_id: T.Buffer((T.int64(3), T.int64(3)), "float32")):
+ for ax0, ax1 in T.grid(T.int64(3), T.int64(3)):
+ with T.block("T_add"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(rxplaceholder[v_ax0, v_ax1])
+ T.writes(T_id[v_ax0, v_ax1])
+ T_id[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1]
+
+ @T.prim_func
+ def multiply(rxplaceholder: T.Buffer((T.int64(3), T.int64(3)), "float32"), T_multiply: T.Buffer((T.int64(3), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for ax0, ax1 in T.grid(T.int64(3), T.int64(3)):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(rxplaceholder[v_ax0, v_ax1])
+ T.writes(T_multiply[v_ax0, v_ax1])
+ T_multiply[v_ax0, v_ax1] = rxplaceholder[v_ax0, v_ax1] * T.float32(2)
+
+ @R.function
+ def main(x1: R.Tensor((3, 3), dtype="float32")) -> R.Tensor((3, 3), dtype="float32"):
+ gv1: R.Tensor((3, 3), dtype="float32") = mul2(x1)
+ gv11 = R.call_tir(identity, gv1, R.Tensor((3, 3), dtype="float32"))
+ gv2 = R.call_tir(multiply, (gv11,), R.Tensor((3, 3), dtype="float32"))
+ return gv2
+ # fmt: on
+
+ After = LegalizeOps()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_can_not_legalize():
+ # case 1: does't have legalization
+ add_legalize = tvm.ir.Op.get("relax.add").get_attr("FLegalize")
+ # reset it for test
+ tvm.ir.Op.get("relax.add").reset_attr("FLegalize")
+
+ # fmt: off
+ @tvm.script.ir_module
+ class Before0:
+ @R.function
+ def main(x: R.Tensor((3, 3), "float32")):
+ gv: R.Tensor((3, 3), "float32") = R.add(x, x)
+ return gv
+ # fmt: on
+ After0 = LegalizeOps()(Before0)
+ tvm.ir.assert_structural_equal(After0, Before0)
+
+ register_legalize("relax.add", add_legalize)
+
+ # case 2: don't know all shape
+ s = relax.Var("s", relax.ShapeStructInfo((3, 3)))
+ x = relax.Var("x", relax.TensorStructInfo((3, 3), "float32"))
+ y = relax.Var("y", relax.TensorStructInfo(s, "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ gv = bb.emit_output(R.add(x, y))
+ bb.emit_func_output(gv)
+ Before1 = bb.get()
+ After1 = LegalizeOps()(Before1)
+ tvm.ir.assert_structural_equal(After1, Before1)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_binary.py b/tests/python/relax/test_transform_legalize_ops_binary.py
new file mode 100644
index 0000000000..c2db7e9ba1
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_binary.py
@@ -0,0 +1,1251 @@
+# 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 tvm
+from tvm.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+##################### Binary arithmetic #####################
+
+
+def test_add():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv: R.Tensor((4, 3, 2, 3), "float32") = R.add(x, y)
+ return gv
+
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv = R.call_tir(add, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def add(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_add: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_add"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_add[ax0, ax1, ax2, ax3])
+ T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_add_with_arg0_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), dtype="float32") = R.add(x, R.const(1, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(add, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def add(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_add"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_add[ax0, ax1])
+ T_add[ax0, ax1] = rxplaceholder[ax0, ax1] + T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_add_with_arg1_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), dtype="float32") = R.add(R.const(1, "float32"), x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(add, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def add(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_add: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_add"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_add[ax0, ax1])
+ T_add[ax0, ax1] = T.float32(1) + rxplaceholder[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_add_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "float32") = R.add(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(add, (x, y), R.Tensor((a, b, c, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def add(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_add: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_add = T.match_buffer(var_T_add, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_add"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_add[ax0, ax1, ax2, ax3])
+ T_add[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] + rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_divide():
+ # fmt: off
+ @tvm.script.ir_module
+ class Divide:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv: R.Tensor((4, 3, 2, 3), "float32") = R.divide(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv = R.call_tir(divide, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def divide(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_divide: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_divide[ax0, ax1, ax2, ax3])
+ T_divide[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Divide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_divide_with_arg0_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Divide:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), dtype="float32") = R.divide(x, R.const(1, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(divide, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_divide[ax0, ax1])
+ T_divide[ax0, ax1] = rxplaceholder[ax0, ax1] / T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Divide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_divide_with_arg1_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Divide:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), dtype="float32") = R.divide(R.const(1, "float32"), x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(divide, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_divide[ax0, ax1])
+ T_divide[ax0, ax1] = T.float32(1) / rxplaceholder[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Divide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_divide_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Divide:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "float32") = R.divide(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(divide, (x, y), R.Tensor((a, b, c, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def divide(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_divide: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_divide = T.match_buffer(var_T_divide, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_divide[ax0, ax1, ax2, ax3])
+ T_divide[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Divide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_floor_divide():
+ # fmt: off
+ @tvm.script.ir_module
+ class FloorDivide:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv: R.Tensor((4, 3, 2, 3), "float32") = R.floor_divide(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv = R.call_tir(floor_divide, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def floor_divide(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_floor_divide: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_floor_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_floor_divide[ax0, ax1, ax2, ax3])
+ T_floor_divide[ax0, ax1, ax2, ax3] = T.floor(rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ # fmt: on
+
+ mod = LegalizeOps()(FloorDivide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_floor_divide_with_arg0_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class FloorDivide:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), dtype="float32") = R.floor_divide(x, R.const(1, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(floor_divide, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def floor_divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_floor_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_floor_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_floor_divide[ax0, ax1])
+ T_floor_divide[ax0, ax1] = T.floor(rxplaceholder[ax0, ax1] / T.float32(1))
+ # fmt: on
+
+ mod = LegalizeOps()(FloorDivide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_floor_divide_with_arg1_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class FloorDivide:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), dtype="float32") = R.floor_divide(R.const(1, "float32"), x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(floor_divide, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def floor_divide(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_floor_divide: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_floor_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_floor_divide[ax0, ax1])
+ T_floor_divide[ax0, ax1] = T.floor(T.float32(1) / rxplaceholder[ax0, ax1])
+ # fmt: on
+
+ mod = LegalizeOps()(FloorDivide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_floor_divide_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class FloorDivide:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "float32") = R.floor_divide(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(floor_divide, (x, y), R.Tensor((a, b, c, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def floor_divide(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_floor_divide: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_floor_divide = T.match_buffer(var_T_floor_divide, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_floor_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_floor_divide[ax0, ax1, ax2, ax3])
+ T_floor_divide[ax0, ax1, ax2, ax3] = T.floor(rxplaceholder[T.int64(0), ax2, ax3] / rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ # fmt: on
+
+ mod = LegalizeOps()(FloorDivide)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_multiply():
+ # fmt: off
+ @tvm.script.ir_module
+ class Multiply:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv: R.Tensor((4, 3, 2, 3), "float32") = R.multiply(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv = R.call_tir(multiply, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def multiply(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_multiply: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_multiply[ax0, ax1, ax2, ax3])
+ T_multiply[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] * rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Multiply)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_multiply_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Multiply:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "float32") = R.multiply(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(multiply, (x, y), R.Tensor((a, b, c, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def multiply(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_multiply: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_multiply = T.match_buffer(var_T_multiply, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_multiply[ax0, ax1, ax2, ax3])
+ T_multiply[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] * rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Multiply)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_subtract():
+ # fmt: off
+ @tvm.script.ir_module
+ class Subtract:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv: R.Tensor((4, 3, 2, 3), "float32") = R.subtract(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "float32"):
+ gv = R.call_tir(subtract, (x, y), R.Tensor((4, 3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def subtract(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_subtract: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] - rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Subtract)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_subtract_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Subtract:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "float32") = R.subtract(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(subtract, (x, y), R.Tensor((a, b, c, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def subtract(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_subtract: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_subtract = T.match_buffer(var_T_subtract, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] - rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Subtract)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+##################### Binary comparison #####################
+
+
+def test_equal():
+ # fmt: off
+ @tvm.script.ir_module
+ class Equal:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv: R.Tensor((4, 3, 2, 3), "bool") = R.equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv = R.call_tir(equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_equal[ax0, ax1, ax2, ax3])
+ T_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] == rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Equal)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_equal_with_arg0_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv: R.Tensor((2, 3), dtype="bool") = R.equal(x, R.const(1, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv = R.call_tir(equal, (x,), R.Tensor((2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_equal"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_equal[ax0, ax1])
+ T_equal[ax0, ax1] = rxplaceholder[ax0, ax1] == T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_equal_with_arg1_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv: R.Tensor((2, 3), dtype="bool") = R.equal(R.const(1, "float32"), x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv = R.call_tir(equal, (x,), R.Tensor((2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_equal"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_equal[ax0, ax1])
+ T_equal[ax0, ax1] = T.float32(1) == rxplaceholder[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_equal_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Equal:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "bool") = R.equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(equal, (x, y), R.Tensor((a, b, c, d), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_equal: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_equal = T.match_buffer(var_T_equal, [a, b, c, d], dtype="bool")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_equal[ax0, ax1, ax2, ax3])
+ T_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] == rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Equal)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_greater():
+ # fmt: off
+ @tvm.script.ir_module
+ class Greater:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv: R.Tensor((4, 3, 2, 3), "bool") = R.greater(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv = R.call_tir(greater, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def greater(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_greater: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_greater"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3])
+ T.writes(T_greater[ax0, ax1, ax2, ax3])
+ T_greater[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] < rxplaceholder[T.int64(0), ax2, ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(Greater)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_greater_with_arg0_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv: R.Tensor((2, 3), dtype="bool") = R.greater(x, R.const(1, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv = R.call_tir(greater, (x,), R.Tensor((2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def greater(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_greater: T.Buffer((T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_greater"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_greater[ax0, ax1])
+ T_greater[ax0, ax1] = T.float32(1) < rxplaceholder[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_greater_with_arg1_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv: R.Tensor((2, 3), dtype="bool") = R.greater(R.const(1, "float32"), x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv = R.call_tir(greater, (x,), R.Tensor((2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def greater(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_greater: T.Buffer((T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_greater"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_greater[ax0, ax1])
+ T_greater[ax0, ax1] = rxplaceholder[ax0, ax1] < T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_greater_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Greater:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "bool") = R.greater(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(greater, (x, y), R.Tensor((a, b, c, d), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def greater(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_greater: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_greater = T.match_buffer(var_T_greater, [a, b, c, d], dtype="bool")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_greater"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3])
+ T.writes(T_greater[ax0, ax1, ax2, ax3])
+ T_greater[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] < rxplaceholder[T.int64(0), ax2, ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(Greater)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_greater_equal():
+ # fmt: off
+ @tvm.script.ir_module
+ class GreaterEqual:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv: R.Tensor((4, 3, 2, 3), "bool") = R.greater_equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv = R.call_tir(greater_equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def greater_equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_greater_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_greater_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3])
+ T.writes(T_greater_equal[ax0, ax1, ax2, ax3])
+ T_greater_equal[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] <= rxplaceholder[T.int64(0), ax2, ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(GreaterEqual)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_greater_equal_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class GreaterEqual:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "bool") = R.greater_equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(greater_equal, (x, y), R.Tensor((a, b, c, d), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def greater_equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_greater_equal: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_greater_equal = T.match_buffer(var_T_greater_equal, [a, b, c, d], dtype="bool")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_greater_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[ax0, ax1, ax2, T.int64(0)], rxplaceholder[T.int64(0), ax2, ax3])
+ T.writes(T_greater_equal[ax0, ax1, ax2, ax3])
+ T_greater_equal[ax0, ax1, ax2, ax3] = rxplaceholder_1[ax0, ax1, ax2, T.int64(0)] <= rxplaceholder[T.int64(0), ax2, ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(GreaterEqual)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_less():
+ # fmt: off
+ @tvm.script.ir_module
+ class Less:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv: R.Tensor((4, 3, 2, 3), "bool") = R.less(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv = R.call_tir(less, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def less(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_less: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_less"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_less[ax0, ax1, ax2, ax3])
+ T_less[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] < rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Less)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_less_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Less:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "bool") = R.less(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(less, (x, y), R.Tensor((a, b, c, d), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def less(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_less: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_less = T.match_buffer(var_T_less, [a, b, c, d], dtype="bool")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_less"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_less[ax0, ax1, ax2, ax3])
+ T_less[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] < rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(Less)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_less_equal():
+ # fmt: off
+ @tvm.script.ir_module
+ class LessEqual:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv: R.Tensor((4, 3, 2, 3), "bool") = R.less_equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv = R.call_tir(less_equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def less_equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_less_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_less_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_less_equal[ax0, ax1, ax2, ax3])
+ T_less_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] <= rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(LessEqual)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_less_equal_with_arg0_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv: R.Tensor((2, 3), dtype="bool") = R.less_equal(x, R.const(1, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv = R.call_tir(less_equal, (x,), R.Tensor((2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def less_equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_less_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_less_equal"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_less_equal[ax0, ax1])
+ T_less_equal[ax0, ax1] = rxplaceholder[ax0, ax1] <= T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_less_equal_with_arg1_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Add:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv: R.Tensor((2, 3), dtype="bool") = R.less_equal(R.const(1, "float32"), x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "bool"):
+ gv = R.call_tir(less_equal, (x,), R.Tensor((2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def less_equal(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_less_equal: T.Buffer((T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_less_equal"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_less_equal[ax0, ax1])
+ T_less_equal[ax0, ax1] = T.float32(1) <= rxplaceholder[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Add)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_less_equal_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class LessEqual:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "bool") = R.less_equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(less_equal, (x, y), R.Tensor((a, b, c, d), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def less_equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_less_equal: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_less_equal = T.match_buffer(var_T_less_equal, [a, b, c, d], dtype="bool")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_less_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_less_equal[ax0, ax1, ax2, ax3])
+ T_less_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] <= rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(LessEqual)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_not_equal():
+ # fmt: off
+ @tvm.script.ir_module
+ class NotEqual:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv: R.Tensor((4, 3, 2, 3), "bool") = R.not_equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3), "float32"), y: R.Tensor((4, 3, 2, 1), "float32")) -> R.Tensor((4, 3, 2, 3), "bool"):
+ gv = R.call_tir(not_equal, (x, y), R.Tensor((4, 3, 2, 3), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def not_equal(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(1)), "float32"), T_not_equal: T.Buffer((T.int64(4), T.int64(3), T.int64(2), T.int64(3)), "bool")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_not_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_not_equal[ax0, ax1, ax2, ax3])
+ T_not_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] != rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(NotEqual)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_not_equal_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class NotEqual:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "bool") = R.not_equal(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, "c", "d"), "float32"), y: R.Tensor(("a", "b", "c", 1), "float32")) -> R.Tensor(("a", "b", "c", "d"), "bool"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(not_equal, (x, y), R.Tensor((a, b, c, d), dtype="bool"))
+ return gv
+
+ @T.prim_func
+ def not_equal(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_not_equal: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(1), c, d], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b, c, T.int64(1)], dtype="float32")
+ T_not_equal = T.match_buffer(var_T_not_equal, [a, b, c, d], dtype="bool")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_not_equal"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[T.int64(0), ax2, ax3], rxplaceholder_1[ax0, ax1, ax2, T.int64(0)])
+ T.writes(T_not_equal[ax0, ax1, ax2, ax3])
+ T_not_equal[ax0, ax1, ax2, ax3] = rxplaceholder[T.int64(0), ax2, ax3] != rxplaceholder_1[ax0, ax1, ax2, T.int64(0)]
+ # fmt: on
+
+ mod = LegalizeOps()(NotEqual)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_create_datatype.py b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
new file mode 100644
index 0000000000..2506e96634
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
@@ -0,0 +1,806 @@
+# 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 tvm
+from tvm.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+##################### Creation #####################
+
+
+def test_full():
+ # fmt: off
+ @tvm.script.ir_module
+ class Full:
+ @R.function
+ def main(v: R.Tensor((), "int32")) -> R.Tensor((2, 3), "int32"):
+ gv: R.Tensor((2, 3), "int32") = R.full((2, 3), v, dtype="int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(v: R.Tensor((), "int32")) -> R.Tensor((2, 3), "int32"):
+ gv = R.call_tir(full, (v,), R.Tensor((2, 3), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def full(rxplaceholder: T.Buffer((), "int32"), T_full: T.Buffer((T.int64(2), T.int64(3)), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[()])
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = rxplaceholder[()]
+ # fmt: on
+
+ mod = LegalizeOps()(Full)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_constant_scalar_fill_value():
+ # fmt: off
+ @tvm.script.ir_module
+ class Full:
+ @R.function
+ def main() -> R.Tensor((2, 3), "int32"):
+ gv: R.Tensor((2, 3), "int32") = R.full((2, 3), R.const(3.5, "float32"), dtype="int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main() -> R.Tensor((2, 3), "int32"):
+ gv = R.call_tir(full, R.tuple(), R.Tensor((2, 3), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def full(T_full: T.Buffer((T.int64(2), T.int64(3)), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = 3
+ # fmt: on
+
+ mod = LegalizeOps()(Full)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_different_dtype():
+ # fmt: off
+ @tvm.script.ir_module
+ class Full:
+ @R.function
+ def main(v: R.Tensor((), "int32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.full((2, 3), v, dtype="float32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(v: R.Tensor((), "int32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(full, (v,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def full(rxplaceholder: T.Buffer((), "int32"), T_full: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[()])
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.Cast("float32", rxplaceholder[()])
+ # fmt: on
+
+ mod = LegalizeOps()(Full)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Full:
+ @R.function
+ def main(dumb_param: R.Tensor(("m", "n")), v: R.Tensor((), "int32")) -> R.Tensor(("m", "n"), "int32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "int32") = R.full((m, n), v, dtype="int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(dumb_param: R.Tensor(("m", "n")), v: R.Tensor((), "int32")) -> R.Tensor(("m", "n"), "int32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(full, (v,), R.Tensor((m, n), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def full(rxplaceholder: T.Buffer((), "int32"), var_T_full: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ T_full = T.match_buffer(var_T_full, [m, n], dtype="int32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[()])
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = rxplaceholder[()]
+ # fmt: on
+
+ mod = LegalizeOps()(Full)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_like():
+ # fmt: off
+ @tvm.script.ir_module
+ class FullLike:
+ @R.function
+ def main(x: R.Tensor((2, 3), "int32"), v: R.Tensor((), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.full_like(x, v)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "int32"), v: R.Tensor((), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(full, (v,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def full(rxplaceholder: T.Buffer((), "float32"), T_full: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[()])
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = rxplaceholder[()]
+ # fmt: on
+
+ mod = LegalizeOps()(FullLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_like_constant_scalar_fill_value():
+ # fmt: off
+ @tvm.script.ir_module
+ class FullLike:
+ @R.function
+ def main(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.full_like(x, R.const(-5, "float32"))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "int32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(full, R.tuple(), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def full(T_full: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(-5)
+ # fmt: on
+
+ mod = LegalizeOps()(FullLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_like_different_dtype():
+ # fmt: off
+ @tvm.script.ir_module
+ class FullLike:
+ @R.function
+ def main(x: R.Tensor((2, 3), "int32"), v: R.Tensor((), "float32")) -> R.Tensor((2, 3), "float64"):
+ gv: R.Tensor((2, 3), "float64") = R.full_like(x, v, dtype="float64")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "int32"), v: R.Tensor((), "float32")) -> R.Tensor((2, 3), "float64"):
+ gv = R.call_tir(full, (v,), R.Tensor((2, 3), dtype="float64"))
+ return gv
+
+ @T.prim_func
+ def full(rxplaceholder: T.Buffer((), "float32"), T_full: T.Buffer((T.int64(2), T.int64(3)), "float64")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[()])
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.Cast("float64", rxplaceholder[()])
+ # fmt: on
+
+ mod = LegalizeOps()(FullLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_full_like_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class FullLike:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "int32"), v: R.Tensor((), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.full_like(x, v)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "int32"), v: R.Tensor((), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(full, (v,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def full(rxplaceholder: T.Buffer((), "float32"), var_T_full: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ T_full = T.match_buffer(var_T_full, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[()])
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = rxplaceholder[()]
+ # fmt: on
+
+ mod = LegalizeOps()(FullLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_ones():
+ # fmt: off
+ @tvm.script.ir_module
+ class Ones:
+ @R.function
+ def main() -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.ones((2, 3), "float32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main() -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(ones, R.tuple(), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def ones(T_full: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Ones)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_ones_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Ones:
+ @R.function
+ def main(dumb_param: R.Tensor(("m", "n"))) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.ones((m, n), "float32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(dumb_param: R.Tensor(("m", "n"))) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(ones, R.tuple(), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def ones(var_T_full: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ T_full = T.match_buffer(var_T_full, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(Ones)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_ones_like():
+ # fmt: off
+ @tvm.script.ir_module
+ class OnesLike:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "int32"):
+ gv: R.Tensor((2, 3), "int32") = R.ones_like(x, "int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "int32"):
+ gv = R.call_tir(ones, R.tuple(), R.Tensor((2, 3), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def ones(T_full: T.Buffer((T.int64(2), T.int64(3)), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = 1
+ # fmt: on
+
+ mod = LegalizeOps()(OnesLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_ones_like_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class OnesLike:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.ones_like(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(ones, R.tuple(), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def ones(var_T_full: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ T_full = T.match_buffer(var_T_full, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(1)
+ # fmt: on
+
+ mod = LegalizeOps()(OnesLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_zeros():
+ # fmt: off
+ @tvm.script.ir_module
+ class Zeros:
+ @R.function
+ def main() -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.zeros((2, 3), "float32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main() -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(zeros, R.tuple(), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def zeros(T_full: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(0)
+ # fmt: on
+
+ mod = LegalizeOps()(Zeros)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_zeros_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Zeros:
+ @R.function
+ def main(dumb_param: R.Tensor(("m", "n"))) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.zeros((m, n), "float32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(dumb_param: R.Tensor(("m", "n"))) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(zeros, R.tuple(), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def zeros(var_T_full: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ T_full = T.match_buffer(var_T_full, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(0)
+ # fmt: on
+
+ mod = LegalizeOps()(Zeros)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_zeros_like():
+ # fmt: off
+ @tvm.script.ir_module
+ class ZerosLike:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "int32"):
+ gv: R.Tensor((2, 3), "int32") = R.zeros_like(x, "int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "int32"):
+ gv = R.call_tir(zeros, R.tuple(), R.Tensor((2, 3), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def zeros(T_full: T.Buffer((T.int64(2), T.int64(3)), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = 0
+ # fmt: on
+
+ mod = LegalizeOps()(ZerosLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_zeros_like_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class ZerosLike:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.zeros_like(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(zeros, R.tuple(), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def zeros(var_T_full: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ T_full = T.match_buffer(var_T_full, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_full"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads()
+ T.writes(T_full[ax0, ax1])
+ T_full[ax0, ax1] = T.float32(0)
+ # fmt: on
+
+ mod = LegalizeOps()(ZerosLike)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_tril():
+ # fmt: off
+ @tvm.script.ir_module
+ class Tril:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 3, 4), "float32"):
+ gv: R.Tensor((2, 3, 4), "float32") = R.tril(x, k=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 3, 4), "float32"):
+ gv = R.call_tir(tril, (x,), R.Tensor((2, 3, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tril(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), trilu: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
+ with T.block("trilu"):
+ i0_1, i1_1, i2_1 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1])
+ T.writes(trilu[i0_1, i1_1, i2_1])
+ trilu[i0_1, i1_1, i2_1] = T.Select(i2_1 - T.int64(1) <= i1_1, rxplaceholder[i0_1, i1_1, i2_1], T.float32(0))
+ # fmt: on
+
+ mod = LegalizeOps()(Tril)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_tril_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Tril:
+ @R.function
+ def main(x: R.Tensor(("m", "n", "k"), "int8")) -> R.Tensor(("m", "n", "k"), "int8"):
+ m = T.var("int64")
+ n = T.var("int64")
+ k = T.var("int64")
+ gv: R.Tensor((m, n, k), "int8") = R.tril(x, k=-2)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n", "k"), "int8")) -> R.Tensor(("m", "n", "k"), "int8"):
+ m = T.var("int64")
+ n = T.var("int64")
+ k = T.var("int64")
+ gv = R.call_tir(tril, (x,), R.Tensor((m, n, k), dtype="int8"))
+ return gv
+
+ @T.prim_func
+ def tril(var_rxplaceholder: T.handle, var_trilu: T.handle):
+ T.func_attr({"tir.noalias": True})
+ k = T.var("int64")
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n, k], dtype="int8")
+ trilu = T.match_buffer(var_trilu, [m, n, k], dtype="int8")
+ for i0, i1, i2 in T.grid(m, n, k):
+ with T.block("trilu"):
+ i0_1, i1_1, i2_1 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1])
+ T.writes(trilu[i0_1, i1_1, i2_1])
+ trilu[i0_1, i1_1, i2_1] = T.Select(i2_1 + T.int64(2) <= i1_1, rxplaceholder[i0_1, i1_1, i2_1], T.int8(0))
+ # fmt: on
+
+ mod = LegalizeOps()(Tril)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_triu():
+ # fmt: off
+ @tvm.script.ir_module
+ class Triu:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 3, 4), "float32"):
+ gv: R.Tensor((2, 3, 4), "float32") = R.triu(x, k=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 3, 4), "float32"):
+ gv = R.call_tir(triu, (x,), R.Tensor((2, 3, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def triu(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), trilu: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
+ with T.block("trilu"):
+ i0_1, i1_1, i2_1 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1])
+ T.writes(trilu[i0_1, i1_1, i2_1])
+ trilu[i0_1, i1_1, i2_1] = T.Select(i1_1 <= i2_1 - T.int64(1), rxplaceholder[i0_1, i1_1, i2_1], T.float32(0))
+ # fmt: on
+
+ mod = LegalizeOps()(Triu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_triu_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Triu:
+ @R.function
+ def main(x: R.Tensor(("m", "n", "k"), "int8")) -> R.Tensor(("m", "n", "k"), "int8"):
+ m = T.var("int64")
+ n = T.var("int64")
+ k = T.var("int64")
+ gv: R.Tensor((m, n, k), "int8") = R.triu(x, k=-2)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n", "k"), "int8")) -> R.Tensor(("m", "n", "k"), "int8"):
+ m = T.var("int64")
+ n = T.var("int64")
+ k = T.var("int64")
+ gv = R.call_tir(triu, (x,), R.Tensor((m, n, k), dtype="int8"))
+ return gv
+
+ @T.prim_func
+ def triu(var_rxplaceholder: T.handle, var_trilu: T.handle):
+ T.func_attr({"tir.noalias": True})
+ k = T.var("int64")
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n, k], dtype="int8")
+ trilu = T.match_buffer(var_trilu, [m, n, k], dtype="int8")
+ for i0, i1, i2 in T.grid(m, n, k):
+ with T.block("trilu"):
+ i0_1, i1_1, i2_1 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1])
+ T.writes(trilu[i0_1, i1_1, i2_1])
+ trilu[i0_1, i1_1, i2_1] = T.Select(i1_1 <= i2_1 + T.int64(2), rxplaceholder[i0_1, i1_1, i2_1], T.int8(0))
+ # fmt: on
+
+ mod = LegalizeOps()(Triu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+##################### Datatype #####################
+
+
+def test_astype():
+ # fmt: off
+ @tvm.script.ir_module
+ class Astype:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 3, 4), "int32"):
+ gv: R.Tensor((2, 3, 4), "int32") = R.astype(x, "int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 3, 4), "int32"):
+ gv = R.call_tir(cast, (x,), R.Tensor((2, 3, 4), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def cast(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "int32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
+ with T.block("compute"):
+ i0_1, i1_1, i2_1 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1])
+ T.writes(compute[i0_1, i1_1, i2_1])
+ compute[i0_1, i1_1, i2_1] = T.Cast("int32", rxplaceholder[i0_1, i1_1, i2_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Astype)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_astype_input_constant_scalar():
+ # fmt: off
+ @tvm.script.ir_module
+ class Astype:
+ @R.function
+ def main() -> R.Tensor((), "int32"):
+ gv: R.Tensor((), "int32") = R.astype(R.const(1.5, "float32"), "int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main() -> R.Tensor((), "int32"):
+ gv: R.Tensor((), "int32") = R.const(1, "int32")
+ return gv
+ # fmt: on
+
+ mod = LegalizeOps()(Astype)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_astype_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Astype:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "int32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "int32") = R.astype(x, "int32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "int32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(cast, (x,), R.Tensor((m, n), dtype="int32"))
+ return gv
+
+ @T.prim_func
+ def cast(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="int32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.Cast("int32", rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Astype)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_image.py b/tests/python/relax/test_transform_legalize_ops_image.py
new file mode 100644
index 0000000000..36c8ecdd7b
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_image.py
@@ -0,0 +1,103 @@
+# 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 tvm
+from tvm.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+def test_image_resize2d():
+ # fmt: off
+ @tvm.script.ir_module
+ class Resize2D:
+ @R.function
+ def main(x: R.Tensor((2, 8, 8, 3), "float32")) -> R.Tensor((2, 16, 16, 3), "float32"):
+ gv: R.Tensor((2, 16, 16, 3), "float32") = R.image.resize2d(x, size=(16, 16), layout="NHWC", method="nearest_neighbor", coordinate_transformation_mode="asymmetric")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 8, 8, 3), "float32")) -> R.Tensor((2, 16, 16, 3), "float32"):
+ gv = R.call_tir(resize2d, (x,), R.Tensor((2, 16, 16, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def resize2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(8), T.int64(8), T.int64(3)), "float32"), resize: T.Buffer((T.int64(2), T.int64(16), T.int64(16), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(16), T.int64(16), T.int64(3)):
+ with T.block("resize"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, T.max(T.min(T.Div(i1_1, T.int64(2)), T.int64(7)), T.int64(0)), T.max(T.min(T.Div(i2_1, T.int64(2)), T.int64(7)), T.int64(0)), i3_1])
+ T.writes(resize[i0_1, i1_1, i2_1, i3_1])
+ resize[i0_1, i1_1, i2_1, i3_1] = rxplaceholder[i0_1, T.max(T.min(T.Div(i1_1, T.int64(2)), T.int64(7)), T.int64(0)), T.max(T.min(T.Div(i2_1, T.int64(2)), T.int64(7)), T.int64(0)), i3_1]
+ # fmt: on
+
+ mod = LegalizeOps()(Resize2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_image_resize2d_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Resize2D:
+ @R.function
+ def main(dumb_param: R.Tensor(("oh", "ow")), x: R.Tensor(("n", "c", "h", "w", 16), "float32")) -> R.Tensor(("n", "c", "oh", "ow", 16), "float32"):
+ n = T.var("int64")
+ c = T.var("int64")
+ oh = T.var("int64")
+ ow = T.var("int64")
+ gv: R.Tensor((n, c, oh, ow, 16), "float32") = R.image.resize2d(x, size=(oh, ow), layout="NCHW16c", method="nearest_neighbor", coordinate_transformation_mode="asymmetric")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(dumb_param: R.Tensor(("oh", "ow")), x: R.Tensor(("n", "c", "h", "w", 16), "float32")) -> R.Tensor(("n", "c", "oh", "ow", 16), "float32"):
+ n = T.var("int64")
+ c = T.var("int64")
+ oh = T.var("int64")
+ ow = T.var("int64")
+ gv = R.call_tir(resize2d, (x,), R.Tensor((n, c, oh, ow, 16), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def resize2d(var_rxplaceholder: T.handle, var_resize: T.handle):
+ T.func_attr({"tir.noalias": True})
+ c = T.var("int64")
+ h = T.var("int64")
+ n = T.var("int64")
+ oh = T.var("int64")
+ ow = T.var("int64")
+ w = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [n, c, h, w, T.int64(16)], dtype="float32")
+ resize = T.match_buffer(var_resize, [n, c, oh, ow, T.int64(16)], dtype="float32")
+ for i0, i1, i2, i3, i4 in T.grid(n, c, oh, ow, T.int64(16)):
+ with T.block("resize"):
+ i0_1, i1_1, i2_1, i3_1, i4_1 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
+ T.reads(rxplaceholder[i0_1, i1_1, T.int64(0) : T.max(h, T.int64(1)), T.int64(0) : T.max(w, T.int64(1)), i4_1])
+ T.writes(resize[i0_1, i1_1, i2_1, i3_1, i4_1])
+ resize[i0_1, i1_1, i2_1, i3_1, i4_1] = rxplaceholder[i0_1, i1_1, T.max(T.min(T.Cast("int64", T.round(T.Cast("float32", h) / T.Cast("float32", oh) * T.Cast("float32", i2_1), dtype="float32")), h - T.int64(1)), T.int64(0)), T.max(T.min(T.Cast("int64", T.round(T.Cast("float32", w) / T.Cast("float32", ow) * T.Cast("float32", i3_1), dtype="float32")), w - T.int64(1)), T.int64(0)), i4_1]
+ # fmt: on
+
+ mod = LegalizeOps()(Resize2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
new file mode 100644
index 0000000000..8b6f9de981
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
@@ -0,0 +1,401 @@
+# 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 tvm
+from tvm.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+##################### Indexing #####################
+
+
+def test_take():
+ # fmt: off
+ @tvm.script.ir_module
+ class Take:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32"), indices: R.Tensor((4,), "int64")) -> R.Tensor((2, 4, 4), "float32"):
+ gv: R.Tensor((2, 4, 4), "float32") = R.take(x, indices, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32"), indices: R.Tensor((4,), "int64")) -> R.Tensor((2, 4, 4), "float32"):
+ gv = R.call_tir(take, (x, indices), R.Tensor((2, 4, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def take(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), rxplaceholder_1: T.Buffer(T.int64(4), "int64"), T_take: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
+ with T.block("T_take"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, rxplaceholder_1[ax1], ax2], rxplaceholder_1[ax1])
+ T.writes(T_take[ax0, ax1, ax2])
+ T_take[ax0, ax1, ax2] = rxplaceholder[ax0, rxplaceholder_1[ax1], ax2]
+ # fmt: on
+
+ mod = LegalizeOps()(Take)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_take_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Take:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32"), indices: R.Tensor(("i",), "int64")) -> R.Tensor(("m", "i"), "float32"):
+ m = T.var("int64")
+ i = T.var("int64")
+ gv: R.Tensor((m, i), "float32") = R.take(x, indices, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32"), indices: R.Tensor(("i",), "int64")) -> R.Tensor(("m", "i"), "float32"):
+ m = T.var("int64")
+ i = T.var("int64")
+ gv = R.call_tir(take, (x, indices), R.Tensor((m, i), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def take(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_T_take: T.handle):
+ T.func_attr({"tir.noalias": True})
+ i = T.var("int64")
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [i], dtype="int64")
+ T_take = T.match_buffer(var_T_take, [m, i], dtype="float32")
+ for i0, i1 in T.grid(m, i):
+ with T.block("T_take"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, rxplaceholder_1[ax1]], rxplaceholder_1[ax1])
+ T.writes(T_take[ax0, ax1])
+ T_take[ax0, ax1] = rxplaceholder[ax0, rxplaceholder_1[ax1]]
+ # fmt: on
+
+ mod = LegalizeOps()(Take)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_strided_slice():
+ # fmt: off
+ @tvm.script.ir_module
+ class StridedSlice:
+ @R.function
+ def main(x: R.Tensor((8, 9, 10, 10), "float32")) -> R.Tensor((4, 9, 10, 3), "float32"):
+ gv: R.Tensor((4, 9, 10, 3), "float32") = R.strided_slice(x, axes=[0, 1, 3], begin=[1, 0, 8], end=[8, 9, 0], strides=[2, 1, -3])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((8, 9, 10, 10), dtype="float32")) -> R.Tensor((4, 9, 10, 3), dtype="float32"):
+ gv = R.call_tir(strided_slice, (x,), R.Tensor((4, 9, 10, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def strided_slice(rxplaceholder: T.Buffer((T.int64(8), T.int64(9), T.int64(10), T.int64(10)), "float32"), T_strided_slice_with_axes: T.Buffer((T.int64(4), T.int64(9), T.int64(10), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(9), T.int64(10), T.int64(3)):
+ with T.block("T_strided_slice_with_axes"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0 * T.int64(2) + T.int64(1), ax1, ax2, T.int64(8) - ax3 * T.int64(3)])
+ T.writes(T_strided_slice_with_axes[ax0, ax1, ax2, ax3])
+ T_strided_slice_with_axes[ax0, ax1, ax2, ax3] = rxplaceholder[ax0 * T.int64(2) + T.int64(1), ax1, ax2, T.int64(8) - ax3 * T.int64(3)]
+ # fmt: on
+
+ mod = LegalizeOps()(StridedSlice)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_strided_slice_no_strides():
+ # fmt: off
+ @tvm.script.ir_module
+ class StridedSlice:
+ @R.function
+ def main(x: R.Tensor((8, 9, 10, 10), "float32")) -> R.Tensor((4, 9, 10, 3), "float32"):
+ gv: R.Tensor((4, 9, 10, 3), "float32") = R.strided_slice(x, axes=[0, 1, 3], begin=[1, 0, 2], end=[8, 9, 4])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((8, 9, 10, 10), dtype="float32")) -> R.Tensor((4, 9, 10, 3), dtype="float32"):
+ gv = R.call_tir(strided_slice, (x,), out_sinfo=R.Tensor((7, 9, 10, 2), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def strided_slice(rxplaceholder: T.Buffer((T.int64(8), T.int64(9), T.int64(10), T.int64(10)), "float32"), T_strided_slice_with_axes: T.Buffer((T.int64(7), T.int64(9), T.int64(10), T.int64(2)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ # with T.block("root"):
+ for ax0, ax1, ax2, ax3 in T.grid(T.int64(7), T.int64(9), T.int64(10), T.int64(2)):
+ with T.block("T_strided_slice_with_axes"):
+ v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
+ T.reads(rxplaceholder[v_ax0 + T.int64(1), v_ax1, v_ax2, v_ax3 + T.int64(2)])
+ T.writes(T_strided_slice_with_axes[v_ax0, v_ax1, v_ax2, v_ax3])
+ T_strided_slice_with_axes[v_ax0, v_ax1, v_ax2, v_ax3] = rxplaceholder[v_ax0 + T.int64(1), v_ax1, v_ax2, v_ax3 + T.int64(2)]
+ # fmt: on
+
+ mod = LegalizeOps()(StridedSlice)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_strided_slice_symbolic_sliced_axis():
+ # fmt: off
+ @tvm.script.ir_module
+ class StridedSlice:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor((2, "n"), "float32"):
+ n = T.var("int64")
+ gv: R.Tensor((2, n), "float32") = R.strided_slice(x, axes=[0], begin=[1], end=[8], strides=[3])
+ return gv
+ # fmt: on
+
+ mod = LegalizeOps()(StridedSlice)
+ tvm.ir.assert_structural_equal(mod, StridedSlice)
+
+
+def test_strided_slice_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class StridedSlice:
+ @R.function
+ def main(x: R.Tensor((10, "n"), "float32")) -> R.Tensor((3, "n"), "float32"):
+ n = T.var("int64")
+ gv: R.Tensor((3, n), "float32") = R.strided_slice(x, axes=[0], begin=[1], end=[8], strides=[3])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((10, "n"), dtype="float32")) -> R.Tensor((3, "n"), dtype="float32"):
+ n = T.var("int64")
+ gv = R.call_tir(strided_slice, (x,), R.Tensor((3, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def strided_slice(var_rxplaceholder: T.handle, var_T_strided_slice_with_axes: T.handle):
+ T.func_attr({"tir.noalias": True})
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [T.int64(10), n], dtype="float32")
+ T_strided_slice_with_axes = T.match_buffer(var_T_strided_slice_with_axes, [T.int64(3), n], dtype="float32")
+ for i0, i1 in T.grid(T.int64(3), n):
+ with T.block("T_strided_slice_with_axes"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0 * T.int64(3) + T.int64(1), ax1])
+ T.writes(T_strided_slice_with_axes[ax0, ax1])
+ T_strided_slice_with_axes[ax0, ax1] = rxplaceholder[ax0 * T.int64(3) + T.int64(1), ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(StridedSlice)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+##################### Linear algebra #####################
+
+
+def test_matmul_1_4():
+ # fmt: off
+ @tvm.script.ir_module
+ class Matmul:
+ @R.function
+ def main(x: R.Tensor((4,), "float32"), y: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 3, 5), "float32"):
+ gv: R.Tensor((2, 3, 5), "float32") = R.matmul(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((4,), "float32"), y: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 3, 5), "float32"):
+ gv = R.call_tir(matmul, (x, y), R.Tensor((2, 3, 5), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def matmul(rxplaceholder: T.Buffer(T.int64(4), "float32"), rxplaceholder_1: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), matmul: T.Buffer((T.int64(2), T.int64(3), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(5), T.int64(4)):
+ with T.block("matmul"):
+ i0_1, i1_1, i2_1, k = T.axis.remap("SSSR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k], rxplaceholder_1[i0_1, i1_1, k, i2_1])
+ T.writes(matmul[i0_1, i1_1, i2_1])
+ with T.init():
+ matmul[i0_1, i1_1, i2_1] = T.float32(0)
+ matmul[i0_1, i1_1, i2_1] = matmul[i0_1, i1_1, i2_1] + rxplaceholder[k] * rxplaceholder_1[i0_1, i1_1, k, i2_1]
+ # fmt: on
+
+ mod = LegalizeOps()(Matmul)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_matmul_4_1():
+ # fmt: off
+ @tvm.script.ir_module
+ class Matmul:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32"), y: R.Tensor((5,), "float32")) -> R.Tensor((2, 3, 4), "float32"):
+ gv: R.Tensor((2, 3, 4), "float32") = R.matmul(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32"), y: R.Tensor((5,), "float32")) -> R.Tensor((2, 3, 4), "float32"):
+ gv = R.call_tir(matmul, (x, y), R.Tensor((2, 3, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def matmul(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_1: T.Buffer(T.int64(5), "float32"), matmul: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("matmul"):
+ i0_1, i1_1, i2_1, k = T.axis.remap("SSSR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1, k], rxplaceholder_1[k])
+ T.writes(matmul[i0_1, i1_1, i2_1])
+ with T.init():
+ matmul[i0_1, i1_1, i2_1] = T.float32(0)
+ matmul[i0_1, i1_1, i2_1] = matmul[i0_1, i1_1, i2_1] + rxplaceholder[i0_1, i1_1, i2_1, k] * rxplaceholder_1[k]
+ # fmt: on
+
+ mod = LegalizeOps()(Matmul)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_matmul_1_1():
+ # fmt: off
+ @tvm.script.ir_module
+ class Matmul:
+ @R.function
+ def main(x: R.Tensor((4,), "float32"), y: R.Tensor((4,), "float32")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.matmul(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((4,), "float32"), y: R.Tensor((4,), "float32")) -> R.Tensor((), "float32"):
+ gv = R.call_tir(matmul, (x, y), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def matmul(rxplaceholder: T.Buffer(T.int64(4), "float32"), rxplaceholder_1: T.Buffer(T.int64(4), "float32"), matmul: T.Buffer((), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0 in T.serial(T.int64(4)):
+ with T.block("matmul"):
+ k = T.axis.reduce(T.int64(4), i0)
+ T.reads(rxplaceholder[k], rxplaceholder_1[k])
+ T.writes(matmul[()])
+ with T.init():
+ matmul[()] = T.float32(0)
+ matmul[()] = matmul[()] + rxplaceholder[k] * rxplaceholder_1[k]
+ # fmt: on
+
+ mod = LegalizeOps()(Matmul)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_matmul_4_5():
+ # fmt: off
+ @tvm.script.ir_module
+ class Matmul:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float16"), y: R.Tensor((6, 2, 3, 5, 7), "float16")) -> R.Tensor((6, 2, 3, 4, 7), "float32"):
+ gv: R.Tensor((6, 2, 3, 4, 7), "float32") = R.matmul(x, y, out_dtype="float32")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float16"), y: R.Tensor((6, 2, 3, 5, 7), "float16")) -> R.Tensor((6, 2, 3, 4, 7), "float32"):
+ gv = R.call_tir(matmul, (x, y), R.Tensor((6, 2, 3, 4, 7), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def matmul(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float16"), rxplaceholder_1: T.Buffer((T.int64(6), T.int64(2), T.int64(3), T.int64(5), T.int64(7)), "float16"), matmul: T.Buffer((T.int64(6), T.int64(2), T.int64(3), T.int64(4), T.int64(7)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(6), T.int64(2), T.int64(3), T.int64(4), T.int64(7), T.int64(5)):
+ with T.block("matmul"):
+ i0_1, i1_1, i2_1, i3_1, i4_1, k = T.axis.remap("SSSSSR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[i1_1, i2_1, i3_1, k], rxplaceholder_1[i0_1, i1_1, i2_1, k, i4_1])
+ T.writes(matmul[i0_1, i1_1, i2_1, i3_1, i4_1])
+ with T.init():
+ matmul[i0_1, i1_1, i2_1, i3_1, i4_1] = T.float32(0)
+ matmul[i0_1, i1_1, i2_1, i3_1, i4_1] = matmul[i0_1, i1_1, i2_1, i3_1, i4_1] + T.Cast("float32", rxplaceholder[i1_1, i2_1, i3_1, k]) * T.Cast("float32", rxplaceholder_1[i0_1, i1_1, i2_1, k, i4_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Matmul)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_matmul_4_5_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Matmul:
+ @R.function
+ def main(x: R.Tensor(("b", 1, "m", "k"), "float32"), y: R.Tensor(("a", 1, "c", "k", "n"), "float32")) -> R.Tensor(("a", "b", "c", "m", "n"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((a, b, c, m, n), "float32") = R.matmul(x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("b", 1, "m", "k"), "float32"), y: R.Tensor(("a", 1, "c", "k", "n"), "float32")) -> R.Tensor(("a", "b", "c", "m", "n"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(matmul, (x, y), R.Tensor((a, b, c, m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def matmul(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_matmul: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ k = T.var("int64")
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [b, T.int64(1), m, k], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, T.int64(1), c, k, n], dtype="float32")
+ matmul = T.match_buffer(var_matmul, [a, b, c, m, n], dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(a, b, c, m, n, k):
+ with T.block("matmul"):
+ i0_1, i1_1, i2_1, i3_1, i4_1, k_1 = T.axis.remap("SSSSSR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[i1_1, T.int64(0), i3_1, k_1], rxplaceholder_1[i0_1, T.int64(0), i2_1, k_1, i4_1])
+ T.writes(matmul[i0_1, i1_1, i2_1, i3_1, i4_1])
+ with T.init():
+ matmul[i0_1, i1_1, i2_1, i3_1, i4_1] = T.float32(0)
+ matmul[i0_1, i1_1, i2_1, i3_1, i4_1] = matmul[i0_1, i1_1, i2_1, i3_1, i4_1] + rxplaceholder[i1_1, T.int64(0), i3_1, k_1] * rxplaceholder_1[i0_1, T.int64(0), i2_1, k_1, i4_1]
+ # fmt: on
+
+ mod = LegalizeOps()(Matmul)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py b/tests/python/relax/test_transform_legalize_ops_manipulate.py
new file mode 100644
index 0000000000..53aa868ffe
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -0,0 +1,789 @@
+# 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.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+##################### Manipulation #####################
+
+
+def test_broadcast_to():
+ # fmt: off
+ @tvm.script.ir_module
+ class BroadcastTo:
+ @R.function
+ def main(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"):
+ gv: R.Tensor((4, 2, 5, 3), "float32") = R.broadcast_to(x, (4, 2, 5, 3))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 1, 3), "float32")) -> R.Tensor((4, 2, 5, 3), "float32"):
+ gv = R.call_tir(broadcast_to, (x,), R.Tensor((4, 2, 5, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def broadcast_to(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3)), "float32"), T_broadcast_to: T.Buffer((T.int64(4), T.int64(2), T.int64(5), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(2), T.int64(5), T.int64(3)):
+ with T.block("T_broadcast_to"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax1, T.int64(0), ax3])
+ T.writes(T_broadcast_to[ax0, ax1, ax2, ax3])
+ T_broadcast_to[ax0, ax1, ax2, ax3] = rxplaceholder[ax1, T.int64(0), ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(BroadcastTo)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_broadcast_to_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class BroadcastTo:
+ @R.function
+ def main(dumb_param: R.Tensor(("a", "c")), x: R.Tensor(("b", 1, "d"), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, b, c, d), "float32") = R.broadcast_to(x, (a, b, c, d))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(dumb_param: R.Tensor(("a", "c")), x: R.Tensor(("b", 1, "d"), "float32")) -> R.Tensor(("a", "b", "c", "d"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(broadcast_to, (x,), R.Tensor((a, b, c, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def broadcast_to(var_rxplaceholder: T.handle, var_T_broadcast_to: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [b, T.int64(1), d], dtype="float32")
+ T_broadcast_to = T.match_buffer(var_T_broadcast_to, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_broadcast_to"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax1, T.int64(0), ax3])
+ T.writes(T_broadcast_to[ax0, ax1, ax2, ax3])
+ T_broadcast_to[ax0, ax1, ax2, ax3] = rxplaceholder[ax1, T.int64(0), ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(BroadcastTo)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_concat():
+ # fmt: off
+ @tvm.script.ir_module
+ class Concat:
+ @R.function
+ def main(x1: R.Tensor((1, 2, 3), "float32"), x2: R.Tensor((1, 3, 3), "float32"), x3: R.Tensor((1, 4, 3), "float32")) -> R.Tensor((1, 9, 3), "float32"):
+ gv: R.Tensor((1, 9, 3), "float32") = R.concat((x1, x2, x3), axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x1: R.Tensor((1, 2, 3), "float32"), x2: R.Tensor((1, 3, 3), "float32"), x3: R.Tensor((1, 4, 3), "float32")) -> R.Tensor((1, 9, 3), "float32"):
+ gv = R.call_tir(concatenate, (x1, x2, x3), R.Tensor((1, 9, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def concatenate(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3)), "float32"), rxplaceholder_1: T.Buffer((T.int64(1), T.int64(3), T.int64(3)), "float32"), rxplaceholder_2: T.Buffer((T.int64(1), T.int64(4), T.int64(3)), "float32"), T_concat: T.Buffer((T.int64(1), T.int64(9), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(9), T.int64(3)):
+ with T.block("T_concat"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder_2[ax0, ax1 - T.int64(5), ax2], rxplaceholder_1[ax0, ax1 - T.int64(2), ax2], rxplaceholder[ax0, ax1, ax2])
+ T.writes(T_concat[ax0, ax1, ax2])
+ T_concat[ax0, ax1, ax2] = T.if_then_else(T.int64(5) <= ax1, rxplaceholder_2[ax0, ax1 - T.int64(5), ax2], T.if_then_else(T.int64(2) <= ax1, rxplaceholder_1[ax0, ax1 - T.int64(2), ax2], rxplaceholder[ax0, ax1, ax2]))
+ # fmt: on
+
+ mod = LegalizeOps()(Concat)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_concat_input_tuple_var():
+ # fmt: off
+ @tvm.script.ir_module
+ class Concat:
+ @R.function
+ def main(t: R.Tuple(R.Tensor((3, 4), "float32"), R.Tensor((3, 5), "float32"))) -> R.Tensor((3, 9), "float32"):
+ gv: R.Tensor((3, 9), "float32") = R.concat(t, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor((3, 4), "float32"), R.Tensor((3, 5), "float32"))) -> R.Tensor((3, 9), "float32"):
+ gv: R.Tensor((3, 4), dtype="float32") = t[0]
+ gv1: R.Tensor((3, 5), dtype="float32") = t[1]
+ gv2 = R.call_tir(concatenate, (gv, gv1), R.Tensor((3, 9), dtype="float32"))
+ return gv2
+
+ @T.prim_func
+ def concatenate(rxplaceholder: T.Buffer((T.int64(3), T.int64(4)), "float32"), rxplaceholder_1: T.Buffer((T.int64(3), T.int64(5)), "float32"), T_concat: T.Buffer((T.int64(3), T.int64(9)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(3), T.int64(9)):
+ with T.block("T_concat"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder_1[ax0, ax1 - T.int64(4)], rxplaceholder[ax0, ax1])
+ T.writes(T_concat[ax0, ax1])
+ T_concat[ax0, ax1] = T.if_then_else(T.int64(4) <= ax1, rxplaceholder_1[ax0, ax1 - T.int64(4)], rxplaceholder[ax0, ax1])
+ # fmt: on
+
+ mod = LegalizeOps()(Concat)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_concat_input_tuple_var_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Concat:
+ @R.function
+ def main(t: R.Tuple(R.Tensor(("a", "b0"), "float32"), R.Tensor(("a", "b1"), "float32"), R.Tensor(("a", "b2"), "float32"))) -> R.Tensor(("a", "b0 + b1 + b2"), "float32"):
+ a = T.var("int64")
+ b0 = T.var("int64")
+ b1 = T.var("int64")
+ b2 = T.var("int64")
+ gv: R.Tensor((a, b0 + b1 + b2), "float32") = R.concat(t, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor(("a", "b0"), "float32"), R.Tensor(("a", "b1"), "float32"), R.Tensor(("a", "b2"), "float32"))) -> R.Tensor(("a", "b0 + b1 + b2"), "float32"):
+ a = T.var("int64")
+ b0 = T.var("int64")
+ b1 = T.var("int64")
+ b2 = T.var("int64")
+ gv: R.Tensor((a, b0), dtype="float32") = t[0]
+ gv1: R.Tensor((a, b1), dtype="float32") = t[1]
+ gv2: R.Tensor((a, b2), dtype="float32") = t[2]
+ gv3 = R.call_tir(concatenate, (gv, gv1, gv2), R.Tensor((a, ((b0 + b1) + b2)), dtype="float32"))
+ return gv3
+
+ @T.prim_func
+ def concatenate(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_T_concat: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b0 = T.var("int64")
+ b1 = T.var("int64")
+ b2 = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b0], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [a, b1], dtype="float32")
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [a, b2], dtype="float32")
+ T_concat = T.match_buffer(var_T_concat, [a, b0 + b1 + b2], dtype="float32")
+ for i0, i1 in T.grid(a, b0 + b1 + b2):
+ with T.block("T_concat"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder_2[ax0, ax1 - b0 - b1], rxplaceholder_1[ax0, ax1 - b0], rxplaceholder[ax0, ax1])
+ T.writes(T_concat[ax0, ax1])
+ T_concat[ax0, ax1] = T.if_then_else(T.int64(0) <= ax1 - b0 - b1, rxplaceholder_2[ax0, ax1 - b0 - b1], T.if_then_else(T.int64(0) <= ax1 - b0, rxplaceholder_1[ax0, ax1 - b0], rxplaceholder[ax0, ax1]))
+ # fmt: on
+
+ mod = LegalizeOps()(Concat)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_expand_dims():
+ # fmt: off
+ @tvm.script.ir_module
+ class ExpandDims:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"):
+ gv: R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32") = R.expand_dims(x, axis=[-1, 1, -6, 3, 5])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), "float32"):
+ gv = R.call_tir(expand_dims, (x,), R.Tensor((2, 1, 1, 1, 3, 1, 4, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def expand_dims(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), expand_dims: T.Buffer((T.int64(2), T.int64(1), T.int64(1), T.int64(1), T.int64(3), T.int64(1), T.int64(4), T.int64(1)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(2), T.int64(1), T.int64(1), T.int64(1), T.int64(3), T.int64(1), T.int64(4), T.int64(1)):
+ with T.block("expand_dims"):
+ i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1 = T.axis.remap("SSSSSSSS", [i0, i1, i2, i3, i4, i5, i6, i7])
+ T.reads(rxplaceholder[i0_1, i4_1, i6_1])
+ T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1])
+ expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1, i6_1, i7_1] = rxplaceholder[i0_1, i4_1, i6_1]
+ # fmt: on
+
+ mod = LegalizeOps()(ExpandDims)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_expand_dims_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class ExpandDims:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", 1, "b", 1, "c", 1), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tensor((a, 1, b, 1, c, 1), "float32") = R.expand_dims(x, axis=[1, 3, 5])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", 1, "b", 1, "c", 1), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(expand_dims, (x,), R.Tensor((a, 1, b, 1, c, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def expand_dims(var_rxplaceholder: T.handle, var_expand_dims: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
+ expand_dims = T.match_buffer(var_expand_dims, [a, T.int64(1), b, T.int64(1), c, T.int64(1)], dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(a, T.int64(1), b, T.int64(1), c, T.int64(1)):
+ with T.block("expand_dims"):
+ i0_1, i1_1, i2_1, i3_1, i4_1, i5_1 = T.axis.remap("SSSSSS", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[i0_1, i2_1, i4_1])
+ T.writes(expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1])
+ expand_dims[i0_1, i1_1, i2_1, i3_1, i4_1, i5_1] = rxplaceholder[i0_1, i2_1, i4_1]
+ # fmt: on
+
+ mod = LegalizeOps()(ExpandDims)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_flatten():
+ # fmt: off
+ @tvm.script.ir_module
+ class Flatten:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((24,), "float32"):
+ gv: R.Tensor((24,), "float32") = R.flatten(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4), "float32")) -> R.Tensor((24,), "float32"):
+ gv = R.call_tir(reshape, (x,), R.Tensor((24,), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def reshape(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), T_reshape: T.Buffer(T.int64(24), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0 in T.serial(T.int64(24)):
+ with T.block("T_reshape"):
+ ax0 = T.axis.spatial(T.int64(24), i0)
+ T.reads(rxplaceholder[ax0 % T.int64(24) // T.int64(12), ax0 % T.int64(12) // T.int64(4), ax0 % T.int64(4)])
+ T.writes(T_reshape[ax0])
+ T_reshape[ax0] = rxplaceholder[ax0 % T.int64(24) // T.int64(12), ax0 % T.int64(12) // T.int64(4), ax0 % T.int64(4)]
+ # fmt: on
+
+ mod = LegalizeOps()(Flatten)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_flatten_zero_rank():
+ # fmt: off
+ @tvm.script.ir_module
+ class Flatten:
+ @R.function
+ def main(x: R.Tensor((), "float32")) -> R.Tensor((1,), "float32"):
+ gv: R.Tensor((1,), "float32") = R.flatten(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((), "float32")) -> R.Tensor((1,), "float32"):
+ gv = R.call_tir(reshape, (x,), R.Tensor((1,), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def reshape(rxplaceholder: T.Buffer((), "float32"), T_reshape: T.Buffer(T.int64(1), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0 in T.serial(T.int64(1)):
+ with T.block("T_reshape"):
+ ax0 = T.axis.spatial(T.int64(1), i0)
+ T.reads(rxplaceholder[()])
+ T.writes(T_reshape[ax0])
+ T_reshape[ax0] = rxplaceholder[()]
+ # fmt: on
+
+ mod = LegalizeOps()(Flatten)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_flatten_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Flatten:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a * b * c",), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tensor((a * b * c,), "float32") = R.flatten(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a * b * c",), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(reshape, (x,), R.Tensor((((a * b) * c),), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
+ T_reshape = T.match_buffer(var_T_reshape, [a * b * c], dtype="float32")
+ for i0 in T.serial(a * b * c):
+ with T.block("T_reshape"):
+ ax0 = T.axis.spatial(a * b * c, i0)
+ T.reads(rxplaceholder[ax0 // c // b % a, ax0 // c % b, ax0 % c])
+ T.writes(T_reshape[ax0])
+ T_reshape[ax0] = rxplaceholder[ax0 // c // b % a, ax0 // c % b, ax0 % c]
+ # fmt: on
+
+ mod = LegalizeOps()(Flatten)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_permute_dims():
+ # fmt: off
+ @tvm.script.ir_module
+ class PermuteDims:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"):
+ gv: R.Tensor((2, 4, 3, 1), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((2, 4, 3, 1), "float32"):
+ gv = R.call_tir(transpose, (x,), R.Tensor((2, 4, 3, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def transpose(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"), T_transpose: T.Buffer((T.int64(2), T.int64(4), T.int64(3), T.int64(1)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(4), T.int64(3), T.int64(1)):
+ with T.block("T_transpose"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax3, ax0, ax2, ax1])
+ T.writes(T_transpose[ax0, ax1, ax2, ax3])
+ T_transpose[ax0, ax1, ax2, ax3] = rxplaceholder[ax3, ax0, ax2, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(PermuteDims)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_permute_dims_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class PermuteDims:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("b", "d", "c", "a"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((b, d, c, a), "float32") = R.permute_dims(x, axes=[1, -1, 2, -4])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor(("b", "d", "c", "a"), dtype="float32"):
+ b = T.var("int64")
+ d = T.var("int64")
+ c = T.var("int64")
+ a = T.var("int64")
+ gv = R.call_tir(transpose, (x,), R.Tensor((b, d, c, a), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def transpose(var_rxplaceholder: T.handle, var_T_transpose: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ T_transpose = T.match_buffer(var_T_transpose, [b, d, c, a], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(b, d, c, a):
+ with T.block("T_transpose"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax3, ax0, ax2, ax1])
+ T.writes(T_transpose[ax0, ax1, ax2, ax3])
+ T_transpose[ax0, ax1, ax2, ax3] = rxplaceholder[ax3, ax0, ax2, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(PermuteDims)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_reshape():
+ # fmt: off
+ @tvm.script.ir_module
+ class Reshape:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
+ gv: R.Tensor((8, 3), "float32") = R.reshape(x, (8, 3))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((1, 2, 3, 4), "float32")) -> R.Tensor((8, 3), "float32"):
+ gv = R.call_tir(reshape, (x,), R.Tensor((8, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def reshape(rxplaceholder: T.Buffer((T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"), T_reshape: T.Buffer((T.int64(8), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(8), T.int64(3)):
+ with T.block("T_reshape"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[T.int64(0), (ax0 * T.int64(3) + ax1) % T.int64(24) // T.int64(12), (ax0 * T.int64(3) + ax1) % T.int64(12) // T.int64(4), (ax0 * T.int64(3) + ax1) % T.int64(4)])
+ T.writes(T_reshape[ax0, ax1])
+ T_reshape[ax0, ax1] = rxplaceholder[T.int64(0), (ax0 * T.int64(3) + ax1) % T.int64(24) // T.int64(12), (ax0 * T.int64(3) + ax1) % T.int64(12) // T.int64(4), (ax0 * T.int64(3) + ax1) % T.int64(4)]
+ # fmt: on
+
+ mod = LegalizeOps()(Reshape)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_reshape_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Reshape:
+ @R.function
+ def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ gv: R.Tensor((a // 2, b * 2), "float32") = R.reshape(x, (a // 2, b * 2))
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b"), "float32")) -> R.Tensor(("a // 2", "b * 2"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ gv = R.call_tir(reshape, (x,), R.Tensor(((a // 2), (b * 2)), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def reshape(var_rxplaceholder: T.handle, var_T_reshape: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b], dtype="float32")
+ T_reshape = T.match_buffer(var_T_reshape, [a // T.int64(2), b * T.int64(2)], dtype="float32")
+ for i0, i1 in T.grid(a // T.int64(2), b * T.int64(2)):
+ with T.block("T_reshape"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[(ax0 * (b * T.int64(2)) + ax1) // b % a, (ax0 * (b * T.int64(2)) + ax1) % b])
+ T.writes(T_reshape[ax0, ax1])
+ T_reshape[ax0, ax1] = rxplaceholder[(ax0 * (b * T.int64(2)) + ax1) // b % a, (ax0 * (b * T.int64(2)) + ax1) % b]
+ # fmt: on
+
+ mod = LegalizeOps()(Reshape)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_split_by_indices():
+ # fmt: off
+ @tvm.script.ir_module
+ class Split:
+ @R.function
+ def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]):
+ gv: R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]) = R.split(x, [3, 7], axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")]):
+ gv = R.call_tir(split, (x,), [R.Tensor((2, 3, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 3, 4), "float32")])
+ return gv
+
+ @T.prim_func
+ def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32"), T_split_1: T.Buffer((T.int64(2), T.int64(4), T.int64(4)), "float32"), T_split_2: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
+ with T.block("T_split"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1, ax2])
+ T.writes(T_split[ax0, ax1, ax2])
+ T_split[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(4), T.int64(4)):
+ with T.block("T_split_1"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1 + T.int64(3), ax2])
+ T.writes(T_split_1[ax0, ax1, ax2])
+ T_split_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(3), ax2]
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
+ with T.block("T_split_2"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1 + T.int64(7), ax2])
+ T.writes(T_split_2[ax0, ax1, ax2])
+ T_split_2[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(7), ax2]
+ # fmt: on
+
+ mod = LegalizeOps()(Split)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_split_by_indices_n_section_indivisible():
+ # fmt: off
+ @tvm.script.ir_module
+ class Split:
+ @R.function
+ def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]):
+ gv: R.Tuple([R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 4, 4), "float32"), R.Tensor((2, 2, 4), "float32")]) = R.split(x, 3, axis=1)
+ return gv
+ # fmt: on
+
+ mod = LegalizeOps()(Split)
+ tvm.ir.assert_structural_equal(mod, Split)
+
+
+def test_split_by_indices_n_section_divisible():
+ # fmt: off
+ @tvm.script.ir_module
+ class Split:
+ @R.function
+ def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]):
+ gv: R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]) = R.split(x, 2, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 10, 4), "float32")) -> R.Tuple([R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")]):
+ gv = R.call_tir(split, (x,), [R.Tensor((2, 5, 4), "float32"), R.Tensor((2, 5, 4), "float32")])
+ return gv
+
+ @T.prim_func
+ def split(rxplaceholder: T.Buffer((T.int64(2), T.int64(10), T.int64(4)), "float32"), T_split_sections: T.Buffer((T.int64(2), T.int64(5), T.int64(4)), "float32"), T_split_sections_1: T.Buffer((T.int64(2), T.int64(5), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(5), T.int64(4)):
+ with T.block("T_split_sections"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1, ax2])
+ T.writes(T_split_sections[ax0, ax1, ax2])
+ T_split_sections[ax0, ax1, ax2] = rxplaceholder[ax0, ax1, ax2]
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(5), T.int64(4)):
+ with T.block("T_split_sections_1"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1 + T.int64(5), ax2])
+ T.writes(T_split_sections_1[ax0, ax1, ax2])
+ T_split_sections_1[ax0, ax1, ax2] = rxplaceholder[ax0, ax1 + T.int64(5), ax2]
+ # fmt: on
+
+ mod = LegalizeOps()(Split)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_split_by_indices_n_section_divisible_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Split:
+ @R.function
+ def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "n * 3"), "float32")) -> R.Tuple([R.Tensor(("m", "n"), "float32"), R.Tensor(("m", "n"), "float32"), R.Tensor(("m", "n"), "float32")]):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tuple([R.Tensor((m, n), "float32"), R.Tensor((m, n), "float32"), R.Tensor((m, n), "float32")]) = R.split(x, 3, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "(n * 3)"), "float32")) -> R.Tuple(R.Tensor(("m", "((n * 3) // 3)"), "float32"), R.Tensor(("m", "((((n * 3) // 3) * 2) - ((n * 3) // 3))"), "float32"), R.Tensor(("m", "((n * 3) - (((n * 3) // 3) * 2))"), "float32")):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(split, (x,), [R.Tensor((m, ((n * 3) // 3)), "float32"), R.Tensor((m, ((((n * 3) // 3) * 2) - ((n * 3) // 3))), "float32"), R.Tensor((m, ((n * 3) - (((n * 3) // 3) * 2))), "float32")], tir_vars=(n,))
+ return gv
+
+ @T.prim_func
+ def split(var_rxplaceholder: T.handle, var_T_split_sections: T.handle, var_T_split_sections_1: T.handle, var_T_split_sections_2: T.handle, n: T.int64):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n * T.int64(3)], dtype="float32")
+ T_split_sections = T.match_buffer(var_T_split_sections, [m, n * T.int64(3) // T.int64(3)], dtype="float32")
+ T_split_sections_1 = T.match_buffer(var_T_split_sections_1, [m, n * T.int64(3) // T.int64(3) * T.int64(2) - n * T.int64(3) // T.int64(3)], dtype="float32")
+ T_split_sections_2 = T.match_buffer(var_T_split_sections_2, [m, n * T.int64(3) - n * T.int64(3) // T.int64(3) * T.int64(2)], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_split_sections"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_split_sections[ax0, ax1])
+ T_split_sections[ax0, ax1] = rxplaceholder[ax0, ax1]
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_split_sections_1"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, n + ax1])
+ T.writes(T_split_sections_1[ax0, ax1])
+ T_split_sections_1[ax0, ax1] = rxplaceholder[ax0, n + ax1]
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_split_sections_2"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, n * T.int64(2) + ax1])
+ T.writes(T_split_sections_2[ax0, ax1])
+ T_split_sections_2[ax0, ax1] = rxplaceholder[ax0, n * T.int64(2) + ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Split)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_squeeze():
+ # fmt: off
+ @tvm.script.ir_module
+ class Squeeze:
+ @R.function
+ def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
+ gv: R.Tensor((2, 3, 1, 4), "float32") = R.squeeze(x, [1, 4])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
+ gv = R.call_tir(squeeze, (x,), R.Tensor((2, 3, 1, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(1), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(1), T.int64(4)):
+ with T.block("T_squeeze"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, T.int64(0), ax1, ax2, T.int64(0), ax3])
+ T.writes(T_squeeze[ax0, ax1, ax2, ax3])
+ T_squeeze[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, T.int64(0), ax1, ax2, T.int64(0), ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(Squeeze)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_squeeze_no_axis():
+ # fmt: off
+ @tvm.script.ir_module
+ class Squeeze:
+ @R.function
+ def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
+ gv: R.Tensor((2, 3, 1, 4), "float32") = R.squeeze(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 1, 3, 1, 1, 4), "float32")) -> R.Tensor((2, 3, 1, 4), "float32"):
+ gv = R.call_tir(squeeze, (x,), R.Tensor((2, 3, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def squeeze(rxplaceholder: T.Buffer((T.int64(2), T.int64(1), T.int64(3), T.int64(1), T.int64(1), T.int64(4)), "float32"), T_squeeze: T.Buffer((T.int64(2), T.int64(3), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(2), T.int64(3), T.int64(4)):
+ with T.block("T_squeeze"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2])
+ T.writes(T_squeeze[ax0, ax1, ax2])
+ T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, T.int64(0), T.int64(0), ax2]
+ # fmt: on
+
+ mod = LegalizeOps()(Squeeze)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_squeeze_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Squeeze:
+ @R.function
+ def main(x: R.Tensor(("a", 1, "b", 1), "float32")) -> R.Tensor(("a", "b", 1), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ gv: R.Tensor((a, b, 1), "float32") = R.squeeze(x, [1])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", 1, "b", 1), "float32")) -> R.Tensor(("a", "b", 1), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ gv = R.call_tir(squeeze, (x,), R.Tensor((a, b, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def squeeze(var_rxplaceholder: T.handle, var_T_squeeze: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, T.int64(1), b, T.int64(1)], dtype="float32")
+ T_squeeze = T.match_buffer(var_T_squeeze, [a, b, T.int64(1)], dtype="float32")
+ for i0, i1, i2 in T.grid(a, b, T.int64(1)):
+ with T.block("T_squeeze"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, T.int64(0), ax1, ax2])
+ T.writes(T_squeeze[ax0, ax1, ax2])
+ T_squeeze[ax0, ax1, ax2] = rxplaceholder[ax0, T.int64(0), ax1, ax2]
+ # fmt: on
+
+ mod = LegalizeOps()(Squeeze)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py b/tests/python/relax/test_transform_legalize_ops_nn.py
new file mode 100644
index 0000000000..3f9f02c410
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -0,0 +1,1188 @@
+# 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.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+##################### Neural network #####################
+
+
+def test_conv2d():
+ # fmt: off
+ @tvm.script.ir_module
+ class Conv2d:
+ @R.function
+ def main(x: R.Tensor((2, 128, 28, 28), "float32"), w: R.Tensor((64, 16, 3, 3), "float32")) -> R.Tensor((2, 64, 13, 13), "float32"):
+ gv: R.Tensor((2, 4, 13, 13), "float32") = R.nn.conv2d(x, w, strides=(2, 2), padding=(1, 1), dilation=(2, 2), groups=8)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 128, 28, 28), "float32"), w: R.Tensor((64, 16, 3, 3), "float32")) -> R.Tensor((2, 64, 13, 13), "float32"):
+ gv = R.call_tir(conv2d, (x, w), R.Tensor((2, 64, 13, 13), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def conv2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(128), T.int64(28), T.int64(28)), "float32"), rxplaceholder_1: T.Buffer((T.int64(64), T.int64(16), T.int64(3), T.int64(3)), "float32"), group_conv2d_nchw: T.Buffer((T.int64(2), T.int64(64), T.int64(13), T.int64(13)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ pad_temp = T.alloc_buffer([T.int64(2), T.int64(128), T.int64(30), T.int64(30)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(128), T.int64(30), T.int64(30)):
+ with T.block("pad_temp"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1 - T.int64(1), i3_1 - T.int64(1)])
+ T.writes(pad_temp[i0_1, i1_1, i2_1, i3_1])
+ pad_temp[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(T.int64(1) <= i2_1 and i2_1 < T.int64(29) and T.int64(1) <= i3_1 and i3_1 < T.int64(29), rxplaceholder[i0_1, i1_1, i2_1 - T.int64(1), i3_1 - T.int64(1)], T.float32(0), dtype="float32")
+ for i0, i1, i2, i3, i4, i5, i6 in T.grid(T.int64(2), T.int64(64), T.int64(13), T.int64(13), T.int64(16), T.int64(3), T.int64(3)):
+ with T.block("group_conv2d_nchw"):
+ nn, ff, yy, xx, rc, ry, rx = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+ T.reads(pad_temp[nn, ff // T.int64(8) * T.int64(16) + rc, yy * T.int64(2) + ry * T.int64(2), xx * T.int64(2) + rx * T.int64(2)], rxplaceholder_1[ff, rc, ry, rx])
+ T.writes(group_conv2d_nchw[nn, ff, yy, xx])
+ with T.init():
+ group_conv2d_nchw[nn, ff, yy, xx] = T.float32(0)
+ group_conv2d_nchw[nn, ff, yy, xx] = group_conv2d_nchw[nn, ff, yy, xx] + pad_temp[nn, ff // T.int64(8) * T.int64(16) + rc, yy * T.int64(2) + ry * T.int64(2), xx * T.int64(2) + rx * T.int64(2)] * rxplaceholder_1[ff, rc, ry, rx]
+ # fmt: on
+
+ mod = LegalizeOps()(Conv2d)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_conv2d_with_out_dtype():
+ # fmt: off
+ @tvm.script.ir_module
+ class Conv2d:
+ @R.function
+ def main(x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3), "float32")) -> R.Tensor((2, 4, 26, 26), "float16"):
+ gv: R.Tensor((2, 4, 26, 26), "float16") = R.nn.conv2d(x, w, out_dtype="float16")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 28, 28), "float32"), w: R.Tensor((4, 3, 3, 3), "float32")) -> R.Tensor((2, 4, 26, 26), "float16"):
+ gv = R.call_tir(conv2d, (x, w), R.Tensor((2, 4, 26, 26), dtype="float16"))
+ return gv
+
+ @T.prim_func
+ def conv2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(28), T.int64(28)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(3), T.int64(3), T.int64(3)), "float32"), conv2d_nchw: T.Buffer((T.int64(2), T.int64(4), T.int64(26), T.int64(26)), "float16")):
+ T.func_attr({"tir.noalias": True})
+ pad_temp = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(28), T.int64(28)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(28), T.int64(28)):
+ with T.block("pad_temp"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1, i3_1])
+ T.writes(pad_temp[i0_1, i1_1, i2_1, i3_1])
+ pad_temp[i0_1, i1_1, i2_1, i3_1] = rxplaceholder[i0_1, i1_1, i2_1, i3_1]
+ for i0, i1, i2, i3, i4, i5, i6 in T.grid(T.int64(2), T.int64(4), T.int64(26), T.int64(26), T.int64(3), T.int64(3), T.int64(3)):
+ with T.block("conv2d_nchw"):
+ nn, ff, yy, xx, rc, ry, rx = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+ T.reads(pad_temp[nn, rc, yy + ry, xx + rx], rxplaceholder_1[ff, rc, ry, rx])
+ T.writes(conv2d_nchw[nn, ff, yy, xx])
+ with T.init():
+ conv2d_nchw[nn, ff, yy, xx] = T.float16(0)
+ conv2d_nchw[nn, ff, yy, xx] = conv2d_nchw[nn, ff, yy, xx] + T.Cast("float16", pad_temp[nn, rc, yy + ry, xx + rx]) * T.Cast("float16", rxplaceholder_1[ff, rc, ry, rx])
+ # fmt: on
+
+ mod = LegalizeOps()(Conv2d)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_conv2d_nhwc():
+ # fmt: off
+ @tvm.script.ir_module
+ class Conv2d:
+ @R.function
+ def main(x: R.Tensor((2, 28, 28, 128), "float32"), w: R.Tensor((64, 128, 3, 3), "float32")) -> R.Tensor((2, 26, 26, 64), "float32"):
+ gv: R.Tensor((2, 26, 26, 64), "float32") = R.nn.conv2d(x, w, data_layout="NHWC")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 28, 28, 128), "float32"), w: R.Tensor((64, 128, 3, 3), "float32")) -> R.Tensor((2, 26, 26, 64), "float32"):
+ gv = R.call_tir(conv2d, (x, w), R.Tensor((2, 26, 26, 64), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def conv2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(28), T.int64(28), T.int64(128)), "float32"), rxplaceholder_1: T.Buffer((T.int64(64), T.int64(128), T.int64(3), T.int64(3)), "float32"), conv2d_nhwc: T.Buffer((T.int64(2), T.int64(26), T.int64(26), T.int64(64)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ pad_temp = T.alloc_buffer([T.int64(2), T.int64(28), T.int64(28), T.int64(128)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(28), T.int64(28), T.int64(128)):
+ with T.block("pad_temp"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1, i3_1])
+ T.writes(pad_temp[i0_1, i1_1, i2_1, i3_1])
+ pad_temp[i0_1, i1_1, i2_1, i3_1] = rxplaceholder[i0_1, i1_1, i2_1, i3_1]
+ for i0, i1, i2, i3, i4, i5, i6 in T.grid(T.int64(2), T.int64(26), T.int64(26), T.int64(64), T.int64(3), T.int64(3), T.int64(128)):
+ with T.block("conv2d_nhwc"):
+ nn, yy, xx, ff, ry, rx, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+ T.reads(pad_temp[nn, yy + ry, xx + rx, rc], rxplaceholder_1[ff, rc, ry, rx])
+ T.writes(conv2d_nhwc[nn, yy, xx, ff])
+ with T.init():
+ conv2d_nhwc[nn, yy, xx, ff] = T.float32(0)
+ conv2d_nhwc[nn, yy, xx, ff] = conv2d_nhwc[nn, yy, xx, ff] + pad_temp[nn, yy + ry, xx + rx, rc] * rxplaceholder_1[ff, rc, ry, rx]
+ # fmt: on
+
+ mod = LegalizeOps()(Conv2d)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_conv2d_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Conv2d:
+ @R.function
+ def main(x: R.Tensor(("n", "c", "h", "w"), "float32"), kernel: R.Tensor(("f", "c", "kh", "kw"), "float32")) -> R.Tensor(("n", "f", "h - kh + 1", "w - kw + 1"), "float32"):
+ n = T.var("int64")
+ h = T.var("int64")
+ w = T.var("int64")
+ f = T.var("int64")
+ kh = T.var("int64")
+ kw = T.var("int64")
+ gv: R.Tensor((n, f, h - kh + 1, w - kw + 1), "float32") = R.nn.conv2d(x, kernel)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("n", "c", "h", "w"), "float32"), kernel: R.Tensor(("f", "c", "kh", "kw"), "float32")) -> R.Tensor(("n", "f", "h - kh + 1", "w - kw + 1"), "float32"):
+ n = T.var("int64")
+ f = T.var("int64")
+ h = T.var("int64")
+ kh = T.var("int64")
+ w = T.var("int64")
+ kw = T.var("int64")
+ gv = R.call_tir(conv2d, (x, kernel), R.Tensor((n, f, ((h - kh) + 1), ((w - kw) + 1)), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def conv2d(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_conv2d_nchw: T.handle):
+ T.func_attr({"tir.noalias": True})
+ c = T.var("int64")
+ f = T.var("int64")
+ h = T.var("int64")
+ kh = T.var("int64")
+ kw = T.var("int64")
+ n = T.var("int64")
+ w = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [n, c, h, w], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [f, c, kh, kw], dtype="float32")
+ conv2d_nchw = T.match_buffer(var_conv2d_nchw, [n, f, h - kh + T.int64(1), w - kw + T.int64(1)], dtype="float32")
+ pad_temp = T.alloc_buffer([n, c, h, w], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(n, c, h, w):
+ with T.block("pad_temp"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, i1_1, i2_1, i3_1])
+ T.writes(pad_temp[i0_1, i1_1, i2_1, i3_1])
+ pad_temp[i0_1, i1_1, i2_1, i3_1] = rxplaceholder[i0_1, i1_1, i2_1, i3_1]
+ for i0, i1, i2, i3, i4, i5, i6 in T.grid(n, f, h + T.int64(1) - kh, w + T.int64(1) - kw, c, kh, kw):
+ with T.block("conv2d_nchw"):
+ nn, ff, yy, xx, rc, ry, rx = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
+ T.reads(pad_temp[nn, rc, yy + ry, xx + rx], rxplaceholder_1[ff, rc, ry, rx])
+ T.writes(conv2d_nchw[nn, ff, yy, xx])
+ with T.init():
+ conv2d_nchw[nn, ff, yy, xx] = T.float32(0)
+ conv2d_nchw[nn, ff, yy, xx] = conv2d_nchw[nn, ff, yy, xx] + pad_temp[nn, rc, yy + ry, xx + rx] * rxplaceholder_1[ff, rc, ry, rx]
+ # fmt: on
+
+ mod = LegalizeOps()(Conv2d)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_max_pool2d():
+ # fmt: off
+ @tvm.script.ir_module
+ class MaxPool2D:
+ @R.function
+ def main(x: R.Tensor((4, 112, 112, 6), "float32")) -> R.Tensor((4, 56, 56, 6), "float32"):
+ gv: R.Tensor((4, 56, 56, 6), "float32") = R.nn.max_pool2d(x, pool_size=[3, 3], strides=[2, 2], dilation=[1, 1], padding=[1, 1, 1, 1], layout="NHWC")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((4, 112, 112, 6), "float32")) -> R.Tensor((4, 56, 56, 6), "float32"):
+ gv = R.call_tir(max_pool2d, (x,), R.Tensor((4, 56, 56, 6), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def max_pool2d(rxplaceholder: T.Buffer((T.int64(4), T.int64(112), T.int64(112), T.int64(6)), "float32"), pool_max: T.Buffer((T.int64(4), T.int64(56), T.int64(56), T.int64(6)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ pad_temp = T.alloc_buffer([T.int64(4), T.int64(114), T.int64(114), T.int64(6)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(114), T.int64(114), T.int64(6)):
+ with T.block("pad_temp"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1 - T.int64(1), ax2 - T.int64(1), ax3])
+ T.writes(pad_temp[ax0, ax1, ax2, ax3])
+ pad_temp[ax0, ax1, ax2, ax3] = T.if_then_else(T.int64(1) <= ax1 and ax1 < T.int64(113) and T.int64(1) <= ax2 and ax2 < T.int64(113), rxplaceholder[ax0, ax1 - T.int64(1), ax2 - T.int64(1), ax3], T.float32(-3.4028234663852886e+38), dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(4), T.int64(56), T.int64(56), T.int64(6), T.int64(3), T.int64(3)):
+ with T.block("pool_max"):
+ ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(pad_temp[ax0, ax1 * T.int64(2) + rv0, ax2 * T.int64(2) + rv1, ax3])
+ T.writes(pool_max[ax0, ax1, ax2, ax3])
+ T.block_attr({"schedule_rule":"meta_schedule.pool_max"})
+ with T.init():
+ pool_max[ax0, ax1, ax2, ax3] = T.float32(-3.4028234663852886e+38)
+ pool_max[ax0, ax1, ax2, ax3] = T.max(pool_max[ax0, ax1, ax2, ax3], pad_temp[ax0, ax1 * T.int64(2) + rv0, ax2 * T.int64(2) + rv1, ax3])
+ # fmt: on
+
+ mod = LegalizeOps()(MaxPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_max_pool2d_NCHW16c():
+ # fmt: off
+ @tvm.script.ir_module
+ class MaxPool2D:
+ @R.function
+ def main(x: R.Tensor((4, 4, 112, 112, 16), "float32")) -> R.Tensor((4, 4, 110, 110, 16), "float32"):
+ gv: R.Tensor((4, 4, 110, 110, 16), "float32") = R.nn.max_pool2d(x, pool_size=[3, 3], layout="NCHW16c")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((4, 4, 112, 112, 16), "float32")) -> R.Tensor((4, 4, 110, 110, 16), "float32"):
+ gv = R.call_tir(max_pool2d, (x,), R.Tensor((4, 4, 110, 110, 16), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def max_pool2d(rxplaceholder: T.Buffer((T.int64(4), T.int64(4), T.int64(112), T.int64(112), T.int64(16)), "float32"), pool_max: T.Buffer((T.int64(4), T.int64(4), T.int64(110), T.int64(110), T.int64(16)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3, i4, i5, i6 in T.grid(T.int64(4), T.int64(4), T.int64(110), T.int64(110), T.int64(16), T.int64(3), T.int64(3)):
+ with T.block("pool_max"):
+ ax0, ax1, ax2, ax3, ax4, rv0, rv1 = T.axis.remap("SSSSSRR", [i0, i1, i2, i3, i4, i5, i6])
+ T.reads(rxplaceholder[ax0, ax1, ax2 + rv0, ax3 + rv1, ax4])
+ T.writes(pool_max[ax0, ax1, ax2, ax3, ax4])
+ T.block_attr({"schedule_rule":"meta_schedule.pool_max"})
+ with T.init():
+ pool_max[ax0, ax1, ax2, ax3, ax4] = T.float32(-3.4028234663852886e+38)
+ pool_max[ax0, ax1, ax2, ax3, ax4] = T.max(pool_max[ax0, ax1, ax2, ax3, ax4], rxplaceholder[ax0, ax1, ax2 + rv0, ax3 + rv1, ax4])
+ # fmt: on
+
+ mod = LegalizeOps()(MaxPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_max_pool2d_ceil_mode():
+ # fmt: off
+ @tvm.script.ir_module
+ class MaxPool2D:
+ @R.function
+ def main(x: R.Tensor((4, 6, 112, 112), "float32")) -> R.Tensor((4, 6, 38, 38), "float32"):
+ gv: R.Tensor((4, 6, 38, 38), "float32") = R.nn.max_pool2d(x, pool_size=[3, 3], strides=[3, 3], dilation=[1, 1], padding=[1, 1, 1, 1], ceil_mode=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((4, 6, 112, 112), dtype="float32")) -> R.Tensor((4, 6, 38, 38), dtype="float32"):
+ gv = R.call_tir(max_pool2d, (x,), R.Tensor((4, 6, 38, 38), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def max_pool2d(rxplaceholder: T.Buffer((T.int64(4), T.int64(6), T.int64(112), T.int64(112)), "float32"), pool_max: T.Buffer((T.int64(4), T.int64(6), T.int64(38), T.int64(38)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ pad_temp = T.alloc_buffer([T.int64(4), T.int64(6), T.int64(116), T.int64(116)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(4), T.int64(6), T.int64(116), T.int64(116)):
+ with T.block("pad_temp"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2 - T.int64(1), ax3 - T.int64(1)])
+ T.writes(pad_temp[ax0, ax1, ax2, ax3])
+ pad_temp[ax0, ax1, ax2, ax3] = T.if_then_else(T.int64(1) <= ax2 and ax2 < T.int64(113) and T.int64(1) <= ax3 and ax3 < T.int64(113), rxplaceholder[ax0, ax1, ax2 - T.int64(1), ax3 - T.int64(1)], T.float32(-3.4028234663852886e+38), dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(4), T.int64(6), T.int64(38), T.int64(38), T.int64(3), T.int64(3)):
+ with T.block("pool_max"):
+ ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(pad_temp[ax0, ax1, ax2 * T.int64(3) + rv0, ax3 * T.int64(3) + rv1])
+ T.writes(pool_max[ax0, ax1, ax2, ax3])
+ T.block_attr({"schedule_rule":"meta_schedule.pool_max"})
+ with T.init():
+ pool_max[ax0, ax1, ax2, ax3] = T.float32(-3.4028234663852886e+38)
+ pool_max[ax0, ax1, ax2, ax3] = T.max(pool_max[ax0, ax1, ax2, ax3], pad_temp[ax0, ax1, ax2 * T.int64(3) + rv0, ax3 * T.int64(3) + rv1])
+ # fmt: on
+
+ mod = LegalizeOps()(MaxPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+@pytest.mark.skip("TOPI pooling casts every shape value to i32.")
+def test_max_pool2d_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class MaxPool2D:
+ @R.function
+ def main(dumb_param: R.Tensor(("kh", "kw")), x: R.Tensor(("n", "c", "h", "w"), "float32")) -> R.Tensor(("n", "c", "h - kh + 1", "w - kw + 1"), "float32"):
+ n = T.var("int64")
+ c = T.var("int64")
+ h = T.var("int64")
+ w = T.var("int64")
+ kh = T.var("int64")
+ kw = T.var("int64")
+ gv: R.Tensor((n, c, h - kh + 1, w - kw + 1), "float32") = R.nn.max_pool2d(x, pool_size=[kh, kw])
+ return gv
+
+ # fmt: on
+
+ mod = LegalizeOps()(MaxPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_adaptive_avg_pool2d():
+ # fmt: off
+ @tvm.script.ir_module
+ class AdaptiveAvgPool2D:
+ @R.function
+ def main(x: R.Tensor((2, 4, 7, 7, 16), "float32")) -> R.Tensor((2, 4, 1, 1, 16), "float32"):
+ gv: R.Tensor((2, 4, 1, 1, 16), "float32") = R.nn.adaptive_avg_pool2d(x, output_size=[1, 1], layout="NCHW16c")
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 4, 7, 7, 16), "float32")) -> R.Tensor((2, 4, 1, 1, 16), "float32"):
+ gv = R.call_tir(adaptive_avg_pool2d, (x,), R.Tensor((2, 4, 1, 1, 16), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def adaptive_avg_pool2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(4), T.int64(7), T.int64(7), T.int64(16)), "float32"), adaptive_pool_avg: T.Buffer((T.int64(2), T.int64(4), T.int64(1), T.int64(1), T.int64(16)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ adaptive_pool_sum = T.alloc_buffer([T.int64(2), T.int64(4), T.int64(1), T.int64(1), T.int64(16)], dtype="float32")
+ for i0, i1, i2, i3, i4, i5, i6 in T.grid(T.int64(2), T.int64(4), T.int64(1), T.int64(1), T.int64(16), T.int64(7), T.int64(7)):
+ with T.block("adaptive_pool_sum"):
+ ax0, ax1, ax2, ax3, ax4, rv0, rv1 = T.axis.remap("SSSSSRR", [i0, i1, i2, i3, i4, i5, i6])
+ T.reads(rxplaceholder[ax0, ax1, ax2 * T.int64(7) + rv0, ax3 * T.int64(7) + rv1, ax4])
+ T.writes(adaptive_pool_sum[ax0, ax1, ax2, ax3, ax4])
+ with T.init():
+ adaptive_pool_sum[ax0, ax1, ax2, ax3, ax4] = T.float32(0)
+ adaptive_pool_sum[ax0, ax1, ax2, ax3, ax4] = adaptive_pool_sum[ax0, ax1, ax2, ax3, ax4] + rxplaceholder[ax0, ax1, ax2 * T.int64(7) + rv0, ax3 * T.int64(7) + rv1, ax4]
+ for i0, i1, i2, i3, i4 in T.grid(T.int64(2), T.int64(4), T.int64(1), T.int64(1), T.int64(16)):
+ with T.block("adaptive_pool_avg"):
+ ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
+ T.reads(adaptive_pool_sum[ax0, ax1, ax2, ax3, ax4])
+ T.writes(adaptive_pool_avg[ax0, ax1, ax2, ax3, ax4])
+ T.block_attr({"schedule_rule":"meta_schedule.adaptive_pool_avg"})
+ adaptive_pool_avg[ax0, ax1, ax2, ax3, ax4] = adaptive_pool_sum[ax0, ax1, ax2, ax3, ax4] * T.float32(0.020408163265306121)
+ # fmt: on
+
+ mod = LegalizeOps()(AdaptiveAvgPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_adaptive_avg_pool2d_without_output_size():
+ # fmt: off
+ @tvm.script.ir_module
+ class AdaptiveAvgPool2D:
+ @R.function
+ def main(x: R.Tensor((2, 16, 7, 7), "float32")) -> R.Tensor((2, 16, 7, 7), "float32"):
+ gv: R.Tensor((2, 16, 7, 7), "float32") = R.nn.adaptive_avg_pool2d(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 16, 7, 7), "float32")) -> R.Tensor((2, 16, 7, 7), "float32"):
+ gv = R.call_tir(adaptive_avg_pool2d, (x,), R.Tensor((2, 16, 7, 7), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def adaptive_avg_pool2d(rxplaceholder: T.Buffer((T.int64(2), T.int64(16), T.int64(7), T.int64(7)), "float32"), adaptive_pool_avg: T.Buffer((T.int64(2), T.int64(16), T.int64(7), T.int64(7)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ adaptive_pool_sum = T.alloc_buffer([T.int64(2), T.int64(16), T.int64(7), T.int64(7)], dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(2), T.int64(16), T.int64(7), T.int64(7), T.int64(1), T.int64(1)):
+ with T.block("adaptive_pool_sum"):
+ ax0, ax1, ax2, ax3, rv0, rv1 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[ax0, ax1, ax2 + rv0, ax3 + rv1])
+ T.writes(adaptive_pool_sum[ax0, ax1, ax2, ax3])
+ with T.init():
+ adaptive_pool_sum[ax0, ax1, ax2, ax3] = T.float32(0)
+ adaptive_pool_sum[ax0, ax1, ax2, ax3] = adaptive_pool_sum[ax0, ax1, ax2, ax3] + rxplaceholder[ax0, ax1, ax2 + rv0, ax3 + rv1]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(16), T.int64(7), T.int64(7)):
+ with T.block("adaptive_pool_avg"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(adaptive_pool_sum[ax0, ax1, ax2, ax3])
+ T.writes(adaptive_pool_avg[ax0, ax1, ax2, ax3])
+ T.block_attr({"schedule_rule":"meta_schedule.adaptive_pool_avg"})
+ adaptive_pool_avg[ax0, ax1, ax2, ax3] = adaptive_pool_sum[ax0, ax1, ax2, ax3]
+ # fmt: on
+
+ mod = LegalizeOps()(AdaptiveAvgPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+@pytest.mark.skip("TOPI pooling casts every shape value to i32.")
+def test_adaptive_avg_pool2d_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class AdaptiveAvgPool2D:
+ @R.function
+ def main(dumb_param: R.Tensor(("oh", "ow")), x: R.Tensor(("n", "c", "h", "w"), "float32")) -> R.Tensor(("n", "c", "oh", "ow"), "float32"):
+ n = T.var("int64")
+ c = T.var("int64")
+ oh = T.var("int64")
+ ow = T.var("int64")
+ gv: R.Tensor((n, c, oh, ow), "float32") = R.nn.adaptive_avg_pool2d(x, (oh, ow))
+ return gv
+ # fmt: on
+
+ mod = LegalizeOps()(AdaptiveAvgPool2D)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_relu():
+ # fmt: off
+ @tvm.script.ir_module
+ class Relu:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.nn.relu(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(relu, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def relu(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.max(rxplaceholder[i0_1, i1_1], T.float32(0))
+ # fmt: on
+
+ mod = LegalizeOps()(Relu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_relu_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Relu:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.nn.relu(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(relu, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def relu(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.max(rxplaceholder[i0_1, i1_1], T.float32(0))
+ # fmt: on
+
+ mod = LegalizeOps()(Relu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_gelu():
+ # fmt: off
+ @tvm.script.ir_module
+ class Gelu:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.nn.gelu(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(gelu, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def gelu(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_multiply: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ T_multiply_1 = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ compute = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ T_multiply_2 = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ T_divide = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_multiply"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_multiply_1[ax0, ax1])
+ T_multiply_1[ax0, ax1] = rxplaceholder[ax0, ax1] * T.float32(0.70710678118654757)
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(T_multiply_1[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.erf(T_multiply_1[i0_1, i1_1], dtype="float32")
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_multiply_1"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(compute[ax0, ax1])
+ T.writes(T_multiply_2[ax0, ax1])
+ T_multiply_2[ax0, ax1] = compute[ax0, ax1] * T.float32(0.5)
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(T_multiply_2[ax0, ax1])
+ T.writes(T_divide[ax0, ax1])
+ T_divide[ax0, ax1] = T.float32(0.5) + T_multiply_2[ax0, ax1]
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_multiply_2"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1], T_divide[ax0, ax1])
+ T.writes(T_multiply[ax0, ax1])
+ T_multiply[ax0, ax1] = rxplaceholder[ax0, ax1] * T_divide[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Gelu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_gelu_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Gelu:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.nn.gelu(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(gelu, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def gelu(var_rxplaceholder: T.handle, var_T_multiply: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ T_multiply = T.match_buffer(var_T_multiply, [m, n], dtype="float32")
+ T_multiply_1 = T.alloc_buffer([m, n], dtype="float32")
+ compute = T.alloc_buffer([m, n], dtype="float32")
+ T_multiply_2 = T.alloc_buffer([m, n], dtype="float32")
+ T_add = T.alloc_buffer([m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_multiply"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1])
+ T.writes(T_multiply_1[ax0, ax1])
+ T_multiply_1[ax0, ax1] = rxplaceholder[ax0, ax1] * T.float32(0.70710678118654757)
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(T_multiply_1[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.erf(T_multiply_1[i0_1, i1_1], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_multiply_1"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(compute[ax0, ax1])
+ T.writes(T_multiply_2[ax0, ax1])
+ T_multiply_2[ax0, ax1] = compute[ax0, ax1] * T.float32(0.5)
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_add"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(T_multiply_2[ax0, ax1])
+ T.writes(T_add[ax0, ax1])
+ T_add[ax0, ax1] = T.float32(0.5) + T_multiply_2[ax0, ax1]
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_multiply_2"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1], T_add[ax0, ax1])
+ T.writes(T_multiply[ax0, ax1])
+ T_multiply[ax0, ax1] = rxplaceholder[ax0, ax1] * T_add[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Gelu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_silu():
+ # fmt: off
+ @tvm.script.ir_module
+ class Silu:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.nn.silu(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(silu, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def silu(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), T_multiply: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ compute = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sigmoid(rxplaceholder[i0_1, i1_1])
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("T_multiply"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1], compute[ax0, ax1])
+ T.writes(T_multiply[ax0, ax1])
+ T_multiply[ax0, ax1] = rxplaceholder[ax0, ax1] * compute[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Silu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_silu_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Silu:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.nn.silu(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(silu, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def silu(var_rxplaceholder: T.handle, var_T_multiply: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ T_multiply = T.match_buffer(var_T_multiply, [m, n], dtype="float32")
+ compute = T.alloc_buffer([m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sigmoid(rxplaceholder[i0_1, i1_1])
+ for i0, i1 in T.grid(m, n):
+ with T.block("T_multiply"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[ax0, ax1], compute[ax0, ax1])
+ T.writes(T_multiply[ax0, ax1])
+ T_multiply[ax0, ax1] = rxplaceholder[ax0, ax1] * compute[ax0, ax1]
+ # fmt: on
+
+ mod = LegalizeOps()(Silu)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_softmax():
+ # fmt: off
+ @tvm.script.ir_module
+ class Softmax:
+ @R.function
+ def main(x: R.Tensor((2, 3, 16, 32), "float32")) -> R.Tensor((2, 3, 16, 32), "float32"):
+ gv: R.Tensor((2, 3, 16, 32), "float32") = R.nn.softmax(x, axis=-2)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 16, 32), "float32")) -> R.Tensor((2, 3, 16, 32), "float32"):
+ gv = R.call_tir(softmax, (x,), R.Tensor((2, 3, 16, 32), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def softmax(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(16), T.int64(32)), "float32"), T_softmax_norm: T.Buffer((T.int64(2), T.int64(3), T.int64(16), T.int64(32)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ T_softmax_maxelem = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(32)], dtype="float32")
+ T_softmax_exp = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(16), T.int64(32)], dtype="float32")
+ T_softmax_expsum = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(32)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(32), T.int64(16)):
+ with T.block("T_softmax_maxelem"):
+ i0_1, i1_1, i2_1, k = T.axis.remap("SSSR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_1, i1_1, k, i2_1])
+ T.writes(T_softmax_maxelem[i0_1, i1_1, i2_1])
+ with T.init():
+ T_softmax_maxelem[i0_1, i1_1, i2_1] = T.float32(-3.4028234663852886e+38)
+ T_softmax_maxelem[i0_1, i1_1, i2_1] = T.max(T_softmax_maxelem[i0_1, i1_1, i2_1], rxplaceholder[i0_1, i1_1, k, i2_1])
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(16), T.int64(32)):
+ with T.block("T_softmax_exp"):
+ i0_2, i1_2, i2_2, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[i0_2, i1_2, i2_2, i3_1], T_softmax_maxelem[i0_2, i1_2, i3_1])
+ T.writes(T_softmax_exp[i0_2, i1_2, i2_2, i3_1])
+ T_softmax_exp[i0_2, i1_2, i2_2, i3_1] = T.exp(rxplaceholder[i0_2, i1_2, i2_2, i3_1] - T_softmax_maxelem[i0_2, i1_2, i3_1], dtype="float32")
+ for i0_3, i1_3, i2_3, i3 in T.grid(T.int64(2), T.int64(3), T.int64(32), T.int64(16)):
+ with T.block("T_softmax_expsum"):
+ i0_4, i1_4, i2_4, k = T.axis.remap("SSSR", [i0_3, i1_3, i2_3, i3])
+ T.reads(T_softmax_exp[i0_4, i1_4, k, i2_4])
+ T.writes(T_softmax_expsum[i0_4, i1_4, i2_4])
+ with T.init():
+ T_softmax_expsum[i0_4, i1_4, i2_4] = T.float32(0)
+ T_softmax_expsum[i0_4, i1_4, i2_4] = T_softmax_expsum[i0_4, i1_4, i2_4] + T_softmax_exp[i0_4, i1_4, k, i2_4]
+ for i0_5, i1_5, i2_5, i3 in T.grid(T.int64(2), T.int64(3), T.int64(16), T.int64(32)):
+ with T.block("T_softmax_norm"):
+ i0_6, i1_6, i2_6, i3_2 = T.axis.remap("SSSS", [i0_5, i1_5, i2_5, i3])
+ T.reads(T_softmax_exp[i0_6, i1_6, i2_6, i3_2], T_softmax_expsum[i0_6, i1_6, i3_2])
+ T.writes(T_softmax_norm[i0_6, i1_6, i2_6, i3_2])
+ T.block_attr({"axis":2})
+ T_softmax_norm[i0_6, i1_6, i2_6, i3_2] = T_softmax_exp[i0_6, i1_6, i2_6, i3_2] / T_softmax_expsum[i0_6, i1_6, i3_2]
+ # fmt: on
+
+ mod = LegalizeOps()(Softmax)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_softmax_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Softmax:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", "b", "c"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tensor((a, b, c), "float32") = R.nn.softmax(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c"), "float32")) -> R.Tensor(("a", "b", "c"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(softmax, (x,), R.Tensor((a, b, c), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def softmax(var_rxplaceholder: T.handle, var_T_softmax_norm: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c], dtype="float32")
+ T_softmax_norm = T.match_buffer(var_T_softmax_norm, [a, b, c], dtype="float32")
+ T_softmax_maxelem = T.alloc_buffer([a, b], dtype="float32")
+ T_softmax_exp = T.alloc_buffer([a, b, c], dtype="float32")
+ T_softmax_expsum = T.alloc_buffer([a, b], dtype="float32")
+ for i0, i1, i2 in T.grid(a, b, c):
+ with T.block("T_softmax_maxelem"):
+ i0_1, i1_1, k = T.axis.remap("SSR", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_1, i1_1, k])
+ T.writes(T_softmax_maxelem[i0_1, i1_1])
+ with T.init():
+ T_softmax_maxelem[i0_1, i1_1] = T.float32(-3.4028234663852886e+38)
+ T_softmax_maxelem[i0_1, i1_1] = T.max(T_softmax_maxelem[i0_1, i1_1], rxplaceholder[i0_1, i1_1, k])
+ for i0, i1, i2 in T.grid(a, b, c):
+ with T.block("T_softmax_exp"):
+ i0_2, i1_2, i2_1 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[i0_2, i1_2, i2_1], T_softmax_maxelem[i0_2, i1_2])
+ T.writes(T_softmax_exp[i0_2, i1_2, i2_1])
+ T_softmax_exp[i0_2, i1_2, i2_1] = T.exp(rxplaceholder[i0_2, i1_2, i2_1] - T_softmax_maxelem[i0_2, i1_2], dtype="float32")
+ for i0_3, i1_3, i2 in T.grid(a, b, c):
+ with T.block("T_softmax_expsum"):
+ i0_4, i1_4, k = T.axis.remap("SSR", [i0_3, i1_3, i2])
+ T.reads(T_softmax_exp[i0_4, i1_4, k])
+ T.writes(T_softmax_expsum[i0_4, i1_4])
+ with T.init():
+ T_softmax_expsum[i0_4, i1_4] = T.float32(0)
+ T_softmax_expsum[i0_4, i1_4] = T_softmax_expsum[i0_4, i1_4] + T_softmax_exp[i0_4, i1_4, k]
+ for i0_5, i1_5, i2 in T.grid(a, b, c):
+ with T.block("T_softmax_norm"):
+ i0_6, i1_6, i2_2 = T.axis.remap("SSS", [i0_5, i1_5, i2])
+ T.reads(T_softmax_exp[i0_6, i1_6, i2_2], T_softmax_expsum[i0_6, i1_6])
+ T.writes(T_softmax_norm[i0_6, i1_6, i2_2])
+ T.block_attr({"axis":2})
+ T_softmax_norm[i0_6, i1_6, i2_2] = T_softmax_exp[i0_6, i1_6, i2_2] / T_softmax_expsum[i0_6, i1_6]
+ # fmt: on
+
+ mod = LegalizeOps()(Softmax)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_batch_norm():
+ # fmt: off
+ @tvm.script.ir_module
+ class BatchNorm:
+ @R.function
+ def main(x: R.Tensor((2, 3, 28, 28), "float32"), gamma: R.Tensor((3,), "float32"), beta: R.Tensor((3,), "float32"), moving_mean: R.Tensor((3,), "float32"), moving_var: R.Tensor((3,), "float32")) -> R.Tuple(R.Tensor((2, 3, 28, 28), "float32"), R.Tensor((3,), "float32"), R.Tensor((3,), "float32")):
+ gv: R.Tuple(R.Tensor((2, 3, 28, 28), "float32"), R.Tensor((3,), "float32"), R.Tensor((3,), "float32")) = R.nn.batch_norm(x, gamma, beta, moving_mean, moving_var, axis=1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 28, 28), "float32"), gamma: R.Tensor((3,), "float32"), beta: R.Tensor((3,), "float32"), moving_mean: R.Tensor((3,), "float32"), moving_var: R.Tensor((3,), "float32")) -> R.Tuple(R.Tensor((2, 3, 28, 28), "float32"), R.Tensor((3,), "float32"), R.Tensor((3,), "float32")):
+ gv = R.call_tir(batch_norm, (x, gamma, beta, moving_mean, moving_var), [R.Tensor((2, 3, 28, 28), "float32"), R.Tensor((3,), "float32"), R.Tensor((3,), "float32")])
+ return gv
+
+ @T.prim_func
+ def batch_norm(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(28), T.int64(28)), "float32"), rxplaceholder_1: T.Buffer(T.int64(3), "float32"), rxplaceholder_2: T.Buffer(T.int64(3), "float32"), rxplaceholder_3: T.Buffer(T.int64(3), "float32"), rxplaceholder_4: T.Buffer(T.int64(3), "float32"), T_add: T.Buffer((T.int64(2), T.int64(3), T.int64(28), T.int64(28)), "float32"), T_multiply: T.Buffer(T.int64(3), "float32"), T_multiply_1: T.Buffer(T.int64(3), "float32")):
+ T.func_attr({"tir.noalias": True})
+ T_reshape = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(1), T.int64(1)], dtype="float32")
+ T_subtract = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(28), T.int64(28)], dtype="float32")
+ T_reshape_1 = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(1), T.int64(1)], dtype="float32")
+ T_add_1 = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(1), T.int64(1)], dtype="float32")
+ compute = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(1), T.int64(1)], dtype="float32")
+ T_divide = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(28), T.int64(28)], dtype="float32")
+ T_reshape_2 = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(1), T.int64(1)], dtype="float32")
+ T_multiply_2 = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(28), T.int64(28)], dtype="float32")
+ T_reshape_3 = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(1), T.int64(1)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(1), T.int64(1)):
+ with T.block("T_reshape"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_3[(ax1 + ax2 + ax3) % T.int64(3)])
+ T.writes(T_reshape[ax0, ax1, ax2, ax3])
+ T_reshape[ax0, ax1, ax2, ax3] = rxplaceholder_3[(ax1 + ax2 + ax3) % T.int64(3)]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(28), T.int64(28)):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_reshape[T.int64(0), ax1, T.int64(0), T.int64(0)])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_reshape[T.int64(0), ax1, T.int64(0), T.int64(0)]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(1), T.int64(1)):
+ with T.block("T_reshape_1"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_4[(ax1 + ax2 + ax3) % T.int64(3)])
+ T.writes(T_reshape_1[ax0, ax1, ax2, ax3])
+ T_reshape_1[ax0, ax1, ax2, ax3] = rxplaceholder_4[(ax1 + ax2 + ax3) % T.int64(3)]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(1), T.int64(1)):
+ with T.block("T_add"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_reshape_1[ax0, ax1, ax2, ax3])
+ T.writes(T_add_1[ax0, ax1, ax2, ax3])
+ T_add_1[ax0, ax1, ax2, ax3] = T_reshape_1[ax0, ax1, ax2, ax3] + T.float32(1.0000000000000001e-05)
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(1), T.int64(1)):
+ with T.block("compute"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_add_1[i0_1, i1_1, i2_1, i3_1])
+ T.writes(compute[i0_1, i1_1, i2_1, i3_1])
+ compute[i0_1, i1_1, i2_1, i3_1] = T.sqrt(T_add_1[i0_1, i1_1, i2_1, i3_1], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(28), T.int64(28)):
+ with T.block("T_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_subtract[ax0, ax1, ax2, ax3], compute[T.int64(0), ax1, T.int64(0), T.int64(0)])
+ T.writes(T_divide[ax0, ax1, ax2, ax3])
+ T_divide[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] / compute[T.int64(0), ax1, T.int64(0), T.int64(0)]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(1), T.int64(1)):
+ with T.block("T_reshape_2"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[(ax1 + ax2 + ax3) % T.int64(3)])
+ T.writes(T_reshape_2[ax0, ax1, ax2, ax3])
+ T_reshape_2[ax0, ax1, ax2, ax3] = rxplaceholder_1[(ax1 + ax2 + ax3) % T.int64(3)]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(28), T.int64(28)):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_divide[ax0, ax1, ax2, ax3], T_reshape_2[T.int64(0), ax1, T.int64(0), T.int64(0)])
+ T.writes(T_multiply_2[ax0, ax1, ax2, ax3])
+ T_multiply_2[ax0, ax1, ax2, ax3] = T_divide[ax0, ax1, ax2, ax3] * T_reshape_2[T.int64(0), ax1, T.int64(0), T.int64(0)]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(1), T.int64(1)):
+ with T.block("T_reshape_3"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_2[(ax1 + ax2 + ax3) % T.int64(3)])
+ T.writes(T_reshape_3[ax0, ax1, ax2, ax3])
+ T_reshape_3[ax0, ax1, ax2, ax3] = rxplaceholder_2[(ax1 + ax2 + ax3) % T.int64(3)]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(28), T.int64(28)):
+ with T.block("T_add_1"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_multiply_2[ax0, ax1, ax2, ax3], T_reshape_3[T.int64(0), ax1, T.int64(0), T.int64(0)])
+ T.writes(T_add[ax0, ax1, ax2, ax3])
+ T_add[ax0, ax1, ax2, ax3] = T_multiply_2[ax0, ax1, ax2, ax3] + T_reshape_3[T.int64(0), ax1, T.int64(0), T.int64(0)]
+ for i0 in T.serial(T.int64(3)):
+ with T.block("T_multiply_1"):
+ ax0 = T.axis.spatial(T.int64(3), i0)
+ T.reads(rxplaceholder_3[ax0])
+ T.writes(T_multiply[ax0])
+ T_multiply[ax0] = rxplaceholder_3[ax0]
+ for i0 in T.serial(T.int64(3)):
+ with T.block("T_multiply_2"):
+ ax0 = T.axis.spatial(T.int64(3), i0)
+ T.reads(rxplaceholder_4[ax0])
+ T.writes(T_multiply_1[ax0])
+ T_multiply_1[ax0] = rxplaceholder_4[ax0]
+ # fmt: on
+
+ mod = LegalizeOps()(BatchNorm)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_batch_norm_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class BatchNorm:
+ @R.function
+ def main(x: R.Tensor(("n", "h", "w", "c"), "float32"), gamma: R.Tensor(("c",), "float32"), beta: R.Tensor(("c",), "float32"), moving_mean: R.Tensor(("c",), "float32"), moving_var: R.Tensor(("c",), "float32")) -> R.Tuple(R.Tensor(("n", "h", "w", "c"), "float32"), R.Tensor(("c",), "float32"), R.Tensor(("c",), "float32")):
+ n = T.var("int64")
+ h = T.var("int64")
+ w = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tuple(R.Tensor((n, h, w, c), "float32"), R.Tensor((c,), "float32"), R.Tensor((c,), "float32")) = R.nn.batch_norm(x, gamma, beta, moving_mean, moving_var, axis=-1)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("n", "h", "w", "c"), "float32"), gamma: R.Tensor(("c",), "float32"), beta: R.Tensor(("c",), "float32"), moving_mean: R.Tensor(("c",), "float32"), moving_var: R.Tensor(("c",), "float32")) -> R.Tuple(R.Tensor(("n", "h", "w", "c"), "float32"), R.Tensor(("c",), "float32"), R.Tensor(("c",), "float32")):
+ n = T.var("int64")
+ h = T.var("int64")
+ w = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(batch_norm, (x, gamma, beta, moving_mean, moving_var), [R.Tensor((n, h, w, c), "float32"), R.Tensor((c,), "float32"), R.Tensor((c,), "float32")])
+ return gv
+
+ @T.prim_func
+ def batch_norm(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_rxplaceholder_3: T.handle, var_rxplaceholder_4: T.handle, var_T_add: T.handle, var_T_multiply: T.handle, var_T_multiply_1: T.handle):
+ T.func_attr({"tir.noalias": True})
+ c = T.var("int64")
+ h = T.var("int64")
+ n = T.var("int64")
+ w = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [n, h, w, c], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [c], dtype="float32")
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [c], dtype="float32")
+ rxplaceholder_3 = T.match_buffer(var_rxplaceholder_3, [c], dtype="float32")
+ rxplaceholder_4 = T.match_buffer(var_rxplaceholder_4, [c], dtype="float32")
+ T_add = T.match_buffer(var_T_add, [n, h, w, c], dtype="float32")
+ T_multiply = T.match_buffer(var_T_multiply, [c], dtype="float32")
+ T_multiply_1 = T.match_buffer(var_T_multiply_1, [c], dtype="float32")
+ T_reshape = T.alloc_buffer([T.int64(1), T.int64(1), T.int64(1), c], dtype="float32")
+ T_subtract = T.alloc_buffer([n, h, w, c], dtype="float32")
+ T_reshape_1 = T.alloc_buffer([T.int64(1), T.int64(1), T.int64(1), c], dtype="float32")
+ T_add_1 = T.alloc_buffer([T.int64(1), T.int64(1), T.int64(1), c], dtype="float32")
+ compute = T.alloc_buffer([T.int64(1), T.int64(1), T.int64(1), c], dtype="float32")
+ T_divide = T.alloc_buffer([n, h, w, c], dtype="float32")
+ T_reshape_2 = T.alloc_buffer([T.int64(1), T.int64(1), T.int64(1), c], dtype="float32")
+ T_multiply_2 = T.alloc_buffer([n, h, w, c], dtype="float32")
+ T_reshape_3 = T.alloc_buffer([T.int64(1), T.int64(1), T.int64(1), c], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(1), c):
+ with T.block("T_reshape"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_3[((ax0 + ax1 + ax2) * c + ax3) % c])
+ T.writes(T_reshape[ax0, ax1, ax2, ax3])
+ T_reshape[ax0, ax1, ax2, ax3] = rxplaceholder_3[((ax0 + ax1 + ax2) * c + ax3) % c]
+ for i0, i1, i2, i3 in T.grid(n, h, w, c):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_reshape[T.int64(0), T.int64(0), T.int64(0), ax3])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_reshape[T.int64(0), T.int64(0), T.int64(0), ax3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(1), c):
+ with T.block("T_reshape_1"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_4[((ax0 + ax1 + ax2) * c + ax3) % c])
+ T.writes(T_reshape_1[ax0, ax1, ax2, ax3])
+ T_reshape_1[ax0, ax1, ax2, ax3] = rxplaceholder_4[((ax0 + ax1 + ax2) * c + ax3) % c]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(1), c):
+ with T.block("T_add"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_reshape_1[ax0, ax1, ax2, ax3])
+ T.writes(T_add_1[ax0, ax1, ax2, ax3])
+ T_add_1[ax0, ax1, ax2, ax3] = T_reshape_1[ax0, ax1, ax2, ax3] + T.float32(1.0000000000000001e-05)
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(1), c):
+ with T.block("compute"):
+ i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_add_1[i0_1, i1_1, i2_1, i3_1])
+ T.writes(compute[i0_1, i1_1, i2_1, i3_1])
+ compute[i0_1, i1_1, i2_1, i3_1] = T.sqrt(T_add_1[i0_1, i1_1, i2_1, i3_1], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(n, h, w, c):
+ with T.block("T_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_subtract[ax0, ax1, ax2, ax3], compute[T.int64(0), T.int64(0), T.int64(0), ax3])
+ T.writes(T_divide[ax0, ax1, ax2, ax3])
+ T_divide[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] / compute[T.int64(0), T.int64(0), T.int64(0), ax3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(1), c):
+ with T.block("T_reshape_2"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_1[((ax0 + ax1 + ax2) * c + ax3) % c])
+ T.writes(T_reshape_2[ax0, ax1, ax2, ax3])
+ T_reshape_2[ax0, ax1, ax2, ax3] = rxplaceholder_1[((ax0 + ax1 + ax2) * c + ax3) % c]
+ for i0, i1, i2, i3 in T.grid(n, h, w, c):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_divide[ax0, ax1, ax2, ax3], T_reshape_2[T.int64(0), T.int64(0), T.int64(0), ax3])
+ T.writes(T_multiply_2[ax0, ax1, ax2, ax3])
+ T_multiply_2[ax0, ax1, ax2, ax3] = T_divide[ax0, ax1, ax2, ax3] * T_reshape_2[T.int64(0), T.int64(0), T.int64(0), ax3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(1), c):
+ with T.block("T_reshape_3"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_2[((ax0 + ax1 + ax2) * c + ax3) % c])
+ T.writes(T_reshape_3[ax0, ax1, ax2, ax3])
+ T_reshape_3[ax0, ax1, ax2, ax3] = rxplaceholder_2[((ax0 + ax1 + ax2) * c + ax3) % c]
+ for i0, i1, i2, i3 in T.grid(n, h, w, c):
+ with T.block("T_add_1"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_multiply_2[ax0, ax1, ax2, ax3], T_reshape_3[T.int64(0), T.int64(0), T.int64(0), ax3])
+ T.writes(T_add[ax0, ax1, ax2, ax3])
+ T_add[ax0, ax1, ax2, ax3] = T_multiply_2[ax0, ax1, ax2, ax3] + T_reshape_3[T.int64(0), T.int64(0), T.int64(0), ax3]
+ for i0 in T.serial(c):
+ with T.block("T_multiply_1"):
+ ax0 = T.axis.spatial(c, i0)
+ T.reads(rxplaceholder_3[ax0])
+ T.writes(T_multiply[ax0])
+ T_multiply[ax0] = rxplaceholder_3[ax0]
+ for i0 in T.serial(c):
+ with T.block("T_multiply_2"):
+ ax0 = T.axis.spatial(c, i0)
+ T.reads(rxplaceholder_4[ax0])
+ T.writes(T_multiply_1[ax0])
+ T_multiply_1[ax0] = rxplaceholder_4[ax0]
+ # fmt: on
+
+ mod = LegalizeOps()(BatchNorm)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_layer_norm():
+ # fmt: off
+ @tvm.script.ir_module
+ class LayerNorm:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32"), gamma: R.Tensor((4, 5), "float32"), beta: R.Tensor((4, 5), "float32")) -> R.Tensor((2, 3, 4, 5), "float32"):
+ gv: R.Tensor((2, 3, 4, 5), "float32") = R.nn.layer_norm(x, gamma, beta, axes=[-2, -1])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32"), gamma: R.Tensor((4, 5), "float32"), beta: R.Tensor((4, 5), "float32")) -> R.Tensor((2, 3, 4, 5), "float32"):
+ gv = R.call_tir(layer_norm, (x, gamma, beta), R.Tensor((2, 3, 4, 5), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def layer_norm(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4), T.int64(5)), "float32"), rxplaceholder_2: T.Buffer((T.int64(4), T.int64(5)), "float32"), T_layer_norm: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ rxplaceholder_red_temp_v0 = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ rxplaceholder_red_temp_v1 = T.alloc_buffer([T.int64(2), T.int64(3)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("rxplaceholder_red_temp"):
+ ax0, ax1, k2, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[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.float32 = rxplaceholder_red_temp_v0[ax0, ax1] + rxplaceholder[ax0, ax1, k2, k3]
+ v_rxplaceholder_red_temp_v1: T.float32 = rxplaceholder_red_temp_v1[ax0, ax1] + rxplaceholder[ax0, ax1, k2, k3] * rxplaceholder[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(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("T_layer_norm"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], rxplaceholder_red_temp_v0[ax0, ax1], rxplaceholder_red_temp_v1[ax0, ax1], rxplaceholder_1[ax2, ax3], rxplaceholder_2[ax2, ax3])
+ T.writes(T_layer_norm[ax0, ax1, ax2, ax3])
+ T_layer_norm[ax0, ax1, ax2, ax3] = (rxplaceholder[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") * rxplaceholder_1[ax2, ax3] + rxplaceholder_2[ax2, ax3]
+ # fmt: on
+ mod = LegalizeOps()(LayerNorm)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_layer_norm_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class LayerNorm:
+ @R.function
+ def main(x: R.Tensor(("n", "s", "f"), "float32"), gamma: R.Tensor(("s", "f"), "float32"), beta: R.Tensor(("s", "f"), "float32")) -> R.Tensor(("n", "s", "f"), "float32"):
+ n = T.var("int64")
+ s = T.var("int64")
+ f = T.var("int64")
+ gv: R.Tensor((n, s, f), "float32") = R.nn.layer_norm(x, gamma, beta, axes=[1, 2])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("n", "s", "f"), "float32"), gamma: R.Tensor(("s", "f"), "float32"), beta: R.Tensor(("s", "f"), "float32")) -> R.Tensor(("n", "s", "f"), "float32"):
+ n = T.var("int64")
+ s = T.var("int64")
+ f = T.var("int64")
+ gv = R.call_tir(layer_norm, (x, gamma, beta), R.Tensor((n, s, f), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def layer_norm(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_T_layer_norm: T.handle):
+ T.func_attr({"tir.noalias": True})
+ f = T.var("int64")
+ n = T.var("int64")
+ s = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [n, s, f], dtype="float32")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [s, f], dtype="float32")
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [s, f], dtype="float32")
+ T_layer_norm = T.match_buffer(var_T_layer_norm, [n, s, f], dtype="float32")
+ rxplaceholder_red_temp_v0 = T.alloc_buffer([n], dtype="float32")
+ rxplaceholder_red_temp_v1 = T.alloc_buffer([n], dtype="float32")
+ for i0, i1, i2 in T.grid(n, s, f):
+ with T.block("rxplaceholder_red_temp"):
+ ax0, k1, k2 = T.axis.remap("SRR", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, k1, k2])
+ T.writes(rxplaceholder_red_temp_v0[ax0], rxplaceholder_red_temp_v1[ax0])
+ with T.init():
+ rxplaceholder_red_temp_v0[ax0] = T.float32(0)
+ rxplaceholder_red_temp_v1[ax0] = T.float32(0)
+ v_rxplaceholder_red_temp_v0: T.float32 = rxplaceholder_red_temp_v0[ax0] + rxplaceholder[ax0, k1, k2]
+ v_rxplaceholder_red_temp_v1: T.float32 = rxplaceholder_red_temp_v1[ax0] + rxplaceholder[ax0, k1, k2] * rxplaceholder[ax0, k1, k2]
+ rxplaceholder_red_temp_v0[ax0] = v_rxplaceholder_red_temp_v0
+ rxplaceholder_red_temp_v1[ax0] = v_rxplaceholder_red_temp_v1
+ for i0, i1, i2 in T.grid(n, s, f):
+ with T.block("T_layer_norm"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1, ax2], rxplaceholder_red_temp_v0[ax0], rxplaceholder_red_temp_v1[ax0], rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, ax2])
+ T.writes(T_layer_norm[ax0, ax1, ax2])
+ T_layer_norm[ax0, ax1, ax2] = (rxplaceholder[ax0, ax1, ax2] - rxplaceholder_red_temp_v0[ax0] / (T.Cast("float32", s) * T.Cast("float32", f))) * T.rsqrt(rxplaceholder_red_temp_v1[ax0] / (T.Cast("float32", s) * T.Cast("float32", f)) - rxplaceholder_red_temp_v0[ax0] / (T.Cast("float32", s) * T.Cast("float32", f)) * (rxplaceholder_red_temp_v0[ax0] / (T.Cast("float32", s) * T.Cast("float32", f))) + T.float32(1e-05), dtype="float32") * rxplaceholder_1[ax1, ax2] + rxplacehol [...]
+ # fmt: on
+ mod = LegalizeOps()(LayerNorm)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_search_statistical.py b/tests/python/relax/test_transform_legalize_ops_search_statistical.py
new file mode 100644
index 0000000000..4c31077d9c
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_search_statistical.py
@@ -0,0 +1,793 @@
+# 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 tvm
+from tvm.relax.transform import LegalizeOps
+from tvm.script import relax as R, tir as T
+import tvm.testing
+
+
+##################### Search #####################
+
+
+def test_where():
+ # fmt: off
+ @tvm.script.ir_module
+ class Where:
+ @R.function
+ def main(condition: R.Tensor((3, 2, 1), "bool"), x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 1), "float32")) -> R.Tensor((3, 2, 3), "float32"):
+ gv: R.Tensor((3, 2, 3), "float32") = R.where(condition, x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(condition: R.Tensor((3, 2, 1), "bool"), x: R.Tensor((2, 3), "float32"), y: R.Tensor((2, 1), "float32")) -> R.Tensor((3, 2, 3), "float32"):
+ gv = R.call_tir(where, (condition, x, y), R.Tensor((3, 2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def where(rxplaceholder: T.Buffer((T.int64(3), T.int64(2), T.int64(1)), "bool"), rxplaceholder_1: T.Buffer((T.int64(2), T.int64(3)), "float32"), rxplaceholder_2: T.Buffer((T.int64(2), T.int64(1)), "float32"), T_where: T.Buffer((T.int64(3), T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2 in T.grid(T.int64(3), T.int64(2), T.int64(3)):
+ with T.block("T_where"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1, T.int64(0)], rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
+ T.writes(T_where[ax0, ax1, ax2])
+ T_where[ax0, ax1, ax2] = T.Select(0 < T.Cast("int32", rxplaceholder[ax0, ax1, T.int64(0)]), rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
+ # fmt: on
+
+ mod = LegalizeOps()(Where)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_where_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Where:
+ @R.function
+ def main(condition: R.Tensor(("a", "b", 1), "bool"), x: R.Tensor(("b", "c"), "float32"), y: R.Tensor(("b", 1), "float32")) -> R.Tensor(("a", "b", "c"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tensor((a, b, c), "float32") = R.where(condition, x, y)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(condition: R.Tensor(("a", "b", 1), "bool"), x: R.Tensor(("b", "c"), "float32"), y: R.Tensor(("b", 1), "float32")) -> R.Tensor(("a", "b", "c"), "float32"):
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(where, (condition, x, y), R.Tensor((a, b, c), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def where(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, var_rxplaceholder_2: T.handle, var_T_where: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, T.int64(1)], dtype="bool")
+ rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [b, c], dtype="float32")
+ rxplaceholder_2 = T.match_buffer(var_rxplaceholder_2, [b, T.int64(1)], dtype="float32")
+ T_where = T.match_buffer(var_T_where, [a, b, c], dtype="float32")
+ for i0, i1, i2 in T.grid(a, b, c):
+ with T.block("T_where"):
+ ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(rxplaceholder[ax0, ax1, T.int64(0)], rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
+ T.writes(T_where[ax0, ax1, ax2])
+ T_where[ax0, ax1, ax2] = T.Select(0 < T.Cast("int32", rxplaceholder[ax0, ax1, T.int64(0)]), rxplaceholder_1[ax1, ax2], rxplaceholder_2[ax1, T.int64(0)])
+ # fmt: on
+
+ mod = LegalizeOps()(Where)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+##################### Statistical #####################
+
+
+def test_max():
+ # fmt: off
+ @tvm.script.ir_module
+ class Max:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 5), "float32"):
+ gv: R.Tensor((2, 5), "float32") = R.max(x, axis=[1, 2])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 5), "float32"):
+ gv = R.call_tir(max, (x,), R.Tensor((2, 5), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def max(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(5), T.int64(3), T.int64(4)):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, k1, k2 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, k1, k2, ax1])
+ T.writes(rxplaceholder_red[ax0, ax1])
+ with T.init():
+ rxplaceholder_red[ax0, ax1] = T.min_value("float32")
+ rxplaceholder_red[ax0, ax1] = T.max(rxplaceholder_red[ax0, ax1], rxplaceholder[ax0, k1, k2, ax1])
+ # fmt: on
+
+ mod = LegalizeOps()(Max)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_max_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Max:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", "d"), "float32"):
+ a = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, d), "float32") = R.max(x, axis=[1, 2])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", "d"), "float32"):
+ a = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(max, (x,), R.Tensor((a, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def max(var_rxplaceholder: T.handle, var_rxplaceholder_red: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ rxplaceholder_red = T.match_buffer(var_rxplaceholder_red, [a, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, d, b, c):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, k1, k2 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, k1, k2, ax1])
+ T.writes(rxplaceholder_red[ax0, ax1])
+ with T.init():
+ rxplaceholder_red[ax0, ax1] = T.min_value("float32")
+ rxplaceholder_red[ax0, ax1] = T.max(rxplaceholder_red[ax0, ax1], rxplaceholder[ax0, k1, k2, ax1])
+ # fmt: on
+
+ mod = LegalizeOps()(Max)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_min():
+ # fmt: off
+ @tvm.script.ir_module
+ class Min:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 1, 1, 5), "float32"):
+ gv: R.Tensor((2, 1, 1, 5), "float32") = R.min(x, axis=[1, 2], keepdims=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((2, 1, 1, 5), "float32"):
+ gv = R.call_tir(min, (x,), R.Tensor((2, 1, 1, 5), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def min(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(2), T.int64(1), T.int64(1), T.int64(5)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(2), T.int64(1), T.int64(1), T.int64(5), T.int64(3), T.int64(4)):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, ax2, ax3, k1, k2 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[ax0, k1, k2, ax3])
+ T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.max_value("float32")
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.min(rxplaceholder_red[ax0, ax1, ax2, ax3], rxplaceholder[ax0, k1, k2, ax3])
+ # fmt: on
+
+ mod = LegalizeOps()(Min)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_min_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Min:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", 1, 1, "d"), "float32"):
+ a = T.var("int64")
+ d = T.var("int64")
+ gv: R.Tensor((a, 1, 1, d), "float32") = R.min(x, axis=[1, 2], keepdims=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("a", 1, 1, "d"), "float32"):
+ a = T.var("int64")
+ d = T.var("int64")
+ gv = R.call_tir(min, (x,), R.Tensor((a, 1, 1, d), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def min(var_rxplaceholder: T.handle, var_rxplaceholder_red: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ rxplaceholder_red = T.match_buffer(var_rxplaceholder_red, [a, T.int64(1), T.int64(1), d], dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(a, T.int64(1), T.int64(1), d, b, c):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, ax2, ax3, k1, k2 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[ax0, k1, k2, ax3])
+ T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.max_value("float32")
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.min(rxplaceholder_red[ax0, ax1, ax2, ax3], rxplaceholder[ax0, k1, k2, ax3])
+ # fmt: on
+
+ mod = LegalizeOps()(Min)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sum():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sum:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.sum(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
+ gv = R.call_tir(sum, (x,), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def sum(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("rxplaceholder_red"):
+ k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k0, k1, k2, k3])
+ T.writes(rxplaceholder_red[()])
+ with T.init():
+ rxplaceholder_red[()] = T.float32(0)
+ rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[k0, k1, k2, k3]
+ # fmt: on
+
+ mod = LegalizeOps()(Sum)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sum_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sum:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.sum(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
+ gv = R.call_tir(sum, (x,), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def sum(var_rxplaceholder: T.handle, rxplaceholder_red: T.Buffer((), "float32")):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("rxplaceholder_red"):
+ k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k0, k1, k2, k3])
+ T.writes(rxplaceholder_red[()])
+ with T.init():
+ rxplaceholder_red[()] = T.float32(0)
+ rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[k0, k1, k2, k3]
+ # fmt: on
+
+ mod = LegalizeOps()(Sum)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_prod():
+ # fmt: off
+ @tvm.script.ir_module
+ class Prod:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
+ gv: R.Tensor((1, 1, 1, 1), "float32") = R.prod(x, keepdims=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
+ gv = R.call_tir(prod, (x,), R.Tensor((1, 1, 1, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def prod(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), rxplaceholder_red: T.Buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, ax2, ax3, k0, k1, k2, k3 = T.axis.remap("SSSSRRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
+ T.reads(rxplaceholder[k0, k1, k2, k3])
+ T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(1)
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] * rxplaceholder[k0, k1, k2, k3]
+ # fmt: on
+
+ mod = LegalizeOps()(Prod)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_prod_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Prod:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
+ gv: R.Tensor((1, 1, 1, 1), "float32") = R.prod(x, keepdims=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, 1, 1, 1), "float32"):
+ gv = R.call_tir(prod, (x,), R.Tensor((1, 1, 1, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def prod(var_rxplaceholder: T.handle, rxplaceholder_red: T.Buffer((T.int64(1), T.int64(1), T.int64(1), T.int64(1)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ for i0, i1, i2, i3, i4, i5, i6, i7 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1), a, b, c, d):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, ax2, ax3, k0, k1, k2, k3 = T.axis.remap("SSSSRRRR", [i0, i1, i2, i3, i4, i5, i6, i7])
+ T.reads(rxplaceholder[k0, k1, k2, k3])
+ T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(1)
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] * rxplaceholder[k0, k1, k2, k3]
+ # fmt: on
+
+ mod = LegalizeOps()(Prod)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_mean():
+ # fmt: off
+ @tvm.script.ir_module
+ class Mean:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((3, 4), "float32"):
+ gv: R.Tensor((3, 4), "float32") = R.mean(x, [0, 3])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((3, 4), "float32"):
+ gv = R.call_tir(mean, (x,), R.Tensor((3, 4), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def mean(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), T_divide: T.Buffer((T.int64(3), T.int64(4)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ rxplaceholder_red = T.alloc_buffer([T.int64(3), T.int64(4)], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(3), T.int64(4), T.int64(2), T.int64(5)):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, k0, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k0, ax0, ax1, k3])
+ T.writes(rxplaceholder_red[ax0, ax1])
+ with T.init():
+ rxplaceholder_red[ax0, ax1] = T.float32(0)
+ rxplaceholder_red[ax0, ax1] = rxplaceholder_red[ax0, ax1] + rxplaceholder[k0, ax0, ax1, k3]
+ for i0, i1 in T.grid(T.int64(3), T.int64(4)):
+ with T.block("T_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder_red[ax0, ax1])
+ T.writes(T_divide[ax0, ax1])
+ T_divide[ax0, ax1] = rxplaceholder_red[ax0, ax1] * T.float32(0.1)
+ # fmt: on
+
+ mod = LegalizeOps()(Mean)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_mean_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Mean:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor(("b", "c"), "float32"):
+ b = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tensor((b, c), "float32") = R.mean(x, [0, 3])
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), dtype="float32")) -> R.Tensor(("b", "c"), dtype="float32"):
+ b = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(mean, (x,), R.Tensor((b, c), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def mean(var_rxplaceholder: T.handle, var_T_divide: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ T_divide = T.match_buffer(var_T_divide, [b, c], dtype="float32")
+ rxplaceholder_red = T.alloc_buffer([b, c], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(b, c, a, d):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, k0, k3 = T.axis.remap("SSRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k0, ax0, ax1, k3])
+ T.writes(rxplaceholder_red[ax0, ax1])
+ with T.init():
+ rxplaceholder_red[ax0, ax1] = T.float32(0)
+ rxplaceholder_red[ax0, ax1] = rxplaceholder_red[ax0, ax1] + rxplaceholder[k0, ax0, ax1, k3]
+ for i0, i1 in T.grid(b, c):
+ with T.block("T_divide"):
+ ax0, ax1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder_red[ax0, ax1])
+ T.writes(T_divide[ax0, ax1])
+ T_divide[ax0, ax1] = rxplaceholder_red[ax0, ax1] / T.Cast("float32", a * d)
+ # fmt: on
+
+ mod = LegalizeOps()(Mean)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_std():
+ # fmt: off
+ @tvm.script.ir_module
+ class Std:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.std(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((), "float32"):
+ gv = R.call_tir(std, (x,), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def std(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), compute: T.Buffer((), "float32")):
+ T.func_attr({"tir.noalias": True})
+ rxplaceholder_red = T.alloc_buffer([], dtype="float32")
+ T_divide = T.alloc_buffer([], dtype="float32")
+ T_subtract = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(4), T.int64(5)], dtype="float32")
+ T_multiply = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(4), T.int64(5)], dtype="float32")
+ T_multiply_red = T.alloc_buffer([], dtype="float32")
+ T_divide_1 = T.alloc_buffer([], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("rxplaceholder_red"):
+ k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k0, k1, k2, k3])
+ T.writes(rxplaceholder_red[()])
+ with T.init():
+ rxplaceholder_red[()] = T.float32(0)
+ rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[k0, k1, k2, k3]
+ with T.block("T_divide"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(rxplaceholder_red[()])
+ T.writes(T_divide[()])
+ T_divide[()] = rxplaceholder_red[()] * T.float32(0.0083333333333333332)
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_divide[()])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_divide[()]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_subtract[ax0, ax1, ax2, ax3])
+ T.writes(T_multiply[ax0, ax1, ax2, ax3])
+ T_multiply[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] * T_subtract[ax0, ax1, ax2, ax3]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("T_multiply_red"):
+ k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
+ T.reads(T_multiply[k0, k1, k2, k3])
+ T.writes(T_multiply_red[()])
+ with T.init():
+ T_multiply_red[()] = T.float32(0)
+ T_multiply_red[()] = T_multiply_red[()] + T_multiply[k0, k1, k2, k3]
+ with T.block("T_divide_1"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(T_multiply_red[()])
+ T.writes(T_divide_1[()])
+ T_divide_1[()] = T_multiply_red[()] * T.float32(0.0083333333333333332)
+ with T.block("compute"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(T_divide_1[()])
+ T.writes(compute[()])
+ compute[()] = T.sqrt(T_divide_1[()])
+ # fmt: on
+
+ mod = LegalizeOps()(Std)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_std_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Std:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
+ gv: R.Tensor((), "float32") = R.std(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((), "float32"):
+ gv = R.call_tir(std, (x,), R.Tensor((), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def std(var_rxplaceholder: T.handle, compute: T.Buffer((), "float32")):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ rxplaceholder_red = T.alloc_buffer([], dtype="float32")
+ T_divide = T.alloc_buffer([], dtype="float32")
+ T_subtract = T.alloc_buffer([a, b, c, d], dtype="float32")
+ T_multiply = T.alloc_buffer([a, b, c, d], dtype="float32")
+ T_multiply_red = T.alloc_buffer([], dtype="float32")
+ T_divide_1 = T.alloc_buffer([], dtype="float32")
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("rxplaceholder_red"):
+ k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[k0, k1, k2, k3])
+ T.writes(rxplaceholder_red[()])
+ with T.init():
+ rxplaceholder_red[()] = T.float32(0)
+ rxplaceholder_red[()] = rxplaceholder_red[()] + rxplaceholder[k0, k1, k2, k3]
+ with T.block("T_divide"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(rxplaceholder_red[()])
+ T.writes(T_divide[()])
+ T_divide[()] = rxplaceholder_red[()] / T.Cast("float32", a * b * c * d)
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_divide[()])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_divide[()]
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_subtract[ax0, ax1, ax2, ax3])
+ T.writes(T_multiply[ax0, ax1, ax2, ax3])
+ T_multiply[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] * T_subtract[ax0, ax1, ax2, ax3]
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_multiply_red"):
+ k0, k1, k2, k3 = T.axis.remap("RRRR", [i0, i1, i2, i3])
+ T.reads(T_multiply[k0, k1, k2, k3])
+ T.writes(T_multiply_red[()])
+ with T.init():
+ T_multiply_red[()] = T.float32(0)
+ T_multiply_red[()] = T_multiply_red[()] + T_multiply[k0, k1, k2, k3]
+ with T.block("T_divide_1"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(T_multiply_red[()])
+ T.writes(T_divide_1[()])
+ T_divide_1[()] = T_multiply_red[()] / T.Cast("float32", a * b * c * d)
+ with T.block("compute"):
+ vi = T.axis.spatial(1, T.int64(0))
+ T.reads(T_divide_1[()])
+ T.writes(compute[()])
+ compute[()] = T.sqrt(T_divide_1[()])
+ # fmt: on
+
+ mod = LegalizeOps()(Std)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_variance():
+ # fmt: off
+ @tvm.script.ir_module
+ class Variance:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), "float32")) -> R.Tensor((1, 3, 4, 1), "float32"):
+ gv: R.Tensor((1, 3, 4, 1), "float32") = R.variance(x, [0, 3], keepdims=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3, 4, 5), dtype="float32")) -> R.Tensor((1, 3, 4, 1), dtype="float32"):
+ gv = R.call_tir(variance, (x,), R.Tensor((1, 3, 4, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def variance(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4), T.int64(5)), "float32"), T_divide: T.Buffer((T.int64(1), T.int64(3), T.int64(4), T.int64(1)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ rxplaceholder_red = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(4), T.int64(1)], dtype="float32")
+ T_divide_1 = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(4), T.int64(1)], dtype="float32")
+ T_subtract = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(4), T.int64(5)], dtype="float32")
+ T_multiply = T.alloc_buffer([T.int64(2), T.int64(3), T.int64(4), T.int64(5)], dtype="float32")
+ T_multiply_red = T.alloc_buffer([T.int64(1), T.int64(3), T.int64(4), T.int64(1)], dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1), T.int64(2), T.int64(5)):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[k0, ax1, ax2, k3])
+ T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(0)
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] + rxplaceholder[k0, ax1, ax2, k3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1)):
+ with T.block("T_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ T.writes(T_divide_1[ax0, ax1, ax2, ax3])
+ T_divide_1[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] * T.float32(0.10000000000000001)
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_divide_1[T.int64(0), ax1, ax2, T.int64(0)])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_divide_1[T.int64(0), ax1, ax2, T.int64(0)]
+ for i0, i1, i2, i3 in T.grid(T.int64(2), T.int64(3), T.int64(4), T.int64(5)):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_subtract[ax0, ax1, ax2, ax3])
+ T.writes(T_multiply[ax0, ax1, ax2, ax3])
+ T_multiply[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] * T_subtract[ax0, ax1, ax2, ax3]
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1), T.int64(2), T.int64(5)):
+ with T.block("T_multiply_red"):
+ ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(T_multiply[k0, ax1, ax2, k3])
+ T.writes(T_multiply_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ T_multiply_red[ax0, ax1, ax2, ax3] = T.float32(0)
+ T_multiply_red[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] + T_multiply[k0, ax1, ax2, k3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(3), T.int64(4), T.int64(1)):
+ with T.block("T_divide_1"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_multiply_red[ax0, ax1, ax2, ax3])
+ T.writes(T_divide[ax0, ax1, ax2, ax3])
+ T_divide[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] * T.float32(0.10000000000000001)
+ # fmt: on
+
+ mod = LegalizeOps()(Variance)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_variance_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Variance:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, "b", "c", 1), "float32"):
+ b = T.var("int64")
+ c = T.var("int64")
+ gv: R.Tensor((1, b, c, 1), "float32") = R.variance(x, [0, 3], keepdims=True)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("a", "b", "c", "d"), "float32")) -> R.Tensor((1, "b", "c", 1), "float32"):
+ b = T.var("int64")
+ c = T.var("int64")
+ gv = R.call_tir(variance, (x,), R.Tensor((1, b, c, 1), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def variance(var_rxplaceholder: T.handle, var_T_divide: T.handle):
+ T.func_attr({"tir.noalias": True})
+ a = T.var("int64")
+ b = T.var("int64")
+ c = T.var("int64")
+ d = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [a, b, c, d], dtype="float32")
+ T_divide = T.match_buffer(var_T_divide, [T.int64(1), b, c, T.int64(1)], dtype="float32")
+ rxplaceholder_red = T.alloc_buffer([T.int64(1), b, c, T.int64(1)], dtype="float32")
+ T_divide_1 = T.alloc_buffer([T.int64(1), b, c, T.int64(1)], dtype="float32")
+ T_subtract = T.alloc_buffer([a, b, c, d], dtype="float32")
+ T_multiply = T.alloc_buffer([a, b, c, d], dtype="float32")
+ T_multiply_red = T.alloc_buffer([T.int64(1), b, c, T.int64(1)], dtype="float32")
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), b, c, T.int64(1), a, d):
+ with T.block("rxplaceholder_red"):
+ ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(rxplaceholder[k0, ax1, ax2, k3])
+ T.writes(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = T.float32(0)
+ rxplaceholder_red[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] + rxplaceholder[k0, ax1, ax2, k3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), b, c, T.int64(1)):
+ with T.block("T_divide"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder_red[ax0, ax1, ax2, ax3])
+ T.writes(T_divide_1[ax0, ax1, ax2, ax3])
+ T_divide_1[ax0, ax1, ax2, ax3] = rxplaceholder_red[ax0, ax1, ax2, ax3] / T.Cast("float32", a * d)
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_subtract"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(rxplaceholder[ax0, ax1, ax2, ax3], T_divide_1[T.int64(0), ax1, ax2, T.int64(0)])
+ T.writes(T_subtract[ax0, ax1, ax2, ax3])
+ T_subtract[ax0, ax1, ax2, ax3] = rxplaceholder[ax0, ax1, ax2, ax3] - T_divide_1[T.int64(0), ax1, ax2, T.int64(0)]
+ for i0, i1, i2, i3 in T.grid(a, b, c, d):
+ with T.block("T_multiply"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_subtract[ax0, ax1, ax2, ax3])
+ T.writes(T_multiply[ax0, ax1, ax2, ax3])
+ T_multiply[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] * T_subtract[ax0, ax1, ax2, ax3]
+ for i0, i1, i2, i3, i4, i5 in T.grid(T.int64(1), b, c, T.int64(1), a, d):
+ with T.block("T_multiply_red"):
+ ax0, ax1, ax2, ax3, k0, k3 = T.axis.remap("SSSSRR", [i0, i1, i2, i3, i4, i5])
+ T.reads(T_multiply[k0, ax1, ax2, k3])
+ T.writes(T_multiply_red[ax0, ax1, ax2, ax3])
+ with T.init():
+ T_multiply_red[ax0, ax1, ax2, ax3] = T.float32(0)
+ T_multiply_red[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] + T_multiply[k0, ax1, ax2, k3]
+ for i0, i1, i2, i3 in T.grid(T.int64(1), b, c, T.int64(1)):
+ with T.block("T_divide_1"):
+ ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
+ T.reads(T_multiply_red[ax0, ax1, ax2, ax3])
+ T.writes(T_divide[ax0, ax1, ax2, ax3])
+ T_divide[ax0, ax1, ax2, ax3] = T_multiply_red[ax0, ax1, ax2, ax3] / T.Cast("float32", a * d)
+ # fmt: on
+
+ mod = LegalizeOps()(Variance)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_unary.py b/tests/python/relax/test_transform_legalize_ops_unary.py
new file mode 100644
index 0000000000..12ae366dcc
--- /dev/null
+++ b/tests/python/relax/test_transform_legalize_ops_unary.py
@@ -0,0 +1,693 @@
+# 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 tvm
+import tvm.testing
+from tvm.relax.transform import LegalizeOps
+from tvm.script import relax as R
+from tvm.script import tir as T
+
+
+def test_abs():
+ # fmt: off
+ @tvm.script.ir_module
+ class Abs:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.abs(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
+ gv = R.call_tir(tir_abs, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_abs(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32"),):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("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.fabs(rxplaceholder[v_i0, v_i1], dtype="float32")
+ # fmt: on
+
+ mod = LegalizeOps()(Abs)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_abs_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Abs:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.abs(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m", "n"), dtype="float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_abs, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_abs(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("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.fabs(rxplaceholder[v_i0, v_i1], dtype="float32")
+ # fmt: on
+
+ mod = LegalizeOps()(Abs)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_cos():
+ # fmt: off
+ @tvm.script.ir_module
+ class Cos:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.cos(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_cos, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_cos(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.cos(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Cos)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_cos_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Cos:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.cos(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_cos, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_cos(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.cos(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Cos)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_exp():
+ # fmt: off
+ @tvm.script.ir_module
+ class Exp:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.exp(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), dtype="float32")) -> R.Tensor((2, 3), dtype="float32"):
+ gv = R.call_tir(tir_exp, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_exp(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32"),):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("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.exp(rxplaceholder[v_i0, v_i1], dtype="float32")
+ # fmt: on
+
+ mod = LegalizeOps()(Exp)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_exp_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Exp:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.exp(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m", "n"), dtype="float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_exp, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_exp(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("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.exp(rxplaceholder[v_i0, v_i1], dtype="float32")
+ # fmt: on
+
+ mod = LegalizeOps()(Exp)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_log():
+ # fmt: off
+ @tvm.script.ir_module
+ class Log:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.log(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_log, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_log(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.log(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Log)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_log_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Log:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.log(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_log, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_log(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.log(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Log)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_negative():
+ # fmt: off
+ @tvm.script.ir_module
+ class Negative:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.negative(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_negative, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_negative(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = rxplaceholder[i0_1, i1_1] * T.float32(-1)
+ # fmt: on
+
+ mod = LegalizeOps()(Negative)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_negative_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Negative:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.negative(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_negative, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_negative(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = rxplaceholder[i0_1, i1_1] * T.float32(-1)
+ # fmt: on
+
+ mod = LegalizeOps()(Negative)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sigmoid():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sigmoid:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.sigmoid(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_sigmoid, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sigmoid(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sigmoid(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sigmoid)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sigmoid_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sigmoid:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.sigmoid(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_sigmoid, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sigmoid(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sigmoid(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sigmoid)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sin():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sin:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.sin(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_sin, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sin(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sin(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sin)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sin_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sin:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.sin(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_sin, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sin(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sin(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sin)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sqrt():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sqrt:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.sqrt(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_sqrt, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sqrt(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sqrt(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sqrt)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_sqrt_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Sqrt:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.sqrt(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_sqrt, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_sqrt(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.sqrt(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Sqrt)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_tanh():
+ # fmt: off
+ @tvm.script.ir_module
+ class Tanh:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv: R.Tensor((2, 3), "float32") = R.tanh(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor((2, 3), "float32")) -> R.Tensor((2, 3), "float32"):
+ gv = R.call_tir(tir_tanh, (x,), R.Tensor((2, 3), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_tanh(rxplaceholder: T.Buffer((T.int64(2), T.int64(3)), "float32"), compute: T.Buffer((T.int64(2), T.int64(3)), "float32")):
+ T.func_attr({"tir.noalias": True})
+ for i0, i1 in T.grid(T.int64(2), T.int64(3)):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.tanh(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Tanh)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_tanh_symbolic():
+ # fmt: off
+ @tvm.script.ir_module
+ class Tanh:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.tanh(x)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_tanh, (x,), R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_tanh(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
+ T.reads(rxplaceholder[i0_1, i1_1])
+ T.writes(compute[i0_1, i1_1])
+ compute[i0_1, i1_1] = T.tanh(rxplaceholder[i0_1, i1_1])
+ # fmt: on
+
+ mod = LegalizeOps()(Tanh)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+def test_clip_symbolic():
+ @tvm.script.ir_module
+ class Clip:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), "float32")) -> R.Tensor(("m", "n"), "float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv: R.Tensor((m, n), "float32") = R.clip(x, 5, 8)
+ return gv
+
+ @tvm.script.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor(("m", "n"), dtype="float32")) -> R.Tensor(("m", "n"), dtype="float32"):
+ m = T.var("int64")
+ n = T.var("int64")
+ gv = R.call_tir(tir_clip, (x,), out_sinfo=R.Tensor((m, n), dtype="float32"))
+ return gv
+
+ @T.prim_func
+ def tir_clip(var_rxplaceholder: T.handle, var_compute: T.handle):
+ T.func_attr({"tir.noalias": True})
+ m = T.var("int64")
+ n = T.var("int64")
+ rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n], dtype="float32")
+ compute = T.match_buffer(var_compute, [m, n], dtype="float32")
+ for i0, i1 in T.grid(m, n):
+ with T.block("compute"):
+ v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
+ compute[v_i0, v_i1] = T.max(
+ T.min(rxplaceholder[v_i0, v_i1], T.float32(8)), T.float32(5)
+ )
+
+ mod = LegalizeOps()(Clip)
+ tvm.ir.assert_structural_equal(mod, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()