chore: import upstream snapshot with attribution

python/tvm/relax/transform/__init__.py

@@ -0,0 +1,102 @@
+# isort: skip_file
+# 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.
+"""Relax transformations."""
+
+from .transform import (
+    AdjustMatmulOrder,
+    AllocateWorkspace,
+    AlterOpImpl,
+    AnnotateTIROpPattern,
+    AttachAttrLayoutFreeBuffers,
+    AttachGlobalSymbol,
+    BindParams,
+    BindSymbolicVars,
+    BundleModelParams,
+    CallTIRRewrite,
+    CanonicalizeBindings,
+    CombineParallelMatmul,
+    ComputePrimValue,
+    ConvertLayout,
+    ConvertToDataflow,
+    DataflowBlockPass,
+    DataflowUseInplaceCalls,
+    DeadCodeElimination,
+    DecomposeOpsForInference,
+    DecomposeOpsForTraining,
+    EliminateCommonSubexpr,
+    ExpandMatmulOfSum,
+    ExpandTupleArguments,
+    FoldConstant,
+    FunctionPass,
+    FuseOps,
+    FuseOpsByPattern,
+    FuseTIR,
+    FusionPattern,
+    Gradient,
+    InlinePrivateFunctions,
+    KillAfterLastUse,
+    LambdaLift,
+    LazyGetInput,
+    LazySetOutput,
+    LegalizeOps,
+    LiftTransformParams,
+    LowerAllocTensor,
+    LowerRuntimeBuiltin,
+    MergeCompositeFunctions,
+    MetaScheduleApplyDatabase,
+    MetaScheduleTuneIRMod,
+    MetaScheduleTuneTIR,
+    Normalize,
+    NormalizeGlobalVar,
+    PatternCheckContext,
+    RealizeVDevice,
+    RemovePurityChecking,
+    RemoveUnusedOutputs,
+    RemoveUnusedParameters,
+    ReorderPermuteDimsAfterConcat,
+    ReorderTakeAfterMatmul,
+    RewriteCUDAGraph,
+    RewriteDataflowReshape,
+    RunCodegen,
+    SplitCallTIRByPattern,
+    SplitLayoutRewritePreproc,
+    StaticPlanBlockMemory,
+    ToMixedPrecision,
+    ToNonDataflow,
+    TopologicalSort,
+    UpdateParamType,
+    UpdateVDevice,
+    VMBuiltinLower,
+    VMShapeLower,
+    SpecializePrimFuncBasedOnCallSite,
+    dataflowblock_pass,
+    function_pass,
+)
+
+from .attach_external_modules import AttachExternModules
+from .fast_math import FastMathTransform
+from .fuse_transpose_matmul import FuseTransposeMatmul
+from .ipc_allreduce_rewrite import IPCAllReduceRewrite
+from .lazy_transform_params import LazyTransformParams
+from .lower_gpu_ipc_alloc_storage import LowerGPUIPCAllocStorage
+from .optimize_layout_transform import OptimizeLayoutTransform
+from .fold_batch_norm_to_conv2d_for_inference import FoldBatchnormToConv2D
+from .remove_redundant_reshape import RemoveRedundantReshape
+
+# Import to register the legalization functions.
+from . import legalize_ops

python/tvm/relax/transform/_ffi_api.py

@@ -0,0 +1,20 @@
+# 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
+"""FFI APIs for tvm.transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.transform", __name__)

python/tvm/relax/transform/attach_external_modules.py

@@ -0,0 +1,53 @@
+# 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.
+"""A pass that attaches external modules to the IRModule.
+
+Note: "external modules" here refers to `relax.frontend.nn.ExternModule`.
+"""
+
+from typing import TYPE_CHECKING
+
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+
+if TYPE_CHECKING:
+    from tvm.relax.frontend.nn import ExternModule
+
+
+@module_pass(opt_level=0, name="AttachExternalModules")
+class AttachExternModules:  # pylint: disable=too-few-public-methods
+    """Attach variable bounds to each Relax function, which primarily helps with memory planning."""
+
+    def __init__(self, extern_modules: list["ExternModule"]):
+        self.extern_modules = extern_modules
+
+    def transform_module(self, mod: IRModule, _ctx: PassContext) -> IRModule:
+        """Entrypoint"""
+        from tvm.relax.frontend.nn import (  # pylint: disable=import-outside-toplevel
+            ExternModule,
+        )
+
+        def _load(ext_mod: ExternModule):
+            assert isinstance(ext_mod, ExternModule), f"Expected ExternModule, but got: {ext_mod}"
+            return ext_mod.load()
+
+        mod_attrs = dict(mod.attrs) if mod.attrs else {}
+        external_mods = mod_attrs.get("external_mods", [])
+        for ext in self.extern_modules:
+            external_mods.append(_load(ext))
+        mod = mod.with_attr("external_mods", external_mods)
+        return mod

python/tvm/relax/transform/fast_math.py

@@ -0,0 +1,65 @@
+# 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, redefined-argument-from-local
+"""Relax Use Fast Math pass."""
+
+import tvm
+from tvm import topi
+from tvm.ir.module import IRModule
+from tvm.relax import Call, Expr, PyExprMutator, expr_functor
+
+
+@expr_functor.mutator
+class FastMathCodeGenerator(PyExprMutator):
+    """
+    Converts the expensive non linear functions to their fast but approximate counterparts.
+
+    Parameters
+    ----------
+    mod: IRModule
+        The module to be transformed
+    """
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_call_(self, call: Call) -> Expr:
+        if call.op.name == "relax.nn.softmax":
+            return self.builder_.call_te(topi.nn.fast_softmax, call.args[0], call.attrs.axis)
+        if call.op.name == "relax.exp":
+            return self.builder_.call_te(topi.fast_exp, call.args[0])
+        if call.op.name == "relax.erf":
+            return self.builder_.call_te(topi.fast_erf, call.args[0])
+        if call.op.name == "relax.tanh":
+            return self.builder_.call_te(topi.fast_tanh, call.args[0])
+
+        return super().visit_call_(call)
+
+
+@tvm.transform.module_pass(opt_level=0, name="FastMathTransform")
+class FastMathTransform:
+    """
+    Pass to convert the expensive non linear functions to their fast but approximate counterparts.
+    """
+
+    def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
+        fast_math_codegen = FastMathCodeGenerator(mod)
+        for gv, func in mod.functions_items():
+            if isinstance(func, tvm.relax.Function):
+                func = fast_math_codegen.visit_expr(func)
+                fast_math_codegen.builder_.update_func(gv, func)
+        return fast_math_codegen.builder_.get()

python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.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.
+# pylint: disable=invalid-name, unused-argument, redefined-argument-from-local
+"""Relax Fold Batchnorm into Conv2D."""
+
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext
+from tvm.relax import Expr
+from tvm.relax.dpl import TupleGetItemPattern, is_const, is_op, rewrite_call, wildcard
+
+from . import function_pass
+
+
+@function_pass(opt_level=0)
+class FoldBatchnormToConv2D:
+    """
+    Fuse Batchnorm to its previous Conv2D
+    This optimization is a special case of FoldScaleAxis that folds scale into conv2d weights.
+    This pass can be removed when FoldScaleAcis enhances to support this case.
+    """
+
+    def __init__(self):
+        self.input = wildcard()
+        self.weight = is_const()
+        self.pattern_conv2d = is_op("relax.nn.conv2d")(self.input, self.weight)
+        self.bn_weight = is_const()
+        self.bias = is_const()
+        self.mean = is_const()
+        self.variance = is_const()
+        self.pattern_bn = is_op("relax.nn.batch_norm")(
+            self.pattern_conv2d, self.bn_weight, self.bias, self.mean, self.variance
+        )
+
+        self.pattern = TupleGetItemPattern(self.pattern_bn, 0)
+
+    def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRModule:
+        """
+        Tranformation function for pattern Conv2D+BatchNorm+TupleGetItem pattern
+        Parameters
+        ----------
+        func: Expr
+            The relax function to be optimized
+        mod: IRModule
+            The ir module
+        ctx: PassContext
+            Relax pass context
+        """
+
+        self.mod = mod
+        updated_call = func
+
+        # Skip primitive functions
+        if "Primitive" in func.attrs.keys() and func.attrs["Primitive"] != 0:
+            return updated_call
+
+        def rewriter(expr, matches):
+            conv_input = matches[self.input]
+            conv_weight = matches[self.weight]
+            bn_weight = matches[self.bn_weight]
+            bn_bias = matches[self.bias]
+            bn_mean = matches[self.mean]
+            bn_variance = matches[self.variance]
+            conv_op = matches[self.pattern_conv2d]
+            bn_op = matches[self.pattern_bn]
+            conv_attrs = conv_op.attrs
+            bn_attrs = bn_op.attrs
+
+            bn_variance = relax.op.add(
+                bn_variance, relax.prim_value(tirx.FloatImm("float32", bn_attrs["epsilon"]))
+            )
+            dino = relax.op.sqrt(bn_variance)
+            wt = relax.op.divide(bn_weight, dino)
+            bs = relax.op.subtract(bn_bias, relax.op.multiply(bn_mean, wt))
+            if conv_attrs["kernel_layout"] == "OIHW":
+                wt = relax.op.reshape(wt, shape=(bn_weight.ty.shape[0], 1, 1, 1))
+            elif conv_attrs["kernel_layout"] == "IOHW":
+                wt = wt.reshape(1, bn_weight.ty.shape[0], 1, 1)
+            else:
+                return expr
+            wt_conv = relax.op.multiply(conv_weight, wt)
+            bs_args = relax.op.reshape(bs, shape=(1, bn_bias.ty.shape[0], 1, 1))
+
+            conv_out = relax.Call(conv_op.op, (conv_input, wt_conv), conv_attrs)
+            return relax.op.add(conv_out, bs_args)
+
+        updated_call = rewrite_call(self.pattern, rewriter, func)
+
+        return updated_call

python/tvm/relax/transform/fuse_transpose_matmul.py

@@ -0,0 +1,174 @@
+# 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.
+
+"""A compiler pass that fuses transpose + matmul and generate TIR function.
+Note that
+1. Please put the pass before LegalizeOps pass.
+2. The pass only works for XW^T but not X^TW
+3. The pass would rewrite the relax ops into TIR functions. If you'd like to dispatch the
+   ops into library (e.g. cuBLAS) calls, please run dispatch pass before this pass.
+"""
+
+import tvm
+from tvm import IRModule, relax, te, tirx
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
+class FuseTransposeMatmul:  # pylint: disable=too-few-public-methods
+    """A compiler pass that fuses transpose + matmul."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        mod = relax.transform.FuseOpsByPattern(
+            [
+                (
+                    "transpose_matmul_fuse",
+                    *_pattern(),
+                ),
+            ],
+            bind_constants=False,
+        )(mod)
+        transpose_matmul_codegen = _TransposeMatmulFuser(mod)
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = transpose_matmul_codegen.visit_expr(func)
+                transpose_matmul_codegen.builder_.update_func(g_var, func)
+        return transpose_matmul_codegen.builder_.get()
+
+
+def _pattern():
+    """Pattern for transpose + matmul."""
+    # pylint: disable=invalid-name
+    w = wildcard()
+    x = wildcard()
+    wT = is_op("relax.permute_dims")(w)
+    o = is_op("relax.matmul")(x, wT)
+    # pylint: enable=invalid-name
+    annotations = {"o": o, "w": w, "x": x, "wT": wT}
+
+    def _check(context: relax.transform.PatternCheckContext) -> bool:
+        transpose_call = context.annotated_expr["wT"]
+        ndim = transpose_call.args[0].ty.ndim
+        if ndim == -1:
+            return False
+        if ndim == 2 and transpose_call.attrs.axes is None:
+            return True
+        axes = list(range(ndim))
+        axes[-1], axes[-2] = axes[-2], axes[-1]
+        return list(transpose_call.attrs.axes) == axes
+
+    return o, annotations, _check
+
+
+# pylint: disable=missing-docstring,invalid-name
+
+
+@mutator
+class _TransposeMatmulFuser(PyExprMutator):  # pylint: disable=abstract-method
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_call_(  # pylint: disable=arguments-renamed
+        self,
+        call: relax.Call,
+    ) -> relax.Expr:
+        out_dtype = None
+
+        def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
+            nonlocal out_dtype
+            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.TensorType(a.shape))
+            bT_shape = list(b.shape)
+            bT_shape[-1], bT_shape[-2] = bT_shape[-2], bT_shape[-1]
+            bT_relax = relax.Var("b", relax.TensorType(bT_shape))
+            output_shape = self.builder_.normalize(relax.op.matmul(a_relax, bT_relax)).ty.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]
+                        dim_equal = a_dim == b_dim
+                        if not isinstance(dim_equal, tirx.IntImm) or dim_equal == 0:
+                            a_dim_is_one = isinstance(a_dim, tirx.IntImm) and a_dim == 1
+                            b_dim_is_one = isinstance(b_dim, tirx.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])
+                        else:
+                            a_indices.append(idx_spatial[i])
+                            b_indices.append(idx_spatial[i])
+
+                    if not a_prepended:
+                        a_indices.append(idx_spatial[-2 + b_appended])
+                    a_indices.append(idx_reduce)
+                    if not b_appended:
+                        b_indices.append(idx_spatial[-1])
+                    b_indices.append(idx_reduce)
+
+                    dtype = out_dtype
+                    if dtype != "":
+                        return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
+                    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="NT_matmul",
+            )
+
+        if isinstance(call.op, relax.GlobalVar):
+            function = self.builder_.get()[call.op]
+            if (
+                "Composite" in function.attrs
+                and function.attrs["Composite"] == "transpose_matmul_fuse"
+            ):
+                out_dtype = function.ret_ty.dtype
+                return self.builder_.call_te(
+                    te_transposed_matmul,
+                    call.args[1],
+                    call.args[0],
+                    primfunc_name_hint="NT_matmul",
+                )
+
+        return super().visit_call_(call)

python/tvm/relax/transform/ipc_allreduce_rewrite.py

@@ -0,0 +1,148 @@
+# 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.
+"""Rewrite all-reduce operation to customized all-reduce impl with IPC memory.
+The pass is written in Python for experiment, fast development.
+"""
+
+import tvm
+from tvm import relax
+from tvm.ir.module import IRModule
+from tvm.relax.expr import Expr, Var
+from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+
+
+@tvm.transform.module_pass(opt_level=0, name="IPCAllReduceRewrite")
+class IPCAllReduceRewrite:
+    """Rewrite all-reduce operation to customized all-reduce impl with IPC memory."""
+
+    def __init__(self, allreduce_strategy: int) -> None:
+        """Constructor
+
+        Parameters
+        ----------
+        allreduce_strategy : int
+            The all-reduce strategy. Only "1" and "2" are supported.
+            "1" stands for one-shot, and "2" stands for two-shot.
+        """
+        self.allreduce_strategy = allreduce_strategy
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        fcustom_allreduce = tvm.get_global_func(
+            "runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True
+        )
+        if fcustom_allreduce is None:
+            # Customized allreduce is not available.
+            return mod
+
+        binding_replacement_map = _Visitor(self.allreduce_strategy).visit(mod)
+        return _Rewriter(mod, binding_replacement_map).transform()
+
+
+@visitor
+class _Visitor(PyExprVisitor):  # pylint: disable=abstract-method
+    def __init__(self, allreduce_strategy: int) -> None:
+        self.allreduce_strategy = allreduce_strategy
+        self.alloc_map: dict[Var, relax.Call] = {}
+        self.binding_replacement_map: dict[relax.Expr, relax.Expr] = {}
+        self.builtin_alloc_tensor_op = tvm.ir.Op.get("relax.builtin.alloc_tensor")
+        self.reshape_op = tvm.ir.Op.get("relax.reshape")
+
+    def visit(self, mod: IRModule) -> dict[relax.Expr, relax.Expr]:
+        """Entry point"""
+        for _, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                self.alloc_map.clear()
+                self.visit_expr(func)
+        return self.binding_replacement_map
+
+    def visit_var_binding_(self, binding: relax.VarBinding):
+        super().visit_var_binding_(binding)
+        if (
+            isinstance(binding.value, relax.Call)
+            and binding.value.op == self.builtin_alloc_tensor_op
+        ):
+            self.alloc_map[binding.var] = binding.value
+        elif isinstance(binding.value, relax.Var) and binding.value in self.alloc_map:
+            self.alloc_map[binding.var] = self.alloc_map[binding.value]
+        elif (
+            isinstance(binding.value, relax.Call)
+            and binding.value.op == self.reshape_op
+            and binding.value.args[0] in self.alloc_map
+        ):
+            self.alloc_map[binding.var] = self.alloc_map[binding.value.args[0]]
+
+    def visit_call_(self, call: relax.Call) -> None:  # pylint: disable=arguments-renamed
+        if (
+            not isinstance(call.op, relax.ExternFunc)
+            or call.op.global_symbol != "runtime.disco.allreduce"
+            or call.args[1].values[0] != 0
+        ):
+            # Return if the call is not a summation all-reduce.
+            return
+
+        assert len(call.args) == 4
+        allreduce_input, _strategy, _ingroup, allreduce_output = call.args
+        alloc_tensor = self.alloc_map.get(allreduce_input, None)
+        if alloc_tensor is None or alloc_tensor.args[3].value != "global":
+            # Return if the allocation of all-reduce input is not recorded,
+            # or the scope of the allocation is not global.
+            return
+
+        # Set the scope of the alloc_tensor to IPC memory.
+        alloc_tensor = self.alloc_map[allreduce_input]
+        self.binding_replacement_map[alloc_tensor] = relax.op.builtin.alloc_tensor(
+            alloc_tensor.args[0],
+            alloc_tensor.args[1],
+            alloc_tensor.args[2],
+            relax.StringImm("ipc_memory"),
+        )
+
+        self.binding_replacement_map[call] = relax.Call(
+            relax.ExternFunc("runtime.disco.cuda_ipc.custom_allreduce"),
+            # The "cuda_ipc.custom_allreduce" implementation does not
+            # yet support num_groups>1, and therefore does not use the
+            # `in_group` argument.
+            [allreduce_input, relax.prim_value(self.allreduce_strategy), allreduce_output],
+        )
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    """Rewrite the IRModule according to the binding replacement provided by the visitor."""
+
+    def __init__(
+        self, mod: IRModule, binding_replacement_map: dict[relax.Expr, relax.Expr]
+    ) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.binding_replacement_map = binding_replacement_map
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                self.builder_.update_func(g_var, updated_func)
+        return self.builder_.get()
+
+    def visit_call_(self, call: relax.Call) -> Expr:  # pylint: disable=arguments-renamed
+        return (
+            super().visit_call_(self.binding_replacement_map[call])
+            if call in self.binding_replacement_map
+            else super().visit_call_(call)
+        )

python/tvm/relax/transform/lazy_transform_params.py

@@ -0,0 +1,399 @@
+# 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, missing-function-docstring, abstract-method, missing-class-docstring
+# ruff: noqa: RUF005
+"""Relax LazyTransformParams pass."""
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+
+
+@visitor
+class ForwardCollector(PyExprVisitor):
+    """
+    Perform a forward pass to collect the following information:
+    out_tuple_map: map from var to its index in the output tuple
+    var_tuple_get_item: list of var that is bound to v = params[i]
+
+    Parameters
+    ----------
+    tuple_var: relax.Var
+        The output tuple var
+
+    input_params: relax.Var
+        The input tuple var
+
+    """
+
+    def __init__(self, tuple_var: relax.Var, input_params: relax.Var) -> None:
+        self.out_tuple_map = {}
+        self.out_tuple_var = tuple_var
+        self.input_params = input_params
+        self.var_tuple_get_item = []
+        self.is_tuple_get_item_input = False
+
+    def visit_tuple_getitem_(self, op: relax.TupleGetItem) -> None:
+        if op.tuple_value == self.input_params:
+            self.is_tuple_get_item_input = True
+        else:
+            self.is_tuple_get_item_input = False
+        super().visit_tuple_getitem_(op)
+
+    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
+        if binding.var == self.out_tuple_var:
+            assert isinstance(binding.value, relax.Tuple)
+            for i, expr in enumerate(binding.value.fields):
+                if expr not in self.out_tuple_map:
+                    self.out_tuple_map[expr] = []
+                self.out_tuple_map[expr].append(relax.prim_value(i))
+        else:
+            self.is_tuple_get_item_input = False
+            super().visit_var_binding_(binding)
+            if self.is_tuple_get_item_input:
+                self.var_tuple_get_item.append(binding.var)
+
+
+@visitor
+class LivenessAnalysis(PyExprVisitor):
+    """
+    Perform a backward pass to collect the following information:
+    var_liveness_end: map from var to the list of var whose liveness is killed by this var binding
+
+    Parameters
+    ----------
+    out_tuple_var: relax.Var
+        The output tuple var
+    input_params: set
+        The set of vars that are bound to v = params[i]
+    """
+
+    def __init__(self, out_tuple_var: relax.Var) -> None:
+        self.last_appear_in_var_binding = []
+        self.out_tuple_var = out_tuple_var
+        self.var_liveness_end = {}
+        self.ended_vars = set()
+
+    def visit_binding_block_(self, block: relax.BindingBlock) -> None:
+        for binding in reversed(block.bindings):
+            self.visit_binding(binding)
+
+    def visit_var_(self, op: relax.Var) -> None:
+        if op not in self.ended_vars:
+            self.last_appear_in_var_binding.append(op)
+            self.ended_vars.add(op)
+
+    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
+        if self.out_tuple_var == binding.var:
+            return
+        self.last_appear_in_var_binding = []
+        super().visit_var_binding_(binding)
+        # param[i] is in output
+        if binding.var not in self.ended_vars:
+            self.last_appear_in_var_binding.append(binding.var)
+            self.ended_vars.add(binding.var)
+        self.var_liveness_end[binding.var] = self.last_appear_in_var_binding
+
+
+class LazyTransformParamsFuncCreator:
+    """
+    Transform transform_params functions into a lazy version.
+
+    Parameters
+    ----------
+    mod: IRModule
+        The module to be transformed
+    """
+
+    def __init__(
+        self,
+        fget_item,
+        fset_item,
+        extra_get_item_params,
+        extra_set_item_params,
+        mod: IRModule = None,
+    ) -> None:
+        self.mod = mod
+        self.fget_item = fget_item
+        self.extra_get_item_params = extra_get_item_params
+        self.fset_item = fset_item
+        self.extra_set_item_params = extra_set_item_params
+        self.input_params_set = None
+        self.out_tuple_map = None
+        self.out_tuple_var = None
+        self.memory_free_insertion = None
+
+    def transform(self, func: relax.Function) -> relax.Function:
+        if "num_input" in func.attrs:
+            num_input = func.attrs["num_input"]
+        else:
+            num_input = 0
+
+        seq_expr = func.body
+        self.out_tuple_var = seq_expr.body
+
+        # Step 1. collect out_tuple_map and input_params_set
+        forward_collector = ForwardCollector(self.out_tuple_var, func.params[num_input])
+        forward_collector.visit_expr(func)
+        self.out_tuple_map = forward_collector.out_tuple_map
+        # input_params_set is the set of binding var for var = params[i]
+        self.input_params_set = set(forward_collector.var_tuple_get_item)
+
+        # Step 2. liveness analysis and get where to insert kill_object instruction
+        liveness = LivenessAnalysis(self.out_tuple_var)
+        liveness.visit_expr(func)
+        self.memory_free_insertion = liveness.var_liveness_end
+
+        # Step 3. rewrite get item and set item
+        if self.fget_item is not None:
+            new_func = LazyInputMutator(self, self.mod).visit_expr(func)
+
+        new_body = new_func.body
+        if self.fset_item is not None:
+            # The LazyOutputMutator only inspects variable bindings
+            # for replacement.  If the output tuple includes elements
+            # that do not have a variable binding, such as
+            # `relax.Const`, these must still produce a call to the
+            # `"set_item"` function.
+            leaf_outputs = {
+                expr: indices
+                for expr, indices in self.out_tuple_map.items()
+                if not isinstance(expr, relax.Var)
+            }
+            if leaf_outputs:
+                new_bindings = [
+                    relax.VarBinding(
+                        relax.Var("_", relax.AnyType()),
+                        relax.Call(
+                            relax.ExternFunc(self.fset_item),
+                            [*self.extra_set_item_params, index, expr],
+                            None,
+                            [relax.AnyType()],
+                        ),
+                    )
+                    for expr, indices in leaf_outputs.items()
+                    for index in indices
+                ]
+                new_body = relax.SeqExpr(
+                    [*new_body.blocks, relax.BindingBlock(new_bindings)], new_body.body
+                )
+
+            new_body = LazyOutputMutator(self, self.mod).visit_expr(new_body)
+
+        # Step 4. Add parameters of get_item and set_item (except index) to the function.
+        params = [
+            *func.params[:num_input],
+            *self.extra_get_item_params,
+            *self.extra_set_item_params,
+        ]
+
+        # Step 5. Find all shape parameters that should be retained as
+        # parameters.
+        symbolic_vars = relax.analysis.defined_symbolic_vars(func)
+        if symbolic_vars:
+
+            def unpack_ty(ty):
+                if isinstance(ty, relax.TupleType):
+                    for field in ty.fields:
+                        yield from unpack_ty(field)
+                else:
+                    yield ty
+
+            # direct iterate over the type annotation
+            for param in func.params[num_input:]:
+                for ty in unpack_ty(param.ty):
+                    if isinstance(ty, tvm.ir.PrimType | relax.ShapeType):
+                        params.append(relax.Var("symbolic_var_holder", ty))
+
+        return relax.Function(
+            params,
+            new_body,
+            relax.AnyType(),
+            attrs=func.attrs,
+            is_pure=False,
+        ).without_attr("relax.force_pure")
+
+
+@mutator
+class LazyInputMutator(PyExprMutator):
+    def __init__(self, func_creator, mod: IRModule | None = None) -> None:
+        self.func_creator = func_creator
+        super().__init__(mod)
+
+    def visit_function_(self, func: relax.Function) -> relax.Expr:
+        if "num_input" in func.attrs:
+            num_input = func.attrs["num_input"]
+        else:
+            num_input = 0
+
+        params = list(func.params)[num_input:]
+        if len(params) == 1 and isinstance(params[0].ty, relax.TupleType):
+            self.tuple_param = params[0]
+            self.params = {}
+        else:
+            self.tuple_param = None
+            self.params = {var: i for i, var in enumerate(params)}
+        func = relax.Function(
+            func.params[:num_input],
+            func.body,
+            func.ret_ty,
+            is_pure=False,
+            attrs=func.attrs,
+            span=func.span,
+        ).without_attr("relax.force_pure")
+        output = super().visit_function_(func)
+        self.tuple_param = None
+        self.params = {}
+        return output
+
+    def visit_var_(self, var: relax.Var) -> relax.Expr:
+        if var in self.params:
+            index = self.params[var]
+            get_item_result = self.builder_.emit(
+                relax.Call(
+                    relax.ExternFunc(self.func_creator.fget_item),
+                    self.func_creator.extra_get_item_params + [relax.prim_value(index)],
+                    None,
+                    [relax.AnyType()],
+                )
+            )
+            match_cast = relax.MatchCast(var, get_item_result, var.ty)
+            self.builder_.emit_normalized(match_cast)
+
+            del self.params[var]
+
+        return super().visit_var_(var)
+
+    def visit_tuple_getitem_(self, node: relax.TupleGetItem) -> relax.Expr:
+        ty = node.ty
+
+        node = super().visit_tuple_getitem_(node)
+
+        if self.tuple_param is not None and node.tuple_value.same_as(self.tuple_param):
+            get_item_result = self.builder_.emit(
+                relax.Call(
+                    relax.ExternFunc(self.func_creator.fget_item),
+                    self.func_creator.extra_get_item_params + [relax.prim_value(node.index)],
+                    None,
+                    [relax.AnyType()],
+                )
+            )
+            return self.builder_.match_cast(get_item_result, ty)
+        else:
+            return node
+
+
+@mutator
+class LazyOutputMutator(PyExprMutator):
+    def __init__(self, func_creator, mod: IRModule | None = None) -> None:
+        self.func_creator = func_creator
+        self.killed_vars = set()
+        super().__init__(mod)
+
+    def visit_var_(self, var: relax.Var) -> None:
+        assert var not in self.killed_vars
+        return super().visit_var_(var)
+
+    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
+        if binding.var == self.func_creator.out_tuple_var:
+            # The function after rewriting returns a empty tuple.
+            func_output = self.builder_.emit(relax.Tuple([]))
+            self.set_var_remap(binding.var, func_output)
+            return
+
+        super().visit_var_binding_(binding)
+
+        if binding.var in self.func_creator.memory_free_insertion:
+            for var in self.func_creator.memory_free_insertion[binding.var]:
+                if var in self.func_creator.out_tuple_map:
+                    self.killed_vars.add(var)
+                    for index in self.func_creator.out_tuple_map[var]:
+                        # rewrite set item
+                        self.builder_.emit(
+                            relax.Call(
+                                relax.ExternFunc(self.func_creator.fset_item),
+                                self.func_creator.extra_set_item_params
+                                + [index, super().visit_var_(var)],
+                                None,
+                                [relax.AnyType()],
+                            ),
+                            name_hint="_",
+                        )
+
+                if var in self.func_creator.input_params_set:
+                    self.builder_.emit(
+                        relax.op.vm.kill_object(super().visit_var_(var)), name_hint="_"
+                    )
+
+
+@tvm.transform.module_pass(opt_level=0, name="LazyTransformParams")
+class LazyTransformParams:
+    """
+    Convert transform_params functions into a lazy version.
+    (Load the input to memory on demand, and immediately free it after the last use.)
+
+    Note: ToNonDataflow() and RemovePurityTracking() should be invoked before this pass.
+
+    Parameters
+    ----------
+    fget_item: str
+        The name of the get_item function.
+    fset_item: str
+        The name of the set_item function.
+    extra_get_item_params: list of relax.Var
+        The parameters of the get_item function except index.
+        The given parameters will be placed before index.
+        For example, if extra_get_item_params is [param1, param2], then the pass will generate
+        call_packed(fget_item, [param1, param2, index])
+    extra_set_item_params: list of relax.Var
+        The parameters of the set_item function except index and value.
+        The given parameters will be placed before index and value.
+        For example, if extra_set_item_params is [param1, param2], then the pass will generate
+        call_packed(fset_item, [param1, param2, index, value])
+    """
+
+    def __init__(
+        self,
+        fget_item="get_item",
+        fset_item="set_item",
+        extra_get_item_params=None,
+        extra_set_item_params=None,
+    ) -> None:
+        self.fget_item = fget_item
+        self.extra_get_item_params = [] if extra_get_item_params is None else extra_get_item_params
+        assert self.fget_item is not None, "transforming set_item only is not supported"
+        self.fset_item = fset_item
+        self.extra_set_item_params = [] if extra_set_item_params is None else extra_set_item_params
+
+    def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
+        lazy_mutator = LazyTransformParamsFuncCreator(
+            self.fget_item,
+            self.fset_item,
+            self.extra_get_item_params,
+            self.extra_set_item_params,
+            mod,
+        )
+        builder = relax.BlockBuilder(mod)
+        for gv, _ in mod.functions_items():
+            if gv.name_hint.endswith("transform_params"):
+                func = mod[gv]
+                if not isinstance(func, relax.Function):
+                    continue
+                func = lazy_mutator.transform(func)
+                builder.update_func(gv, func)
+
+        return builder.get()

python/tvm/relax/transform/legalize_ops/__init__.py

@@ -0,0 +1,39 @@
+# isort: skip_file
+# 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.
+"""Legalize high-level operator calls in Relax functions to call_tir."""
+
+from . import binary
+from . import ccl
+from . import create
+from . import datatype
+from . import distributed
+from . import grad
+from . import image
+from . import index
+from . import inspect_op
+from . import linear_algebra
+from . import manipulate
+from . import nn
+from . import qdq
+from . import search
+from . import statistical
+from . import unary
+from . import vision
+
+# Device specific legalizations
+from . import adreno

python/tvm/relax/transform/legalize_ops/adreno/__init__.py

@@ -0,0 +1,20 @@
+# isort: skip_file
+# 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.
+"""Legalize high-level operator calls in Relax functions to call_tir."""
+
+from .convolution import conv2d_NCHWc_OIHWo

python/tvm/relax/transform/legalize_ops/adreno/convolution.py

@@ -0,0 +1,36 @@
+# 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=missing-docstring, invalid-name
+"""A Convolution impl for Adreno GPU."""
+
+from tvm import relax, topi
+
+
+def conv2d_NCHWc_OIHWo(bb: relax.BlockBuilder, call: relax.Call) -> relax.Expr:
+    return bb.call_te(
+        topi.nn.conv2d_NCHWc_OIHWo,
+        data=call.args[0],
+        kernel=call.args[1],
+        stride=call.attrs.strides,
+        padding=call.attrs.padding,
+        dilation=call.attrs.dilation,
+        layout=call.attrs.data_layout,
+        out_layout=call.attrs.out_layout,
+        # out_dtype=call.attrs.out_dtype,
+        ty_args=call.ty_args,
+        primfunc_name_hint="conv2d_NCHWc_OIHWo",
+    )

python/tvm/relax/transform/legalize_ops/binary.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 binary operators."""
+
+from tvm import topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import (
+    LegalizeFunc,
+    TEFunc,
+    _is_relax_expr,
+    _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 _is_relax_expr(arg1):
+            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.log_add_exp", _binary(topi.log_add_exp))
+register_legalize("relax.multiply", _binary(topi.multiply))
+register_legalize("relax.power", _binary(topi.power))
+register_legalize("relax.atan2", _binary(topi.atan2))
+register_legalize("relax.subtract", _binary(topi.subtract))
+register_legalize("relax.equal", _binary(topi.equal))
+register_legalize("relax.mod", _binary(topi.mod))
+register_legalize("relax.floor_mod", _binary(topi.floor_mod))
+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))
+
+register_legalize("relax.maximum", _binary(topi.maximum))
+register_legalize("relax.minimum", _binary(topi.minimum))
+
+# bitwise
+register_legalize("relax.bitwise_and", _binary(topi.bitwise_and))
+register_legalize("relax.bitwise_or", _binary(topi.bitwise_or))
+register_legalize("relax.bitwise_xor", _binary(topi.bitwise_xor))
+register_legalize("relax.left_shift", _binary(topi.left_shift))
+register_legalize("relax.right_shift", _binary(topi.right_shift))
+
+# logical
+register_legalize("relax.logical_and", _binary(topi.logical_and))
+register_legalize("relax.logical_or", _binary(topi.logical_or))
+register_legalize("relax.logical_xor", _binary(topi.logical_xor))

python/tvm/relax/transform/legalize_ops/ccl.py

@@ -0,0 +1,125 @@
+# 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
+# ruff: noqa: RUF005
+"""Default legalization function for ccl operators."""
+
+from tvm import arith, tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, ShapeExpr
+from ...op import call_dps_packed
+from ...type import ShapeType, TensorType
+from .common import register_legalize
+
+
+@register_legalize("relax.ccl.allreduce")
+def _allreduce(_bb: BlockBuilder, call: Call) -> Expr:
+    op_type_str = call.attrs.op_type
+    op_type_map = {
+        "sum": 0,
+        "prod": 1,
+        "min": 2,
+        "max": 3,
+        "avg": 4,
+    }
+    if op_type_str not in op_type_map:
+        raise ValueError(
+            f"Unsupported reduction operation: {op_type_str}. "
+            f"Supported operations are {op_type_map.keys()}."
+        )
+    return call_dps_packed(
+        "runtime.disco.allreduce",
+        [call.args[0], ShapeExpr([op_type_map[op_type_str]]), call.attrs.in_group],
+        out_ty=call.args[0].ty,
+    )
+
+
+@register_legalize("relax.ccl.allgather")
+def _allgather(_bb: BlockBuilder, call: Call) -> Expr:
+    output_shape = []
+    arg_ty = call.args[0].ty
+    assert isinstance(arg_ty, TensorType), "The input type of allgather should be TensorType."
+    assert isinstance(arg_ty.shape.ty, ShapeType)
+    arg_shape = arg_ty.shape.ty
+    for i, shape_value in enumerate(arg_shape.values):
+        if i == 0:
+            output_shape.append(shape_value * call.attrs.num_workers)
+        else:
+            output_shape.append(shape_value)
+    return call_dps_packed(
+        "runtime.disco.allgather",
+        [call.args[0], call.attrs.in_group],
+        out_ty=TensorType(
+            shape=output_shape,
+            dtype=arg_ty.dtype,
+            vdevice=arg_ty.vdevice,
+        ),
+    )
+
+
+@register_legalize("relax.ccl.broadcast_from_worker0")
+def _broadcast_from_worker0(_bb: BlockBuilder, call: Call) -> Expr:
+    return call_dps_packed(
+        "runtime.disco.broadcast_from_worker0",
+        [call.args[0], False],
+        out_ty=call.args[0].ty,
+    )
+
+
+# Since collective communication ops are performed on contiguous memory,
+# we need to reshape and transpose the input tensor to make sharding dimension in the highest order
+def _transpose_for_ccl(_bb: BlockBuilder, expr: Expr, axis: int, num_workers: int):
+    assert isinstance(expr.ty, TensorType), "The input type should be TensorType."
+    assert isinstance(expr.ty.shape.ty, ShapeType)
+    arg_shape = expr.ty.shape.ty
+    new_shape = []
+    for i, shape_value in enumerate(arg_shape.values):
+        if i == axis:
+            modulo = arith.Analyzer().simplify(shape_value % num_workers)
+            assert modulo == 0, (
+                f"scatter_from_worker0 expects the size of axis {axis} of input tensor "
+                "to be divisible by num_workers. However, the axis 0 of input tensor "
+                f"is {shape_value} while num_workers is {num_workers}"
+            )
+            new_shape.append(num_workers)
+            new_shape.append(tirx.div(shape_value, num_workers))
+        else:
+            new_shape.append(shape_value)
+    reshape_var = _bb.emit_te(topi.reshape, expr, new_shape)
+    if axis == 0:
+        return reshape_var
+    permute_order = [axis] + list(range(axis)) + list(range(axis + 1, len(new_shape)))
+    transpose_var = _bb.emit_te(topi.transpose, reshape_var, permute_order)
+    return transpose_var
+
+
+@register_legalize("relax.ccl.scatter_from_worker0")
+def _scatter_from_worker0(_bb: BlockBuilder, call: Call) -> Expr:
+    transpose_var = _transpose_for_ccl(_bb, call.args[0], call.attrs.axis, call.attrs.num_workers)
+    output_shape = transpose_var.ty.shape.ty.values
+    output_shape = output_shape[1:]
+    return call_dps_packed(
+        "runtime.disco.scatter_from_worker0",
+        [transpose_var, False],
+        out_ty=TensorType(
+            shape=output_shape,
+            dtype=call.args[0].ty.dtype,
+            vdevice=call.args[0].ty.vdevice,
+        ),
+    )

python/tvm/relax/transform/legalize_ops/common.py

@@ -0,0 +1,125 @@
+# 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 collections.abc import Callable
+
+import tvm
+from tvm import te
+from tvm.ir import Call
+from tvm.runtime import DataTypeCode
+from tvm.tirx import FloatImm, IntImm
+
+from ...block_builder import BlockBuilder
+from ...expr import Constant, Expr
+
+##################### 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]
+
+
+def _is_relax_expr(expr: object) -> bool:
+    return isinstance(expr, Expr) and not tvm.ir.is_prim_expr(expr)
+
+
+def _try_convert_to_scalar_const(
+    expr: Expr, python_native: bool = False
+) -> Expr | FloatImm | IntImm | bool | float | int:
+    """Check if the input Expr is a scalar constant.
+    If it is, return its plain value with the same data type or in native python type.
+    If it is not, return the input expr.
+
+    Note that if the python_native flag is True, the returned value will be in native python type,
+    this might cause loss of data type for example, a float16 constant will be converted to float32
+    and a int64 constant will be converted to int32.
+
+    Parameters
+    ----------
+    expr : Expr
+        The expr to be checked and converted.
+
+    Returns
+    -------
+    ret : Union[Expr, FloatImm, IntImm, bool, float, int]
+        Return a FloatImm or IntImm if the given expr is a scalar integer or float constant, and the
+        python native flag is False. Or return the plain value of the constant in native python type
+        if the python native flag is True.
+        Or return the input itself if it is not a scalar constant.
+    """
+    if isinstance(expr, Constant) and expr.ty.ndim == 0:
+        # get the value of the scalar constant
+        value = expr.data.numpy()[()].item()
+        dtype = expr.ty.dtype
+        dtype_str = str(dtype.dtype)
+        if python_native:
+            return value
+        # preserve the data type of the constant
+        if dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT):
+            return tvm.tirx.FloatImm(dtype_str, value)
+        elif dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT, DataTypeCode.BOOL):
+            return tvm.tirx.IntImm(dtype_str, value)
+    return expr
+
+
+def _call_topi_without_attr(te_func: TEFunc, primfunc_name: str | None = 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)

python/tvm/relax/transform/legalize_ops/create.py

@@ -0,0 +1,166 @@
+# 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
+# ruff: noqa: E731
+"""Default legalization function for creation operators."""
+
+import numpy as np
+
+import tvm
+from tvm import te, tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, ShapeExpr, const
+from ...type import ShapeType
+from .common import LegalizeFunc, _try_convert_to_scalar_const, register_legalize
+
+
+def _full(is_like: bool, fill_value: float | None, primfunc_name: str) -> LegalizeFunc:
+    def full_call_te(bb: BlockBuilder, call: Call) -> Expr:
+        _fill_value = (
+            _try_convert_to_scalar_const(call.args[1], python_native=True)
+            if fill_value is None
+            else fill_value
+        )
+        shape = call.args[0].ty.shape if is_like else call.args[0]
+
+        if isinstance(shape, ShapeExpr):
+            output_shape = shape.values
+        else:
+            assert isinstance(shape.ty, ShapeType)
+            assert shape.ty.ndim >= 0
+
+            shape = bb.emit(shape)
+            output_shape = [tirx.Var(f"s{i}", "int64") for i in range(shape.ty.ndim)]
+            bb.match_cast(shape, ShapeType(output_shape))
+
+        return bb.call_te(
+            topi.full,
+            output_shape,
+            call.ty.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:
+        data, k = call.args
+        return bb.call_te(
+            topi.trilu,
+            data,
+            k,
+            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"))
+
+
+def _eye(is_like: bool, primfunc_name: str) -> LegalizeFunc:
+    def eye_call_te(bb: BlockBuilder, call: Call) -> Expr:
+        _convert_to_scalar_const = lambda x: _try_convert_to_scalar_const(x, python_native=True)
+        if is_like:
+            x = call.args[0]
+            k = _convert_to_scalar_const(call.args[1]) if len(call.args) > 1 else 0
+            n, m = x.ty.shape
+            dtype = x.ty.dtype
+        else:
+            n = _convert_to_scalar_const(call.args[0])
+            m = _convert_to_scalar_const(call.args[1]) if len(call.args) > 1 else n
+            k = _convert_to_scalar_const(call.args[2]) if len(call.args) > 2 else 0
+            dtype = call.attrs.dtype
+
+        return bb.call_te(
+            topi.eye,
+            n,
+            m,
+            k,
+            dtype,
+            primfunc_name_hint=primfunc_name,
+        )
+
+    return eye_call_te
+
+
+register_legalize("relax.eye", _eye(is_like=False, primfunc_name="eye"))
+register_legalize("relax.eye_like", _eye(is_like=True, primfunc_name="eye_like"))
+
+
+@register_legalize("relax.arange")
+def _arange(bb: BlockBuilder, call: Call) -> Expr:
+    assert len(call.args) == 3
+    assert all(tvm.ir.is_prim_expr(x) for x in call.args)
+    start, end, step = call.args
+    dtype = call.attrs.dtype
+
+    def is_const_scalar(x: tirx.Expr):
+        return isinstance(x, tirx.IntImm | tirx.FloatImm)
+
+    if all([is_const_scalar(x) for x in call.args]):
+        return const(np.arange(start.value, end.value, step.value, dtype=dtype), dtype=dtype)
+    else:
+        return bb.call_te(topi.arange, start, end, step, dtype)
+
+
+@register_legalize("relax.shape_to_tensor")
+def _shape_to_tensor(bb: BlockBuilder, call: Call) -> Expr:
+    shape = call.args[0]
+    values = shape.values if isinstance(shape, ShapeExpr) else shape.ty.values
+    if values is None:
+        return call
+    values = list(values)
+    n = len(values)
+    symbolic = [v for v in values if not isinstance(v, tirx.IntImm)]
+
+    def te_shape_to_tensor(*sym):
+        sym = list(sym)
+        resolved = [v if isinstance(v, tirx.IntImm) else sym.pop(0) for v in values]
+
+        def fcompute(i):
+            result = tirx.const(0, "int64")
+            for idx in range(n - 1, -1, -1):
+                result = tirx.if_then_else(i == idx, tirx.Cast("int64", resolved[idx]), result)
+            return result
+
+        return te.compute((n,), fcompute, name="shape_to_tensor")
+
+    return bb.call_te(te_shape_to_tensor, *symbolic, primfunc_name_hint="shape_to_tensor")
+
+
+@register_legalize("relax.hamming_window")
+def _hamming_window(bb: BlockBuilder, call: Call) -> Expr:
+    assert len(call.args) == 4
+    dtype = call.attrs.dtype
+    window_size = call.args[0]
+    periodic = call.args[1]
+    alpha = call.args[2]
+    beta = call.args[3]
+    return bb.call_te(topi.hamming_window, window_size, periodic, alpha, beta, dtype)

python/tvm/relax/transform/legalize_ops/datatype.py

@@ -0,0 +1,34 @@
+# 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 datatype operators."""
+
+from tvm import relax, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import _is_relax_expr, _try_convert_to_scalar_const, register_legalize
+
+
+@register_legalize("relax.astype")
+def _astype(bb: BlockBuilder, call: Call) -> Expr:
+    arg = _try_convert_to_scalar_const(call.args[0], python_native=True)
+    if _is_relax_expr(arg):
+        return bb.call_te(topi.cast, arg, call.attrs.dtype)
+    else:
+        return relax.const(arg, call.attrs.dtype)

python/tvm/relax/transform/legalize_ops/distributed.py

@@ -0,0 +1,45 @@
+# 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 distir-related operators."""
+
+from tvm import relax, tirx
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from ...op import call_pure_packed
+from ...type import ShapeType
+from .common import register_legalize
+
+
+@register_legalize("relax.dist.redistribute_replica_to_shard")
+def _redistribute_replica_to_shard(_bb: BlockBuilder, call: Call) -> Expr:
+    num_workers = call.attrs.num_workers
+    axis = call.attrs.axis
+    worker_id_symbol = tirx.Var("worker_id", "int64")
+    worker_id_var = _bb.emit(call_pure_packed("runtime.disco.worker_id", ty_args=[ShapeType(None)]))
+    _bb.match_cast(worker_id_var, ShapeType([worker_id_symbol]))
+
+    split_axis_size = call.args[0].ty.shape[axis]
+    return relax.op.strided_slice(
+        call.args[0],
+        axes=[axis],
+        begin=[worker_id_symbol * split_axis_size // num_workers],
+        end=[(worker_id_symbol + 1) * split_axis_size // num_workers],
+        assume_inbound=True,
+    )

python/tvm/relax/transform/legalize_ops/grad.py

@@ -0,0 +1,241 @@
+# 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 perators to implement operaor gradients."""
+
+import logging
+
+from tvm import te, tirx, topi
+from tvm.ir import Call
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+from tvm.tirx.script.builder.utils import buffer_proxy
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import register_legalize
+
+
+@register_legalize("relax.grad.no_grad")
+def _no_grad(bb: BlockBuilder, call: Call) -> Expr:
+    return call.args[0]
+
+
+@register_legalize("relax.grad.start_checkpoint")
+def _start_checkpoint(bb: BlockBuilder, call: Call) -> Expr:
+    return call.args[0]
+
+
+@register_legalize("relax.grad.end_checkpoint")
+def _end_checkpoint(bb: BlockBuilder, call: Call) -> Expr:
+    return call.args[0]
+
+
+@register_legalize("relax.grad.nll_loss_backward")
+def _grad_nll_loss_backward(bb: BlockBuilder, call: Call) -> Expr:
+    # topi.sum don't support zero-dim x
+    # we add support for that
+    def topi_sum_extend(x):
+        return x if x.ndim == 0 else topi.sum(x)
+
+    def te_nll_loss_backward(output_grad, predictions, targets, weights, reduction, ignore_index):
+        # handle ignore_index
+        if ignore_index >= 0:
+            weights = te.compute(
+                weights.shape,
+                lambda i: tirx.Select(i == ignore_index, tirx.const(0, weights.dtype), weights(i)),
+                "weights_new",
+            )
+
+        all_weights = te.compute(targets.shape, lambda *i: weights(targets(*i)), "all_weights")
+
+        # handle reduction
+        if reduction == "sum":
+            output_grad = topi.broadcast_to(output_grad, targets.shape)
+        elif reduction == "mean":
+            weight_sum = topi_sum_extend(all_weights)
+            output_grad = topi.divide(topi.broadcast_to(output_grad, targets.shape), weight_sum)
+
+        # handle no batch
+        if predictions.ndim == 1:
+            return te.compute(
+                predictions.shape,
+                lambda i: tirx.Select(
+                    i == targets(), -all_weights() * output_grad(), tirx.const(0, predictions.dtype)
+                ),
+                "pred_grad",
+            )
+
+        return te.compute(
+            predictions.shape,
+            lambda *i: tirx.Select(
+                i[1] == targets(*i[:1], *i[2:]),
+                -all_weights(*i[:1], *i[2:]) * output_grad(*i[:1], *i[2:]),
+                tirx.const(0, predictions.dtype),
+            ),
+            "pred_grad",
+        )
+
+    def te_nll_loss_backward_no_weight(output_grad, predictions, targets, reduction, ignore_index):
+        weight = topi.full(
+            (predictions.shape[1] if len(predictions.shape) > 1 else predictions.shape[0],),
+            predictions.dtype,
+            1.0,
+        )
+        return te_nll_loss_backward(
+            output_grad, predictions, targets, weight, reduction, ignore_index
+        )
+
+    if len(call.args) == 3:
+        return bb.call_te(
+            te_nll_loss_backward_no_weight,
+            *call.args,
+            reduction=call.attrs.reduction,
+            ignore_index=call.attrs.ignore_index,
+        )
+
+    return bb.call_te(
+        te_nll_loss_backward,
+        *call.args,
+        reduction=call.attrs.reduction,
+        ignore_index=call.attrs.ignore_index,
+        primfunc_name_hint="nll_loss_backward",
+    )
+
+
+@register_legalize("relax.grad.max_pool2d_backward")
+def _grad_max_pool2d_backward(bb: BlockBuilder, call: Call) -> Expr:
+    if not (len(call.attrs.dilation) == 2 and all(i == 1 for i in call.attrs.dilation)):
+        logging.info("Dilation is not supported in TOPI pool_grad and is not legalized.")
+        return call
+    return bb.call_te(
+        topi.nn.pool_grad,
+        call.args[0],
+        call.args[1],
+        kernel=call.attrs.pool_size,
+        stride=call.attrs.strides,
+        padding=call.attrs.padding,
+        pool_type="max",
+        ceil_mode=call.attrs.ceil_mode,
+        count_include_pad=call.attrs.count_include_pad,
+        layout=call.attrs.layout,
+        primfunc_name_hint="max_pool2d_backward",
+    )
+
+
+@register_legalize("relax.grad.avg_pool2d_backward")
+def _grad_avg_pool2d_backward(bb: BlockBuilder, call: Call) -> Expr:
+    if not (len(call.attrs.dilation) == 2 and all(i == 1 for i in call.attrs.dilation)):
+        logging.info("Dilation is not supported in TOPI pool_grad and is not legalized.")
+        return call
+    return bb.call_te(
+        topi.nn.pool_grad,
+        call.args[0],
+        call.args[1],
+        kernel=call.attrs.pool_size,
+        stride=call.attrs.strides,
+        padding=call.attrs.padding,
+        pool_type="avg",
+        ceil_mode=call.attrs.ceil_mode,
+        count_include_pad=call.attrs.count_include_pad,
+        layout=call.attrs.layout,
+        primfunc_name_hint="avg_pool2d_backward",
+    )
+
+
+@register_legalize("relax.grad.take_backward")
+def _grad_take_backward(bb: BlockBuilder, call: Call) -> Expr:
+    axis = call.attrs.axis
+    if axis is not None:
+        axis = int(axis)
+
+    def te_take_backward(output_grad, x, indices):
+        def gen_ir(output_grad_ptr, x_ptr, indices_ptr, out_ptr):
+            # pylint: disable=invalid-name
+            # Use buffer_proxy for flat indexing on multi-dimensional buffers
+            out = buffer_proxy(out_ptr)
+            grad = buffer_proxy(output_grad_ptr)
+            idx = buffer_proxy(indices_ptr)
+
+            fused_shape = 1
+            for i in x_ptr.shape:
+                fused_shape *= i
+
+            assert len(indices_ptr.shape) == 1  # indices in take must be 1-dim Tensor
+            indices_len = indices_ptr.shape[0]
+
+            with IRBuilder() as ib:
+                with T.seq_scope():
+                    # Init loop (zero-fill output buffer)
+                    with T.serial(fused_shape) as i:
+                        out[i] = tirx.const(0, dtype=x_ptr.dtype)
+
+                    # Accumulation loop
+                    if axis is not None:
+                        fused_output_grad_shape_pre = 1
+                        fused_output_grad_shape_nxt = 1
+                        for i in range(len(output_grad_ptr.shape)):
+                            if i < axis:
+                                fused_output_grad_shape_pre *= output_grad_ptr.shape[i]
+                            elif i > axis:
+                                fused_output_grad_shape_nxt *= output_grad_ptr.shape[i]
+
+                        x_axis_len = x_ptr.shape[axis]
+
+                        with T.serial(
+                            fused_output_grad_shape_pre * fused_output_grad_shape_nxt
+                        ) as fused:
+                            i = fused // fused_output_grad_shape_nxt
+                            j = fused % fused_output_grad_shape_nxt
+                            with T.serial(indices_len) as loop_l:
+                                out_idx = (
+                                    i * fused_output_grad_shape_nxt * x_axis_len
+                                    + idx[loop_l] * fused_output_grad_shape_nxt
+                                    + j
+                                )
+                                grad_idx = (
+                                    i * fused_output_grad_shape_nxt * indices_len
+                                    + loop_l * fused_output_grad_shape_nxt
+                                    + j
+                                )
+                                out[out_idx] = out[out_idx] + grad[grad_idx]
+                    else:
+                        with T.serial(indices_len) as loop_l:
+                            out[idx[loop_l]] = out[idx[loop_l]] + grad[loop_l]
+
+                return ib.get()
+
+        shape = x.shape
+        out_buf = tirx.decl_buffer(shape, x.dtype, "out_buf", layout=None)
+
+        return te.extern(
+            [shape],
+            [output_grad, x, indices],
+            lambda ins, outs: gen_ir(ins[0], ins[1], ins[2], outs[0]),
+            dtype=x.dtype,
+            out_buffers=[out_buf],
+            name="take_backward",
+            tag="take_backward",
+        )
+
+    return bb.call_te(
+        te_take_backward,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        primfunc_name_hint="take_backward",
+    )

python/tvm/relax/transform/legalize_ops/image.py

@@ -0,0 +1,88 @@
+# 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 tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import 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,
+    )
+
+
+@register_legalize("relax.image.grid_sample")
+def _image_grid_sample(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.image.grid_sample,
+        call.args[0],
+        call.args[1],
+        method=call.attrs.method,
+        layout=call.attrs.layout,
+        padding_mode=call.attrs.padding_mode,
+        align_corners=call.attrs.align_corners,
+    )
+
+
+@register_legalize("relax.image.affine_grid")
+def _image_affine_grid(bb: BlockBuilder, call: Call) -> Expr:
+    for v in call.args[1].values:
+        if not isinstance(v, int | tirx.IntImm):
+            raise ValueError(
+                f"affine_grid legalization requires static target_shape, got symbolic value: {v}"
+            )
+    target_shape = [int(v) for v in call.args[1].values]
+    return bb.call_te(
+        topi.image.affine_grid,
+        call.args[0],
+        target_shape=target_shape,
+        align_corners=call.attrs.align_corners,
+    )
+
+
+@register_legalize("relax.image.resize3d")
+def _image_resize3d(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.image.resize3d,
+        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,
+    )

python/tvm/relax/transform/legalize_ops/index.py

@@ -0,0 +1,137 @@
+# 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 tvm
+from tvm import te, tirx, topi
+from tvm.ir import Call, PrimType
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, Tuple
+from ...op import tensor_to_shape
+from ...type import ShapeType
+from .common import register_legalize
+
+
+@register_legalize("relax.take")
+def _take(bb: BlockBuilder, call: Call) -> Expr:
+    # Currently "fast" is the default mode, which leads to segmentation faults
+    # when there are out-of-bounds indices.
+    return bb.call_te(topi.take, call.args[0], call.args[1], call.attrs.axis, mode=call.attrs.mode)
+
+
+@register_legalize("relax.strided_slice")
+def _strided_slice(bb: BlockBuilder, call: Call) -> Expr:
+    def _relax_tuple_to_tir(relax_tuple):
+        if isinstance(relax_tuple, Tuple):
+            output = []
+            for field in relax_tuple.fields:
+                assert tvm.ir.is_prim_expr(field)
+                output.append(field)
+            return output
+
+        output = []
+        for field in relax_tuple.ty.fields:
+            assert isinstance(field, PrimType)
+            return None
+        return output
+
+    if len(call.args) == 4:
+        data, axes, begin, end = call.args
+        strides = [tirx.IntImm("int64", 1)] * len(axes.ty.fields)
+    elif len(call.args) == 5:
+        data, axes, begin, end, strides = call.args
+        strides = _relax_tuple_to_tir(strides)
+    else:
+        raise ValueError(
+            f"Expression {call} provides {len(call.args)} arguments, "
+            f"but {call.op} requires either 4 or 5 arguments."
+        )
+
+    axes = _relax_tuple_to_tir(axes)
+    begin = _relax_tuple_to_tir(begin)
+    end = _relax_tuple_to_tir(end)
+
+    return bb.call_te(
+        topi.strided_slice,
+        data,
+        begin,
+        end,
+        strides,
+        axes,
+        slice_mode="end",
+        assume_inbound=call.attrs.assume_inbound,
+    )
+
+
+@register_legalize("relax.dynamic_strided_slice")
+def _dynamic_strided_slice(bb: BlockBuilder, call: Call) -> Expr:
+    assert len(call.args) == 4
+    data, begin, end, strides = call.args
+
+    # 1. Insert shape function
+    def shape_func(data, begin, end, strides):
+        def _compute(i):
+            def canonicalize_index(index, extent, strides):
+                begin_range = tirx.Select(
+                    strides < 0, tirx.const(-1, "int64"), tirx.const(0, "int64")
+                )
+                end_range = tirx.Select(strides < 0, extent - 1, extent)
+                index = tirx.Select(index < 0, index + extent, index)
+                return tirx.Min(tirx.Max(index, begin_range), end_range)
+
+            def get_length(begin, end, strides, length):
+                begin = canonicalize_index(begin, length, strides)
+                end = canonicalize_index(end, length, strides)
+                len1 = tirx.ceildiv(begin - end, -strides)
+                len2 = tirx.ceildiv(end - begin, strides)
+                return tirx.Select(strides < 0, len1, len2)
+
+            length = tirx.const(-1, "int64")
+            for idx in range(data.ndim):
+                length = tirx.Select(i == tirx.const(idx, "int64"), data.shape[idx], length)
+
+            return get_length(begin[i], end[i], strides[i], length)
+
+        return te.compute((begin.shape[0],), _compute, name="T_shape_func_strided_slice_dynamic")
+
+    output_shape = bb.normalize(
+        bb.call_te(
+            shape_func,
+            data,
+            begin,
+            end,
+            strides,
+        )
+    )
+
+    # 2. Convert tensor to shape and match cast with new symbolic vars
+    ndim = int(output_shape.ty.shape[0])
+    output_shape = bb.emit(tensor_to_shape(output_shape))
+    output_shape_vars = [tirx.Var("s", "int64") for i in range(ndim)]
+    bb.match_cast(output_shape, ShapeType(output_shape_vars))
+
+    # 3. Pass the output shape vars to TOPI
+    return bb.call_te(
+        topi.dynamic_strided_slice,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        call.args[3],
+        output_shape=output_shape_vars,
+    )

python/tvm/relax/transform/legalize_ops/inspect_op.py

@@ -0,0 +1,137 @@
+# 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
+"""Legalization functions for DLTensor inspection."""
+
+import enum
+
+from tvm.ir import Call
+from tvm.script import tirx as T
+
+from ... import op
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import register_legalize
+
+
+class TVMStructFieldKind(enum.IntEnum):
+    """Equivalent to tvm::tirx::builtin::TVMStructFieldKind
+
+    This does not use `enum.auto()` to define the values, because
+    `enum.auto()` starts from 1, and this must match the C++
+    definition which starts from 0.
+    """
+
+    kDLTensorAddr = 0
+    kDLTensorData = 1
+    kDLTensorShape = 2
+    kDLTensorStrides = 3
+    kDLTensorNDim = 4
+    kDLTensorTypeCode = 5
+    kDLTensorTypeBits = 6
+    kDLTensorTypeLanes = 7
+    kDLTensorByteOffset = 8
+    kDLTensorDeviceId = 9
+    kDLTensorDeviceType = 10
+    kDLTensorKindBound_ = 11
+    kTVMValueContent = 12
+    kTVMValueKindBound_ = 13
+
+
+@register_legalize("relax.inspect.tensor_stride_i")
+def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
+    @T.prim_func(private=True, s_tir=True)
+    def _get_tensor_stride_i(dlpack_handle: T.handle, axis: T.int64) -> T.int64:
+        T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
+        assert T.int64(0) <= axis, "Specified axis may not be negative"
+        ndim: T.let[T.int32] = T.tvm_struct_get(
+            dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorNDim), "int32"
+        )
+        assert axis < T.Cast("int64", ndim), (
+            "Specified axis may not be larger than the tensor's dimensionality"
+        )
+        stride_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
+            dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorStrides), T.handle("int64").ty
+        )
+
+        if T.isnullptr(stride_ptr):
+            shape_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
+                dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), T.handle("int64").ty
+            )
+            shape = T.decl_buffer(ndim, "int64", data=shape_ptr)
+
+            product = T.decl_buffer([], "int64")
+            product[()] = 1
+
+            # TODO(Lunderberg): Add a TIR lowering pass to allow
+            # ranges to start somewhere other than zero.  This loop
+            # could then iterate on `range(axis+1, ndim)`.
+            for dim_offset in range(ndim - (axis + 1)):
+                dim: T.let[T.int64] = dim_offset + (axis + 1)
+                product[()] = product[()] * shape[dim]
+
+            return product[()]
+        else:
+            strides = T.decl_buffer(ndim, "int64", data=stride_ptr)
+            stride: T.let[T.int64] = strides[axis]
+            return stride
+
+    gvar = bb.add_func(_get_tensor_stride_i, "_get_tensor_stride_i")
+    return Call(gvar, call.args)
+
+
+@register_legalize("relax.inspect.tensor_byte_offset")
+def _tensor_byte_offset(bb: BlockBuilder, call: Call) -> Expr:
+    @T.prim_func(private=True, s_tir=True)
+    def _get_tensor_byte_offset(dlpack_handle: T.handle) -> T.int64:
+        T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
+        byte_offset: T.let[T.uint64] = T.tvm_struct_get(
+            dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
+        )
+        return byte_offset
+
+    gvar = bb.add_func(_get_tensor_byte_offset, "_get_tensor_byte_offset")
+    return Call(gvar, call.args)
+
+
+@register_legalize("relax.inspect.tensor_elem_offset")
+def _tensor_elem_offset(bb: BlockBuilder, call: Call) -> Expr:
+    @T.prim_func(private=True, s_tir=True)
+    def _get_tensor_elem_offset(dlpack_handle: T.handle) -> T.int64:
+        T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
+        byte_offset: T.let[T.uint64] = T.tvm_struct_get(
+            dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
+        )
+        scalar_bits: T.let[T.uint8] = T.tvm_struct_get(
+            dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeBits), "uint8"
+        )
+        lanes: T.let[T.uint16] = T.tvm_struct_get(
+            dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeLanes), "uint16"
+        )
+        bytes_per_element: T.let[T.uint64] = T.ceildiv(
+            scalar_bits.astype("uint64") * lanes.astype("uint64"), 8
+        )
+        elem_offset: T.let[T.uint64] = byte_offset // bytes_per_element
+        return elem_offset
+
+    gvar = bb.add_func(_get_tensor_elem_offset, "_get_tensor_elem_offset")
+    return Call(gvar, call.args)
+
+
+@register_legalize("relax.size")
+def _size(_bb: BlockBuilder, call: Call) -> Expr:
+    return op.prod(op.shape_to_tensor(op.shape_of(call.args[0])))

python/tvm/relax/transform/legalize_ops/linear_algebra.py

@@ -0,0 +1,158 @@
+# 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 DataTypeCode, relax, te, tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, Tuple, TupleGetItem, Var
+from .common import register_legalize
+
+
+@register_legalize("relax.matmul")
+def _matmul(bb: BlockBuilder, call: Call) -> Expr:
+    def is_known_tensor_dtype(dtype) -> bool:
+        raw_dtype = dtype.dtype
+        return not (
+            raw_dtype.type_code == int(DataTypeCode.HANDLE)
+            and raw_dtype.bits == 0
+            and raw_dtype.lanes == 0
+        )
+
+    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.TensorType(a.shape))
+        b_relax = relax.Var("b", relax.TensorType(b.shape))
+        f_infer_ty = call.op.get_attr("FInferType")
+        output_shape = f_infer_ty(relax.op.matmul(a_relax, b_relax), bb).shape
+        if isinstance(a_shape[-1], tirx.IntImm) and a_shape[-1] == 0:
+            return te.compute(
+                output_shape,
+                lambda *_: tirx.const(0, call.ty.dtype),
+                name="matmul",
+            )
+
+        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]
+                    dim_equal = a_dim == b_dim
+                    if not isinstance(dim_equal, tirx.IntImm) or dim_equal == 0:
+                        a_dim_is_one = isinstance(a_dim, tirx.IntImm) and a_dim == 1
+                        b_dim_is_one = isinstance(b_dim, tirx.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])
+                    else:
+                        a_indices.append(idx_spatial[i])
+                        b_indices.append(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 is not None and dtype != "":
+                    return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
+                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",
+        )
+
+    lhs, rhs = call.args
+    lhs_ty = call.args[0].ty
+    rhs_ty = call.args[1].ty
+    assert (
+        lhs_ty.dtype
+        and rhs_ty.dtype
+        and is_known_tensor_dtype(lhs_ty.dtype)
+        and is_known_tensor_dtype(rhs_ty.dtype)
+    ), (
+        f"To legalize R.matmul into R.call_tir, the dtype of both operands must be known.  "
+        f"However, the LHS {lhs} has type {lhs_ty} (dtype='{lhs_ty.dtype}') "
+        f"and the RHS {rhs} has type {rhs_ty} (dtype='{rhs_ty.dtype}')."
+    )
+    return bb.call_te(te_matmul, call.args[0], call.args[1], primfunc_name_hint="matmul")
+
+
+@register_legalize("relax.einsum")
+def _einsum(bb: BlockBuilder, call: Call) -> Expr:
+    t = call.args[0]
+    n_field = len(t.ty.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.einsum, call.attrs.subscripts, *fields)
+
+
+@register_legalize("relax.outer")
+def _outer(bb: BlockBuilder, call: Call) -> Expr:
+    def te_outer(a: te.Tensor, b: te.Tensor) -> te.Tensor:
+        a_shape = list(a.shape)
+        b_shape = list(b.shape)
+        assert len(a_shape) == 1 and len(b_shape) == 1, "outer requires 1D tensors"
+
+        n = a_shape[0]
+        m = b_shape[0]
+
+        def compute_fn(i, j):
+            return a[i] * b[j]
+
+        return te.compute((n, m), compute_fn, name="outer")
+
+    lhs, rhs = call.args
+    return bb.call_te(te_outer, lhs, rhs, primfunc_name_hint="outer")

python/tvm/relax/transform/legalize_ops/manipulate.py

@@ -0,0 +1,369 @@
+# 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
+# ruff: noqa: RUF005
+"""Default legalization function for manipulate operators."""
+
+import tvm
+from tvm import DataTypeCode, relax, s_tir, te, tirx, topi
+from tvm.ir import Call
+from tvm.relax.op.base import call_tir
+from tvm.relax.type import TensorType
+from tvm.relax.utils import gen_call_tir_inputs
+from tvm.tirx.expr import IntImm
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, ShapeExpr, Tuple, TupleGetItem, Var
+from .common import LegalizeFunc, TEFunc, 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].ty.shape if is_collapse_sum_like else call.args[1]
+        # If target shape is Var, pass its bound expr only when it is ShapeExpr
+        if isinstance(tgt_shape, Var):
+            tgt_shape = bb.lookup_binding(tgt_shape)
+            assert isinstance(tgt_shape, ShapeExpr)
+        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.collapse_sum_like",
+    _reshape(topi.collapse_sum, "collapse_sum", is_collapse_sum_like=True),
+)
+
+register_legalize("relax.collapse_sum_to", _reshape(topi.collapse_sum, "collapse_sum"))
+
+
+@register_legalize("relax.concat")
+def _concat(bb: BlockBuilder, call: Call) -> Expr:
+    t = call.args[0]
+    n_field = len(t.ty.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
+    )
+
+
+@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.TensorType(data.shape))
+        f_infer_ty = call.op.get_attr("FInferType")
+        output_shape = f_infer_ty(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.ty.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, tirx.IntImm):
+        indices_or_sections = call.attrs.indices_or_sections.value
+    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)
+
+
+@register_legalize("relax.stack")
+def _stack(bb: BlockBuilder, call: Call) -> Expr:
+    t = call.args[0]
+    n_field = len(t.ty.fields)
+
+    # Follow bindings to find the actual tuple
+    while isinstance(t, Var):
+        binding = bb.lookup_binding(t)
+        if not isinstance(binding, Tuple | Var):
+            break
+        t = binding
+
+    assert isinstance(t, Tuple | Var)
+
+    # Extract fields from either Tuple or bound 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.stack, fields, 0 if call.attrs.axis is None else call.attrs.axis)
+
+
+@register_legalize("relax.repeat")
+def _repeat(bb: BlockBuilder, call: Call) -> Expr:
+    def te_repeat(data: te.Tensor, repeats: IntImm, axis: IntImm | None):
+        if axis is None:
+            # flatten data
+            out_shape = data.shape[0]
+            for i in data.shape[1:]:
+                out_shape *= i
+            data = topi.reshape(data, (out_shape,))
+            axis = 0
+        # topi only receives int repeats and axis
+        return topi.repeat(data, int(repeats), int(axis))
+
+    return bb.call_te(
+        te_repeat, call.args[0], call.attrs.repeats, call.attrs.axis, primfunc_name_hint="repeat"
+    )
+
+
+@register_legalize("relax.tile")
+def _tile(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(topi.tile, call.args[0], call.attrs.repeats)
+
+
+@register_legalize("relax.flip")
+def _flip(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(topi.flip, call.args[0], int(call.attrs.axis))
+
+
+@register_legalize("relax.reverse_sequence")
+def _reverse_sequence(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.reverse_sequence,
+        call.args[0],
+        call.args[1],
+        int(call.attrs.seq_axis),
+        int(call.attrs.batch_axis),
+        primfunc_name_hint="reverse_sequence",
+    )
+
+
+@register_legalize("relax.gather_elements")
+def _gather_elements(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(topi.gather, call.args[0], int(call.attrs.axis), call.args[1])
+
+
+@register_legalize("relax.gather_nd")
+def _gather_nd(bb: BlockBuilder, call: Call) -> Expr:
+    def te_gather_nd(data, indices, batch_dims):
+        indices_ndim = len(indices.shape)
+        axes = [indices_ndim - 1] + list(range(indices_ndim - 1))
+        indices = topi.transpose(indices, axes)
+        return topi.gather_nd(data, indices, batch_dims)
+
+    return bb.call_te(te_gather_nd, call.args[0], call.args[1], int(call.attrs.batch_dims))
+
+
+@register_legalize("relax.index_tensor")
+def _index_tensor(bb: BlockBuilder, call: Call) -> Expr:
+    t = call.args[1]
+    n_field = len(t.ty.fields)
+    fields = [bb.emit(TupleGetItem(t, i)) for i in range(n_field)]
+    return bb.call_te(topi.index_tensor, call.args[0], fields)
+
+
+@register_legalize("relax.index_put")
+def _index_put(bb: BlockBuilder, call: Call) -> Expr:
+    data = call.args[0]
+    indices = call.args[1]
+    values = call.args[2]
+    accumulate = call.attrs.accumulate
+
+    # If indices is a Tuple, unpack it into individual tensors
+    if isinstance(indices, relax.Tuple):
+        indices_list = [indices.fields[i] for i in range(len(indices.fields))]
+    else:
+        indices_list = [indices]
+
+    return bb.call_te(
+        topi.index_put,
+        data,
+        indices_list,
+        values,
+        accumulate=accumulate,
+    )
+
+
+@register_legalize("relax.meshgrid")
+def _meshgrid(bb: BlockBuilder, call: Call) -> Expr:
+    t = call.args[0]
+    n_field = len(t.ty.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.meshgrid, fields, "ij" if call.attrs.indexing is None else call.attrs.indexing
+    )
+
+
+def _is_gpu_target():
+    target = tvm.target.Target.current(allow_none=True)
+    return target is not None and "gpu" in target.keys
+
+
+@register_legalize("relax.scatter_elements")
+def _scatter_elements(bb: BlockBuilder, call: Call) -> Expr:
+    te_func = topi.gpu.scatter_elements if _is_gpu_target() else topi.scatter_elements
+    return bb.call_te(
+        te_func,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        call.attrs.axis,
+        call.attrs.reduction,
+    )
+
+
+@register_legalize("relax.scatter_nd")
+def _scatter_nd(bb: BlockBuilder, call: Call) -> Expr:
+    # TODO(relax-team): Support native scatter_nd without te extern
+    base_te = topi.gpu.scatter_nd if _is_gpu_target() else topi.scatter_nd
+
+    def scatter_nd(data, indices, updates, reduction):
+        axes = list(range(len(indices.shape)))
+        indices = topi.transpose(indices, axes[-1:] + axes[:-1])
+        return base_te(data, indices, updates, reduction)
+
+    return bb.call_te(
+        scatter_nd,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        call.attrs.reduction,
+    )
+
+
+@register_legalize("relax.slice_scatter")
+def _slice_scatter(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.slice_scatter,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        call.args[3],
+        call.args[4],
+        call.attrs.axis,
+    )
+
+
+@register_legalize("relax.one_hot")
+def _one_hot(bb: BlockBuilder, call: Call) -> Expr:
+    indices, on_value, off_value = call.args
+    if not (tvm.ir.is_prim_expr(on_value) and tvm.ir.is_prim_expr(off_value)):
+        raise ValueError("on_value and off_value must be Expr")
+    if on_value.ty != off_value.ty:
+        raise ValueError("on_value and off_value must have the same dtype")
+    return bb.call_te(
+        topi.one_hot,
+        indices,
+        on_value,
+        off_value,
+        call.attrs.depth,
+        call.attrs.axis,
+        on_value.ty,
+    )
+
+
+@register_legalize("relax.layout_transform")
+def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
+    def te_layout_transform(data, name):
+        """
+        Returns a passthrough TE compute with appropriate name. This is needed to generate
+        TIR function, output shape info, TIR vars from gen_call_tir_inputs function.
+        """
+        return te.compute(
+            data.shape,
+            data,
+            name=name,
+        )
+
+    def set_axis_sep(axis_sep: list, sch: s_tir.schedule, buffer_type: str):
+        sch.set_axis_separator(primfunc_name, (buffer_type, 0), axis_separators=axis_sep)
+
+    index_map: tvm.tirx.IndexMap = call.attrs.index_map
+    pad_value = call.attrs.pad_value
+    if pad_value is not None:
+        pad_value = pad_value.value
+    else:
+        if call.args[0].ty.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT):
+            pad_value = 0
+        else:
+            pad_value = 0.0
+
+    axis_separators: tvm.tirx.IndexMap.AXIS_SEPARATOR = call.attrs.axis_separators
+    input_axis_separators: tvm.tirx.IndexMap.AXIS_SEPARATOR = call.attrs.input_axis_separators
+
+    # Convert to list from array
+    axis_separators = [int(sep) for sep in axis_separators]
+    primfunc_name = "te_layout_transform"
+    _, padding_predicate = index_map.non_surjective_inverse(call.args[0].ty.shape)
+    if not isinstance(padding_predicate, tvm.tirx.expr.IntImm):
+        primfunc_name += "_with_pad"
+    if len(axis_separators) != 0:
+        primfunc_name += "_axis_separator"
+    tir_func, call_args, _, tir_vars = gen_call_tir_inputs(
+        te_layout_transform, call.args[0], primfunc_name
+    )
+    # Create TIR schedule to apply layout changes with axis separators
+    sch = tvm.s_tir.Schedule(tir_func)
+    sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value)
+    set_axis_sep(axis_separators, sch, "write")
+    if input_axis_separators is not None:
+        set_axis_sep(input_axis_separators, sch, "read")
+    gvar = bb.add_func(sch.mod["main"], primfunc_name)
+    output_shape = index_map.map_shape(list(call_args[0].ty.shape))
+    output_dtype = call_args[0].ty.dtype
+    output_ty = [TensorType(output_shape, output_dtype)]
+    return call_tir(gvar, call_args, output_ty, tir_vars)

python/tvm/relax/transform/legalize_ops/nn.py

@@ -0,0 +1,826 @@
+# 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
+import math
+
+from tvm import s_tir, te, tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import _call_topi_without_attr, register_legalize
+
+
+@register_legalize("relax.nn.conv1d")
+def _nn_conv1d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.data_layout:
+        logging.info(
+            "TOPI conv1d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+    if len(call.attrs.data_layout) != 3 or len(call.attrs.kernel_layout) != 3:
+        logging.info(
+            "Conv1D 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 = s_tir.slayout(call.attrs.data_layout)
+        kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
+        ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
+        oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
+        if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
+            logging.info(
+                "Conv1D 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.conv1d,
+        data=call.args[0],
+        kernel=call.args[1],
+        strides=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="conv1d",
+    )
+
+
+@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 = s_tir.slayout(call.attrs.data_layout)
+        kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
+        ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
+        oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
+        if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.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.conv3d")
+def _nn_conv3d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.data_layout:
+        logging.info(
+            "TOPI conv3d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+    if len(call.attrs.data_layout) != 5 or len(call.attrs.kernel_layout) != 5:
+        logging.info(
+            "Conv3D 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 = s_tir.slayout(call.attrs.data_layout)
+        kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
+        ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
+        oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
+        if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
+            logging.info(
+                "Conv3D 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="conv3d",
+    )
+
+
+@register_legalize("relax.nn.conv1d_transpose")
+def _nn_conv1d_transpose(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.data_layout:
+        logging.info(
+            "TOPI conv1d_transpose does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+    if call.attrs.data_layout != "NCW" or call.attrs.kernel_layout != "IOW":
+        logging.info(
+            "TOPI conv1d_transpose does not support input layout other than NCW, "
+            "and kernel layout other than IOW, so cannot be legalized by TOPI"
+        )
+        return call
+    strides = [int(s) for s in call.attrs.strides]
+    padding = [int(p) for p in call.attrs.padding]
+    output_padding = [int(o) for o in call.attrs.output_padding]
+    groups = int(call.attrs.groups)
+    out_dtype = call.ty.dtype
+    dilation = [int(d) for d in call.attrs.dilation]
+
+    def te_conv1d_transpose(data, kernel):
+        # Dilated transposed conv == transposed conv with a spatially dilated (zero-filled) kernel.
+        if any(d != 1 for d in dilation):
+            kernel = topi.nn.dilate(kernel, [1, 1, dilation[0]], name="kernel_dilate")
+        return topi.nn.group_conv1d_transpose_ncw(
+            data, kernel, strides, padding, out_dtype, output_padding, groups
+        )
+
+    return bb.call_te(
+        te_conv1d_transpose, call.args[0], call.args[1], primfunc_name_hint="conv1d_transpose"
+    )
+
+
+@register_legalize("relax.nn.conv2d_transpose")
+def _nn_conv2d_transpose(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.data_layout:
+        logging.info(
+            "TOPI conv2d_transpose does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+    if call.attrs.data_layout != "NCHW" or call.attrs.kernel_layout != "IOHW":
+        logging.info(
+            "TOPI conv2d_transpose does not support input layout other than NCHW, "
+            "and kernel layout other than IOHW, so cannot be legalized by TOPI"
+        )
+        return call
+    strides = [int(s) for s in call.attrs.strides]
+    padding = [int(p) for p in call.attrs.padding]
+    output_padding = [int(o) for o in call.attrs.output_padding]
+    groups = int(call.attrs.groups)
+    out_dtype = call.ty.dtype
+    dilation = [int(d) for d in call.attrs.dilation]
+
+    def te_conv2d_transpose(data, kernel):
+        # Dilated transposed conv == transposed conv with a spatially dilated (zero-filled) kernel.
+        if any(d != 1 for d in dilation):
+            kernel = topi.nn.dilate(kernel, [1, 1, dilation[0], dilation[1]], name="kernel_dilate")
+        return topi.nn.group_conv2d_transpose_nchw(
+            data, kernel, strides, padding, out_dtype, output_padding, groups
+        )
+
+    return bb.call_te(
+        te_conv2d_transpose, call.args[0], call.args[1], primfunc_name_hint="conv2d_transpose"
+    )
+
+
+@register_legalize("relax.nn.conv3d_transpose")
+def _nn_conv3d_transpose(bb: BlockBuilder, call: Call) -> Expr:
+    # Keep policy in sync with _nn_conv2d_transpose: only lower when TOPI supports
+    # the layout/dilation.
+    if call.attrs.out_layout != call.attrs.data_layout:
+        logging.info(
+            "TOPI conv3d_transpose does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+    if call.attrs.data_layout != "NCDHW" or call.attrs.kernel_layout != "IODHW":
+        logging.info(
+            "TOPI conv3d_transpose does not support input layout other than NCDHW, "
+            "and kernel layout other than IODHW, so cannot be legalized by TOPI"
+        )
+        return call
+    strides = [int(s) for s in call.attrs.strides]
+    padding = [int(p) for p in call.attrs.padding]
+    output_padding = [int(o) for o in call.attrs.output_padding]
+    groups = int(call.attrs.groups)
+    out_dtype = call.ty.dtype
+    dilation = [int(d) for d in call.attrs.dilation]
+
+    def te_conv3d_transpose(data, kernel):
+        # Dilated transposed conv == transposed conv with a spatially dilated (zero-filled) kernel.
+        if any(d != 1 for d in dilation):
+            kernel = topi.nn.dilate(
+                kernel, [1, 1, dilation[0], dilation[1], dilation[2]], name="kernel_dilate"
+            )
+        return topi.nn.group_conv3d_transpose_ncdhw(
+            data, kernel, strides, padding, out_dtype, output_padding, groups
+        )
+
+    return bb.call_te(
+        te_conv3d_transpose, call.args[0], call.args[1], primfunc_name_hint="conv3d_transpose"
+    )
+
+
+@register_legalize("relax.nn.pad")
+def _nn_pad(bb: BlockBuilder, call: Call) -> Expr:
+    pad_mode = call.attrs.pad_mode
+    pad_widths = call.attrs.pad_width
+    pad_before = pad_widths[::2]
+    pad_after = pad_widths[1::2]
+    if pad_mode == "reflect":
+        return bb.call_te(
+            topi.nn.reflect_pad, call.args[0], pad_before=pad_before, pad_after=pad_after
+        )
+    elif pad_mode == "replicate":
+        return bb.call_te(
+            topi.nn.replicate_pad, call.args[0], pad_before=pad_before, pad_after=pad_after
+        )
+    elif pad_mode == "circular":
+        return bb.call_te(
+            topi.nn.circular_pad, call.args[0], pad_before=pad_before, pad_after=pad_after
+        )
+    else:
+        return bb.call_te(
+            topi.nn.pad,
+            call.args[0],
+            pad_before=pad_before,
+            pad_after=pad_after,
+            pad_value=call.attrs.pad_value,
+            primfunc_name_hint="pad",
+        )
+
+
+@register_legalize("relax.nn.pixel_shuffle")
+def _nn_pixel_shuffle(bb: BlockBuilder, call: Call) -> Expr:
+    upscale_factor = call.attrs.upscale_factor
+    return bb.call_te(topi.nn.pixel_shuffle, call.args[0], upscale_factor=upscale_factor)
+
+
+@register_legalize("relax.nn.max_pool1d")
+def _nn_max_pool1d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI max_pool1d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    return bb.call_te(
+        topi.nn.pool1d,
+        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_pool1d",
+    )
+
+
+@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.max_pool3d")
+def _nn_max_pool3d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI max_pool3d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    return bb.call_te(
+        topi.nn.pool3d,
+        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_pool3d",
+    )
+
+
+@register_legalize("relax.nn.avg_pool1d")
+def _nn_avg_pool1d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI avg_pool1d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    return bb.call_te(
+        topi.nn.pool1d,
+        call.args[0],
+        kernel=call.attrs.pool_size,
+        stride=call.attrs.strides,
+        dilation=call.attrs.dilation,
+        padding=call.attrs.padding,
+        pool_type="avg",
+        ceil_mode=call.attrs.ceil_mode,
+        layout=call.attrs.layout,
+        count_include_pad=call.attrs.count_include_pad,
+        primfunc_name_hint="avg_pool1d",
+    )
+
+
+@register_legalize("relax.nn.avg_pool2d")
+def _nn_avg_pool2d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI avg_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="avg",
+        ceil_mode=call.attrs.ceil_mode,
+        layout=call.attrs.layout,
+        count_include_pad=call.attrs.count_include_pad,
+        primfunc_name_hint="avg_pool2d",
+    )
+
+
+@register_legalize("relax.nn.avg_pool3d")
+def _nn_avg_pool3d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI avg_pool3d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    return bb.call_te(
+        topi.nn.pool3d,
+        call.args[0],
+        kernel=call.attrs.pool_size,
+        stride=call.attrs.strides,
+        dilation=call.attrs.dilation,
+        padding=call.attrs.padding,
+        pool_type="avg",
+        ceil_mode=call.attrs.ceil_mode,
+        layout=call.attrs.layout,
+        count_include_pad=call.attrs.count_include_pad,
+        primfunc_name_hint="avg_pool3d",
+    )
+
+
+@register_legalize("relax.nn.adaptive_avg_pool1d")
+def _nn_adaptive_avg_pool1d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI adaptive_avg_pool1d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    def te_adaptive_avg_pool1d(data, output_size, layout_str):
+        if output_size is None:
+            layout = s_tir.slayout(layout_str)
+            idx_W = layout.index_of("W")
+            assert idx_W != -1
+            output_size = data.shape[idx_W]
+
+        return topi.nn.adaptive_pool1d(data, output_size, "avg", layout_str)
+
+    return bb.call_te(
+        te_adaptive_avg_pool1d,
+        call.args[0],
+        call.attrs.output_size,
+        call.attrs.layout,
+        primfunc_name_hint="adaptive_avg_pool1d",
+    )
+
+
+@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 = s_tir.slayout(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.adaptive_avg_pool3d")
+def _nn_adaptive_avg_pool3d(bb: BlockBuilder, call: Call) -> Expr:
+    if call.attrs.out_layout != call.attrs.layout:
+        logging.info(
+            "TOPI adaptive_avg_pool3d does not support different input-output "
+            "layouts, and thus cannot be legalized by TOPI"
+        )
+        return call
+
+    def te_adaptive_avg_pool3d(data, output_size, layout_str):
+        if output_size is None:
+            layout = s_tir.slayout(layout_str)
+            idx_D = layout.index_of("D")
+            idx_H = layout.index_of("H")
+            idx_W = layout.index_of("W")
+            assert idx_D != -1 and idx_H != -1 and idx_W != -1
+            output_size = (data.shape[idx_D], data.shape[idx_H], data.shape[idx_W])
+
+        return topi.nn.adaptive_pool3d(data, output_size, "avg", layout_str)
+
+    return bb.call_te(
+        te_adaptive_avg_pool3d,
+        call.args[0],
+        call.attrs.output_size,
+        call.attrs.layout,
+        primfunc_name_hint="adaptive_avg_pool3d",
+    )
+
+
+register_legalize("relax.nn.relu", _call_topi_without_attr(topi.nn.relu))
+
+
+@register_legalize("relax.nn.leakyrelu")
+def _nn_leakyrelu(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(topi.nn.leaky_relu, call.args[0], call.attrs.alpha)
+
+
+@register_legalize("relax.nn.prelu")
+def _nn_prelu(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(topi.nn.prelu, call.args[0], call.args[1], call.attrs.axis)
+
+
+@register_legalize("relax.nn.gelu")
+def _nn_gelu(bb: BlockBuilder, call: Call) -> Expr:
+    def te_gelu(x: te.Tensor):
+        dtype = x.dtype
+        erf_inp = x * tirx.const(0.5**0.5, dtype)
+
+        if dtype == "float16":
+            erf = topi.math.cast(topi.erf(topi.math.cast(erf_inp, "float32")), "float16")
+        else:
+            erf = topi.erf(erf_inp)
+
+        return x * (tirx.const(0.5, dtype) + erf * tirx.const(0.5, dtype))
+
+    return bb.call_te(te_gelu, call.args[0], primfunc_name_hint="gelu")
+
+
+@register_legalize("relax.nn.gelu_tanh")
+def _nn_gelu_tanh(bb: BlockBuilder, call: Call) -> Expr:
+    def te_gelu_tanh(x: te.Tensor):
+        dtype = x.dtype
+        return (
+            tirx.const(0.5, dtype)
+            * x
+            * (
+                tirx.const(1.0, dtype)
+                + topi.tanh(
+                    tirx.const(math.sqrt(2.0 / math.pi), dtype)
+                    * x
+                    * (1 + tirx.const(0.044715, dtype) * x * x)
+                )
+            )
+        )
+
+    return bb.call_te(te_gelu_tanh, call.args[0], primfunc_name_hint="gelu_tanh")
+
+
+@register_legalize("relax.nn.selu")
+def _nn_selu(bb: BlockBuilder, call: Call) -> Expr:
+    def te_selu(x: te.Tensor):
+        dtype = x.dtype
+        alpha = tirx.const(1.6732632423543772848170429916717, dtype)
+        scale = tirx.const(1.0507009873554804934193349852946, dtype)
+
+        # Compute SELU
+        # SELU(x) = scale*(max(0,x)+min(0,a*(exp(x)-1)))
+        positive_part = topi.maximum(x, tirx.const(0, dtype))
+        negative_part = topi.minimum(
+            tirx.const(0, dtype), alpha * (topi.exp(x) - tirx.const(1, dtype))
+        )
+        return scale * (positive_part + negative_part)
+
+    return bb.call_te(te_selu, call.args[0], primfunc_name_hint="selu")
+
+
+@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.softplus")
+def _nn_softplus(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.nn.softplus,
+        call.args[0],
+        call.attrs.beta,
+        call.attrs.threshold,
+    )
+
+
+@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.log_softmax")
+def _nn_log_softmax(bb: BlockBuilder, call: Call):
+    return bb.call_te(topi.nn.log_softmax, call.args[0], call.attrs.axis)
+
+
+@register_legalize("relax.nn.cross_entropy_with_logits")
+def _nn_cross_entropy_with_logits(bb: BlockBuilder, call: Call):
+    def te_cross_entropy_with_logits(x, y):
+        if len(x.shape) > 1:
+            return -topi.sum(x * y) / x.shape[0]
+        return -topi.sum(x * y)
+
+    return bb.call_te(
+        te_cross_entropy_with_logits,
+        call.args[0],
+        call.args[1],
+        primfunc_name_hint="cross_entropy_with_logits",
+    )
+
+
+@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,
+        training=call.attrs.training,
+        momentum=call.attrs.momentum,
+    )
+
+
+@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.group_norm")
+def _nn_group_norm(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.nn.group_norm,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        call.attrs.num_groups,
+        call.attrs.channel_axis,
+        call.attrs.axes,
+        call.attrs.epsilon,
+    )
+
+
+@register_legalize("relax.nn.instance_norm")
+def _nn_instance_norm(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.nn.instance_norm,
+        data=call.args[0],
+        gamma=call.args[1],
+        beta=call.args[2],
+        channel_axis=call.attrs.channel_axis,
+        axis=call.attrs.axes,
+        epsilon=call.attrs.epsilon,
+    )
+
+
+@register_legalize("relax.nn.rms_norm")
+def _nn_rms_norm(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.nn.rms_norm,
+        call.args[0],
+        call.args[1],
+        axis=call.attrs.axes,
+        epsilon=call.attrs.epsilon,
+    )
+
+
+@register_legalize("relax.nn.dropout")
+def _nn_dropout(bb: BlockBuilder, call: Call) -> Expr:
+    # Dropout is a no-op at inference: pass the input through and return an all-ones mask.
+    return bb.call_te(
+        lambda x: [topi.identity(x), topi.full_like(x, 1.0)],
+        call.args[0],
+        primfunc_name_hint="dropout",
+    )
+
+
+def _te_attention(
+    q: te.Tensor,
+    k: te.Tensor,
+    v: te.Tensor,
+    bias: te.Tensor,
+    scale: tirx.FloatImm,
+    causal_mask: str | None,
+) -> te.Tensor:
+    batch_size, seq_len, num_head, head_dim = q.shape
+    _, seq_len_kv, _, head_dim_v = v.shape
+    q = topi.transpose(q, [0, 2, 1, 3])
+    k = topi.transpose(k, [0, 2, 1, 3])
+    v = topi.transpose(v, [0, 2, 1, 3])
+    bs = batch_size * num_head
+    q = topi.reshape(q, [bs, seq_len, head_dim])
+    k = topi.reshape(k, [bs, seq_len_kv, head_dim])
+    v = topi.reshape(v, [bs, seq_len_kv, head_dim_v])
+    p = topi.nn.batch_matmul(q, k, oshape=[bs, seq_len, seq_len_kv])
+    if scale is not None:
+        p = topi.multiply(p, scale)
+    else:
+        p = topi.divide(p, tirx.sqrt(tirx.Cast(p.dtype, head_dim)))
+    if bias is not None:
+        p = topi.reshape(p, [batch_size, num_head, seq_len, seq_len_kv])
+        p = topi.add(p, bias)
+        p = topi.reshape(p, [bs, seq_len, seq_len_kv])
+    if causal_mask is None:
+        s = topi.nn.softmax(p)
+    else:
+        if causal_mask == "TopLeft":
+            offset = tirx.IntImm("int32", 0)
+        elif causal_mask == "BottomRight":
+            offset = tirx.abs(seq_len - seq_len_kv).astype("int32")
+        else:
+            raise NotImplementedError()
+        p_masked = topi.trilu(p, k=offset, upper=False)
+        p_masked_exp = topi.trilu(
+            topi.exp(p_masked - topi.max(p_masked, axis=-1, keepdims=True)), k=offset, upper=False
+        )
+        p_masked_sum = topi.sum(p_masked_exp, axis=-1, keepdims=True)
+        s = topi.divide(p_masked_exp, p_masked_sum)
+    o = topi.nn.batch_matmul(s, v, transpose_b=False, oshape=[bs, seq_len, head_dim_v])
+    o = topi.reshape(o, [batch_size, num_head, seq_len, head_dim_v])
+    return topi.transpose(o, [0, 2, 1, 3])
+
+
+@register_legalize("relax.nn.attention")
+def _nn_attention(bb: BlockBuilder, call: Call) -> Expr:
+    assert call.attrs.window_size is None, (
+        "Legalization for sliding-window attention is not supported yet."
+    )
+    return bb.call_te(
+        _te_attention,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        None,
+        call.attrs.scale,
+        call.attrs.causal_mask,
+        primfunc_name_hint="attention",
+    )
+
+
+@register_legalize("relax.nn.attention_bias")
+def _nn_attention_bias(bb: BlockBuilder, call: Call) -> Expr:
+    assert call.attrs.window_size is None, (
+        "Legalization for sliding-window attention is not supported yet."
+    )
+    return bb.call_te(
+        _te_attention,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        call.args[3],
+        call.attrs.scale,
+        call.attrs.causal_mask,
+        primfunc_name_hint="attention_bias",
+    )
+
+
+@register_legalize("relax.nn.attention_var_len")
+def _nn_attention_var_len(bb: BlockBuilder, call: Call) -> Expr:
+    raise RuntimeError("Legalization of attention_var_len op is not supported yet.")
+
+
+@register_legalize("relax.nn.nll_loss")
+def _nn_nll_loss(bb: BlockBuilder, call: Call) -> Expr:
+    def nll_loss_without_weight(predictions, targets, reduction, ignore_index):
+        weight = topi.full(
+            (predictions.shape[1] if len(predictions.shape) > 1 else predictions.shape[0],),
+            predictions.dtype,
+            1.0,
+        )
+        return topi.nn.nll_loss(predictions, targets, weight, reduction, ignore_index)
+
+    if len(call.args) == 2:
+        return bb.call_te(
+            nll_loss_without_weight,
+            call.args[0],
+            call.args[1],
+            reduction=call.attrs.reduction,
+            ignore_index=call.attrs.ignore_index,
+        )
+
+    return bb.call_te(
+        topi.nn.nll_loss,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        reduction=call.attrs.reduction,
+        ignore_index=call.attrs.ignore_index,
+    )
+
+
+@register_legalize("relax.nn.batch_flatten")
+def _nn_batch_flatten(bb: BlockBuilder, call: Call) -> Expr:
+    if call.ty.shape is None:
+        return call
+    return bb.call_te(topi.reshape, call.args[0], call.ty.shape.values)

python/tvm/relax/transform/legalize_ops/qdq.py

@@ -0,0 +1,165 @@
+# 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 quantize/dequantize operators."""
+
+import tvm
+from tvm import te, tirx
+from tvm.ir import Call
+from tvm.runtime import DataTypeCode
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import _try_convert_to_scalar_const, register_legalize
+
+
+def clip_cast(val, dtype):
+    const_min = tvm.tirx.min_value(dtype)
+    const_max = tvm.tirx.max_value(dtype)
+    return te.max(te.min(val, const_max), const_min).astype(dtype)
+
+
+def is_const_scalar(x):
+    return isinstance(x, tvm.tirx.IntImm | tvm.tirx.FloatImm)
+
+
+def _is_singleton_qparam(qparam: te.Tensor) -> bool:
+    """Return True if qparam is a tensor with all dimensions equal to 1."""
+    if not isinstance(qparam, te.Tensor):
+        return False
+    if len(qparam.shape) == 0:
+        return True
+    for dim in qparam.shape:
+        if not isinstance(dim, tirx.IntImm) or dim.value != 1:
+            return False
+    return True
+
+
+@register_legalize("relax.quantize")
+def _quantize(bb: BlockBuilder, call: Call) -> Expr:
+    """
+    Lower relax.quantize into the sequence of simple operations.
+    Quantization formula is defined as: out = clip(round(input / scale) + zp, min_val, max_val)
+    """
+    axis = call.attrs.axis
+    out_dtype = call.attrs.out_dtype
+
+    def te_quantize(
+        data: te.Tensor,
+        scale: te.Tensor | tirx.IntImm | tirx.FloatImm,
+        zp: te.Tensor | tirx.IntImm | tirx.FloatImm,
+    ):
+        scale_singleton = _is_singleton_qparam(scale) if isinstance(scale, te.Tensor) else False
+        zp_singleton = _is_singleton_qparam(zp) if isinstance(zp, te.Tensor) else False
+
+        def quantize_compute(*indices):
+            if is_const_scalar(scale):
+                scale_value = scale
+            elif scale_singleton:
+                scale_value = scale[(0,) * len(scale.shape)]
+            else:
+                scale_value = scale[indices[axis]]
+
+            if is_const_scalar(zp):
+                zp_value = zp
+            elif zp_singleton:
+                zp_value = zp[(0,) * len(zp.shape)]
+            else:
+                zp_value = zp[indices[axis]]
+            scaled = data[indices] / scale_value
+            round_val = (te.round(scaled) if "int" in out_dtype else scaled) + zp_value
+            return clip_cast(round_val, out_dtype)
+
+        output_shape = data.shape
+        return te.compute(output_shape, quantize_compute, name="quantized")
+
+    return bb.call_te(
+        te_quantize,
+        call.args[0],
+        _try_convert_to_scalar_const(call.args[1]),
+        _try_convert_to_scalar_const(call.args[2]),
+        primfunc_name_hint="quantize",
+    )
+
+
+@register_legalize("relax.dequantize")
+def _dequantize(bb: BlockBuilder, call: Call) -> Expr:
+    """
+    Lower relax.dequantize into the sequence of simple operations.
+    Dequantization formula is defined as: out = scale * (input - zp)
+    Compute datatype: float32
+
+    Example of lowering:
+
+        dtype = ["int32"|"float32"]
+
+        qnn.dequantize(data, scale, zp, "float32") -->
+            sub = subtract(cast(data, dtype), zp)
+            out = multiply(cast(sub, "float32"), scale)
+
+        qnn.dequantize(data, scale, zp, "float16") -->
+            sub = subtract(cast(data, dtype), zp)
+            mul = multiply(cast(sub, "float32"), cast(scale, "float32"))
+            clipped_out = clip(mul, float32(-65504.0), float32(65504.0))
+            out = cast(clipped_out, "float16")
+    """
+    axis = call.attrs.axis
+    out_dtype = call.attrs.out_dtype
+
+    def te_dequantize(
+        data: te.Tensor,
+        scale: te.Tensor | tirx.IntImm | tirx.FloatImm,
+        zp: te.Tensor | tirx.IntImm | tirx.FloatImm,
+    ):
+        scale_singleton = _is_singleton_qparam(scale) if isinstance(scale, te.Tensor) else False
+        zp_singleton = _is_singleton_qparam(zp) if isinstance(zp, te.Tensor) else False
+
+        def dequantize_compute(*indices):
+            if is_const_scalar(scale):
+                scale_value = scale
+            elif scale_singleton:
+                scale_value = scale[(0,) * len(scale.shape)]
+            else:
+                scale_value = scale[indices[axis]]
+
+            if is_const_scalar(zp):
+                zp_value = zp
+            elif zp_singleton:
+                zp_value = zp[(0,) * len(zp.shape)]
+            else:
+                zp_value = zp[indices[axis]]
+            dtype = (
+                "float32"
+                if data.dtype.matches_code(DataTypeCode.FLOAT, DataTypeCode.BFLOAT)
+                else "int32"
+            )
+            sub = data[indices].astype(dtype) - zp_value
+            out = sub * scale_value.astype("float32")
+            if out_dtype == "float32":
+                return out
+            return clip_cast(out, out_dtype)
+
+        output_shape = data.shape
+        return te.compute(output_shape, dequantize_compute, name="dequantized")
+
+    return bb.call_te(
+        te_dequantize,
+        call.args[0],
+        _try_convert_to_scalar_const(call.args[1]),
+        _try_convert_to_scalar_const(call.args[2]),
+        primfunc_name_hint="dequantize",
+    )

python/tvm/relax/transform/legalize_ops/search.py

@@ -0,0 +1,52 @@
+# 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 search operators."""
+
+from tvm import topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+from .common import LegalizeFunc, TEFunc, _call_topi_without_attr, register_legalize
+
+register_legalize("relax.where", _call_topi_without_attr(topi.where))
+
+
+def _argmax_argmin(te_func: TEFunc) -> LegalizeFunc:
+    def argmax_argmin_call_te(bb: BlockBuilder, call: Call) -> Expr:
+        return bb.call_te(
+            te_func,
+            call.args[0],
+            None if call.attrs.axis is None else call.attrs.axis,
+            call.attrs.keepdims,
+        )
+
+    return argmax_argmin_call_te
+
+
+register_legalize("relax.argmax", _argmax_argmin(topi.argmax))
+register_legalize("relax.argmin", _argmax_argmin(topi.argmin))
+
+
+@register_legalize("relax.bucketize")
+def _bucketize(bb, call):
+    input_tensor = call.args[0]
+    boundaries = call.args[1]
+    right = call.attrs.right
+    out_dtype = "int32" if call.attrs.out_int32 else "int64"
+    return bb.call_te(topi.searchsorted, boundaries, input_tensor, right, out_dtype)

python/tvm/relax/transform/legalize_ops/statistical.py

@@ -0,0 +1,179 @@
+# 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 collections.abc import Callable
+
+from tvm import te, tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, ShapeExpr
+from .common import LegalizeFunc, TEFunc, register_legalize
+
+
+def _normalize_reduction_axes(axis: list[int] | None, ndim: int) -> list[int]:
+    if axis is None:
+        return list(range(ndim))
+
+    axes = []
+    for dim in axis:
+        if isinstance(dim, tirx.IntImm):
+            dim = dim.value
+        dim = int(dim)
+        axes.append(dim + ndim if dim < 0 else dim)
+    return axes
+
+
+def _has_const_zero_reduction_dim(call: Call) -> bool:
+    input_shape = call.args[0].ty.shape
+    if not isinstance(input_shape, ShapeExpr):
+        return False
+
+    axes = _normalize_reduction_axes(call.attrs.axis, len(input_shape.values))
+    return any(
+        isinstance(input_shape.values[dim], tirx.IntImm) and input_shape.values[dim] == 0
+        for dim in axes
+    )
+
+
+def _statistical(
+    te_func: TEFunc,
+    zero_dim_identity: int | float | bool | Callable[[str], int | float | bool] | None = None,
+) -> LegalizeFunc:
+    def statistical_call_te(bb: BlockBuilder, call: Call) -> Expr:
+        if zero_dim_identity is not None and _has_const_zero_reduction_dim(call):
+            fill_value = (
+                zero_dim_identity(call.ty.dtype)
+                if callable(zero_dim_identity)
+                else zero_dim_identity
+            )
+            return bb.call_te(
+                topi.full,
+                call.ty.shape.values,
+                call.ty.dtype,
+                fill_value,
+            )
+        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[int]) -> tirx.Expr:
+    shape_prod = tirx.const(1, "int32")
+    axes = list(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[int], 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[int], keepdims: bool) -> te.Tensor:
+    dev = x - _te_mean(x, axis, True)
+    return _te_mean(dev * dev, axis, keepdims)
+    # This version has better memory locality and performance
+    # But may trigger some precision problems, so we will use the previous version now
+    # mean = _te_mean(x, axis, keepdims)
+    # return _te_mean(x * x, axis, keepdims) - mean * mean
+
+
+def _te_median(
+    x: te.Tensor, axis: list[int], keepdims: bool
+) -> te.Tensor | tuple[te.Tensor, te.Tensor]:
+    # currently only supports one axis or no axis ~ same pytorch
+    # todo: support multiple axis ~ same numpy
+    shape_prod = _compute_shape_prod(x, axis)
+    mid_index = (shape_prod - 1) // 2
+
+    if axis is None or len(axis) == 0:
+        x = topi.reshape(x, [shape_prod])
+        ax = -1
+    else:
+        ax = axis[0]
+    index_sorted = topi.argsort(x, axis=ax, is_ascend=True, dtype="int64")
+    x_sorted = topi.gather(x, axis=ax, indices=index_sorted)
+
+    new_shape = list(x.shape)
+    new_shape[ax] = 1
+    indices = topi.full(new_shape, fill_value=mid_index, dtype="int64")
+
+    median_val = topi.gather(x_sorted, axis=ax, indices=indices)
+    median_idx = topi.gather(index_sorted, axis=ax, indices=indices)
+
+    if axis is None or len(axis) == 0:
+        return median_val if keepdims else topi.squeeze(median_val, axis=axis)
+
+    val = median_val
+    idx = median_idx
+    if not keepdims:
+        val = topi.squeeze(val, axis=axis)
+        idx = topi.squeeze(idx, axis=axis)
+    return val, idx
+
+
+@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[int], 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.median")
+def _median(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        _te_median,
+        call.args[0],
+        call.attrs.axis,
+        call.attrs.keepdims,
+        primfunc_name_hint="median",
+    )
+
+
+register_legalize("relax.max", _statistical(topi.max))
+register_legalize("relax.min", _statistical(topi.min))
+register_legalize(
+    "relax.prod",
+    _statistical(topi.prod, zero_dim_identity=lambda dtype: True if dtype == "bool" else 1),
+)
+register_legalize("relax.sum", _statistical(topi.sum, zero_dim_identity=0))

python/tvm/relax/transform/legalize_ops/unary.py

@@ -0,0 +1,72 @@
+# 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 unary operators."""
+
+from tvm import te, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr
+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.acos", _call_topi_without_attr(topi.acos, "tir_acos"))
+register_legalize("relax.acosh", _call_topi_without_attr(topi.acosh, "tir_acosh"))
+register_legalize("relax.asin", _call_topi_without_attr(topi.asin, "tir_asin"))
+register_legalize("relax.asinh", _call_topi_without_attr(topi.asinh, "tir_asinh"))
+register_legalize("relax.atan", _call_topi_without_attr(topi.atan, "tir_atan"))
+register_legalize("relax.atanh", _call_topi_without_attr(topi.atanh, "tir_atanh"))
+register_legalize("relax.bitwise_not", _call_topi_without_attr(topi.bitwise_not, "tir_bitwise_not"))
+register_legalize("relax.ceil", _call_topi_without_attr(topi.ceil, "tir_ceil"))
+register_legalize("relax.cos", _call_topi_without_attr(topi.cos, "tir_cos"))
+register_legalize("relax.cosh", _call_topi_without_attr(topi.cosh, "tir_cosh"))
+register_legalize("relax.exp", _call_topi_without_attr(topi.exp, "tir_exp"))
+register_legalize("relax.floor", _call_topi_without_attr(topi.floor, "tir_floor"))
+register_legalize("relax.isfinite", _call_topi_without_attr(topi.isfinite, "tir_isfinite"))
+register_legalize("relax.isinf", _call_topi_without_attr(topi.isinf, "tir_isinf"))
+register_legalize("relax.isnan", _call_topi_without_attr(topi.isnan, "tir_isnan"))
+register_legalize("relax.log", _call_topi_without_attr(topi.log, "tir_log"))
+register_legalize("relax.logical_not", _call_topi_without_attr(topi.logical_not, "tir_logical_not"))
+register_legalize("relax.negative", _call_topi_without_attr(topi.negative, "tir_negative"))
+register_legalize("relax.round", _call_topi_without_attr(topi.round, "tir_round"))
+register_legalize("relax.rsqrt", _call_topi_without_attr(topi.rsqrt, "tir_rsqrt"))
+register_legalize("relax.sigmoid", _call_topi_without_attr(topi.sigmoid, "tir_sigmoid"))
+register_legalize("relax.sign", _call_topi_without_attr(topi.sign, "tir_sign"))
+register_legalize("relax.sin", _call_topi_without_attr(topi.sin, "tir_sin"))
+register_legalize("relax.sinh", _call_topi_without_attr(topi.sinh, "tir_sinh"))
+register_legalize("relax.square", _call_topi_without_attr(lambda x: x * x, "tir_square"))
+register_legalize("relax.sqrt", _call_topi_without_attr(topi.sqrt, "tir_sqrt"))
+register_legalize("relax.tan", _call_topi_without_attr(topi.tan, "tir_tan"))
+register_legalize("relax.tanh", _call_topi_without_attr(topi.tanh, "tir_tanh"))
+register_legalize("relax.trunc", _call_topi_without_attr(topi.trunc, "tir_trunc"))
+register_legalize("relax.clip", _call_topi_without_attr(topi.clip, "tir_clip"))
+
+
+@register_legalize("relax.erf")
+def _erf(bb: BlockBuilder, call: Call) -> Expr:
+    def te_erf(x: te.Tensor):
+        dtype = x.dtype
+        if dtype == "float16":
+            erf = topi.math.cast(topi.erf(topi.math.cast(x, "float32")), "float16")
+        else:
+            erf = topi.erf(x)
+        return erf
+
+    return bb.call_te(te_erf, call.args[0], primfunc_name_hint="erf")

python/tvm/relax/transform/legalize_ops/vision.py

@@ -0,0 +1,194 @@
+# 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 vision network related operators."""
+
+from tvm import relax, te, tirx, topi
+from tvm.ir import Call
+
+from ...block_builder import BlockBuilder
+from ...expr import Expr, TupleGetItem
+from .common import register_legalize
+
+
+@register_legalize("relax.vision.all_class_non_max_suppression")
+def _all_class_non_max_suppression(block_builder: BlockBuilder, call: Call) -> Expr:
+    """Legalize all_class_non_max_suppression with dynamic output trimming.
+
+    This implementation uses dynamic_strided_slice to trim the NMS output to only
+    contain valid detections, improving memory efficiency and ONNX compatibility.
+
+    Returns
+    -------
+    result : Expr
+        The legalized NMS result.
+
+        - For ONNX output format, returns a tuple of
+          `(trimmed_indices, num_total_detections)`, where `trimmed_indices`
+          contains only valid detection indices.
+        - For TensorFlow output format, returns the TOPI result directly to
+          preserve the `(selected_indices, selected_scores, num_detections)`
+          layout expected by the Relax op.
+    """
+    boxes = call.args[0]
+    scores = call.args[1]
+    max_output_boxes_per_class = call.args[2]
+    iou_threshold = call.args[3]
+    score_threshold = call.args[4]
+    output_format = call.attrs.output_format
+
+    scores_shape = scores.ty.shape
+    if len(scores_shape) == 3:
+        _, _, num_boxes = scores_shape
+    elif len(scores_shape) == 2:
+        _, num_boxes = scores_shape
+    else:
+        raise ValueError(f"Unexpected scores shape: {scores_shape}")
+
+    if isinstance(max_output_boxes_per_class, relax.Constant):
+        max_boxes_val = int(max_output_boxes_per_class.data.numpy())
+    else:
+        max_boxes_val = int(num_boxes)
+
+    # Get NMS result with fixed shape from TOPI
+    nms_result = block_builder.call_te(
+        topi.vision.all_class_non_max_suppression,
+        boxes,
+        scores,
+        max_boxes_val,
+        iou_threshold,
+        score_threshold,
+        output_format,
+    )
+
+    if output_format == "tensorflow":
+        return nms_result
+
+    selected_indices = block_builder.emit(TupleGetItem(nms_result, 0))
+    num_total_detections = block_builder.emit(TupleGetItem(nms_result, 1))
+
+    # Build slicing parameters using TE to avoid high-level Relax ops during legalization
+    def build_begin():
+        return te.compute((2,), lambda i: tirx.const(0, "int64"), name="begin")
+
+    def build_strides():
+        return te.compute((2,), lambda i: tirx.const(1, "int64"), name="strides")
+
+    def build_end(count_tensor):
+        # end = [count_tensor[0], 3]
+        def compute_end(i):
+            return tirx.if_then_else(
+                i == 0,
+                tirx.Cast("int64", count_tensor[0]),
+                tirx.const(3, "int64"),
+            )
+
+        return te.compute((2,), compute_end, name="end")
+
+    begin = block_builder.call_te(build_begin)
+    strides = block_builder.call_te(build_strides)
+    end = block_builder.call_te(build_end, num_total_detections)
+
+    # Apply dynamic strided slice to trim to valid detections only
+    trimmed_indices = block_builder.emit(
+        relax.op.dynamic_strided_slice(selected_indices, begin, end, strides)
+    )
+
+    # Return trimmed indices along with num_total_detections for compatibility
+    return relax.Tuple([trimmed_indices, num_total_detections])
+
+
+@register_legalize("relax.vision.roi_align")
+def _roi_align(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.vision.roi_align,
+        call.args[0],
+        call.args[1],
+        pooled_size=call.attrs.pooled_size,
+        spatial_scale=call.attrs.spatial_scale,
+        mode=call.attrs.mode,
+        sample_ratio=call.attrs.sample_ratio,
+        aligned=call.attrs.aligned,
+        layout=call.attrs.layout,
+    )
+
+
+@register_legalize("relax.vision.get_valid_counts")
+def _get_valid_counts(block_builder: BlockBuilder, call: Call) -> Expr:
+    return block_builder.call_te(
+        topi.vision.get_valid_counts,
+        call.args[0],
+        score_threshold=call.attrs.score_threshold,
+        id_index=call.attrs.id_index,
+        score_index=call.attrs.score_index,
+    )
+
+
+@register_legalize("relax.vision.non_max_suppression")
+def _non_max_suppression(block_builder: BlockBuilder, call: Call) -> Expr:
+    return block_builder.call_te(
+        topi.vision.non_max_suppression,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        max_output_size=call.attrs.max_output_size,
+        iou_threshold=call.attrs.iou_threshold,
+        force_suppress=call.attrs.force_suppress,
+        top_k=call.attrs.top_k,
+        coord_start=call.attrs.coord_start,
+        score_index=call.attrs.score_index,
+        id_index=call.attrs.id_index,
+        return_indices=call.attrs.return_indices,
+        invalid_to_bottom=call.attrs.invalid_to_bottom,
+        soft_nms_sigma=call.attrs.soft_nms_sigma,
+        score_threshold=call.attrs.score_threshold,
+    )
+
+
+@register_legalize("relax.vision.roi_pool")
+def _roi_pool(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(
+        topi.vision.roi_pool,
+        call.args[0],
+        call.args[1],
+        pooled_size=call.attrs.pooled_size,
+        spatial_scale=call.attrs.spatial_scale,
+        layout=call.attrs.layout,
+    )
+
+
+@register_legalize("relax.vision.multibox_transform_loc")
+def _multibox_transform_loc(bb: BlockBuilder, call: Call) -> Expr:
+    variances = tuple(float(x) for x in call.attrs.variances)
+
+    def _te(cls_pred, loc_pred, anchor):
+        return topi.vision.multibox_transform_loc(
+            cls_pred,
+            loc_pred,
+            anchor,
+            variances,
+            clip=call.attrs.clip,
+            threshold=call.attrs.threshold,
+            keep_background=call.attrs.keep_background,
+        )
+
+    return bb.call_te(
+        _te,
+        call.args[0],
+        call.args[1],
+        call.args[2],
+        primfunc_name_hint="multibox_transform_loc",
+    )

python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.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.
+"""Lower the storage/tensor allocation on IPC memory.
+The pass is written in Python for experiment, fast development.
+"""
+
+import tvm
+from tvm import relax
+from tvm.ir.module import IRModule
+from tvm.relax.expr import Expr
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="LowerGPUIPCAllocStorage")
+class LowerGPUIPCAllocStorage:
+    """Lower the storage/tensor allocation on IPC memory."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _Rewriter(mod).transform()
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    def __init__(self, mod: IRModule) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.memory_alloc_storage_op = tvm.ir.Op.get("relax.memory.alloc_storage")
+        self.memory_alloc_tensor_op = tvm.ir.Op.get("relax.memory.alloc_tensor")
+        self.builtin_alloc_tensor_op = tvm.ir.Op.get("relax.builtin.alloc_tensor")
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                self.builder_.update_func(g_var, updated_func)
+        return self.builder_.get()
+
+    def visit_call_(self, call: relax.Call) -> Expr:  # pylint: disable=arguments-renamed
+        if call.op == self.memory_alloc_storage_op and call.args[2].value == "ipc_memory":
+            return self.rewrite_alloc_storage(call)
+        elif call.op == self.builtin_alloc_tensor_op and call.args[3].value == "ipc_memory":
+            return self.rewrite_alloc_tensor(call)
+        else:
+            return call
+
+    def rewrite_alloc_storage(self, call: relax.Call) -> relax.Call:
+        shape = call.args[0]
+        dtype = call.args[3]
+        return relax.Call(
+            relax.ExternFunc("runtime.disco.cuda_ipc.alloc_storage"),
+            args=[shape, dtype],
+            ty_args=[call.ty],
+        )
+
+    def rewrite_alloc_tensor(self, call: relax.Call) -> relax.Call:
+        shape = call.args[0]
+        dtype = call.args[1]
+        ipc_alloc_storage = relax.Call(
+            relax.ExternFunc("runtime.disco.cuda_ipc.alloc_storage"),
+            args=[shape, dtype],
+            ty_args=[relax.AnyType()],
+        )
+        return relax.Call(
+            self.memory_alloc_tensor_op,
+            args=[ipc_alloc_storage, call.args[2], shape, dtype, relax.prim_value(0)],
+            ty_args=call.ty_args,
+        )

python/tvm/relax/transform/optimize_layout_transform.py

@@ -0,0 +1,88 @@
+# 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, redefined-argument-from-local
+"""Relax Optimize Layout Transform pass."""
+
+import tvm_ffi
+
+from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext
+from tvm.relax import Expr
+from tvm.relax.dpl import TuplePattern, is_op, rewrite_call, wildcard
+
+from . import function_pass
+
+
+@function_pass(opt_level=0)
+class OptimizeLayoutTransform:
+    """
+    Pass to remove redundant transform layout operators
+    introduced by AlterOpImpl pass.
+    """
+
+    def __init__(self):
+        self.input = wildcard()
+        pattern_transform_layout = is_op("relax.layout_transform")(self.input)
+        pattern_1 = is_op("relax.layout_transform")(pattern_transform_layout)
+
+        self.gv_ = wildcard()
+        args = TuplePattern([pattern_transform_layout])
+        pattern_2 = is_op("relax.call_tir")(self.gv_, args)
+        self.pattern_2 = is_op("relax.layout_transform")(pattern_2)
+
+        self.pattern = pattern_1 | self.pattern_2
+
+    def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRModule:
+        """
+        Tranformation function to pattern match layout_transform -> layout_transform
+        pattern
+
+        Parameters
+        ----------
+        func: Expr
+            The relax function to be optimized
+
+        mod: IRModule
+            The ir module
+
+        ctx: PassContext
+            Relax pass context
+        """
+
+        self.mod = mod
+        updated_func = func
+
+        # Skip primitive functions
+        if "Primitive" in func.attrs.keys() and func.attrs["Primitive"] != 0:
+            return updated_func
+
+        def rewriter(expr, matches):
+            arg1 = matches[self.pattern]
+            if self.pattern_2 not in matches.keys():
+                arg2 = matches[self.input]
+            else:
+                arg2 = matches[self.gv_]
+                if "remove_pad" == self.mod[arg2].attrs["operator_name"]:
+                    arg2 = matches[self.input]
+            if hasattr(arg1.ty, "shape") and hasattr(arg2.ty, "shape"):
+                if tvm_ffi.structural_equal(arg1.ty.shape, arg2.ty.shape):
+                    return arg2
+            return expr
+
+        updated_func = rewrite_call(self.pattern, rewriter, func)
+
+        return updated_func

python/tvm/relax/transform/remove_redundant_reshape.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, unused-argument, missing-function-docstring, abstract-method
+"""Relax Remove Redundant Reshape ops"""
+
+import tvm_ffi
+
+from tvm import IRModule, relax
+from tvm.ir.transform import PassContext
+from tvm.relax import Expr
+from tvm.relax.dpl import is_op, rewrite_call, wildcard
+
+from . import function_pass
+
+
+@function_pass(opt_level=0)
+class RemoveRedundantReshape:
+    """
+    Transformation pass to remove redundant reshape operator
+    """
+
+    def __init__(self):
+        self.input1 = wildcard()
+        shape1 = wildcard()
+        pattern_redundant_reshape = is_op("relax.reshape")(self.input1, shape1)
+        self.no_op_reshape = pattern_redundant_reshape
+        shape2 = wildcard()
+        self.repeated_reshape = is_op("relax.reshape")(pattern_redundant_reshape, shape2)
+        self.pattern = self.repeated_reshape | self.no_op_reshape
+
+    def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRModule:
+        """
+        Tarnsformation function to remove redundant reshape
+        where tensors before and after reshape are of same dimentions.
+
+        Parameters
+        --------------
+        func: Expr
+            The relax function to be optimized
+
+        mod: IRModule
+            The IR module
+
+        ctx: PassContext
+            Relax pass context
+        """
+
+        updated_func = func
+
+        # Skip primitive functions
+        if "Primitive" in func.attrs.keys() and func.attrs["Primitive"] != 0:
+            return updated_func
+
+        def rewriter(expr, matches):
+            arg = matches[self.input1]
+
+            if self.repeated_reshape in matches:
+                output_shape = matches[self.repeated_reshape].args[1]
+                return relax.op.reshape(arg, output_shape)
+
+            elif self.no_op_reshape in matches:
+                output_shape = matches[self.no_op_reshape].args[1]
+                if arg.ty.shape and tvm_ffi.structural_equal(arg.ty.shape, output_shape):
+                    return arg
+            return expr
+
+        updated_func = rewrite_call(self.pattern, rewriter, func)
+
+        return updated_func