chore: import upstream snapshot with attribution

python/tvm/relax/backend/__init__.py

@@ -0,0 +1,23 @@
+# 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 backends"""
+
+from . import contrib, cpu_generic, cuda, gpu_generic, metal, rocm, adreno
+from .dispatch_sampling import DispatchSampling
+from .dispatch_sort_scan import DispatchSortScan
+from .pattern_registry import get_pattern, get_patterns_with_prefix

python/tvm/relax/backend/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI API for Relax backend."""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.backend", __name__)

python/tvm/relax/backend/adreno/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+"""The Relax Adreno backend compilation pipeline and other passes."""
+
+from . import transform
+
+from .pipeline import (
+    finalize_passes,
+    get_default_pipeline,
+    dataflow_lower_passes,
+    legalize_passes,
+    library_dispatch_passes,
+)

python/tvm/relax/backend/adreno/clml.py

@@ -0,0 +1,727 @@
+# 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, pointless-exception-statement
+"""Pattern table for CLML backend"""
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.ir.transform import PassContext, module_pass
+from tvm.relax import transform
+from tvm.relax.dpl.pattern import (
+    GlobalVarPattern,
+    TuplePattern,
+    is_const,
+    is_op,
+    is_tuple_get_item,
+    wildcard,
+)
+from tvm.relax.expr import TupleGetItem, VarBinding
+from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import register_patterns
+
+
+def _dtype_str(dtype):
+    return str(dtype.dtype) if isinstance(dtype, tvm.ir.PrimType) else str(dtype)
+
+
+@mutator
+class AppendReshapeToBNRewriter(PyExprMutator):
+    """
+    Append Reshape Operator to BatchNorm Pass Rewriter Pass
+
+    - Automatically appends a reshape operation after BatchNorm operators
+    - Resolves fusion issues for custom backends where BatchNorm output
+      might explicitly access the first elment of the Tuple
+
+    Algo:
+    Identifies BatchNorm operators in the computational graph
+    When BatchNorm's first output is accessed via TupleGetItem
+    Automatically inserts a reshape operation to match input shape
+
+    """
+
+    def __init__(self, mod):
+        super().__init__(mod)
+        self.bn_vars = {}
+
+    def visit_tuple_getitem_(self, op: TupleGetItem):
+        tuple_value = op.tuple_value
+        reshape_op = tvm.ir.Op.get("relax.reshape")
+
+        if isinstance(tuple_value, relax.Var) and tuple_value in self.bn_vars:
+            bn_call = self.bn_vars[tuple_value]
+            if op.index == 0:
+                bn_out = relax.TupleGetItem(bn_call, 0)
+                input_shape = bn_call.args[0].ty.shape
+                return relax.Call(reshape_op, [bn_out, input_shape])
+
+        return super().visit_tuple_getitem_(op)
+
+    def visit_var_binding_(self, binding: VarBinding):
+        if isinstance(binding.value, relax.Call) and binding.value.op.name == "relax.nn.batch_norm":
+            self.bn_vars[binding.var] = binding.value
+        return super().visit_var_binding_(binding)
+
+
+@transform.function_pass(opt_level=0, name="AppendReshapeToBN")
+class AppendReshapeToBNRewriterPass:
+    def transform_function(
+        self, func: relax.Function, mod: IRModule, _ctx: tvm.transform.PassContext
+    ) -> relax.Function:
+        updated_func = AppendReshapeToBNRewriter(mod).visit_expr(func)
+        updated_func = relax.analysis.remove_all_unused(updated_func)
+        return updated_func
+
+
+def clml_sdk_version():
+    """Utility function to get clml version.
+
+    Probes the FFI registry for the OpenCLML version registered by the
+    CLML backend at build time.  Returns 2 when CLML is not present.
+    """
+    # Registry: "relax.get_openclml_version" — returns the CLML SDK version
+    # that TVM was built against; registered unconditionally in codegen.cc.
+    # Grep hint: grep -rn 'relax.get_openclml_version' src/
+    get_version = tvm.get_global_func("relax.get_openclml_version", allow_missing=True)
+    if get_version is None:
+        return 2
+    return int(get_version())
+
+
+def is_clml_runtime_enabled():
+    """Check if the CLML graph runtime is present.
+
+    Returns
+    -------
+    ret: bool
+        True if present, False if not.
+    """
+    check_enabled = tvm.get_global_func("relax.op.is_openclml_runtime_enabled", True)
+    if check_enabled:
+        return check_enabled()
+    return False
+
+
+def _check_default(context: PatternCheckContext) -> bool:
+    return True
+
+
+def clml_pattern_table():
+    """Get the CLML pattern table."""
+
+    def _check_conv2d(context: PatternCheckContext) -> bool:
+        if "root" in context.annotated_expr:
+            root_call = context.annotated_expr["root"]
+            if root_call.op.name == "relax.nn.conv2d":
+                input_layout = root_call.attrs.data_layout
+                weight_layout = root_call.attrs.kernel_layout
+                if input_layout != "NCHW" or weight_layout != "OIHW":
+                    return False
+            if root_call.op.name == "relax.nn.conv2d_transpose":
+                input_layout = root_call.attrs.data_layout
+                weight_layout = root_call.attrs.kernel_layout
+                if input_layout != "NCHW" or weight_layout != "OIHW":
+                    return False
+
+        if "data" in context.annotated_expr:
+            input_expr = context.annotated_expr["data"]
+            input_dtype = _dtype_str(input_expr.ty.dtype)
+            if input_dtype not in ["float32", "float16"]:
+                return False
+
+        if "weight" in context.annotated_expr:
+            weight_expr = context.annotated_expr["weight"]
+            weight_dtype = _dtype_str(weight_expr.ty.dtype)
+            if weight_dtype not in ["float32", "float16"]:
+                return False
+
+        return True
+
+    def populate_patterns(patterns, name, op, annotations, *args):
+        ret = {}
+        for k, v in patterns.items():
+            ret_ann = v["annotation"].copy()
+            ret_ann.update(annotations)
+            ret[name + "." + k] = {"pattern": op(v["pattern"], *args), "annotation": ret_ann.copy()}
+
+        return ret
+
+    def conv_pattern():
+        """Create a convolution pattern."""
+        data = wildcard()
+        weight = wildcard()
+        bias = is_const()
+        bn_scale = is_const()
+        bn_bias = is_const()
+        bn_mean = is_const()
+        bn_var = is_const()
+
+        annotations = {
+            "data": data,
+            "weight": weight,
+        }
+
+        patterns = {}
+        patterns["nn.conv2d"] = {
+            "pattern": is_op("relax.nn.conv2d")(data, weight),
+            "annotation": annotations.copy(),
+        }
+
+        pad_annotations = annotations.copy()
+        patterns["pad.nn.conv2d"] = {
+            "pattern": is_op("relax.nn.conv2d")(is_op("relax.nn.pad")(data), weight),
+            "annotation": pad_annotations,
+        }
+
+        patterns["nn.conv2d_transpose"] = {
+            "pattern": is_op("relax.nn.conv2d_transpose")(data, weight),
+            "annotation": annotations.copy(),
+        }
+        patterns.update(
+            populate_patterns(patterns, "bias", is_op("relax.add"), {"bias": bias}, bias)
+        )
+        patterns.update(
+            populate_patterns(
+                patterns,
+                "bn",
+                is_op("relax.nn.batch_norm"),
+                {
+                    "bn_scale": bn_scale,
+                    "bn_bias": bn_bias,
+                    "bn_mean": bn_mean,
+                    "bn_var": bn_var,
+                },
+                bn_scale,
+                bn_bias,
+                bn_mean,
+                bn_var,
+            )
+        )
+        tuple_patterns = {}
+        for k, v in patterns.items():
+            tuple_annotation = v["annotation"].copy()
+            tuple_patterns["tuple" + "." + k] = {
+                "pattern": is_tuple_get_item(v["pattern"], 0),
+                "annotation": tuple_annotation,
+            }
+        patterns.update(tuple_patterns)
+
+        relu_patterns = populate_patterns(patterns, "relu", is_op("relax.nn.relu"), {})
+        clip_patterns = populate_patterns(patterns, "clip", is_op("relax.clip"), {})
+        patterns.update(relu_patterns)
+        patterns.update(clip_patterns)
+
+        conv_patterns = []
+        for k, v in patterns.items():
+            ret_annotations = v["annotation"]
+            ret_annotations["root"] = v["pattern"]
+            conv_patterns.append(
+                ("openclml." + (k), v["pattern"], ret_annotations.copy(), _check_conv2d)
+            )
+        return conv_patterns[::-1]
+
+    def _check_maxpool2d(context: PatternCheckContext) -> bool:
+        root = context.annotated_expr.get("root")
+        if root is None or not isinstance(root, relax.Call):
+            return False
+
+        if root.op.name != "relax.nn.max_pool2d":
+            return False
+
+        if "data" not in context.annotated_expr:
+            return False
+
+        data = context.annotated_expr["data"]
+        input_shape = data.ty.shape
+
+        if len(input_shape) != 4:
+            return False
+
+        if any(dim <= 0 for dim in input_shape):
+            return False
+
+        pool_size = root.attrs.pool_size
+        if len(pool_size) != 2:
+            return False
+        if any(size <= 0 for size in pool_size):
+            return False
+
+        strides = root.attrs.strides
+        if len(strides) != 2:
+            return False
+        if any(stride <= 0 for stride in strides):
+            return False
+
+        dilation = root.attrs.dilation
+        if len(dilation) != 2:
+            return False
+        if any(d <= 0 for d in dilation):
+            return False
+
+        padding = root.attrs.padding
+        if len(padding) != 4:
+            return False
+        if any(p < 0 for p in padding):
+            return False
+
+        return True
+
+    def maxpool_pattern():
+        """Create Pool Pattern"""
+        data = wildcard()
+        annotations = {
+            "data": data,
+        }
+        patterns = {}
+        patterns["nn.max_pool2d"] = {
+            "pattern": is_op("relax.nn.max_pool2d")(data),
+            "annotation": annotations.copy(),
+        }
+
+        pool_patterns = []
+        for k, v in patterns.items():
+            ret_annotations = v["annotation"]
+            ret_annotations["root"] = v["pattern"]
+            pool_patterns.append(
+                ("openclml." + (k), v["pattern"], ret_annotations.copy(), _check_maxpool2d)
+            )
+        return pool_patterns
+
+    def _check_avgpool2d(context: PatternCheckContext) -> bool:
+        root = context.annotated_expr.get("root")
+        if root is None or not isinstance(root, relax.Call):
+            return False
+
+        if root.op.name != "relax.nn.avg_pool2d":
+            return False
+
+        if "data" not in context.annotated_expr:
+            return False
+
+        data = context.annotated_expr["data"]
+        input_shape = data.ty.shape
+
+        if len(input_shape) != 4:
+            return False
+
+        if any(dim <= 0 for dim in input_shape):
+            return False
+
+        pool_size = root.attrs.pool_size
+        if len(pool_size) != 2:
+            return False
+        if any(size <= 0 for size in pool_size):
+            return False
+
+        strides = root.attrs.strides
+        if len(strides) != 2:
+            return False
+        if any(stride <= 0 for stride in strides):
+            return False
+
+        padding = root.attrs.padding
+        if len(padding) != 4:
+            return False
+        if any(p < 0 for p in padding):
+            return False
+
+        return True
+
+    def avgpool_pattern():
+        data = wildcard()
+        annotations = {
+            "data": data,
+        }
+        patterns = {}
+        patterns["nn.avg_pool2d"] = {
+            "pattern": is_op("relax.nn.avg_pool2d")(data),
+            "annotation": annotations.copy(),
+        }
+
+        pool_patterns = []
+        for k, v in patterns.items():
+            ret_annotations = v["annotation"]
+            ret_annotations["root"] = v["pattern"]
+            pool_patterns.append(
+                ("openclml." + (k), v["pattern"], ret_annotations.copy(), _check_avgpool2d)
+            )
+        return pool_patterns
+
+    def _check_global_avgpool(context: PatternCheckContext) -> bool:
+        root = context.annotated_expr.get("root")
+        if root is None or not isinstance(root, relax.Call):
+            return False
+
+        if root.op.name != "relax.mean":
+            return False
+
+        if "data" not in context.annotated_expr:
+            return False
+
+        data = context.annotated_expr["data"]
+        input_shape = data.ty.shape
+
+        if len(input_shape) != 4:
+            return False
+
+        if input_shape[1] <= 0 or input_shape[2] <= 0 or input_shape[3] <= 0:
+            return False
+
+        if not hasattr(root.attrs, "axis"):
+            return False
+
+        axis = root.attrs.axis
+        if not (len(axis) == 2 and axis[0] == 2 and axis[1] == 3):
+            return False
+
+        return True
+
+    def global_avgpool_pattern():
+        """Create Pool Pattern"""
+        data = wildcard()
+        pattern = is_op("relax.mean")(data).has_attr({"axis": [2, 3]})
+
+        annotations = {
+            "data": data,
+            "root": pattern,
+        }
+
+        return [
+            ("openclml.nn.global_avg_pool2d", pattern, annotations, _check_global_avgpool),
+        ]
+
+    def _check_reshape(context: PatternCheckContext) -> bool:
+        root = context.annotated_expr.get("root")
+        if root is None or not isinstance(root, relax.Call):
+            return False
+
+        if root.op.name != "relax.reshape":
+            return False
+
+        shape_arg = root.args[1]
+        if not isinstance(shape_arg, relax.Expr):
+            return False
+
+        return True
+
+    def reshape_pattern():
+        """Create Reshape Pattern"""
+
+        pattern = is_op("relax.reshape")(wildcard(), wildcard())
+        annotations = {
+            "root": pattern,
+        }
+        return [("openclml.reshape", pattern, annotations, _check_reshape)]
+
+    def _check_batchnorm(context: PatternCheckContext) -> bool:
+        root = context.annotated_expr.get("root")
+        if root is None or not isinstance(root, relax.Call):
+            return False
+
+        if root.op.name != "relax.reshape":
+            return False
+
+        required_params = ["moving_var", "gamma", "moving_mean", "beta"]
+        for param in required_params:
+            if param not in context.annotated_expr:
+                return False
+
+        params = {
+            "moving_var": context.annotated_expr["moving_var"],
+            "gamma": context.annotated_expr["gamma"],
+            "moving_mean": context.annotated_expr["moving_mean"],
+            "beta": context.annotated_expr["beta"],
+        }
+
+        for param in params.values():
+            if not isinstance(param, relax.expr.Constant):
+                return False
+
+        base_shape = None
+        for param in params.values():
+            shape = param.ty.shape
+            dtype = _dtype_str(param.ty.dtype)
+
+            if dtype not in {"float32"}:
+                return False
+
+            # Initialize base_shape if not set
+            if base_shape is None:
+                base_shape = shape
+                continue
+
+            # All parameters should have same shape
+            if len(shape) != len(base_shape):
+                return False
+            if any(s1 != s2 for s1, s2 in zip(shape, base_shape)):
+                return False
+
+        return True
+
+    def batch_norm_pattern():
+        """Create a batch norm pattern."""
+        data = wildcard()
+        bn_scale = is_const()
+        bn_bias = is_const()
+        bn_mean = is_const()
+        bn_var = is_const()
+
+        pattern = is_op("relax.nn.batch_norm")(data, bn_scale, bn_bias, bn_mean, bn_var)
+        pattern = is_tuple_get_item(pattern, 0)
+        pattern = is_op("relax.reshape")(pattern, wildcard())
+
+        annotations = {
+            "gamma": bn_scale,
+            "beta": bn_bias,
+            "moving_mean": bn_mean,
+            "moving_var": bn_var,
+            "root": pattern,
+        }
+
+        return [
+            ("openclml.nn.batch_norm", pattern, annotations, _check_batchnorm),
+        ]
+
+    def _check_binary_op(context: PatternCheckContext) -> bool:
+        def _check_arg(input_expr):
+            input_dtype = _dtype_str(input_expr.ty.dtype)
+            input_shape = input_expr.ty.shape
+            if len(input_shape) == 0:
+                return False
+
+            # Avoid any operators with dtype Int64
+            if input_dtype == "int64":
+                return False
+
+            # No support for batch> 1
+            if input_shape[0] > 1:
+                return False
+
+            return True
+
+        def compare_shapes(lhs_shape, rhs_shape):
+            if len(lhs_shape) != len(rhs_shape):
+                return False
+            for lhs_dim, rhs_dim in zip(lhs_shape, rhs_shape):
+                if lhs_dim != rhs_dim:
+                    return False
+            return True
+
+        lhs_shape = None
+        rhs_shape = None
+        if "lhs" in context.annotated_expr:
+            lhs = context.annotated_expr["lhs"]
+            lhs_shape = lhs.ty.shape
+            if not _check_arg(lhs):
+                return False
+
+        if "rhs" in context.annotated_expr:
+            rhs = context.annotated_expr["rhs"]
+            rhs_shape = rhs.ty.shape
+            if not _check_arg(rhs):
+                return False
+
+        # Checking for BinaryOps ( False for unaryOp )
+        if (
+            "lhs" in context.annotated_expr
+            and "rhs" in context.annotated_expr
+            and not compare_shapes(lhs_shape, rhs_shape)
+        ):
+            return False
+
+        return True
+
+    def binary_op_pattern():
+        """Create a binary op pattern."""
+
+        def make_pattern(op):
+            lhs = wildcard()
+            rhs = wildcard()
+            pattern = is_op(op)(lhs, rhs)
+            annotations = {"lhs": lhs, "rhs": rhs}
+            return ("openclml." + op, pattern, annotations, _check_binary_op)
+
+        binary_ops = [
+            "relax.add",
+            "relax.subtract",
+            "relax.multiply",
+            "relax.divide",
+            "relax.maximum",
+            "relax.minimum",
+        ]
+
+        return [make_pattern(op) for op in binary_ops]
+
+    def unary_op_pattern():
+        """Create a unary op pattern."""
+
+        def make_pattern(op):
+            lhs = wildcard()
+            pattern = is_op(op)(lhs)
+            annotations = {"lhs": lhs}
+            return ("openclml." + op, pattern, annotations, _check_binary_op)
+
+        unary_ops = [
+            "relax.nn.softmax",
+            "relax.nn.relu",
+            "relax.clip",
+        ]
+
+        return [make_pattern(op) for op in unary_ops]
+
+    return [
+        *conv_pattern(),
+        *batch_norm_pattern(),
+        *binary_op_pattern(),
+        *unary_op_pattern(),
+        *maxpool_pattern(),
+        *avgpool_pattern(),
+        *global_avgpool_pattern(),
+        *reshape_pattern(),
+    ]
+
+
+clml_patterns = clml_pattern_table()
+register_patterns(clml_patterns)
+
+
+@module_pass(opt_level=0, name="OpenCLMLOffLoad")
+class OpenCLMLOffLoad:
+    """The pass sequence used for CLML offload"""
+
+    def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule:
+        """The transform"""
+
+        clml_layouts = {
+            "relax.nn.conv2d": ["NCHW", "OIHW"],
+            "relax.nn.conv2d_transpose": ["NCHW", "OIHW"],
+        }
+        seq = tvm.transform.Sequential(
+            [
+                transform.ConvertLayout(clml_layouts),
+                transform.Normalize(),
+                transform.FoldBatchnormToConv2D(),
+                AppendReshapeToBNRewriterPass(),
+                transform.FoldConstant(),
+                transform.FuseOpsByPattern(clml_pattern_table()),
+                transform.MergeCompositeFunctions(),
+                transform.RunCodegen(),
+            ],
+        )
+        mod = seq(mod)
+        return mod
+
+
+def _check_dequantize_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
+    _input = ctx.annotated_expr["lhs"]
+    root = ctx.annotated_expr["root"]
+    wdq = ctx.annotated_expr["w_decoded"]
+    w_pack = ctx.annotated_expr["w_encoded"]
+
+    if _dtype_str(ctx.annotated_expr["lhs"].ty.dtype) != "float16":
+        return False
+    if not isinstance(wdq, relax.Call):
+        return False
+    g_var = wdq.args[0]
+    if not (isinstance(g_var, relax.GlobalVar) and "dequantize" in g_var.name_hint):
+        return False
+
+    if not (
+        (len(root.ty.shape) == 3)
+        and isinstance(root.ty.shape[0], tirx.IntImm)
+        and (_dtype_str(root.ty.dtype) == "float16")
+        and (root.ty.shape[0] == 1)
+    ):
+        return False
+
+    if not (
+        (len(wdq.ty.shape) == 2)
+        and (w_pack.ty.shape[-1] == root.ty.shape[-1])
+        and (wdq.ty.shape[-2] == _input.ty.shape[-1])
+    ):
+        return False
+
+    return True
+
+
+def dequantize_matmul_patterns():
+    """Returns a list of supported decode -> matmul patterns."""
+
+    def _dequantize_matmul_pattern(name):
+        scales = wildcard()
+        x = wildcard()
+        w_packed = wildcard()
+
+        w_decoded = is_op("relax.call_tir")(
+            GlobalVarPattern(),
+            TuplePattern([w_packed, scales]),
+        )
+        matmul = is_op("relax.matmul")(x, w_decoded)
+
+        annotations = {
+            "root": matmul,
+            "lhs": x,
+            "w_encoded": w_packed,
+            "w_decoded": w_decoded,
+            "scales": scales,
+        }
+
+        return name, matmul, annotations, _check_dequantize_matmul
+
+    return [
+        _dequantize_matmul_pattern("openclml.dequant_matmul"),
+    ]
+
+
+clml_llm_patterns = [
+    *dequantize_matmul_patterns(),
+]
+register_patterns(clml_llm_patterns)
+
+
+@tvm.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
+class OpenCLMLOffLoadForLLM:
+    """A compiler pass that partition the graph with dequant Matmul to CLML backend offload."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        """Initializer.
+        Parameters
+        ----------
+        target : tvm.target.Target
+            Target device.
+        """
+        self.target = target
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """Apply required passed to transform"""
+
+        if "adreno" in self.target.keys and (clml_sdk_version() >= 5):
+            mod = tvm.transform.Sequential(
+                [
+                    transform.Normalize(),
+                    transform.FuseOpsByPattern(clml_llm_patterns, annotate_codegen=True),
+                    transform.RunCodegen(),
+                ]
+            )(mod)
+
+        return mod

python/tvm/relax/backend/adreno/pipeline.py

@@ -0,0 +1,138 @@
+# 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.
+"""The Relax Adreno GPU backend compilation pipeline and other passes."""
+
+import tvm
+from tvm import relax
+from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
+
+
+def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default library dispatch passes for Adreno GPU backend."""
+    if "clml" in target.keys:
+        return [
+            relax.backend.adreno.clml.OpenCLMLOffLoadForLLM(target),
+            relax.backend.adreno.clml.OpenCLMLOffLoad(),
+        ]
+    else:
+        return []
+
+
+def legalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default legalization passes for Adreno GPU backend."""
+    desired_layouts = {"relax.nn.conv2d": ["NCHW4c", "OIHW4o", "NCHW4c"]}
+    skip_ops = [
+        "relax.nn.conv2d",
+        "relax.nn.max_pool2d",
+        "relax.nn.adaptive_avg_pool2d",
+    ]
+    pass_list = []
+
+    pass_list.extend(
+        [
+            tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)),
+            relax.transform.DecomposeOpsForInference(),
+        ]
+    )
+    if "texture" in target.keys:
+        pass_list.extend(
+            [
+                relax.transform.ConvertLayout(desired_layouts),
+                relax.transform.Normalize(),
+                relax.transform.FoldConstant(),
+                relax.transform.LegalizeOps(skip_ops=skip_ops),
+                relax.transform.AnnotateTIROpPattern(),
+                relax.backend.adreno.transform.AnnotateCustomMemoryScope(target),
+            ]
+        )
+    pass_list.extend([tvm.relax.transform.LegalizeOps()])
+    if "texture" in target.keys:
+        pass_list.extend(
+            [
+                relax.transform.LegalizeOps(
+                    {"relax.nn.conv2d": legalize_adreno.conv2d_NCHWc_OIHWo},
+                )
+            ]
+        )
+
+    pass_list.extend(
+        [
+            relax.transform.AnnotateTIROpPattern(),
+            relax.transform.FoldConstant(),
+            relax.transform.FuseOps(),
+            relax.transform.FuseTIR(),
+            relax.transform.DeadCodeElimination(),
+        ]
+    )
+    if "texture" in target.keys:
+        pass_list.extend(
+            [
+                relax.backend.adreno.transform.FoldVDeviceScopeChange(),
+                relax.transform.DeadCodeElimination(),
+                relax.transform.SpecializePrimFuncBasedOnCallSite(),
+            ]
+        )
+    from tvm.s_tir import dlight as dl  # pylint: disable=import-outside-toplevel
+
+    pass_list.extend([relax.transform.Normalize()])
+    pass_list.extend(
+        [
+            dl.ApplyDefaultSchedule(
+                dl.adreno.Conv2d(),
+                dl.adreno.LayoutTransform(),
+                dl.adreno.Pool2D(),
+            )
+        ]
+    )
+    pass_list.extend(
+        [
+            dl.ApplyDefaultSchedule(
+                dl.gpu.Reduction(),
+                dl.gpu.GeneralReduction(),
+                dl.gpu.Fallback(),
+            )
+        ]
+    )
+    return pass_list
+
+
+def dataflow_lower_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default dataflow lowering passes for Adreno GPU backend."""
+    return relax.backend.gpu_generic.dataflow_lower_passes(target)
+
+
+def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default finalization passes for Adreno GPU backend."""
+    return relax.backend.gpu_generic.finalize_passes(target)
+
+
+def get_default_pipeline(target: tvm.target.Target):
+    """Return the default compilation pipeline for Adreno GPU."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+        with target:
+            seq = tvm.transform.Sequential(
+                library_dispatch_passes(target)
+                + legalize_passes(target)
+                + dataflow_lower_passes(target)
+                + finalize_passes(target)
+            )
+            mod = seq(mod)
+        return mod
+
+    return _pipeline

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

@@ -0,0 +1,23 @@
+# 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.
+"""Adreno Relax transformations."""
+
+from .transform import (
+    AnnotateCustomMemoryScope,
+    FoldVDeviceScopeChange,
+)

python/tvm/relax/backend/adreno/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 Adreno transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.backend.adreno.transform", __name__)

python/tvm/relax/backend/adreno/transform/transform.py

@@ -0,0 +1,49 @@
+# 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
+"""Adreno Relax transformation passes."""
+
+import tvm.ir
+from tvm.target import Target
+
+from . import _ffi_api
+
+
+def AnnotateCustomMemoryScope(target: Target | None = None) -> tvm.ir.transform.Pass:
+    """Allocate the memory scope information. This is Adreno specific pass to annotate
+    The memory scope information and realize the same with RealizeVDevice pass followed by
+    updating the Prim Function var_buffer mapping using SpecializePrimFuncBasedOnCallSite.
+
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+        The registered pass for allocating workspace.
+    """
+    return _ffi_api.AnnotateCustomMemoryScope(target)  # type: ignore
+
+
+def FoldVDeviceScopeChange() -> tvm.ir.transform.Pass:
+    """This pass is a texture specific pass that can optimize unnecessary to_device copies.
+    Like texture_scope -> ToVDevice -> global scope. In this case the producer can directly
+    store into global scope avoiding unnecessary device copy.
+
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+        The registered pass for allocating workspace.
+    """
+    return _ffi_api.FoldVDeviceScopeChange()  # type: ignore

python/tvm/relax/backend/contrib/__init__.py

@@ -0,0 +1,17 @@
+# 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 backends contrib"""

python/tvm/relax/backend/contrib/example_npu/README.md

@@ -0,0 +1,243 @@
+<!--- 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. -->
+
+# Example NPU Backend
+
+A hands-on example showing how to build a Neural Processing Unit (NPU) backend for TVM's Relax framework using Bring Your Own Codegen (BYOC).
+
+## Context
+
+NPUs are purpose-built accelerators designed around a fixed set of operations common in neural network inference, such as matrix multiplication, convolution, and activation functions. This example shows the architectural patterns you will encounter when building real NPU backends, making it easier to adapt to specific hardware like:
+
+- Mobile NPUs (AMD XDNA, Google Edge TPU, Samsung NPU)
+- Dedicated AI chips (Intel Movidius, Qualcomm Hexagon, MediaTek APU)
+- Cloud AI accelerators (AWS Inferentia, Google TPU, Microsoft Azure Maia)
+- Custom ASIC designs and embedded AI processors
+
+## What This Is
+
+This is an educational template that demonstrates real NPU concepts without requiring actual NPU hardware. It shows developers how to:
+
+- **Pattern-based partitioning**: Identify and group operations that should run on specialized hardware
+- **Memory hierarchy management**: Handle different memory tiers (L0/L1/L2/L3) common in NPUs
+- **Automatic tiling**: Break large tensors into smaller chunks that fit in on-chip memory
+- **Quantization support**: Handle different data precisions efficiently
+- **BYOC integration**: Connect custom backends to TVM's compilation pipeline
+
+## Building TVM with Example NPU Support
+
+Add the following flags when configuring TVM with CMake:
+
+```bash
+cmake -DUSE_EXAMPLE_NPU_CODEGEN=ON -DUSE_EXAMPLE_NPU_RUNTIME=ON ..
+```
+
+Or set them in your `config.cmake`:
+
+```cmake
+set(USE_EXAMPLE_NPU_CODEGEN ON)
+set(USE_EXAMPLE_NPU_RUNTIME ON)
+```
+
+## Quick Start
+
+```python
+import tvm
+from tvm import relax
+from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
+from tvm.relax.transform import FuseOpsByPattern, RunCodegen
+
+# Import to register patterns
+import tvm.relax.backend.contrib.example_npu
+
+# Get available patterns
+patterns = get_patterns_with_prefix("example_npu")
+print(f"Available patterns: {[p.name for p in patterns]}")
+
+# Your model gets automatically partitioned
+# Operations matching patterns get fused into "Composite" functions
+# Those get lowered to the example NPU backend
+```
+
+The snippet above shows how to discover registered patterns. A minimal runnable example that demonstrates the BYOC flow (partition -> merge -> codegen) looks like this:
+
+```python
+import tvm
+from tvm import relax
+from tvm.script import relax as R
+from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
+import tvm.relax.backend.contrib.example_npu  # registers patterns
+
+
+@tvm.script.ir_module
+class MatmulReLU:
+    @R.function
+    def main(
+        x: R.Tensor((2, 4), "float32"),
+        w: R.Tensor((4, 8), "float32"),
+    ) -> R.Tensor((2, 8), "float32"):
+        with R.dataflow():
+            y = relax.op.matmul(x, w)
+            z = relax.op.nn.relu(y)
+            R.output(z)
+        return z
+
+
+mod = MatmulReLU
+patterns = get_patterns_with_prefix("example_npu")
+
+# Apply partitioning and codegen annotation
+mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
+mod = MergeCompositeFunctions()(mod)
+mod = RunCodegen()(mod)
+
+print(mod)
+```
+
+A compact visualization of the BYOC flow:
+
+```
+Model source (Relax)
+	│
+	▼
+Pattern-based partition (FuseOpsByPattern)
+	│
+	▼
+Composite functions (MergeCompositeFunctions)
+	│
+	▼
+Lower/Codegen for example NPU (RunCodegen / relax.ext.example_npu)
+	│
+	▼
+Runtime dispatch to NPU runtime (runtime.ExampleNPUJSONRuntimeCreate)
+```
+
+## Supported Operations
+
+The backend recognizes these common neural network patterns:
+
+### Core Operations
+- `example_npu.dense` - Dense/fully connected layers
+- `example_npu.matmul` - Matrix multiplication operations
+- `example_npu.conv1d` - 1D convolution for sequence processing
+- `example_npu.conv2d` - 2D convolution for image processing
+- `example_npu.depthwise_conv2d` - Depthwise separable convolutions
+- `example_npu.max_pool2d` - 2D max pooling
+- `example_npu.avg_pool2d` - 2D average pooling
+- `example_npu.batch_norm` - Batch normalization
+- `example_npu.softmax` - Softmax
+- `example_npu.add` - Element-wise addition
+- `example_npu.multiply` - Element-wise multiplication
+- `example_npu.subtract` - Element-wise subtraction
+- `example_npu.divide` - Element-wise division
+- `example_npu.relu` - ReLU activation
+- `example_npu.gelu` - Gaussian Error Linear Unit
+- `example_npu.quantize` - Quantization
+- `example_npu.dequantize` - Dequantization
+
+### Build-dependent Operations
+These patterns are registered only when the corresponding Relax op is present
+in the TVM build:
+- `example_npu.relu6` - ReLU6 activation (`relax.nn.relu6`)
+- `example_npu.sigmoid` - Sigmoid activation (`relax.nn.sigmoid`)
+- `example_npu.tanh` - Hyperbolic tangent (`relax.nn.tanh`)
+
+### Fused Patterns
+- `example_npu.conv2d_relu_fused` - Optimized Conv2D+ReLU fusion
+
+
+## Files
+
+### Backend Implementation
+- `patterns.py` - Defines which operations get fused together, along with pattern metadata and architectural annotations used by the partitioner. Includes operator availability checking and NPU-specific constraints.
+- `__init__.py` - Registers the backend and its BYOC entry points with TVM so the compiler can discover and use the example NPU.
+
+### Runtime Implementation
+- `src/runtime/extra/contrib/example_npu/example_npu_runtime.cc` - C++ runtime implementation that handles JSON-based graph execution for the NPU backend.
+
+### Tests and Examples
+- `tests/python/contrib/test_example_npu.py` - Comprehensive test suite containing example IRModules (e.g. `MatmulReLU`, `Conv2dReLU`) and demonstrating the complete BYOC flow from pattern registration to runtime execution.
+
+## Status / Build
+
+- The example backend is an educational, CPU-backed emulation. It does not require real NPU hardware.
+- Tests are skipped automatically when the example codegen/runtime are not built into TVM. The test checks for the presence of these global functions before running:
+
+```python
+import tvm
+has_codegen = tvm.get_global_func("relax.ext.example_npu", True)
+has_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
+has_example_npu = has_codegen and has_runtime
+```
+
+If `has_example_npu` is False, tests are skipped. This ensures compatibility across different TVM build configurations.
+
+## Testing
+
+Run the tests to see it in action:
+
+```bash
+pytest tests/python/contrib/test_example_npu.py -v
+```
+
+Tests are skipped if the backend isn't built — see the test file for the exact runtime/codegen checks.
+
+The test suite includes:
+- Pattern registration verification (checks that core patterns are available)
+- Graph partitioning validation (ensures operations get grouped correctly)
+- End-to-end execution testing (verifies runtime integration)
+- Build-dependent pattern verification (confirms build-dependent ops register when present)
+
+### Example output
+
+When you run the quick-start snippet or the test, you should see output similar to the following (truncated for brevity):
+
+```
+Available patterns: ['example_npu.dense', 'example_npu.matmul', 'example_npu.conv1d', 'example_npu.conv2d', 'example_npu.depthwise_conv2d', 'example_npu.max_pool2d', 'example_npu.avg_pool2d', 'example_npu.batch_norm', 'example_npu.relu', 'example_npu.add', 'example_npu.multiply', 'example_npu.conv2d_relu_fused']
+
+Relax IRModule
+def @main(...) -> ...
+	%0 = call_extern("relax.ext.example_npu", ...)
+
+# composite functions
+def @composite_0(...) /* Composite */ = ...
+```
+
+This shows the registered patterns and that matched subgraphs were turned into composite functions and lowered to the example NPU codegen/runtime.
+
+## Key Features Demonstrated
+
+### NPU Architectural Concepts
+- **Multi-tier memory hierarchy**: SRAM (256KB), CMX (512KB), and DRAM management
+- **Tiling constraints**: 32x32 tiles with 16-element vectors for optimal NPU utilization
+- **Quantization support**: INT8/INT16 for inference acceleration, mixed precision handling
+- **Specialized execution units**: Matrix engines (16x16), vector units (64-wide), pooling units
+- **Power management**: Support for different power modes (high_performance, balanced, low_power)
+
+### Pattern Matching Features
+- **Memory constraint hooks**: Placeholder checks where a real backend would reject tensors that exceed on-chip memory; the example accepts all
+- **Fusion opportunities**: Identifies conv+activation and other beneficial fusions
+- **Layout preferences**: NHWC channel-last layouts preferred by NPUs
+
+### Error Handling
+- **Robust exception handling**: Catches specific exception types instead of generic exceptions
+- **Comprehensive testing**: Validates both successful cases and error conditions
+
+## Learn More
+
+This backend serves as both a working example and educational resource for understanding NPU integration patterns. The implementation demonstrates vendor-neutral concepts that apply across different NPU architectures, making it a valuable starting point for real NPU backend development.

python/tvm/relax/backend/contrib/example_npu/__init__.py

@@ -0,0 +1,31 @@
+# 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.
+"""
+Example NPU Backend for BYOC Integration
+
+This module provides an educational example of how to implement
+a custom NPU backend in TVM using the Bring Your Own Codegen (BYOC)
+framework. It demonstrates key NPU architectural concepts including
+memory hierarchy, tiling, quantization, and operation fusion.
+
+The patterns module registers all supported NPU operations and their
+constraints, making them available for graph partitioning.
+"""
+
+from . import patterns
+
+__all__ = ["patterns"]

python/tvm/relax/backend/contrib/example_npu/patterns.py

@@ -0,0 +1,543 @@
+# 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.
+"""
+Example NPU Pattern Table with Architectural Concepts
+
+This module demonstrates NPU-specific architectural patterns that are common
+across different NPU vendors, including memory hierarchy, quantization,
+tiling, and fusion strategies.
+"""
+
+from typing import ClassVar
+
+from tvm.ir import Op
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.transform import PatternCheckContext
+
+from ...pattern_registry import register_patterns
+
+
+# NPU-specific configuration constants (vendor-neutral)
+class NPUConfig:
+    """NPU architectural parameters common across vendors"""
+
+    # Memory hierarchy sizes (in KB) - typical NPU values
+    SRAM_SIZE_KB = 256  # On-chip SRAM/scratchpad
+    CMX_SIZE_KB = 512  # Compute memory (near compute units)
+
+    # Tiling constraints
+    TILE_HEIGHT = 32
+    TILE_WIDTH = 32
+    VECTOR_SIZE = 16
+
+    # Supported data types for NPU acceleration
+    SUPPORTED_DTYPES: ClassVar[list[str]] = ["int8", "int16", "float16", "float32"]
+    QUANTIZED_DTYPES: ClassVar[list[str]] = ["int8", "int16"]
+
+    # NPU execution units
+    MATRIX_ENGINE_SIZE = 16  # MxN matrix engine
+    VECTOR_ENGINE_WIDTH = 64  # Vector processing width
+
+    # Power modes
+    POWER_MODES: ClassVar[list[str]] = ["high_performance", "balanced", "low_power"]
+
+
+def _check_npu_memory_constraints(
+    context: PatternCheckContext,  # pylint: disable=unused-argument
+) -> bool:
+    """
+    Placeholder for NPU memory hierarchy constraint checking.
+
+    A real implementation would inspect the annotated expression's
+    TensorType to verify the tensor fits within the NPU's
+    on-chip SRAM (L1) or compute memory (L2/CMX). Tensors that
+    exceed on-chip capacity require tiling before offload.
+    """
+    return True
+
+
+def _check_npu_quantization(
+    context: PatternCheckContext,  # pylint: disable=unused-argument
+) -> bool:
+    """
+    Placeholder for NPU quantization requirement checking.
+
+    A real implementation would verify the op's dtype falls within
+    the set supported by the NPU (e.g. int8, int16, float16, float32)
+    and reject ops with unsupported dtypes so they fall back to CPU.
+    """
+    return True
+
+
+def conv2d_relu_fused_pattern():
+    """
+    NPU-optimized Conv2D+ReLU fusion pattern.
+
+    This is a key NPU optimization - fusing convolution with activation
+    avoids memory traffic between operations.
+    """
+
+    def _make_conv2d_relu_pattern():
+        input_tensor = wildcard()
+        weight = wildcard()
+        conv = is_op("relax.nn.conv2d")(input_tensor, weight)
+        relu = is_op("relax.nn.relu")(conv)
+
+        annotations = {
+            "input": input_tensor,
+            "weight": weight,
+            "conv": conv,
+            "root": relu,
+        }
+        return relu, annotations
+
+    def _check_conv2d_relu(context: PatternCheckContext) -> bool:
+        """Check if Conv2D+ReLU fusion is beneficial for NPU"""
+        if not _check_npu_memory_constraints(context):
+            return False
+        if not _check_npu_quantization(context):
+            return False
+        return True
+
+    return ("example_npu.conv2d_relu_fused", *_make_conv2d_relu_pattern(), _check_conv2d_relu)
+
+
+def matmul_relu_fused_pattern():
+    """
+    NPU-optimized MatMul+ReLU fusion pattern.
+
+    Fusing the matrix engine output with the activation unit avoids a
+    write/read round-trip through L1 SRAM, mirroring the conv2d+relu
+    fusion below.
+    """
+
+    def _make_matmul_relu_pattern():
+        input_tensor = wildcard()
+        weight = wildcard()
+        matmul = is_op("relax.matmul")(input_tensor, weight)
+        relu = is_op("relax.nn.relu")(matmul)
+
+        annotations = {
+            "input": input_tensor,
+            "weight": weight,
+            "matmul": matmul,
+            "root": relu,
+        }
+        return relu, annotations
+
+    def _check_matmul_relu(context: PatternCheckContext) -> bool:
+        """Check if MatMul+ReLU fusion is beneficial for NPU"""
+        if not _check_npu_memory_constraints(context):
+            return False
+        if not _check_npu_quantization(context):
+            return False
+        return True
+
+    return ("example_npu.matmul_relu_fused", *_make_matmul_relu_pattern(), _check_matmul_relu)
+
+
+def matmul_patterns():
+    """
+    NPU-optimized matrix multiplication patterns.
+
+    NPUs typically have dedicated matrix engines (systolic arrays,
+    tensor cores) that require specific layouts and sizes.
+    """
+
+    def _make_matmul_pattern():
+        input_tensor = wildcard()
+        weight = wildcard()
+        output = is_op("relax.matmul")(input_tensor, weight)
+
+        annotations = {
+            "input": input_tensor,
+            "weight": weight,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_matmul(context: PatternCheckContext) -> bool:
+        """Check if matmul can use NPU matrix engine"""
+        return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
+
+    def _matmul_pattern(pattern_name):
+        return (pattern_name, *_make_matmul_pattern(), _check_matmul)
+
+    # Register both common names used for matrix multiplication in patterns/tests
+    return [
+        _matmul_pattern("example_npu.dense"),
+        _matmul_pattern("example_npu.matmul"),
+    ]
+
+
+def conv1d_patterns():
+    """
+    1D Convolution patterns optimized for NPU execution.
+
+    NPUs handle 1D convolution by mapping to 2D operations
+    or using specialized 1D processing units.
+    """
+
+    def _make_conv1d_pattern():
+        input_tensor = wildcard()
+        weight = wildcard()
+        output = is_op("relax.nn.conv1d")(input_tensor, weight)
+
+        annotations = {
+            "input": input_tensor,
+            "weight": weight,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_conv1d(context: PatternCheckContext) -> bool:
+        """Check if conv1d can use NPU vector engine"""
+        return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
+
+    def _conv1d_pattern(pattern_name):
+        return (pattern_name, *_make_conv1d_pattern(), _check_conv1d)
+
+    return [_conv1d_pattern("example_npu.conv1d")]
+
+
+def conv2d_patterns():
+    """
+    2D Convolution patterns with NPU tiling and memory management.
+
+    2D convolution is the most important NPU operation, with
+    dedicated hardware for efficient processing.
+    """
+
+    def _make_conv2d_pattern():
+        input_tensor = wildcard()
+        weight = wildcard()
+        output = is_op("relax.nn.conv2d")(input_tensor, weight)
+
+        annotations = {
+            "input": input_tensor,
+            "weight": weight,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_conv2d(context: PatternCheckContext) -> bool:
+        """Check conv2d NPU constraints"""
+        return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
+
+    def _conv2d_pattern(pattern_name):
+        return (pattern_name, *_make_conv2d_pattern(), _check_conv2d)
+
+    return [_conv2d_pattern("example_npu.conv2d")]
+
+
+def depthwise_conv2d_patterns():
+    """
+    Depthwise convolution - critical for mobile NPUs.
+
+    Many NPUs have specialized units for depthwise operations
+    used in MobileNet-style architectures.
+    """
+
+    def _make_depthwise_pattern():
+        input_tensor = wildcard()
+        weight = wildcard()
+        output = is_op("relax.nn.conv2d")(input_tensor, weight)
+
+        annotations = {
+            "input": input_tensor,
+            "weight": weight,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_depthwise(context: PatternCheckContext) -> bool:
+        """Check if this is a depthwise conv that NPU can accelerate"""
+        conv_call = context.annotated_expr["root"]
+        # groups > 1 distinguishes depthwise/grouped conv from standard conv2d.
+        # True depthwise has groups == in_channels; we accept any grouped variant
+        # here since the NPU's depthwise unit handles all grouped convolutions.
+        if conv_call.attrs.groups <= 1:
+            return False
+        return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
+
+    return [("example_npu.depthwise_conv2d", *_make_depthwise_pattern(), _check_depthwise)]
+
+
+def pooling_patterns():
+    """
+    Pooling operations with NPU memory streaming.
+
+    NPUs often process pooling with the convolution engine
+    or dedicated pooling units.
+    """
+
+    def _make_maxpool2d_pattern():
+        input_tensor = wildcard()
+        output = is_op("relax.nn.max_pool2d")(input_tensor)
+
+        annotations = {
+            "input": input_tensor,
+            "root": output,
+        }
+        return output, annotations
+
+    def _make_avgpool2d_pattern():
+        input_tensor = wildcard()
+        output = is_op("relax.nn.avg_pool2d")(input_tensor)
+
+        annotations = {
+            "input": input_tensor,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_pooling(context: PatternCheckContext) -> bool:
+        """Check pooling NPU constraints"""
+        return _check_npu_memory_constraints(context)
+
+    return [
+        ("example_npu.max_pool2d", *_make_maxpool2d_pattern(), _check_pooling),
+        ("example_npu.avg_pool2d", *_make_avgpool2d_pattern(), _check_pooling),
+    ]
+
+
+def batch_norm_patterns():
+    """
+    Batch normalization - often fused with conv on NPUs.
+
+    NPUs typically fuse BN into convolution to avoid
+    separate memory passes.
+    """
+
+    def _make_batch_norm_pattern():
+        input_tensor = wildcard()
+        gamma = wildcard()
+        beta = wildcard()
+        moving_mean = wildcard()
+        moving_var = wildcard()
+
+        output = is_op("relax.nn.batch_norm")(input_tensor, gamma, beta, moving_mean, moving_var)
+
+        annotations = {
+            "input": input_tensor,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_batch_norm(context: PatternCheckContext) -> bool:
+        """Check if batch norm should be offloaded or fused"""
+        return _check_npu_quantization(context)
+
+    return [("example_npu.batch_norm", *_make_batch_norm_pattern(), _check_batch_norm)]
+
+
+def softmax_patterns():
+    """
+    Softmax - used in classification heads and attention mechanisms.
+
+    NPUs typically implement softmax via dedicated hardware or
+    a combination of exp, sum, and divide operations.
+    """
+
+    def _make_softmax_pattern():
+        input_tensor = wildcard()
+        output = is_op("relax.nn.softmax")(input_tensor)
+
+        annotations = {
+            "input": input_tensor,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_softmax(context: PatternCheckContext) -> bool:
+        """Check if softmax can use NPU activation unit"""
+        return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
+
+    patterns = []
+    try:
+        Op.get("relax.nn.softmax")
+        patterns.append(("example_npu.softmax", *_make_softmax_pattern(), _check_softmax))
+    except (KeyError, AttributeError):
+        pass
+
+    return patterns
+
+
+def activation_patterns():
+    """
+    NPU activation functions with specialized hardware.
+
+    NPUs have dedicated activation units that can handle
+    various functions efficiently.
+    """
+
+    def _make_activation_pattern(op_name: str):
+        def _pattern():
+            input_tensor = wildcard()
+            output = is_op(op_name)(input_tensor)
+
+            annotations = {
+                "input": input_tensor,
+                "root": output,
+            }
+            return output, annotations
+
+        return _pattern
+
+    def _check_activation(context: PatternCheckContext) -> bool:
+        """Check if activation can use NPU activation unit"""
+        return _check_npu_quantization(context)
+
+    activations = [
+        ("example_npu.relu", "relax.nn.relu"),
+        ("example_npu.relu6", "relax.nn.relu6"),
+        ("example_npu.sigmoid", "relax.nn.sigmoid"),
+        ("example_npu.tanh", "relax.nn.tanh"),
+        ("example_npu.gelu", "relax.nn.gelu"),
+    ]
+
+    patterns = []
+    for pattern_name, op_name in activations:
+        try:
+            Op.get(op_name)
+        except (KeyError, AttributeError):
+            continue
+
+        pattern_fn = _make_activation_pattern(op_name)
+        patterns.append((pattern_name, *pattern_fn(), _check_activation))
+
+    return patterns
+
+
+def elementwise_patterns():
+    """
+    Element-wise operations that NPUs can vectorize.
+
+    NPUs process element-wise ops using vector units
+    with SIMD capabilities.
+    """
+
+    def _make_elementwise_pattern(op_name: str):
+        def _pattern():
+            input1 = wildcard()
+            input2 = wildcard()
+            output = is_op(op_name)(input1, input2)
+
+            annotations = {
+                "input1": input1,
+                "input2": input2,
+                "root": output,
+            }
+            return output, annotations
+
+        return _pattern
+
+    def _check_elementwise(context: PatternCheckContext) -> bool:
+        """Check if elementwise op can use NPU vector unit"""
+        return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
+
+    ops = ["relax.add", "relax.multiply", "relax.subtract", "relax.divide"]
+    patterns = []
+    for op in ops:
+        try:
+            Op.get(op)
+        except (KeyError, AttributeError):
+            continue
+
+        op_short = op.split(".")[-1]
+        pattern_fn = _make_elementwise_pattern(op)
+        patterns.append((f"example_npu.{op_short}", *pattern_fn(), _check_elementwise))
+
+    return patterns
+
+
+def quantization_patterns():
+    """
+    Quantization/dequantization patterns for NPU.
+
+    NPUs need explicit quantization boundaries to switch
+    between precision levels.
+    """
+
+    def _make_quantize_pattern():
+        input_tensor = wildcard()
+        output = is_op("relax.quantize")(input_tensor)
+
+        annotations = {
+            "input": input_tensor,
+            "root": output,
+        }
+        return output, annotations
+
+    def _make_dequantize_pattern():
+        input_tensor = wildcard()
+        output = is_op("relax.dequantize")(input_tensor)
+
+        annotations = {
+            "input": input_tensor,
+            "root": output,
+        }
+        return output, annotations
+
+    def _check_quantization(
+        context: PatternCheckContext,  # pylint: disable=unused-argument
+    ) -> bool:
+        """Check quantization operations"""
+        return True
+
+    patterns = []
+
+    try:
+        Op.get("relax.quantize")
+        patterns.append(("example_npu.quantize", *_make_quantize_pattern(), _check_quantization))
+    except (KeyError, AttributeError):
+        pass
+
+    try:
+        Op.get("relax.dequantize")
+        patterns.append(
+            ("example_npu.dequantize", *_make_dequantize_pattern(), _check_quantization)
+        )
+    except (KeyError, AttributeError):
+        pass
+
+    return patterns
+
+
+# Register all NPU patterns with architectural awareness
+# register_patterns priority: patterns that appear LATER in the list win.
+# So we place general / standalone patterns first, and fused (more
+# specific) patterns last so they take precedence over their constituents.
+register_patterns(
+    [
+        *quantization_patterns(),
+        *elementwise_patterns(),
+        *activation_patterns(),
+        *softmax_patterns(),
+        *batch_norm_patterns(),
+        *pooling_patterns(),
+        *matmul_patterns(),
+        *conv1d_patterns(),
+        # Plain conv2d is more general than depthwise (groups>1); list
+        # plain first so depthwise wins on grouped convs.
+        *conv2d_patterns(),
+        *depthwise_conv2d_patterns(),
+        # Fused patterns last (highest priority).
+        matmul_relu_fused_pattern(),
+        conv2d_relu_fused_pattern(),
+    ]
+)

python/tvm/relax/backend/contrib/nnapi.py

@@ -0,0 +1,321 @@
+# 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.
+
+"""Pattern table for NNAPI backend"""
+
+from collections.abc import Mapping
+
+from tvm.ir import IRModule
+from tvm.relax.dpl.pattern import (
+    DFPattern,
+    is_op,
+    wildcard,
+)
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+
+
+def elementwise_binary_patterns() -> list[tuple[str, DFPattern, Mapping[str, DFPattern]]]:
+    """
+    Returns a list of tuples representing elementwise binary operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+
+    def _elementwise_binary_pattern(
+        pattern_name: str,
+        op_name: str,
+    ) -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+        input0 = wildcard()
+        input1 = wildcard()
+
+        pattern = is_op(op_name)(input0, input1)
+
+        return (pattern_name, pattern, {})
+
+    return [
+        _elementwise_binary_pattern("nnapi.add", "relax.add"),
+        _elementwise_binary_pattern("nnapi.mul", "relax.multiply"),
+        _elementwise_binary_pattern("nnapi.div", "relax.divide"),
+        _elementwise_binary_pattern("nnapi.sub", "relax.subtract"),
+        _elementwise_binary_pattern("nnapi.pow", "relax.power"),
+        _elementwise_binary_pattern("nnapi.equal", "relax.equal"),
+        _elementwise_binary_pattern("nnapi.greater", "relax.greater"),
+        _elementwise_binary_pattern("nnapi.greater_equal", "relax.greater_equal"),
+        _elementwise_binary_pattern("nnapi.less", "relax.less"),
+        _elementwise_binary_pattern("nnapi.less_equal", "relax.less_equal"),
+        _elementwise_binary_pattern("nnapi.not_equal", "relax.not_equal"),
+        _elementwise_binary_pattern("nnapi.maximum", "relax.maximum"),
+        _elementwise_binary_pattern("nnapi.minimum", "relax.minimum"),
+    ]
+
+
+def unary_patterns() -> list[tuple[str, DFPattern, Mapping[str, DFPattern]]]:
+    """
+    Returns a list of tuples representing unary operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+
+    def _unary_pattern(
+        pattern_name: str, op_name: str
+    ) -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+        input0 = wildcard()
+        pattern = is_op(op_name)(input0)
+        return (pattern_name, pattern, {})
+
+    return [
+        _unary_pattern("nnapi.floor", "relax.floor"),
+        _unary_pattern("nnapi.relu", "relax.nn.relu"),
+        _unary_pattern("nnapi.logistic", "relax.sigmoid"),
+        _unary_pattern("nnapi.softmax", "relax.nn.softmax"),
+        _unary_pattern("nnapi.tanh", "relax.tanh"),
+        _unary_pattern("nnapi.abs", "relax.abs"),
+        _unary_pattern("nnapi.exp", "relax.exp"),
+        _unary_pattern("nnapi.log", "relax.log"),
+        _unary_pattern("nnapi.neg", "relax.negative"),
+        _unary_pattern("nnapi.cast", "relax.astype"),
+        _unary_pattern("nnapi.sqrt", "relax.sqrt"),
+        _unary_pattern("nnapi.rsqrt", "relax.rsqrt"),
+    ]
+
+
+def matmul_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing matmul operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    input1 = wildcard()
+    pattern = is_op("relax.matmul")(input0, input1)
+    return ("nnapi.batch_matmul", pattern, {})
+
+
+def permute_dims_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing permute operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    pattern = is_op("relax.permute_dims")(input0)
+    return ("nnapi.transpose", pattern, {})
+
+
+def astype_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing astype operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard().has_dtype("float16") | wildcard().has_dtype("float32")
+    pattern = is_op("relax.astype")(input0).has_dtype("float16") | is_op("relax.astype")(
+        input0
+    ).has_dtype("float32")
+
+    return ("nnapi.cast", pattern, {})
+
+
+def mean_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing mean operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    pattern = is_op("relax.mean")(input0)
+
+    return ("nnapi.mean", pattern, {})
+
+
+def conv2d_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing conv2d operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    input1 = wildcard()
+    input2 = wildcard()
+    conv = is_op("relax.nn.conv2d")(input0, input1)
+    pattern = is_op("relax.add")(conv, input2)
+    return ("nnapi.conv2d", pattern, {})
+
+
+def max_pool2d_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing max_pool2d operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    pattern = is_op("relax.nn.max_pool2d")(input0)
+    return ("nnapi.max_pool_2d", pattern, {})
+
+
+register_patterns(
+    [
+        *elementwise_binary_patterns(),
+        *unary_patterns(),
+        matmul_pattern(),
+        permute_dims_pattern(),
+        astype_pattern(),
+        mean_pattern(),
+        conv2d_pattern(),
+        max_pool2d_pattern(),
+    ]
+)
+
+
+def min_feature_level(pattern_name: str) -> int:
+    """
+    Returns the minimum feature level required to support a given NNAPI operation pattern.
+
+    Args:
+        pattern_name (str): The name of the NNAPI operation pattern
+        (e.g., "nnapi.add", "nnapi.conv2d").
+
+    Returns:
+        int: The minimum feature level for the specified pattern, or 1 if the pattern is not found.
+    """
+
+    levels = {
+        "nnapi.add": 1,
+        "nnapi.average_pool_2d": 1,
+        "nnapi.concatenation": 1,
+        "nnapi.conv2d": 1,
+        "nnapi.depthwise_conv_2d": 1,
+        "nnapi.depth_to_space": 1,
+        "nnapi.dequantize": 1,
+        "nnapi.embedding_lookup": 1,
+        "nnapi.floor": 1,
+        "nnapi.fully_connected": 1,
+        "nnapi.hashtable_lookup": 1,
+        "nnapi.l2_normalization": 1,
+        "nnapi.l2_pool_2d": 1,
+        "nnapi.local_response_normalization": 1,
+        "nnapi.logistic": 1,
+        "nnapi.lsh_projection": 1,
+        "nnapi.lstm": 1,
+        "nnapi.max_pool_2d": 1,
+        "nnapi.mul": 1,
+        "nnapi.relu": 1,
+        "nnapi.relu1": 1,
+        "nnapi.relu6": 1,
+        "nnapi.reshape": 1,
+        "nnapi.resize_bilinear": 1,
+        "nnapi.rnn": 1,
+        "nnapi.softmax": 1,
+        "nnapi.space_to_depth": 1,
+        "nnapi.svdf": 1,
+        "nnapi.tanh": 1,
+        "nnapi.batch_to_space_nd": 2,
+        "nnapi.div": 2,
+        "nnapi.mean": 2,
+        "nnapi.pad": 2,
+        "nnapi.space_to_batch_nd": 2,
+        "nnapi.squeeze": 2,
+        "nnapi.strided_slice": 2,
+        "nnapi.sub": 2,
+        "nnapi.transpose": 2,
+        "nnapi.abs": 3,
+        "nnapi.argmax": 3,
+        "nnapi.argmin": 3,
+        "nnapi.axis_aligned_bbox_transform": 3,
+        "nnapi.bidirectional_sequence_lstm": 3,
+        "nnapi.bidirectional_sequence_rnn": 3,
+        "nnapi.box_with_nms_limit": 3,
+        "nnapi.cast": 3,
+        "nnapi.channel_shuffle": 3,
+        "nnapi.detection_postprocessing": 3,
+        "nnapi.equal": 3,
+        "nnapi.exp": 3,
+        "nnapi.expand_dims": 3,
+        "nnapi.gather": 3,
+        "nnapi.generate_proposals": 3,
+        "nnapi.greater": 3,
+        "nnapi.greater_equal": 3,
+        "nnapi.grouped_conv_2d": 3,
+        "nnapi.heatmap_max_keypoint": 3,
+        "nnapi.instance_normalization": 3,
+        "nnapi.less": 3,
+        "nnapi.less_equal": 3,
+        "nnapi.log": 3,
+        "nnapi.logical_and": 3,
+        "nnapi.logical_not": 3,
+        "nnapi.logical_or": 3,
+        "nnapi.log_softmax": 3,
+        "nnapi.maximum": 3,
+        "nnapi.minimum": 3,
+        "nnapi.neg": 3,
+        "nnapi.not_equal": 3,
+        "nnapi.pad_v2": 3,
+        "nnapi.pow": 3,
+        "nnapi.prelu": 3,
+        "nnapi.quantize": 3,
+        "nnapi.quantized_16bit_lstm": 3,
+        "nnapi.random_multinomial": 3,
+        "nnapi.reduce_all": 3,
+        "nnapi.reduce_any": 3,
+        "nnapi.reduce_max": 3,
+        "nnapi.reduce_min": 3,
+        "nnapi.reduce_prod": 3,
+        "nnapi.reduce_sum": 3,
+        "nnapi.roi_align": 3,
+        "nnapi.roi_pooling": 3,
+        "nnapi.rsqrt": 3,
+        "nnapi.select": 3,
+        "nnapi.sin": 3,
+        "nnapi.slice": 3,
+        "nnapi.split": 3,
+        "nnapi.sqrt": 3,
+        "nnapi.tile": 3,
+        "nnapi.topk_v2": 3,
+        "nnapi.transpose_conv_2d": 3,
+        "nnapi.unidirectional_sequence_lstm": 3,
+        "nnapi.unidirectional_sequence_rnn": 3,
+        "nnapi.resize_nearest_neighbor": 3,
+        "nnapi.quantized_lstm": 4,
+        "nnapi.if": 4,
+        "nnapi.while": 4,
+        "nnapi.elu": 4,
+        "nnapi.hard_swish": 4,
+        "nnapi.fill": 4,
+        "nnapi.rank": 4,
+        "nnapi.batch_matmul": 6,
+        "nnapi.pack": 6,
+        "nnapi.mirror_pad": 7,
+        "nnapi.reverse": 7,
+    }
+    return levels[pattern_name]
+
+
+def partition_for_nnapi(mod: IRModule, feature_level: int | None = None) -> IRModule:
+    """Partition the graph greedily offloading supported operators to NNAPI.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The module to run passes on.
+    feature_level : Optional[int]
+        The maximum NNAPI feature level.
+
+    Returns
+    -------
+    mod : tvm.ir.IRModule
+        Annotated and partitioned module.
+    """
+    patterns = get_patterns_with_prefix("nnapi")
+    if feature_level is not None:
+        patterns = [pat for pat in patterns if feature_level >= min_feature_level(pat.name)]
+    mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    mod = MergeCompositeFunctions()(mod)
+    return mod

python/tvm/relax/backend/contrib/tensorrt.py

@@ -0,0 +1,140 @@
+# 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.
+
+"""Pattern table and partitioning for the TensorRT BYOC backend.
+
+The composite name of each pattern is "tensorrt.<op>", matching the runtime
+converter registered under the same name (the converters are keyed by
+"tensorrt." + op_name). ``partition_for_tensorrt`` carves the matched subgraphs
+out of the module and annotates them for the ``tensorrt`` codegen.
+"""
+
+from collections.abc import Mapping
+
+from tvm.ir import IRModule
+from tvm.relax.dpl.pattern import DFPattern, is_op, wildcard
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+
+Pattern = tuple[str, DFPattern, Mapping[str, DFPattern]]
+
+
+def _op_pattern(composite_name: str, op_name: str, num_args: int) -> Pattern:
+    """A pattern matching a single op called with ``num_args`` wildcard arguments."""
+    args = [wildcard() for _ in range(num_args)]
+    return (composite_name, is_op(op_name)(*args), {})
+
+
+def _tensorrt_patterns() -> list[Pattern]:
+    patterns: list[Pattern] = []
+
+    # Activations and unary elementwise ops (single tensor argument).
+    for composite, op in [
+        ("tensorrt.nn.relu", "relax.nn.relu"),
+        ("tensorrt.sigmoid", "relax.sigmoid"),
+        ("tensorrt.tanh", "relax.tanh"),
+        ("tensorrt.exp", "relax.exp"),
+        ("tensorrt.log", "relax.log"),
+        ("tensorrt.sqrt", "relax.sqrt"),
+        ("tensorrt.abs", "relax.abs"),
+        ("tensorrt.negative", "relax.negative"),
+        ("tensorrt.sin", "relax.sin"),
+        ("tensorrt.cos", "relax.cos"),
+        ("tensorrt.atan", "relax.atan"),
+        ("tensorrt.ceil", "relax.ceil"),
+        ("tensorrt.floor", "relax.floor"),
+        ("tensorrt.erf", "relax.erf"),
+        ("tensorrt.nn.softmax", "relax.nn.softmax"),
+        ("tensorrt.nn.batch_flatten", "relax.nn.batch_flatten"),
+        ("tensorrt.expand_dims", "relax.expand_dims"),
+        ("tensorrt.squeeze", "relax.squeeze"),
+        ("tensorrt.transpose", "relax.permute_dims"),
+        ("tensorrt.layout_transform", "relax.layout_transform"),
+        ("tensorrt.nn.max_pool2d", "relax.nn.max_pool2d"),
+        ("tensorrt.nn.avg_pool2d", "relax.nn.avg_pool2d"),
+        ("tensorrt.nn.max_pool3d", "relax.nn.max_pool3d"),
+        ("tensorrt.nn.avg_pool3d", "relax.nn.avg_pool3d"),
+        ("tensorrt.nn.adaptive_avg_pool2d", "relax.nn.adaptive_avg_pool2d"),
+        ("tensorrt.sum", "relax.sum"),
+        ("tensorrt.prod", "relax.prod"),
+        ("tensorrt.max", "relax.max"),
+        ("tensorrt.min", "relax.min"),
+        ("tensorrt.mean", "relax.mean"),
+        ("tensorrt.concatenate", "relax.concat"),
+        ("tensorrt.split", "relax.split"),
+    ]:
+        patterns.append(_op_pattern(composite, op, 1))
+
+    # Binary elementwise ops (two tensor arguments).
+    for composite, op in [
+        ("tensorrt.add", "relax.add"),
+        ("tensorrt.subtract", "relax.subtract"),
+        ("tensorrt.multiply", "relax.multiply"),
+        ("tensorrt.divide", "relax.divide"),
+        ("tensorrt.power", "relax.power"),
+        ("tensorrt.maximum", "relax.maximum"),
+        ("tensorrt.minimum", "relax.minimum"),
+    ]:
+        patterns.append(_op_pattern(composite, op, 2))
+
+    # Convolutions and matmul (data + weight).
+    for composite, op in [
+        ("tensorrt.nn.conv1d", "relax.nn.conv1d"),
+        ("tensorrt.nn.conv2d", "relax.nn.conv2d"),
+        ("tensorrt.nn.conv3d", "relax.nn.conv3d"),
+        ("tensorrt.nn.conv2d_transpose", "relax.nn.conv2d_transpose"),
+        ("tensorrt.nn.conv3d_transpose", "relax.nn.conv3d_transpose"),
+        ("tensorrt.nn.batch_matmul", "relax.matmul"),
+        ("tensorrt.reshape", "relax.reshape"),
+    ]:
+        patterns.append(_op_pattern(composite, op, 2))
+
+    # layer_norm (data, gamma, beta) and clip (data, min, max).
+    patterns.append(_op_pattern("tensorrt.nn.layer_norm", "relax.nn.layer_norm", 3))
+    patterns.append(_op_pattern("tensorrt.clip", "relax.clip", 3))
+
+    # strided_slice is called either with or without the optional strides argument.
+    patterns.append(_op_pattern("tensorrt.strided_slice", "relax.strided_slice", 5))
+    patterns.append(_op_pattern("tensorrt.strided_slice", "relax.strided_slice", 4))
+
+    return patterns
+
+
+register_patterns(_tensorrt_patterns())
+
+
+def partition_for_tensorrt(mod: IRModule) -> IRModule:
+    """Partition the module, offloading TensorRT-supported subgraphs.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The module to partition. Bind model parameters (e.g. via
+        ``relax.transform.BindParams``) before calling this so that weights are
+        available to TensorRT as constants.
+
+    Returns
+    -------
+    mod : tvm.ir.IRModule
+        The module with TensorRT-supported subgraphs grouped into composite
+        functions annotated for the ``tensorrt`` codegen.
+    """
+    patterns = get_patterns_with_prefix("tensorrt")
+    mod = FuseOpsByPattern(patterns, bind_constants=True, annotate_codegen=False)(mod)
+    mod = MergeCompositeFunctions()(mod)
+    return mod

python/tvm/relax/backend/cpu_generic/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+"""The Relax CPU backend compilation pipeline and other passes."""
+
+from .pipeline import (
+    finalize_passes,
+    get_default_pipeline,
+    legalize_passes,
+    library_dispatch_passes,
+)

python/tvm/relax/backend/cpu_generic/pipeline.py

@@ -0,0 +1,80 @@
+# 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.
+"""The Relax CPU backend compilation pipeline and other passes."""
+
+import tvm
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default library dispatch passes for CPU backend."""
+    return [
+        relax.backend.DispatchSampling(),
+        relax.backend.DispatchSortScan(),
+    ]
+
+
+def legalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default legalization passes for CPU backend."""
+    return [
+        tvm.relax.transform.LegalizeOps(),
+        tvm.relax.transform.AnnotateTIROpPattern(),
+        tvm.relax.transform.FoldConstant(),
+        tvm.relax.transform.FuseOps(),
+        tvm.relax.transform.FuseTIR(),
+    ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default dataflow lowering passes for CPU backend."""
+    return [
+        relax.transform.RewriteDataflowReshape(),
+        relax.transform.ToNonDataflow(),
+        relax.transform.RemovePurityChecking(),
+        relax.transform.CallTIRRewrite(),
+    ]
+
+
+def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default finalization passes for CPU backend."""
+    return [
+        relax.transform.StaticPlanBlockMemory(),
+        relax.transform.LowerAllocTensor(),
+        relax.transform.KillAfterLastUse(),
+        relax.transform.LowerRuntimeBuiltin(),
+        relax.transform.ComputePrimValue(),
+        relax.transform.VMShapeLower(),
+        relax.transform.AttachGlobalSymbol(),
+    ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+    """Return the default compilation pipeline for CPU."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+        with target:
+            seq = tvm.transform.Sequential(
+                library_dispatch_passes(target)
+                + legalize_passes(target)
+                + dataflow_lower_passes(target)
+                + finalize_passes(target)
+            )
+            mod = seq(mod)
+        return mod
+
+    return _pipeline

python/tvm/relax/backend/cuda/__init__.py

@@ -0,0 +1,26 @@
+# 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.
+"""The Relax CUDA backend compilation pipeline and other passes."""
+
+from . import flashinfer
+from .pipeline import (
+    finalize_passes,
+    get_default_pipeline,
+    legalize_passes,
+    library_dispatch_passes,
+)

python/tvm/relax/backend/cuda/cublas.py

@@ -0,0 +1,245 @@
+# 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.
+
+"""Pattern table for cuBLAS backend"""
+
+import operator
+from functools import reduce
+
+import tvm
+from tvm import DataType
+from tvm.arith import Analyzer
+from tvm.relax import transform
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import (
+    make_matmul_dequantize_pattern,
+    make_matmul_multiply_pattern,
+    make_matmul_pattern,
+)
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
+    """Check if dtypes in the given workload are supported by cuBLAS BYOC."""
+    if lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
+        # The output cannot be 'float8_e5m2' if inputs are 'float8_e4m3fn'
+        return out_dtype != "float8_e5m2"
+    return (
+        (lhs_dtype == "float16" and rhs_dtype == "float16")
+        or (lhs_dtype == "float32" and rhs_dtype == "float32")
+        or (lhs_dtype == "int8" and rhs_dtype == "int8")
+        or (lhs_dtype == "bfloat16" and rhs_dtype == "bfloat16")
+    )
+
+
+def _check_matmul(context: PatternCheckContext) -> bool:
+    if has_leaking_intermediate_variables(context):
+        return False
+    lhs = context.annotated_expr["lhs"]
+    rhs = context.annotated_expr["rhs"]
+    matmul_call = context.annotated_expr["root"]
+
+    if "scale" in context.annotated_expr and "zp" in context.annotated_expr:
+        scale = context.annotated_expr["scale"]
+        zero_point = context.annotated_expr["zp"]
+        # Only scalar values for scale and zero_point are supported.
+        if scale.ty.ndim != 0 or zero_point.ty.ndim != 0:
+            return False
+        # Only zero_point == 0.0 is supported.
+        if zero_point.data.numpy()[()].item() != 0.0:
+            return False
+
+    lhs_dtype = lhs.ty.dtype
+    rhs_dtype = rhs.ty.dtype
+    out_dtype = matmul_call.ty.dtype
+    if not _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
+        return False
+
+    lhs_shape = lhs.ty.shape.values
+    rhs_shape = rhs.ty.shape.values
+
+    if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+        # Reduction axis must be constant
+        return False
+
+    if lhs_dtype == "int8" and rhs_dtype == "int8":
+        if lhs_shape[-1] % 4 != 0:
+            # Reduction axis must be multiples of 4 for IGEMM
+            return False
+        if not isinstance(rhs_shape[-1], tvm.tirx.expr.IntImm | int) or rhs_shape[-1] % 4 != 0:
+            # Rows number must be multiples of 4 for IGEMM
+            return False
+    elif lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
+        matmul_rhs_var = matmul_call.args[1]
+        rhs_transposed = False
+        if matmul_rhs_var in context.matched_bindings:
+            matmul_rhs_call = context.matched_bindings[matmul_rhs_var]
+            assert (
+                isinstance(matmul_rhs_call, tvm.relax.Call)
+                and matmul_rhs_call.op.name == "relax.permute_dims"
+            )
+            rhs_transposed = True
+
+        if not rhs_transposed:
+            # cuBLAS FP8 operations require rhs being transposed
+            return False
+
+        # cuBLAS FP8 operations require all tensors being aligned to 16 bytes.
+        if (
+            not isinstance(rhs_shape[-1], tvm.tirx.expr.IntImm | int)
+            or rhs_shape[-1] % (16 // DataType(lhs_dtype).itemsize) != 0
+        ):
+            return False
+        if (
+            not isinstance(rhs_shape[-2], tvm.tirx.expr.IntImm | int)
+            or rhs_shape[-2] % (16 // DataType(out_dtype).itemsize) != 0
+        ):
+            return False
+
+    lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
+    rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
+
+    if "bias" in context.annotated_expr:
+        if lhs_dtype == "int8" and rhs_dtype == "int8":
+            # Non-default epilogue not supported for IGEMM
+            return False
+        bias = context.annotated_expr["bias"]
+        bias_shape = bias.ty.shape.values
+        bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
+        if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
+            # cuBLAS only supports bias vector
+            return False
+
+    analyzer = Analyzer()
+
+    # cuBLASLt does not seem to support batched GEMM with one of matrices having
+    # one batch (with batch_stride 0). So for batched GEMM, the two batch counts
+    # must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
+    # flattening all batch axes into the M axis.
+    return (
+        isinstance(lhs_batches, tvm.tirx.Var)
+        or isinstance(rhs_batches, tvm.tirx.Var)
+        or (analyzer.can_prove_equal(lhs_batches, rhs_batches))
+        or (analyzer.can_prove(lhs_batches >= 1) and analyzer.can_prove(rhs_batches == 1))
+    )
+
+
+register_patterns(
+    [
+        (
+            "cublas.matmul",
+            *make_matmul_pattern(
+                with_bias=False,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_bias",
+            *make_matmul_pattern(
+                with_bias=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_bias_relu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_bias_gelu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.gelu",
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed",
+            *make_matmul_pattern(
+                with_bias=False,
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_bias",
+            *make_matmul_pattern(
+                with_bias=True,
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_bias_relu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_bias_gelu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.gelu",
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_dequantize",
+            *make_matmul_dequantize_pattern(transposed_rhs=True),
+            _check_matmul,
+        ),
+        (
+            "cublas.matmul_transposed_multiply",
+            *make_matmul_multiply_pattern(transposed_rhs=True),
+            _check_matmul,
+        ),
+    ]
+)
+
+
+def partition_for_cublas(mod, bind_constants=False):
+    """
+    Partition the input module into cuBLAS-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    bind_constants : bool
+        Whether or not to keep bound constants in the grouped function.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the cuBLAS backend.
+    """
+
+    patterns = get_patterns_with_prefix("cublas")
+    return transform.FuseOpsByPattern(
+        patterns, bind_constants=bind_constants, annotate_codegen=True
+    )(mod)

python/tvm/relax/backend/cuda/cudnn.py

@@ -0,0 +1,202 @@
+# 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.
+
+"""Pattern table for cuDNN backend"""
+
+import operator
+from functools import partial, reduce
+
+import tvm
+from tvm import relax
+from tvm.relax import PyExprMutator, expr_functor, transform
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import make_conv2d_pattern, make_stacked_attention_pattern
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+    """Check if dtypes in the given workload are supported by cuDNN BYOC."""
+    return (lhs_dtype == "float16" and rhs_dtype == "float16") or (
+        lhs_dtype == "float32" and rhs_dtype == "float32"
+    )
+
+
+def _is_supported_format(data_layout, kernel_layout):
+    """Check if layouts in the given workload are supported by cuDNN BYOC."""
+    return (data_layout == "NHWC" and kernel_layout == "OHWI") or (
+        data_layout == "NCHW" and kernel_layout == "OIHW"
+    )
+
+
+def _check_conv2d(context: PatternCheckContext) -> bool:
+    if has_leaking_intermediate_variables(context):
+        return False
+    # Retrieve the annotated expression from context
+    conv2d_call = context.annotated_expr["root"]
+    input_expr = context.annotated_expr["input"]
+    weight_expr = context.annotated_expr["weight"]
+
+    # Check if the data types of input and weights are supported by cuDNN BYOC
+    input_dtype = input_expr.ty.dtype
+    weight_dtype = weight_expr.ty.dtype
+    if not _is_supported_dtype(input_dtype, weight_dtype):
+        return False
+
+    input_layout = conv2d_call.attrs.data_layout
+    weight_layout = conv2d_call.attrs.kernel_layout
+    if not _is_supported_format(input_layout, weight_layout):
+        return False
+
+    return True
+
+
+def _check_stacked_attention(context: PatternCheckContext, layout: str) -> bool:
+    """Check if the given stacked attention workload can be offloaded to cuDNN."""
+    if has_leaking_intermediate_variables(context):
+        return False
+    if layout == "BS3NH":
+        if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
+            return False
+        if "split" in context.annotated_expr:
+            split_op = context.annotated_expr["split"]
+            if not split_op.attrs.axis == 2:
+                return False
+    elif layout == "SBN3H":
+        if not context.annotated_expr["stacked_qkv"].ty.ndim == 4:
+            return False
+        if "split" in context.annotated_expr:
+            split_op = context.annotated_expr["split"]
+            if not split_op.attrs.axis == 3:
+                return False
+    else:
+        raise NotImplementedError(f"Unsupported layout: {layout}")
+    return True
+
+
+register_patterns(
+    [
+        (
+            "cudnn.conv2d.nhwc_ohwi",
+            *make_conv2d_pattern(
+                with_bias=False,
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.conv2d.nhwc_ohwi_bias",
+            *make_conv2d_pattern(
+                with_bias=True,
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.conv2d.nhwc_ohwi_bias_relu",
+            *make_conv2d_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+            ),
+            _check_conv2d,
+        ),
+        (
+            "cudnn.attention.BS3NH",
+            *make_stacked_attention_pattern(start_op="split", layout="BS3NH"),
+            partial(_check_stacked_attention, layout="BS3NH"),
+        ),
+        (
+            "cudnn.attention.SBN3H",
+            *make_stacked_attention_pattern(start_op="split", layout="SBN3H"),
+            partial(_check_stacked_attention, layout="SBN3H"),
+        ),
+    ]
+)
+
+
+def partition_for_cudnn(mod):
+    """
+    Partition the input module into cuDNN-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the cuDNN backend.
+    """
+
+    patterns = get_patterns_with_prefix("cudnn")
+    return tvm.transform.Sequential(
+        [
+            transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True),
+            annotate_workspace,
+            transform.AllocateWorkspace(),
+        ]
+    )(mod)
+
+
+def _shape_1d(shape):
+    return reduce(operator.mul, shape, 1)
+
+
+@expr_functor.mutator
+class WorkspaceAnnotator(PyExprMutator):
+    """Annotate a workspace requirement for each cuDNN-offloaded function."""
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_function_(self, f):
+        if "Composite" not in f.attrs:
+            body = super().visit_expr(f.body)
+            new_f = relax.Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
+
+            if "global_symbol" in f.attrs and "cudnn" in f.attrs["global_symbol"]:
+                composite_func = body.blocks[0].bindings[0].value
+                if "WorkspaceSize" in composite_func.attrs:
+                    return new_f.with_attr("WorkspaceSize", composite_func.attrs["WorkspaceSize"])
+
+            return new_f
+
+        if "attention" in f.attrs["Composite"] and "cudnn" in f.attrs["Composite"]:
+            # Workspace is needed only for larger head sizes, but for simplicity we always allocate.
+            out_dtype = f.ret_ty.dtype
+            out_size_1d = _shape_1d(f.ret_ty.shape)
+            # This needs to be in sync with the actual value that the kernel expects.
+            workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
+            if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
+                # Tempororay workaround for dynamic shape workload. Will be removed when
+                # workspace for dynamic shape workload is implemented.
+                workspace_size_bytes = 8
+            return f.with_attr("WorkspaceSize", workspace_size_bytes)
+
+        return f
+
+
+@tvm.transform.module_pass(opt_level=0)
+def annotate_workspace(mod, _):
+    """Pass to annotate a workspace requirement for each cuDNN-offloaded function."""
+    annotator = WorkspaceAnnotator(mod)
+    for name, f in mod.functions_items():
+        if isinstance(f, relax.Function):
+            new_f = annotator.visit_expr(f)
+            mod.update_func(name, new_f)
+    return mod

python/tvm/relax/backend/cuda/cutlass.py

@@ -0,0 +1,620 @@
+# 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
+"""Pattern table for CUTLASS backend"""
+
+import operator
+from collections.abc import Mapping, Sequence
+from functools import reduce
+
+import tvm
+from tvm.contrib.cutlass.build import is_shape_valid_for_cutlass_matmul
+from tvm.relax import (
+    Call,
+    ExternFunc,
+    Function,
+    PyExprMutator,
+    Var,
+    expr_functor,
+    transform,
+)
+from tvm.relax.dpl import rewrite_call
+from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import (
+    make_attention_pattern,
+    make_attention_rewrite_pattern,
+    make_fused_bias_activation_pattern,
+    make_layer_norm_pattern,
+    make_matmul_pattern,
+    make_residual_block_pattern,
+    make_rms_norm_pattern,
+    make_stacked_attention_pattern,
+)
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+    """Check if dtypes in the given workload are supported by CUTLASS."""
+    return (
+        (lhs_dtype == "float16" and rhs_dtype == "float16")
+        or (lhs_dtype == "float32" and rhs_dtype == "float32")
+        or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
+    )
+
+
+def _shape_1d(shape):
+    return reduce(operator.mul, shape, 1)
+
+
+def _has_dependency(from_var: Var, to_var: Var, var_usages: Mapping[Var, Sequence[Var]]):
+    if from_var == to_var:
+        return True
+
+    checked = set()
+    vars_to_check = [to_var]
+    while vars_to_check:
+        current_var = vars_to_check.pop()
+        for user in var_usages.get(current_var, []):
+            if user == from_var:
+                return True
+            if user not in checked:
+                checked.add(user)
+                vars_to_check.append(user)
+
+    return False
+
+
+def _is_same_shape(shape1, shape2):
+    analyzer = tvm.arith.Analyzer()
+    return all([analyzer.can_prove_equal(s1, s2) for s1, s2 in zip(shape1, shape2)])
+
+
+def _is_bias_like(shape, out_channel):
+    return shape[-1] == out_channel and _shape_1d(shape) == out_channel
+
+
+def _check_residual(root_call: Call, context: PatternCheckContext) -> bool:
+    if "residual" in context.annotated_expr:
+        residual = context.annotated_expr["residual"]
+        if not isinstance(residual, Var):
+            if residual not in context.value_to_bound_var:
+                return False
+
+            residual = context.value_to_bound_var[residual]
+
+        root_var = context.value_to_bound_var[root_call]
+        if _has_dependency(from_var=residual, to_var=root_var, var_usages=context.var_usages):
+            # If residual depends on the result of the root call, this cannot be handled by cutlass.
+            return False
+
+        shape1 = root_var.ty.shape
+        shape2 = residual.ty.shape
+        out_channel = shape1[-1]
+
+        if not _is_same_shape(shape1, shape2) and not _is_bias_like(shape2, out_channel):
+            return False
+
+    return True
+
+
+def _check_conv2d(context: PatternCheckContext) -> bool:
+    """Check if the given conv2d workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(context):
+        return False
+
+    conv2d_call = context.annotated_expr["root"]
+    data_layout = conv2d_call.attrs.data_layout
+    kernel_layout = conv2d_call.attrs.kernel_layout
+    data, weight, *_ = conv2d_call.args
+    if (
+        data_layout != "NHWC"
+        or kernel_layout != "OHWI"
+        or not _is_supported_dtype(data.ty.dtype, weight.ty.dtype)
+    ):
+        return False
+
+    if not _check_residual(conv2d_call, context):
+        return False
+
+    # Check if any dimensions are symbolic.
+    for dim in data.ty.shape.values:
+        if isinstance(dim, tvm.tirx.Var):
+            return False
+
+    # pylint: disable=invalid-name
+    IC = data.ty.shape.values[3]
+    OC = weight.ty.shape.values[0]
+    # not depthwise conv2d
+    return not IC == OC == conv2d_call.attrs.groups
+
+
+def _check_matmul(context: PatternCheckContext) -> bool:
+    """Check if the given matmul workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(context):
+        return False
+
+    lhs = context.annotated_expr["lhs"]
+    rhs = context.annotated_expr["rhs"]
+
+    lhs_dtype = lhs.ty.dtype
+    rhs_dtype = rhs.ty.dtype
+    if not _is_supported_dtype(lhs_dtype, rhs_dtype):
+        return False
+
+    if not _check_residual(context.annotated_expr["root"], context):
+        return False
+
+    lhs_shape = lhs.ty.shape.values
+    rhs_shape = rhs.ty.shape.values
+    return is_shape_valid_for_cutlass_matmul(lhs_shape, rhs_shape)
+
+
+def _get_activation_from_name(pattern_name):
+    if "_relu" in pattern_name:
+        return "relax.nn.relu"
+    elif "_gelu_tanh" in pattern_name:
+        return "relax.nn.gelu_tanh"
+    elif "_gelu" in pattern_name:
+        return "relax.nn.gelu"
+    elif "_silu" in pattern_name:
+        return "relax.nn.silu"
+    else:
+        return None
+
+
+def matmul_patterns():
+    """
+    Returns a list of all matmul patterns in cutlass BYOC backend.
+    """
+
+    def _matmul_pattern(pattern_name):
+        transposed_rhs = "_transposed" in pattern_name
+        with_bias = "_bias" in pattern_name
+        activation = _get_activation_from_name(pattern_name)
+
+        return (
+            pattern_name,
+            *make_matmul_pattern(
+                transposed_rhs=transposed_rhs,
+                with_bias=with_bias,
+                activation=activation,
+            ),
+            _check_matmul,
+        )
+
+    return [
+        _matmul_pattern("cutlass.matmul"),
+        _matmul_pattern("cutlass.matmul_bias"),
+        _matmul_pattern("cutlass.matmul_bias_relu"),
+        _matmul_pattern("cutlass.matmul_bias_gelu"),
+        _matmul_pattern("cutlass.matmul_transposed"),
+        _matmul_pattern("cutlass.matmul_transposed_bias"),
+        _matmul_pattern("cutlass.matmul_transposed_bias_relu"),
+        _matmul_pattern("cutlass.matmul_transposed_bias_gelu"),
+    ]
+
+
+def _check_decode_matmul(ctx):
+    """Check if the given decode -> matmul workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(ctx):
+        return False
+
+    root = ctx.annotated_expr["root"]
+
+    if not _check_residual(root, ctx):
+        return False
+
+    # out_dtype = "float32" not supported unless matmul is followed by cast to fp16.
+    if root.ty.dtype == "float32":
+        return False
+
+    call_tir_decode = ctx.annotated_expr["w_decoded"]
+    if "decode" not in call_tir_decode.args[0].name_hint:
+        return False
+
+    N = root.ty.shape[-1]
+
+    if ctx.annotated_expr["lhs"].ty.dtype != "float16":
+        return False
+
+    # weight needs to be packed to int8.
+    packed_weight = ctx.annotated_expr["w_encoded"]
+
+    if packed_weight.ty.dtype != "int8":
+        return False
+
+    # The kernel expects the weight to be preprocessed by this packed function.
+    if (
+        isinstance(packed_weight, Call)
+        and isinstance(packed_weight.args[0], ExternFunc)
+        and packed_weight.args[0].global_symbol != "cutlass.ft_preprocess_weight"
+    ):
+        return False
+
+    scales = ctx.annotated_expr["scales"]
+
+    if scales.ty.dtype != "float16":
+        return False
+
+    # scale shape needs to be (N,) or (1, N) or (K // group_size, N)
+    if len(scales.ty.shape) > 2 or scales.ty.shape[-1] != N:
+        return False
+
+    if "bias" in ctx.annotated_expr:
+        out_shape = root.ty.shape
+        bias_shape = ctx.annotated_expr["bias"].ty.shape
+
+        # bias shape needs to be (N,), possibly with additional axes on the front.
+        # It can also have the same shape as the output.
+        if not _is_bias_like(bias_shape, N) and not _is_same_shape(out_shape, bias_shape):
+            return False
+
+    return True
+
+
+def decode_matmul_patterns():
+    """Returns a list of supported decode -> matmul patterns."""
+
+    def _decode_matmul_pattern(name):
+        scales = wildcard()
+        x = wildcard()
+        w_packed = wildcard()
+
+        w = is_op("relax.call_tir")(
+            GlobalVarPattern(),
+            TuplePattern([w_packed, scales]),
+        )
+        matmul = is_op("relax.matmul")(x, w)
+
+        if "cast" in name:
+            matmul = is_op("relax.astype")(matmul)
+
+        annotations = {
+            "root": matmul,
+            "lhs": x,
+            "w_encoded": w_packed,
+            "w_decoded": w,
+            "scales": scales,
+        }
+
+        if "bias" in name:
+            annotations["bias"] = bias = wildcard()
+            out = is_op("relax.add")(matmul, bias)
+        else:
+            out = matmul
+
+        if "gelu" in name:
+            out = is_op("relax.nn.gelu")(out)
+
+        return name, out, annotations, _check_decode_matmul
+
+    return [
+        _decode_matmul_pattern("cutlass.decode_matmul"),
+        _decode_matmul_pattern("cutlass.decode_matmul_bias"),
+        _decode_matmul_pattern("cutlass.decode_matmul_cast"),
+        _decode_matmul_pattern("cutlass.decode_matmul_cast_bias"),
+        _decode_matmul_pattern("cutlass.decode_matmul_bias_gelu"),
+        _decode_matmul_pattern("cutlass.decode_matmul_cast_bias_gelu"),
+    ]
+
+
+def conv2d_patterns():
+    """
+    Returns a list of all conv2d patterns in cutlass BYOC backend.
+    """
+
+    def _conv2d_pattern(pattern_name):
+        with_bias = "_bias" in pattern_name
+        activation = _get_activation_from_name(pattern_name)
+
+        return (
+            pattern_name,
+            *make_fused_bias_activation_pattern(
+                "relax.nn.conv2d",
+                with_bias=with_bias,
+                activation=activation,
+            ),
+            _check_conv2d,
+        )
+
+    return [
+        _conv2d_pattern("cutlass.conv2d"),
+        _conv2d_pattern("cutlass.conv2d_bias"),
+        _conv2d_pattern("cutlass.conv2d_bias_relu"),
+        _conv2d_pattern("cutlass.conv2d_bias_silu"),
+    ]
+
+
+def residual_block_patterns():
+    """
+    Returns a list of all residual block patterns in cutlass BYOC backend.
+    """
+    patterns = []
+
+    for activation, name_postfix in [(None, ""), ("relax.nn.relu", "_relu")]:
+        for check, base_patterns in [
+            (_check_conv2d, conv2d_patterns()),
+            (_check_matmul, matmul_patterns()),
+            (_check_decode_matmul, decode_matmul_patterns()),
+        ]:
+            for name, pat, arg_pat, _ in base_patterns:
+                # Append residual patterns only to those base patterns with bias add,
+                # since conv2d or matmul + residual add without bias is already supported
+                # via conv2d or matmul + bias patterns (the residual input is treated as "bias").
+                if "bias" in name:
+                    for bin_op in ["relax.add", "relax.multiply"]:
+                        patterns.append(
+                            (
+                                name + "_residual_" + bin_op.split(".")[-1] + name_postfix,
+                                *make_residual_block_pattern(
+                                    (pat, arg_pat), binary_op=bin_op, activation=activation
+                                ),
+                                check,
+                            )
+                        )
+
+    return patterns
+
+
+def _check_stacked_attention(context: PatternCheckContext) -> bool:
+    """Check if the given stacked attention workload can be offloaded to CUTLASS."""
+    if has_leaking_intermediate_variables(context):
+        return False
+    if not context.annotated_expr["stacked_qkv"].ty.ndim == 3:
+        return False
+    if "split" in context.annotated_expr:
+        split_op = context.annotated_expr["split"]
+        if not split_op.attrs.axis == 2:
+            return False
+    else:
+        get_const_int_list = lambda tup: [int(e.value) for e in tup]
+        last_end = 0
+        for name in ["query", "key", "value"]:
+            assert f"strided_slice_{name}" in context.annotated_expr
+            strided_slice_op = context.annotated_expr[f"strided_slice_{name}"]
+            axes = get_const_int_list(strided_slice_op.args[1])
+            begins = get_const_int_list(strided_slice_op.args[2])
+            ends = get_const_int_list(strided_slice_op.args[3])
+            strides = get_const_int_list(strided_slice_op.args[4])
+
+            if axes != [2]:
+                return False
+            if begins != [last_end]:
+                return False
+            if not len(ends) == 1:
+                return False
+            if strides != [1]:
+                return False
+            last_end = ends[0]
+    return True
+
+
+def attention_patterns():
+    """
+    Returns a list of all attention patterns in cutlass BYOC backend.
+    """
+    return [
+        (
+            "cutlass.attention",
+            *make_attention_pattern(),
+        ),
+        (
+            "cutlass.attention_bias",
+            *make_attention_pattern(with_bias=True),
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="split"),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="split", with_bias=True),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="strided_slice"),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.stacked_attention",
+            *make_stacked_attention_pattern(start_op="strided_slice", with_bias=True),
+            _check_stacked_attention,
+        ),
+        (
+            "cutlass.attention_var_len",
+            *make_attention_pattern(var_len=True),
+        ),
+    ]
+
+
+def _check_layer_norm(context: PatternCheckContext) -> bool:
+    attrs = context.matched_expr.attrs
+
+    if not attrs.center or not attrs.scale:
+        return False
+
+    if len(attrs.axes) != 1:
+        # Contiguous inner-most axes can be supported, but reject it for now for simplicity.
+        return False
+
+    axis = int(attrs.axes[0])
+    rank = len(context.matched_expr.ty.shape)
+
+    if axis < 0:
+        axis += rank
+
+    return axis == rank - 1
+
+
+def layer_norm_pattern():
+    """Create a layer norm pattern for CUTLASS."""
+    return [
+        (
+            "cutlass.layer_norm",
+            *make_layer_norm_pattern(),
+            _check_layer_norm,
+        ),
+    ]
+
+
+def _check_rms_norm(ctx: PatternCheckContext) -> bool:
+    rms_norm = ctx.annotated_expr["rms_norm"]
+    if "rms_norm" not in rms_norm.args[0].name_hint:
+        return False
+
+    return True
+
+
+def rms_norm_pattern():
+    """Create a RMS norm pattern for CUTLASS."""
+    return [
+        (
+            "cutlass.rms_norm",
+            *make_rms_norm_pattern(),
+            _check_rms_norm,
+        ),
+    ]
+
+
+def attention_rewrite_patterns():
+    """
+    Returns a list of all attention rewriting patterns in cutlass BYOC backend.
+    """
+    patterns = []
+    for qkv_layout in ["BSNH", "BSH"]:
+        for out_layout in ["BSNH", "BSH"]:
+            for with_bias in [True, False]:
+                for with_cast in [True, False]:
+                    patterns.append(
+                        make_attention_rewrite_pattern(qkv_layout, out_layout, with_bias, with_cast)
+                    )
+    return patterns
+
+
+register_patterns(
+    [
+        *conv2d_patterns(),
+        *matmul_patterns(),
+        *decode_matmul_patterns(),
+        *residual_block_patterns(),
+        *attention_patterns(),
+        *layer_norm_pattern(),
+        *rms_norm_pattern(),
+    ]
+)
+
+_REWRITE_PATTERNS = [*attention_rewrite_patterns()]
+
+
+@expr_functor.mutator
+class WorkspaceAnnotator(PyExprMutator):
+    """Annotate a workspace requirement for each CUTLASS-offloaded function."""
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_function_(self, f):
+        if "Composite" not in f.attrs:
+            body = super().visit_expr(f.body)
+            new_f = Function(f.params, body, f.ret_ty, f.is_pure, f.attrs, f.span)
+
+            if "global_symbol" in f.attrs and "cutlass" in f.attrs["global_symbol"]:
+                composite_func = body.blocks[0].bindings[0].value
+                if "WorkspaceSize" in composite_func.attrs:
+                    return new_f.with_attr("WorkspaceSize", composite_func.attrs["WorkspaceSize"])
+
+            return new_f
+
+        if "attention" in f.attrs["Composite"] and "cutlass" in f.attrs["Composite"]:
+            # Workspace is needed only for larger head sizes, but for simplicity we always allocate.
+            out_dtype = f.ret_ty.dtype
+            out_size_1d = _shape_1d(f.ret_ty.shape)
+            # This needs to be in sync with the actual value that the kernel expects.
+            workspace_size_bytes = out_size_1d * {"float16": 2, "float32": 4}[out_dtype]
+            if not isinstance(workspace_size_bytes, int | tvm.tirx.expr.IntImm):
+                # Tempororay workaround for dynamic shape workload. Will be removed when
+                # workspace for dynamic shape workload is implemented.
+                workspace_size_bytes = 8
+            return f.with_attr("WorkspaceSize", workspace_size_bytes)
+
+        return f
+
+
+@tvm.transform.module_pass(opt_level=0)
+def annotate_workspace(mod, _):
+    """Pass to annotate a workspace requirement for each CUTLASS-offloaded function."""
+    annotator = WorkspaceAnnotator(mod)
+    for name, f in mod.functions_items():
+        if isinstance(f, Function):
+            new_f = annotator.visit_expr(f)
+            mod.update_func(name, new_f)
+    return mod
+
+
+def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
+    """
+    Partition the input module into CUTLASS-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    annotate_codegen: bool
+        Whether to wrap each created composite function with another function, whose
+        body consists only of a call to the composite function. See the doc of FuseOpsByPattern
+        for more detail.
+
+    use_flash_mqa: bool
+        Whether to consider a rewrite pattern for multi-query attention, which is supported by
+        the Flash Attention kernel.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        compiled by the CUTLASS backend.
+    """
+    for func_name, func in mod.functions_items():
+        if isinstance(func, Function):
+            if use_flash_mqa:
+                mqa_pattern, rewriter = make_attention_rewrite_pattern(
+                    "BSNH", "BSNH", with_bias=False, with_cast=True, with_kv_repeat=True
+                )
+                func = rewrite_call(mqa_pattern, rewriter, func)
+
+            for pattern, rewriter in _REWRITE_PATTERNS:
+                func = rewrite_call(pattern, rewriter, func)
+
+        mod[func_name] = func
+
+    patterns = get_patterns_with_prefix("cutlass")
+    return tvm.transform.Sequential(
+        [
+            transform.FuseOpsByPattern(
+                patterns, bind_constants=False, annotate_codegen=annotate_codegen
+            ),
+            annotate_workspace,
+            transform.AllocateWorkspace(),
+        ]
+    )(mod)

python/tvm/relax/backend/cuda/flashinfer.py

@@ -0,0 +1,345 @@
+# 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.
+
+"""FlashInfer JIT compilation module for CUDA backend"""
+
+import re
+from pathlib import Path
+
+import tvm
+from tvm.target import Target
+
+
+def _rename_exported_func_names(source_paths: list[Path], prefix: str):
+    """Rename the ffi-exported function names in the source files to the given prefix."""
+    pattern = re.compile(r"^(\s*TVM_FFI_DLL_EXPORT_TYPED_FUNC\()([A-Za-z0-9_]+)(,.*)$")
+    for source_path in source_paths:
+        if not source_path.name.endswith("_binding.cu"):
+            continue
+
+        original_text = source_path.read_text(encoding="utf-8")
+        lines = original_text.splitlines(keepends=True)
+        updated = False
+        for idx, line in enumerate(lines):
+            line_body = line.rstrip("\r\n")
+            line_ending = line[len(line_body) :]
+            match = pattern.match(line_body)
+            if not match:
+                continue
+            new_body = f"{match.group(1)}{prefix}_{match.group(2)}{match.group(3)}"
+            lines[idx] = new_body + line_ending
+            updated = True
+
+        if updated:
+            source_path.write_text("".join(lines), encoding="utf-8")
+
+
+def _load_flashinfer_modules(object_files: list[Path]) -> list[tvm.runtime.Module]:
+    return [
+        tvm.runtime.load_static_library(str(obj_path.absolute()), func_names=[])
+        for obj_path in object_files
+    ]
+
+
+def gen_flashinfer_prefill_module(
+    dtype_q: str,
+    dtype_kv: str,
+    dtype_o: str,
+    qk_head_dim: int,
+    v_head_dim: int,
+    enable_inline_rope: bool,
+    return_static_libs: bool = False,
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for prefill.
+
+    Parameters
+    ----------
+    dtype_q : str
+        The data type of the query tensor.
+    dtype_kv : str
+        The data type of the key/value tensors.
+    dtype_o : str
+        The data type of the output tensor.
+    qk_head_dim : int
+        The head dimension of the query and key tensors.
+    v_head_dim : int
+        The head dimension of the value tensor.
+    enable_inline_rope : bool
+        Whether to enable inline rotary positional embedding.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    A list of compiled static library modules for FlashInfer prefill kernels.
+    """
+    try:
+        from flashinfer.jit import (  # pylint: disable=import-outside-toplevel
+            gen_customize_batch_prefill_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+    try:
+        import torch  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
+
+    if enable_inline_rope and qk_head_dim != v_head_dim:
+        raise ValueError("Inline rope mode is not supported when qk_head_dim == v_head_dim")
+
+    torch_dtype_q = getattr(torch, dtype_q)
+    torch_dtype_kv = getattr(torch, dtype_kv)
+    torch_dtype_o = getattr(torch, dtype_o)
+    # Todo(tvm-team): decide which backend ("fa2/fa3") to use
+    backend = "fa2"
+    variant_name = (
+        "DefaultAttention<false, false, false, false>"
+        if backend == "fa2"
+        else "DefaultAttention<false>"
+    )
+    variant_decl = (
+        "#include <flashinfer/attention/variants.cuh>"
+        if backend == "fa2"
+        else "#include <flashinfer/attention/hopper/variants.cuh>"
+    )
+    jit_spec = gen_customize_batch_prefill_module(
+        backend=backend,
+        uri=f"batch_prefill_tvm_dtype_q_{dtype_q}_"
+        + f"dtype_kv_{dtype_kv}_"
+        + f"dtype_o_{dtype_o}_"
+        + f"qk_head_dim_{qk_head_dim}_"
+        + f"v_head_dim_{v_head_dim}_"
+        + f"enable_inline_rope_{enable_inline_rope}",
+        dtype_q=torch_dtype_q,
+        dtype_kv=torch_dtype_kv,
+        dtype_o=torch_dtype_o,
+        idtype=torch.int32,
+        head_dim_qk=qk_head_dim,
+        head_dim_vo=v_head_dim,
+        pos_encoding_mode=int(enable_inline_rope),
+        additional_tensor_names=[],
+        additional_tensor_dtypes=[],
+        additional_scalar_names=["sm_scale", "rope_rcp_scale", "rope_rcp_theta"],
+        additional_scalar_dtypes=["double", "double", "double"],
+        variant_name=variant_name,
+        variant_decl=variant_decl,
+    )
+    _rename_exported_func_names(jit_spec.sources, "batch_prefill")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]
+
+
+def gen_flashinfer_decode_module(
+    dtype_q: str,
+    dtype_kv: str,
+    dtype_o: str,
+    qk_head_dim: int,
+    v_head_dim: int,
+    enable_inline_rope: bool,
+    return_static_libs: bool = False,
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for decode.
+
+    Parameters
+    ----------
+    dtype_q : str
+        The data type of the query tensor.
+    dtype_kv : str
+        The data type of the key/value tensors.
+    dtype_o : str
+        The data type of the output tensor.
+    qk_head_dim : int
+        The head dimension of the query and key tensors.
+    v_head_dim : int
+        The head dimension of the value tensor.
+    enable_inline_rope : bool
+        Whether to enable inline rotary positional embedding.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    A list of compiled static library modules for FlashInfer decode kernels.
+    """
+    try:
+        from flashinfer.jit import (  # pylint: disable=import-outside-toplevel
+            gen_customize_batch_decode_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+    try:
+        import torch  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
+
+    torch_dtype_q = getattr(torch, dtype_q)
+    torch_dtype_kv = getattr(torch, dtype_kv)
+    torch_dtype_o = getattr(torch, dtype_o)
+    jit_spec = gen_customize_batch_decode_module(
+        uri=f"batch_decode_tvm_dtype_q_{dtype_q}_"
+        + f"dtype_kv_{dtype_kv}_"
+        + f"dtype_o_{dtype_o}_"
+        + f"qk_head_dim_{qk_head_dim}_"
+        + f"v_head_dim_{v_head_dim}_"
+        + f"enable_inline_rope_{enable_inline_rope}",
+        dtype_q=torch_dtype_q,
+        dtype_kv=torch_dtype_kv,
+        dtype_o=torch_dtype_o,
+        idtype=torch.int32,
+        head_dim_qk=qk_head_dim,
+        head_dim_vo=v_head_dim,
+        pos_encoding_mode=int(enable_inline_rope),
+        additional_tensor_names=[],
+        additional_tensor_dtypes=[],
+        additional_scalar_names=["sm_scale", "rope_rcp_scale", "rope_rcp_theta"],
+        additional_scalar_dtypes=["double", "double", "double"],
+        variant_name="DefaultAttention<false, false, false, false>",
+        variant_decl="#include <flashinfer/attention/variants.cuh>",
+    )
+    _rename_exported_func_names(jit_spec.sources, "batch_decode")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]
+
+
+def gen_flashinfer_mla_module(
+    dtype_q: str,
+    dtype_kv: str,
+    dtype_o: str,
+    head_dim_ckv: int,
+    head_dim_kpe: int,
+    return_static_libs: bool = False,
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for MLA.
+
+    Parameters
+    ----------
+    dtype_q : str
+        The data type of the query tensor.
+    dtype_kv : str
+        The data type of the key/value tensors.
+    dtype_o : str
+        The data type of the output tensor.
+    head_dim_ckv : int
+        The head dimension of the compressed key/value tensors.
+    head_dim_kpe : int
+        The head dimension of the query/key positional embedding.
+    target : Target
+        The target device to compile for.
+    num_threads : int
+        The number of threads to use for compilation.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    A list of compiled static library modules for FlashInfer MLA kernels.
+    """
+    try:
+        from flashinfer.jit import (  # pylint: disable=import-outside-toplevel
+            gen_batch_mla_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+    try:
+        import torch  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        raise ImportError("PyTorch is not installed. Please install PyTorch to use FlashInfer.")
+
+    torch_dtype_q = getattr(torch, dtype_q)
+    torch_dtype_kv = getattr(torch, dtype_kv)
+    torch_dtype_o = getattr(torch, dtype_o)
+    jit_spec = gen_batch_mla_module(
+        backend="fa2",
+        dtype_q=torch_dtype_q,
+        dtype_kv=torch_dtype_kv,
+        dtype_o=torch_dtype_o,
+        dtype_idx=torch.int32,
+        head_dim_ckv=head_dim_ckv,
+        head_dim_kpe=head_dim_kpe,
+        use_profiler=False,
+    )
+    _rename_exported_func_names(jit_spec.sources, "batch_mla")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]
+
+
+def gen_grouped_gemm_module(
+    target: Target, return_static_libs: bool = False
+) -> list[tvm.runtime.Module]:
+    """Generate a FlashInfer module for FP8 grouped GEMM.
+
+    Parameters
+    ----------
+    target : Target
+        The target device to compile for.
+    return_static_libs : bool
+        Whether to return static library modules instead of compiled modules.
+        When it is False, it returns the loaded shared library that links all the object files.
+        When it is True, it returns the static libraries of each compiled object files.
+
+    Returns
+    -------
+    List[tvm.runtime.Module]
+        A list of compiled static library modules for FlashInfer FP8 grouped GEMM kernels.
+
+    Note
+    _____
+    when apply grouped gemm on A: (total_m, k), B: (batch_size, n, k), m_indptr: (batch_size, )
+    requires all m in m_indptr to be multiple of 4
+    """
+    # NOTE: This function is still under development,
+    # and we currently only support SM100 grouped gemm
+    try:
+        from flashinfer.gemm import (  # pylint: disable=import-outside-toplevel
+            gen_gemm_sm100_module,
+        )
+    except ImportError:
+        raise ImportError(
+            "FlashInfer is not installed. Please follow instructions "
+            "in https://docs.flashinfer.ai to install FlashInfer."
+        )
+
+    compute_version = "".join(tvm.support.nvcc.get_target_compute_version(target).split("."))
+    if compute_version == "100":
+        jit_spec = gen_gemm_sm100_module()
+    else:
+        raise ValueError(f"Unsupported compute version: {compute_version}")
+    if return_static_libs:
+        jit_spec.build(verbose=False)
+        return _load_flashinfer_modules(jit_spec.get_object_paths())
+    return [jit_spec.build_and_load()]

python/tvm/relax/backend/cuda/pipeline.py

@@ -0,0 +1,90 @@
+# 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.
+"""The Relax CUDA backend compilation pipeline and other passes."""
+
+import tvm
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default library dispatch passes for CUDA backend."""
+    return [
+        relax.backend.DispatchSampling(),
+        relax.backend.DispatchSortScan(),
+    ]
+
+
+def legalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default legalization passes for CUDA backend."""
+    from tvm.s_tir import dlight as dl  # pylint: disable=import-outside-toplevel
+
+    return [
+        tvm.relax.transform.LegalizeOps(),
+        tvm.relax.transform.AnnotateTIROpPattern(),
+        tvm.relax.transform.FoldConstant(),
+        tvm.relax.transform.FuseOps(),
+        tvm.relax.transform.FuseTIR(),
+        dl.ApplyDefaultSchedule(
+            dl.gpu.Matmul(),
+            dl.gpu.GEMV(),
+            dl.gpu.Reduction(),
+            dl.gpu.GeneralReduction(),
+            dl.gpu.Fallback(),
+        ),
+    ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default dataflow lowering passes for CUDA backend."""
+    return [
+        relax.transform.RewriteDataflowReshape(),
+        relax.transform.ToNonDataflow(),
+        relax.transform.RemovePurityChecking(),
+        relax.transform.CallTIRRewrite(),
+    ]
+
+
+def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default finalization passes for CUDA backend."""
+    return [
+        relax.transform.StaticPlanBlockMemory(),
+        relax.transform.RewriteCUDAGraph(),
+        relax.transform.LowerAllocTensor(),
+        relax.transform.KillAfterLastUse(),
+        relax.transform.LowerRuntimeBuiltin(),
+        relax.transform.ComputePrimValue(),
+        relax.transform.VMShapeLower(),
+        relax.transform.AttachGlobalSymbol(),
+    ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+    """Return the default compilation pipeline for CUDA."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+        with target:
+            seq = tvm.transform.Sequential(
+                library_dispatch_passes(target)
+                + legalize_passes(target)
+                + dataflow_lower_passes(target)
+                + finalize_passes(target)
+            )
+            mod = seq(mod)
+        return mod
+
+    return _pipeline

python/tvm/relax/backend/dispatch_sampling.py

@@ -0,0 +1,93 @@
+# 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
+"""Dispatch sampling operators to platform dependent implementation."""
+
+from tvm import relax
+from tvm.ir import Op
+from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.relax import expr_functor
+
+from .utils import BackendDispatcher
+
+
+@expr_functor.mutator
+class SamplingDispatcher(BackendDispatcher):
+    """Dispatcher to dispatch sampling op."""
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        if not isinstance(call.op, Op):
+            return super().visit_call_(call)
+
+        if call.op.name == "relax.multinomial_from_uniform":
+            from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
+                generic_get_sample_index,
+                gpu_multinomial_from_uniform,
+            )
+
+            prob, uniform_sample, sample_indices = call.args
+            tgt = self._get_target(call.ty)
+            dtype = call.attrs.dtype
+            _, prob_dtype = self.get_shape_dtype(prob)
+            sample_shape, sample_dtype = self.get_shape_dtype(uniform_sample)
+            sample_indices_shape, sample_indices_dtype = self.get_shape_dtype(sample_indices)
+
+            if len(sample_shape) != 2 or sample_shape[1] != 1:
+                raise ValueError("uniform_sample should be a 2D tensor with shape (N, 1)")
+
+            if len(sample_indices_shape) != 2 or sample_indices_shape[1] != 1:
+                raise ValueError("sample_indices should be a 2D tensor with shape (N, 1)")
+
+            if self.is_gpu_target(tgt):
+                gv = self.builder_.add_func(
+                    gpu_multinomial_from_uniform(
+                        prob_dtype, sample_dtype, sample_indices_dtype, dtype
+                    ),
+                    "gpu_multinomial_from_uniform",
+                )
+                return relax.call_tir(
+                    gv,
+                    [prob, uniform_sample, sample_indices],
+                    out_ty=call.ty,
+                )
+            else:
+                cumsum_prob = relax.op.cumsum(prob, axis=1, dtype=prob_dtype.dtype, exclusive=False)
+                gv = self.builder_.add_func(
+                    generic_get_sample_index(prob_dtype, sample_dtype, sample_indices_dtype, dtype),
+                    "get_sample_index",
+                )
+                return relax.call_tir(
+                    gv,
+                    [cumsum_prob, uniform_sample, sample_indices],
+                    out_ty=call.ty,
+                )
+
+        return super().visit_call_(call)
+
+
+@module_pass(opt_level=0, name="DispatchSampling")
+class DispatchSampling:
+    """Pass to dispatch scan and sort operators to platform dependent implementation."""
+
+    def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule:
+        sampling_dispatcher = SamplingDispatcher(mod)
+        for gv, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = sampling_dispatcher.visit_expr(func)
+                sampling_dispatcher.builder_.update_func(gv, func)
+        return sampling_dispatcher.builder_.finalize()

python/tvm/relax/backend/dispatch_sort_scan.py

@@ -0,0 +1,256 @@
+# 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
+"""Dispatch sort and scan operators to platform dependent implementation."""
+
+from functools import reduce
+from operator import mul
+
+from tvm import DataType, relax, topi
+from tvm.contrib.thrust import can_use_thrust
+from tvm.ir import GlobalVar, Op
+from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.relax import expr_functor
+from tvm.target import Target
+
+from .utils import BackendDispatcher
+
+
+@expr_functor.mutator
+class SortScanDispatcher(BackendDispatcher):
+    """Dispatcher to dispatch sort and scan."""
+
+    calls_to_update: dict[GlobalVar, Target]
+
+    def __init__(self, mod):
+        super().__init__(mod)
+        self.calls_to_update = {}
+
+    def apply_dlight_gpu_fallback(
+        self,
+    ) -> None:
+        """Apply DLight rules for all the calls that need to be updated."""
+        from tvm.s_tir import dlight  # pylint: disable=import-outside-toplevel
+
+        for gvar, target in self.calls_to_update.items():
+            func = self.builder_.get()[gvar]
+            sch = dlight.base.transform._apply_rules(
+                func,
+                target,
+                rules=[dlight.gpu.Fallback()],
+                tunable=False,
+            )
+            if sch is not None:
+                assert len(sch) == 1
+                self.builder_.update_func(
+                    gvar, sch[0].mod["main"].with_attr("tirx.is_scheduled", True)
+                )
+
+    def _append_calls_to_update(self, tir_call: relax.Call, target: Target) -> None:
+        gvar = tir_call.args[0]
+        assert isinstance(gvar, GlobalVar)
+        existing_tgt = self.calls_to_update.get(gvar, None)
+        if existing_tgt is not None and existing_tgt != target:
+            raise ValueError(
+                f"Multiple targets detected for function {gvar}. "
+                f"Existing target: {existing_tgt}, new target: {target}"
+            )
+        self.calls_to_update[gvar] = target
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        if not isinstance(call.op, Op):
+            return super().visit_call_(call)
+
+        if call.op.name == "relax.bucketize":
+            input_tensor = call.args[0]
+            boundaries = call.args[1]
+            right = call.attrs.right
+            tgt = self._get_target(call.ty)
+            te_func = topi.searchsorted
+            with tgt:
+                if self.is_gpu_target(tgt):
+                    te_func = topi.gpu.searchsorted
+            out_dtype = "int32" if call.attrs.out_int32 else "int64"
+            return self.builder_.call_te(te_func, boundaries, input_tensor, right, out_dtype)
+        if call.op.name == "relax.sort":
+            tgt = self._get_target(call.ty)
+            te_func = topi.sort
+            kwargs = {}
+            with tgt:
+                if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+                    te_func = topi.gpu.sort_thrust
+                    kwargs["workspace"] = self.allocate_workspace(call)
+                elif self.is_gpu_target(tgt):
+                    te_func = topi.gpu.sort
+            return self.builder_.call_te(
+                te_func, call.args[0], call.attrs.axis, not call.attrs.descending, **kwargs
+            )
+        if call.op.name == "relax.argsort":
+            tgt = self._get_target(call.ty)
+            te_func = topi.argsort
+            kwargs = {}
+            with tgt:
+                if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+                    te_func = topi.gpu.argsort_thrust
+                    kwargs["workspace"] = self.allocate_workspace(call)
+                elif self.is_gpu_target(tgt):
+                    te_func = topi.gpu.argsort
+            return self.builder_.call_te(
+                te_func,
+                call.args[0],
+                axis=call.attrs.axis,
+                is_ascend=not call.attrs.descending,
+                dtype=call.attrs.dtype,
+                **kwargs,
+            )
+        if call.op.name == "relax.topk":
+            tgt = self._get_target(call.ty)
+            te_func = topi.topk
+            kwargs = {}
+            if can_use_thrust(tgt, "tvm.contrib.thrust.sort"):
+                te_func = topi.gpu.topk_thrust
+                kwargs["workspace"] = self.allocate_workspace(call)
+            elif self.is_gpu_target(tgt):
+                te_func = topi.gpu.topk
+            tir_call = self.builder_.call_te(
+                te_func,
+                call.args[0],
+                k=call.attrs.k,
+                axis=call.attrs.axis,
+                ret_type=call.attrs.ret_type,
+                is_ascend=not call.attrs.largest,
+                dtype=call.attrs.dtype,
+                **kwargs,
+            )
+            self._append_calls_to_update(tir_call, tgt)
+            return tir_call
+        if call.op.name in ("relax.cumprod", "relax.cumsum"):
+            tgt = self._get_target(call.ty)
+            axis = int(call.attrs.axis) if call.attrs.axis is not None else call.attrs.axis
+            shape = call.ty.shape
+            # TODO(tvm-team): Support fully dynamic case with `shape=None`
+            if shape is None:
+                raise ValueError("non-symbolic shape is not supported for now")
+            kwargs = {}
+            if (
+                shape is not None
+                and (axis == -1 or axis == len(shape) - 1)
+                and self.is_gpu_target(tgt)
+                and not can_use_thrust(tgt, "tvm.contrib.thrust.sum_scan")
+                and call.op.name == "relax.cumsum"
+                and call.attrs.exclusive == 0
+            ):
+                from tvm.relax.backend.gpu_generic import (  # pylint: disable=import-outside-toplevel
+                    gpu_2d_continuous_cumsum,
+                )
+
+                dim = 1
+                for i in range(len(shape) - 1):
+                    dim *= shape[i]
+                in_dtype = call.args[0].ty.dtype
+                out_dtype = call.attrs.dtype
+                out_dtype = out_dtype or in_dtype
+                cumsum_2d_shape = relax.ShapeExpr([dim, shape[-1]])
+                reshape = relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    call.args[0],
+                    cumsum_2d_shape,
+                    ty_args=relax.TensorType(cumsum_2d_shape, out_dtype),
+                )
+                gv = self.builder_.add_func(
+                    gpu_2d_continuous_cumsum(in_dtype=in_dtype, out_dtype=out_dtype),
+                    "gpu_2d_continuous_cumsum",
+                )
+                cumsum = relax.call_tir(
+                    gv,
+                    reshape,
+                    out_ty=relax.TensorType(cumsum_2d_shape, out_dtype),
+                )
+                return relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    cumsum,
+                    shape,
+                    ty_args=call.ty,
+                )
+
+            with tgt:
+                if call.op.name == "relax.cumsum":
+                    te_func = topi.gpu.cumsum if self.is_gpu_target(tgt) else topi.cumsum
+                    if can_use_thrust(tgt, "tvm.contrib.thrust.sum_scan"):
+                        kwargs["workspace"] = self.allocate_workspace(call)
+                elif call.op.name == "relax.cumprod":
+                    te_func = topi.gpu.cumprod if self.is_gpu_target(tgt) else topi.cumprod
+                else:
+                    raise ValueError(f"Unsupported op: {call.op.name}")
+                tir_call = self.builder_.call_te(
+                    te_func,
+                    call.args[0],
+                    axis,
+                    call.attrs.dtype,
+                    call.attrs.exclusive,
+                    **kwargs,
+                )
+            self._append_calls_to_update(tir_call, tgt)
+            return tir_call
+        return super().visit_call_(call)
+
+    def estimate_thrust_workspace_size(self, call: relax.Call) -> int:
+        """
+        Estimate the workspace size for thrust sort/argsort/topk/cumsum
+        """
+        input_shape = call.args[0].ty.shape
+        input_byte_per_elem = DataType(call.args[0].ty.dtype.dtype).bits // 8
+        int64_byte_per_elem = DataType("int64").bits // 8
+        int32_byte_per_elem = DataType("int32").bits // 8
+        num_elem = reduce(mul, input_shape, 1)
+        input_size = num_elem * input_byte_per_elem
+        # Most GPU algorithms take O(n) space or less, we choose 8N + 8MB as a safe estimation
+        # for algorithm workspace.
+        # The current thrust sort implementation may need extra int64 and int32 arrays
+        # for temporary data, so we further add this part to the workspace.
+        return (
+            8 * input_size
+            + 8 * 1024 * 1024
+            + num_elem * (int64_byte_per_elem + int32_byte_per_elem)
+        )
+
+    def allocate_workspace(self, call: relax.Call) -> relax.Var:
+        """
+        Allocate workspace for thrust sort/argsort/topk.
+        """
+        workspace_size = self.estimate_thrust_workspace_size(call)
+        alloc = relax.op.builtin.alloc_tensor(
+            relax.ShapeExpr((workspace_size,)), "uint8", runtime_device_index=0
+        )
+        return self.builder_.emit(alloc)
+
+
+@module_pass(opt_level=0, name="DispatchSortScan")
+class DispatchSortScan:
+    """
+    Pass to dispatch scan and sort operators to platform dependent implementation.
+    """
+
+    def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule:
+        sort_scan_dispater = SortScanDispatcher(mod)
+        for gv, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = sort_scan_dispater.visit_expr(func)
+                sort_scan_dispater.builder_.update_func(gv, func)
+        sort_scan_dispater.apply_dlight_gpu_fallback()
+        return sort_scan_dispater.builder_.finalize()

python/tvm/relax/backend/gpu_generic/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+"""The Relax Metal backend compilation pipeline and other passes."""
+
+from .cumsum import gpu_2d_continuous_cumsum
+from .pipeline import (
+    dataflow_lower_passes,
+    finalize_passes,
+    get_default_pipeline,
+    legalize_passes,
+    library_dispatch_passes,
+)
+from .sampling import generic_get_sample_index, gpu_multinomial_from_uniform

python/tvm/relax/backend/gpu_generic/cumsum.py

@@ -0,0 +1,195 @@
+# 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, too-many-nested-blocks
+"""Backend kernels for cumsum operator."""
+
+import math
+
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+
+
+def _is_power_of_two(n: int):
+    """Check if n is a power of 2."""
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def gpu_2d_continuous_cumsum(
+    ty_len: int = 4,
+    tx_len: int = 32,
+    thread_elem: int = 4,
+    in_dtype: str = "int32",
+    out_dtype: str | None = None,
+) -> PrimFunc:
+    """Generate GPU kernel for 2D continuous cumsum, i.e. The cumsum axis is -1
+
+    Parameters
+    ----------
+    ty_len : int
+        The length of `threadIdx.y`
+
+    tx_len : int
+        The length of `threadIdx.x`
+
+    thread_elem : int
+        The number of elements processed by single thread
+
+    in_dtype : str
+        The input data type
+
+    out_dtype : Optional[str]
+        The output data type, if None, it will be the same as in_dtype
+
+    Returns
+    -------
+    cumsum : PrimFunc
+        The generated cumsum kernel
+    """
+
+    out_dtype = out_dtype or in_dtype
+
+    # Configuration for GPU kernel
+    TX = T.int64(tx_len)  # threadIdx.x
+    TY = T.int64(ty_len)  # threadIdx.y
+    N = T.int64(thread_elem)  # number of elements in single thread
+
+    if not _is_power_of_two(TX) or not _is_power_of_two(TY) or not _is_power_of_two(N):
+        raise ValueError("Configuration of TX, TY, N must be power of 2")
+
+    # number of elements to be processed by single warp
+    warp_elem = T.int64(tx_len * thread_elem)
+    # number of elements to be processed by single block(SM)
+    block_elem = T.int64(tx_len * ty_len * thread_elem)
+
+    LOG_TX = T.int64(int(math.log2(tx_len)))
+    LOG_BLOCK_N = T.int64(int(math.log2(tx_len * ty_len * thread_elem)))
+
+    @T.macro
+    def block_inclusive_inside_block(
+        batch: T.int64,
+        cur_len: T.int64,
+        source: T.Buffer,
+        output: T.Buffer,
+        tmp_buf: T.Buffer,
+        src_offset: T.int64,
+        tmp_offset: T.int64,
+    ):
+        for by in T.thread_binding(batch, thread="blockIdx.y"):
+            for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"):
+                with T.sblock():
+                    local_buf = T.sblock_alloc_buffer((thread_elem,), out_dtype, scope="local")
+                    shared_buf = T.sblock_alloc_buffer((block_elem,), out_dtype, scope="shared")
+                    for ty in T.thread_binding(TY, thread="threadIdx.y"):
+                        for tx in T.thread_binding(TX, thread="threadIdx.x"):
+                            tx_idx: T.let[T.int64] = (
+                                bx * block_elem + ty * warp_elem + tx * thread_elem
+                            )
+                            # Load data from global memory
+                            for i in T.vectorized(N):
+                                local_buf[i] = T.if_then_else(
+                                    tx_idx + i < cur_len,
+                                    T.Cast(out_dtype, source[by, src_offset + tx_idx + i]),
+                                    T.Cast(out_dtype, 0),
+                                )
+                            # Inclusive scan inside thread
+                            for i in T.unroll(1, N):
+                                local_buf[i] += local_buf[i - 1]
+                            # Store data to shared memory
+                            for i in T.vectorized(N):
+                                shared_buf[ty * warp_elem + tx * thread_elem + i] = local_buf[i]
+                            # Inclusive scan inside warp
+                            for i in T.unroll(LOG_TX):
+                                for j in T.vectorized(N):
+                                    idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem
+                                    if tx >= (1 << i):
+                                        shared_buf[idx + j] += shared_buf[
+                                            idx - (1 << i) * thread_elem + N - 1
+                                        ]
+                            # Inclusive scan inside block
+                            for i in T.unroll(1, TY):
+                                for j in T.vectorized(N):
+                                    if ty == 0:
+                                        idx: T.let[T.int64] = i * warp_elem + tx * thread_elem
+                                        shared_buf[idx + j] += shared_buf[i * warp_elem - 1]
+                            # Write sum of block to global memory
+                            for i in T.vectorized(N):
+                                idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem + i
+                                if bx * block_elem + idx < cur_len:
+                                    output[by, src_offset + bx * block_elem + idx] = shared_buf[idx]
+                            if tx == 0 and ty == 0:
+                                for i in T.vectorized(N):
+                                    tmp_buf[by, tmp_offset + bx] = shared_buf[block_elem - 1]
+
+    @T.macro
+    def update_cross_block(
+        batch: T.int64,
+        cur_len: T.int64,
+        source: T.Buffer,
+        output: T.Buffer,
+        src_offset: T.int64,
+        out_offset: T.int64,
+    ):
+        for by in T.thread_binding(batch, thread="blockIdx.y"):
+            for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"):
+                for ty in T.thread_binding(TY, thread="threadIdx.y"):
+                    for tx in T.thread_binding(TX, thread="threadIdx.x"):
+                        for i in T.serial(N):
+                            idx: T.let[T.int64] = bx * block_elem + ty * warp_elem + i * TX + tx
+                            if idx < cur_len:
+                                output[by, out_offset + idx] += T.if_then_else(
+                                    bx > 0, source[by, src_offset + bx - 1], 0
+                                )
+
+    @T.prim_func(private=True, s_tir=True)
+    def cumsum(var_a: T.handle, var_out: T.handle):
+        T.func_attr({"tirx.is_scheduled": True})  # prevent further scheduling
+        m, n = T.int64(), T.int64()
+        A = T.match_buffer(var_a, [m, n], dtype=in_dtype)
+        Out = T.match_buffer(var_out, [m, n], dtype=out_dtype)
+        Tmp = T.alloc_buffer([m, n], dtype=out_dtype)
+        total_rounds: T.let[T.int64] = (
+            T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n)))) // LOG_BLOCK_N
+        )
+
+        block_inclusive_inside_block(
+            m, n, A, Out, Tmp, src_offset=T.int64(0), tmp_offset=T.int64(0)
+        )
+        for i in range(total_rounds):
+            cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (i + 1)))
+            block_inclusive_inside_block(
+                m,
+                cur_len,
+                Tmp,
+                Tmp,
+                Tmp,
+                src_offset=i * T.ceildiv(n, block_elem),
+                tmp_offset=(i + 1) * T.ceildiv(n, block_elem),
+            )
+        for i in range(total_rounds - 1):
+            real_idx: T.let[T.int64] = total_rounds - 1 - i - 1
+            cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (real_idx + 1)))
+            update_cross_block(
+                m,
+                cur_len,
+                Tmp,
+                Tmp,
+                src_offset=(real_idx + 1) * T.ceildiv(n, block_elem),
+                out_offset=real_idx * T.ceildiv(n, block_elem),
+            )
+        update_cross_block(m, n, Tmp, Out, src_offset=0, out_offset=0)
+
+    return cumsum

python/tvm/relax/backend/gpu_generic/pipeline.py

@@ -0,0 +1,89 @@
+# 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.
+"""The Relax generic GPU backend compilation pipeline and other passes."""
+
+import tvm
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default library dispatch passes for generic GPU backend."""
+    return [
+        relax.backend.DispatchSampling(),
+        relax.backend.DispatchSortScan(),
+    ]
+
+
+def legalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default legalization passes for generic GPU backend."""
+    from tvm.s_tir import dlight as dl  # pylint: disable=import-outside-toplevel
+
+    return [
+        tvm.relax.transform.LegalizeOps(),
+        tvm.relax.transform.AnnotateTIROpPattern(),
+        tvm.relax.transform.FoldConstant(),
+        tvm.relax.transform.FuseOps(),
+        tvm.relax.transform.FuseTIR(),
+        dl.ApplyDefaultSchedule(
+            dl.gpu.Matmul(),
+            dl.gpu.GEMV(),
+            dl.gpu.Reduction(),
+            dl.gpu.GeneralReduction(),
+            dl.gpu.Fallback(),
+        ),
+    ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default dataflow lowering passes for generic GPU backend."""
+    return [
+        relax.transform.RewriteDataflowReshape(),
+        relax.transform.ToNonDataflow(),
+        relax.transform.RemovePurityChecking(),
+        relax.transform.CallTIRRewrite(),
+    ]
+
+
+def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default finalization passes for generic GPU backend."""
+    return [
+        relax.transform.StaticPlanBlockMemory(),
+        relax.transform.LowerAllocTensor(),
+        relax.transform.KillAfterLastUse(),
+        relax.transform.LowerRuntimeBuiltin(),
+        relax.transform.ComputePrimValue(),
+        relax.transform.VMShapeLower(),
+        relax.transform.AttachGlobalSymbol(),
+    ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+    """Return the default compilation pipeline for generic GPU."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+        with target:
+            seq = tvm.transform.Sequential(
+                library_dispatch_passes(target)
+                + legalize_passes(target)
+                + dataflow_lower_passes(target)
+                + finalize_passes(target)
+            )
+            mod = seq(mod)
+        return mod
+
+    return _pipeline

python/tvm/relax/backend/gpu_generic/sampling.py

@@ -0,0 +1,345 @@
+# 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, too-many-nested-blocks
+"""Backend kernels for sampling operator."""
+
+import math
+from collections.abc import Callable
+
+import tvm
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+
+
+def _is_power_of_two(n: int):
+    """Check if n is a power of 2."""
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def gpu_multinomial_from_uniform(
+    prob_dtype: str = "float32",
+    sample_dtype: str = "float32",
+    sample_indices_dtype: str = "int64",
+    dtype: str = "int64",
+    ty_len: int = 4,
+    tx_len: int = 32,
+    thread_elem: int = 4,
+    eps: float = 1e-6,
+) -> PrimFunc:
+    """Generate GPU kernel for multinomial_from_uniform operator.
+
+    Parameters
+    ----------
+    ty_len : int
+        The length of `threadIdx.y`
+
+    tx_len : int
+        The length of `threadIdx.x`
+
+    thread_elem : int
+        The number of elements processed by single thread
+
+    prob_dtype : str
+        The probability data type
+
+    sample_dtype : str
+        The sample data type
+
+    sample_indices_dtype : str
+        The sample indices data type
+
+    dtype : str
+        The output data type
+
+    Returns
+    -------
+    func : PrimFunc
+        The generated function
+    """
+
+    target = tvm.target.Target.current()
+    target_dtype = "int32" if "webgpu" in str(target) else "int64"
+
+    TX = T.int64(tx_len)  # threadIdx.x
+    TY = T.int64(ty_len)  # threadIdx.y
+
+    # number of elements to be processed by single thread
+    thread_elem = T.int64(thread_elem)
+    # number of elements to be processed by single warp
+    warp_elem = T.int64(tx_len * thread_elem)
+    # number of elements to be processed by single block(SM)
+    block_elem = T.int64(tx_len * ty_len * thread_elem)
+
+    LOG_TX = T.int64(int(math.log2(tx_len)))
+    LOG_TY = T.int64(int(math.log2(ty_len)))
+
+    if (
+        not _is_power_of_two(tx_len)
+        or not _is_power_of_two(ty_len)
+        or not _is_power_of_two(thread_elem)
+    ):
+        raise ValueError(
+            "Configuration of tx_len, ty_len, thread_elem must be power of 2,"
+            f"but got {tx_len}, {ty_len}, {thread_elem}"
+        )
+
+    @T.macro
+    def block_cumsum(
+        ty: T.int64,
+        tx: T.int64,
+        source_local: T.Buffer,
+        output_shared: T.Buffer,
+    ):
+        """cumsum inside block (SM)"""
+        # Inclusive scan inside thread
+        for i in T.unroll(1, thread_elem):
+            source_local[i] += source_local[i - 1]
+        # Store data to shared memory
+        for i in T.vectorized(thread_elem):
+            output_shared[ty * warp_elem + tx * thread_elem + i] = source_local[i]
+        # Inclusive scan inside warp
+        for i in T.unroll(LOG_TX):
+            for j in T.vectorized(thread_elem):
+                idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem
+                if tx >= (1 << i):
+                    output_shared[idx + j] += output_shared[
+                        idx - (1 << i) * thread_elem + thread_elem - 1
+                    ]
+        # Inclusive scan inside block
+        for i in T.unroll(1, TY):
+            for j in T.vectorized(thread_elem):
+                if ty == 0:
+                    idx: T.let[T.int64] = i * warp_elem + tx * thread_elem
+                    output_shared[idx + j] += output_shared[i * warp_elem - 1]
+
+    def compare_bool_not_equal(a: T.bool, b: T.bool) -> T.bool:
+        # Vulkan does not support compare two bool value direct
+        # return a != b
+        return T.Cast("int8", a) != T.Cast("int8", b)
+
+    @T.macro
+    def block_adjacent_difference_left(
+        ty: T.int64,
+        tx: T.int64,
+        source_local: T.Buffer,
+        output_local: T.Buffer,
+    ):
+        with T.sblock():
+            shared_buf = T.sblock_alloc_buffer((TX * TY,), "bool", scope="shared")
+            tx_idx: T.let[T.int64] = ty * TX + tx
+            shared_buf[tx_idx] = source_local[thread_elem - 1]
+            output_local[0] = T.if_then_else(
+                tx_idx != 0,
+                compare_bool_not_equal(source_local[0], shared_buf[tx_idx - 1]),
+                source_local[0],
+            )
+            for i in T.unroll(1, thread_elem):
+                output_local[i] = compare_bool_not_equal(source_local[i], source_local[i - 1])
+
+    def op_reduce_min(a, b):
+        return T.min(a, b)
+
+    def op_reduce_sum(a, b):
+        return a + b
+
+    @T.macro
+    def block_reduce_with_mask(
+        ty: T.int64,
+        tx: T.int64,
+        init_value,
+        data_local: T.Buffer,
+        output_local: T.Buffer,
+        dtype: str,
+        reduce_op: Callable,  # T.macro
+        mask_local: T.Buffer | None = None,
+    ):
+        with T.sblock():
+            local_sum = T.sblock_alloc_buffer((), dtype, scope="local")
+            shared_buf = T.sblock_alloc_buffer((TX * TY,), dtype, scope="shared")
+            idx: T.let[T.int64] = ty * TX + tx
+
+            local_sum[()] = T.Cast(dtype, init_value)
+            for i in T.unroll(thread_elem):
+                if mask_local is not None:
+                    if mask_local[i]:
+                        local_sum[()] = reduce_op(local_sum[()], data_local[i])
+                else:
+                    local_sum[()] = reduce_op(local_sum[()], data_local[i])
+            shared_buf[idx] = local_sum[()]
+
+            for i in T.unroll(LOG_TX + LOG_TY):
+                if idx % (1 << (i + 1)) == 0:
+                    shared_buf[idx] = reduce_op(shared_buf[idx], shared_buf[idx + (1 << i)])
+            output_local[()] = shared_buf[0]
+
+    @T.macro
+    def single_batch_sampling(
+        prob,
+        row_idx,
+        vocab_size,
+        ty,
+        tx,
+        step_iter,
+        threshold,
+        aggregate,
+        uniform_sample,
+        sample_id_local,
+    ):
+        with T.sblock():
+            prob_gt_threshold = T.sblock_alloc_buffer((thread_elem,), prob_dtype, scope="local")
+            cumsum = T.sblock_alloc_buffer((block_elem,), prob_dtype, scope="shared")
+            greater_than_u = T.sblock_alloc_buffer((thread_elem,), "bool", scope="local")
+            mask = T.sblock_alloc_buffer((thread_elem,), "bool", scope="local")
+            valid = T.sblock_alloc_buffer((thread_elem,), "bool", scope="local")
+            indices = T.sblock_alloc_buffer((thread_elem), dtype, scope="local")
+            step_aggregate = T.sblock_alloc_buffer((), prob_dtype, scope="local")
+            # Load prob data from global memory to local memory
+            for v in T.unroll(thread_elem):
+                idx: T.let[T.int64] = step_iter * block_elem + ty * warp_elem + tx * thread_elem + v
+                prob_local: T.let = T.if_then_else(
+                    idx < vocab_size,
+                    prob[row_idx, idx],
+                    T.Cast(prob_dtype, 0),
+                )
+                prob_gt_threshold[v] = T.if_then_else(
+                    prob_local > threshold, prob_local, T.Cast(prob_dtype, 0)
+                )
+                valid[v] = prob_local > threshold and idx < vocab_size
+
+            block_reduce_with_mask(
+                ty,
+                tx,
+                init_value=0,
+                data_local=prob_gt_threshold,
+                output_local=step_aggregate,
+                dtype=prob_dtype,
+                reduce_op=op_reduce_sum,
+                mask_local=None,
+            )
+            if T.tvm_thread_invariant(aggregate[()] + step_aggregate[()] >= uniform_sample - eps):
+                block_cumsum(ty, tx, prob_gt_threshold, cumsum)
+                # Note: it should be `T.vectorized` instead of `T.unroll`
+                # However, it will cause vulkan codegen error
+                for v in T.unroll(thread_elem):
+                    greater_than_u[v] = (
+                        cumsum[ty * warp_elem + tx * thread_elem + v] + aggregate[()]
+                        >= uniform_sample - eps
+                    )
+
+                block_adjacent_difference_left(ty, tx, greater_than_u, mask)
+                # Same as above, it should be `T.vectorized`
+                for v in T.unroll(thread_elem):
+                    mask[v] = mask[v] and valid[v]
+                    indices[v] = step_iter * block_elem + ty * warp_elem + tx * thread_elem + v
+                block_reduce_with_mask(
+                    ty,
+                    tx,
+                    init_value=vocab_size - 1,
+                    data_local=indices,
+                    output_local=sample_id_local,
+                    dtype=dtype,
+                    reduce_op=op_reduce_min,
+                    mask_local=mask,
+                )
+
+            aggregate[()] += step_aggregate[()]
+
+    @T.prim_func(s_tir=True)
+    def parallel_sampling_from_prob(
+        var_prob: T.handle,
+        var_uniform_samples: T.handle,
+        var_row_indices: T.handle,
+        var_sampled_token_ids: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        n, vocab_size, batch_size = T.int64(), T.int64(), T.int64()
+        # match buffers
+        prob = T.match_buffer(var_prob, (n, vocab_size), prob_dtype)
+        uniform_samples = T.match_buffer(var_uniform_samples, (batch_size, 1), sample_dtype)
+        row_indices = T.match_buffer(var_row_indices, (batch_size, 1), sample_indices_dtype)
+        token_ids = T.match_buffer(var_sampled_token_ids, (batch_size, 1), dtype)
+        # local buffers
+        aggregate = T.sblock_alloc_buffer((), prob_dtype, scope="local")
+        sample_id_local = T.sblock_alloc_buffer((), dtype, scope="local")
+        step_iter = T.sblock_alloc_buffer((), "int32", scope="local")
+
+        for bx in T.thread_binding(batch_size, thread="blockIdx.x"):
+            row_idx: T.let[T.int64] = T.Cast("int64", row_indices[bx, 0])
+            for ty in T.thread_binding(TY, thread="threadIdx.y"):
+                for tx in T.thread_binding(TX, thread="threadIdx.x"):
+                    u: T.let[T.float32] = uniform_samples[bx, 0]
+                    aggregate[()] = T.Cast(prob_dtype, 0)
+                    step_iter[()] = T.int32(0)
+                    # at least one iteration
+                    while T.tvm_thread_invariant(
+                        (step_iter[()] == 0 or aggregate[()] < u - eps)
+                        and T.Cast(target_dtype, step_iter[()])
+                        < T.Cast(target_dtype, T.ceildiv(vocab_size, block_elem))
+                    ):
+                        single_batch_sampling(
+                            prob,
+                            row_idx,
+                            vocab_size,
+                            ty,
+                            tx,
+                            T.Cast(target_dtype, step_iter[()]),
+                            0.0,
+                            aggregate,
+                            u,
+                            sample_id_local,
+                        )
+                        step_iter[()] += 1
+                    if tx == 0 and ty == 0:
+                        token_ids[bx, 0] = sample_id_local[()]
+
+    return parallel_sampling_from_prob
+
+
+def generic_get_sample_index(
+    prob_dtype: str = "float32",
+    sample_dtype: str = "float32",
+    sample_indices_dtype: str = "int64",
+    dtype: str = "int64",
+):
+    """Generate a generic get_sample_index kernel."""
+
+    @T.prim_func(private=True, s_tir=True)
+    def _get_sample_index(A: T.handle, B: T.handle, C: T.handle, D: T.handle):
+        batch, vocab_size = T.int64(), T.int64()
+        prob = T.match_buffer(A, (batch, vocab_size), prob_dtype)
+        out_batch = T.int64()
+        usample = T.match_buffer(B, (out_batch, 1), sample_dtype)
+        sample_indices = T.match_buffer(C, (out_batch, 1), sample_indices_dtype)
+        output_index = T.match_buffer(D, (out_batch, 1), dtype)
+
+        for ax0, ax1 in T.grid(out_batch, vocab_size):
+            with T.sblock("T_get_sample_index"):
+                v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+                T.writes(output_index[v_ax0, 0])
+                if (
+                    usample[v_ax0, T.int64(0)] < prob[sample_indices[v_ax0, T.int64(0)], v_ax1]
+                    or v_ax1 + 1 == vocab_size
+                ):
+                    if v_ax1 == 0:
+                        output_index[v_ax0, 0] = 0
+                    elif (
+                        usample[v_ax0, T.int64(0)]
+                        >= prob[sample_indices[v_ax0, T.int64(0)], v_ax1 - 1]
+                    ):
+                        output_index[v_ax0, 0] = v_ax1
+
+    return _get_sample_index

python/tvm/relax/backend/metal/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+"""The Relax Metal backend compilation pipeline and other passes."""

python/tvm/relax/backend/metal/coreml.py

@@ -0,0 +1,496 @@
+# 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, import-outside-toplevel
+"""Pattern table and codegen for CoreML"""
+
+import os
+import shutil
+
+import tvm_ffi
+
+import tvm
+from tvm.contrib import coreml_runtime
+from tvm.ir import Call, PrimType
+from tvm.relax import transform
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.expr import (
+    BindingBlock,
+    Constant,
+    Function,
+    SeqExpr,
+    Var,
+    VarBinding,
+)
+from tvm.relax.transform import PatternCheckContext
+from tvm.relax.type import TensorType
+from tvm.support.xcode import compile_coreml
+
+from ...expr_functor import PyExprVisitor, visitor
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import make_matmul_pattern
+
+
+def _check_default(context: PatternCheckContext) -> bool:
+    return True
+
+
+def default_binary_patterns(op_name: str):
+    """
+    Returns a list of binary op patterns in coreML BYOC backend.
+    """
+
+    def _make_binary_pattern():
+        lhs = wildcard()
+        rhs = wildcard()
+        out = is_op("relax." + op_name)(lhs, rhs)
+        annotations = {"lhs": lhs, "rhs": rhs, "root": out}
+        return out, annotations
+
+    def _binary_pattern(pattern_name):
+        return (pattern_name, *_make_binary_pattern(), _check_default)
+
+    return [_binary_pattern("coreml." + op_name)]
+
+
+def default_unary_patterns(op_name: str):
+    """
+    Returns a list of unary op patterns in coreML BYOC backend.
+    """
+
+    def _make_unary_pattern():
+        lhs = wildcard()
+        out = is_op("relax." + op_name)(lhs)
+        annotations = {"lhs": lhs, "root": out}
+        return out, annotations
+
+    def _unary_pattern(pattern_name):
+        return (pattern_name, *_make_unary_pattern(), _check_default)
+
+    return [_unary_pattern("coreml." + op_name)]
+
+
+def conv2d_patterns():
+    """
+    Returns a list of conv2d patterns in coreML BYOC backend.
+    """
+
+    def _make_conv2d_pattern():
+        lhs = wildcard()
+        rhs = wildcard()
+        out = is_op("relax.nn.conv2d")(lhs, rhs)
+        annotations = {"lhs": lhs, "rhs": rhs, "root": out}
+        return out, annotations
+
+    def _conv2d_pattern(pattern_name):
+        return (pattern_name, *_make_conv2d_pattern(), _check_default)
+
+    return [_conv2d_pattern("coreml.nn.conv2d")]
+
+
+def matmul_patterns():
+    """
+    Returns a list of all matmul patterns in coreML BYOC backend.
+    """
+
+    def _matmul_pattern(pattern_name):
+        return (
+            pattern_name,
+            *make_matmul_pattern(),
+            _check_default,
+        )
+
+    return [_matmul_pattern("coreml.matmul")]
+
+
+def clip_patterns():
+    """
+    Returns a list of clip patterns in coreML BYOC backend.
+    """
+
+    def _make_clip_pattern():
+        arg0 = wildcard()
+        arg1 = wildcard()
+        arg2 = wildcard()
+        out = is_op("relax.clip")(arg0, arg1, arg2)
+        annotations = {"arg0": arg0, "arg1": arg1, "arg2": arg2, "root": out}
+        return out, annotations
+
+    def _conv2d_pattern(pattern_name):
+        return (pattern_name, *_make_clip_pattern(), _check_default)
+
+    return [_conv2d_pattern("coreml.clip")]
+
+
+register_patterns(
+    [
+        *default_binary_patterns(op_name="add"),
+        *default_binary_patterns(op_name="multiply"),
+        *default_unary_patterns(op_name="nn.softmax"),
+        *default_unary_patterns(op_name="nn.relu"),
+        *default_unary_patterns(op_name="expand_dims"),
+        *default_unary_patterns(op_name="nn.avg_pool2d"),
+        *default_unary_patterns(op_name="nn.batch_flatten"),
+        *conv2d_patterns(),
+        *clip_patterns(),
+        *matmul_patterns(),
+    ]
+)
+
+
+def partition_for_coreml(mod):
+    """
+    Partition the input module into coreml-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the coreml backend.
+    """
+
+    patterns = get_patterns_with_prefix("coreml")
+    mod = transform.CanonicalizeBindings()(mod)
+    mod = transform.FuseOpsByPattern(patterns, bind_constants=True, annotate_codegen=False)(mod)
+    mod = transform.MergeCompositeFunctions()(mod)
+    return mod
+
+
+# Codegen for coreml API reference:
+# https://apple.github.io/coremltools/source/coremltools.models.neural_network.html
+def _convert_add(builder, name, inputs, outputs, args, attrs):
+    builder.add_elementwise(name=name, input_names=inputs, output_name=outputs[0], mode="ADD")
+
+
+def _convert_multiply(builder, name, inputs, outputs, args, attrs):
+    builder.add_elementwise(name=name, input_names=inputs, output_name=outputs[0], mode="MULTIPLY")
+
+
+def _convert_matmul(builder, name, inputs, outputs, args, attrs):
+    builder.add_batched_mat_mul(
+        name=name,
+        input_names=inputs,
+        output_name=outputs[0],
+    )
+
+
+def _convert_clip(builder, name, inputs, outputs, args, attrs):
+    builder.add_clip(
+        name=name,
+        input_name=inputs[0],
+        output_name=outputs[0],
+        min_value=inputs[1],
+        max_value=inputs[2],
+    )
+
+
+def _convert_batch_flatten(builder, name, inputs, outputs, args, attrs):
+    builder.add_flatten_to_2d(name=name, input_name=inputs[0], output_name=outputs[0])
+
+
+def _convert_expand_dims(builder, name, inputs, outputs, args, attrs):
+    axes = [int(v) for v in attrs["axis"]]
+    builder.add_expand_dims(name=name, input_name=inputs[0], output_name=outputs[0], axes=axes)
+
+
+def _convert_relu(builder, name, inputs, outputs, args, attrs):
+    builder.add_activation(
+        name=name, non_linearity="RELU", input_name=inputs[0], output_name=outputs[0]
+    )
+
+
+def _convert_softmax(builder, name, inputs, outputs, args, attrs):
+    builder.add_softmax_nd(
+        name=name, input_name=inputs[0], output_name=outputs[0], axis=int(attrs["axis"])
+    )
+
+
+def _convert_conv2d(builder, name, inputs, outputs, args, attrs):
+    weight = args[1].data.numpy()
+    oc, kc, kh, kw = weight.shape
+
+    builder.add_convolution(
+        name=name,
+        kernel_channels=kc,
+        output_channels=oc,
+        height=kh,
+        width=kw,
+        stride_height=int(attrs["strides"][0]),
+        stride_width=int(attrs["strides"][0]),
+        border_mode="valid",
+        groups=int(attrs["groups"]),
+        W=weight,
+        b=None,
+        has_bias=False,
+        input_name=inputs[0],
+        output_name=outputs[0],
+        dilation_factors=[int(v) for v in attrs["dilation"]],
+        padding_top=int(attrs["padding"][0]),
+        padding_bottom=int(attrs["padding"][2]),
+        padding_left=int(attrs["padding"][1]),
+        padding_right=int(attrs["padding"][3]),
+    )
+
+
+def _convert_avg_pool2d(builder, name, inputs, outputs, args, attrs):
+    builder.add_pooling(
+        name=name,
+        height=1,
+        width=1,
+        stride_height=1,
+        stride_width=1,
+        layer_type="AVERAGE",
+        padding_type="VALID",
+        input_name=inputs[0],
+        output_name=outputs[0],
+    )
+
+
+_convert_map = {
+    "add": _convert_add,
+    "multiply": _convert_multiply,
+    "matmul": _convert_matmul,
+    "clip": _convert_clip,
+    "expand_dims": _convert_expand_dims,
+    "nn.relu": _convert_relu,
+    "nn.batch_flatten": _convert_batch_flatten,
+    "nn.softmax": _convert_softmax,
+    "nn.conv2d": _convert_conv2d,
+    "nn.avg_pool2d": _convert_avg_pool2d,
+}
+
+
+@visitor
+class CallNodeInfoCollector(PyExprVisitor):
+    """
+    Collect Expr, Constant and attributes in the inner function
+    """
+
+    def __init__(self, op_name):
+        self.primvals = []
+        self.attrs = []
+        self.consts = []
+        self.op_name = op_name
+
+    def visit_call_(self, call: Call) -> None:
+        self.attrs.append(call.attrs)
+        for arg in call.args:
+            if tvm.ir.is_prim_expr(arg):
+                self.primvals.append(arg)
+            if isinstance(arg, Constant):
+                self.consts.append(arg)
+
+    def collect(self, expr):
+        self.visit_expr(expr)
+        return self.primvals, self.attrs, self.consts
+
+
+@visitor
+class CodegenCoreML(PyExprVisitor):
+    """
+    A visitor to traverse subgraphs and build Core ML models.
+    """
+
+    def __init__(self, model_name, function):
+        try:
+            import coremltools
+            from coremltools.models.neural_network import NeuralNetworkBuilder
+        except ImportError as err:
+            raise ImportError(
+                "coremltools is required by the CoreML backend. "
+                "Install it with: pip install coremltools"
+            ) from err
+
+        self.model_name = model_name
+        self.function = function
+        self.out_map = {}
+        self.const_map = {}  # (buffer name, object)
+        self.model_inputs_ = []
+        self.buf_idx_ = 0
+
+        getter = tvm.get_global_func("relax.analysis.get_var2val")
+        assert getter, "Cannot find `relax.analysis.get_var2val` function."
+
+        self.var2val = getter(function)
+        self.cur_binding_var = None
+
+        inputs = [
+            (
+                "",
+                coremltools.models.datatypes.Array(
+                    1,
+                ),
+            )
+            for _ in self.function.params
+        ]
+        outputs = [
+            (
+                "",
+                coremltools.models.datatypes.Array(
+                    1,
+                ),
+            )
+        ]
+        self.builder = NeuralNetworkBuilder(inputs, outputs, disable_rank5_shape_mapping=True)
+
+    def visit_function_(self, op) -> None:
+        for var in op.params:
+            name = var.name_hint
+            ty = var.ty
+            if isinstance(ty, TensorType):
+                shape = [int(v) for v in list(ty.shape)]
+            elif isinstance(ty, PrimType):
+                shape = []
+            else:
+                raise Exception("Currently not supported: ", type(ty))
+            dtype = ty.dtype
+            self.model_inputs_.append((name, shape, dtype))
+
+        self.visit_expr(op.body)
+
+    def visit_var_(self, var):
+        self.out_map[var] = [var.name_hint]
+        prev_binding_var = self.cur_binding_var
+        self.cur_binding_var = var
+        if var in self.var2val:
+            self.visit_expr(self.var2val[var])
+        self.cur_binding_var = prev_binding_var
+
+    def visit_call_(self, call: Call) -> None:
+        assert isinstance(call.op, Var)
+        assert call.op in self.var2val
+        func = self.var2val[call.op]
+
+        assert "Composite" in func.attrs, "Only composite functions are supported."
+        composite_name = func.attrs["Composite"]
+
+        # Get the op name and remove "relax." prefix.
+        op_name = composite_name[7:]
+
+        inputs = []
+        args = []
+        for arg in call.args:
+            args.append(arg)
+            super().visit_expr(arg)
+            for out in self.out_map[arg]:
+                inputs.append(out)
+
+        primvals, attrs, consts = CallNodeInfoCollector(op_name).collect(func.body)
+        for arg in primvals:
+            args.append(arg)
+            inputs.append(arg.value.value)
+
+        for arg in consts:
+            output = "buf_" + str(self.buf_idx_)
+            self.builder.add_load_constant_nd(
+                name=output,
+                output_name=output,
+                constant_value=arg.data.numpy(),
+                shape=arg.data.shape,
+            )
+            self.buf_idx_ = self.buf_idx_ + 1
+            self.out_map[arg] = [output]
+            inputs.append(output)
+            args.append(arg)
+
+        layer_name = op_name + "_" + str(self.buf_idx_)
+
+        assert op_name in _convert_map, f"{op_name} is not supported"
+        outputs = ["buf_" + str(self.buf_idx_)]
+        _convert_map[op_name](self.builder, layer_name, inputs, outputs, args, attrs[0])
+        self.buf_idx_ = self.buf_idx_ + 1
+        self.out_map[self.cur_binding_var] = outputs
+
+    def visit_var_binding_(self, binding: VarBinding) -> None:
+        # Visit var of the last binding
+        self.visit_expr(binding.var)
+
+    def visit_binding_block_(self, block: BindingBlock) -> None:
+        # We only visit the last VarBinding to retrieve
+        # target composite function
+        self.visit_binding(block.bindings[-1])
+
+    def visit_seq_expr_(self, op: SeqExpr) -> None:
+        for bb in op.blocks:
+            self.visit_binding_block_(bb)
+
+    def serialize(self, func: Function):
+        self.visit_expr(func)
+
+    def compile(self, out_dir):
+        """
+        Build a Core ML model and compile it with Xcode toolchain.
+        """
+        import coremltools
+        from coremltools.proto.Model_pb2 import ArrayFeatureType
+
+        FEATURE_TYPE_MAP = {
+            "float32": ArrayFeatureType.FLOAT32,
+            "float64": ArrayFeatureType.DOUBLE,
+            "int32": ArrayFeatureType.INT32,
+        }
+        input_names, input_dims, input_dtypes = zip(*self.model_inputs_)
+        self.builder.set_input(input_names, input_dims)
+
+        for i, dtype in enumerate(input_dtypes):
+            assert dtype in FEATURE_TYPE_MAP
+            input_desc = self.builder.spec.description.input
+            input_desc[i].type.multiArrayType.dataType = FEATURE_TYPE_MAP[dtype]
+
+        output_dim = [int(n) for n in self.function.ty.ret.shape]
+
+        last_binding_var = self.function.body.blocks[0].bindings[-1].var
+        self.builder.set_output(self.out_map[last_binding_var], [output_dim])
+
+        for i, dtype in enumerate([self.function.ty.ret.dtype]):
+            assert dtype in FEATURE_TYPE_MAP
+            output_desc = self.builder.spec.description.output
+            output_desc[i].type.multiArrayType.dataType = FEATURE_TYPE_MAP[dtype]
+
+        model = coremltools.models.MLModel(self.builder.spec)
+        compile_coreml(model, self.model_name, out_dir)
+
+
+@tvm_ffi.register_global_func("relax.ext.coreml")
+def coreml_compiler(funcs, options, constant_names):
+    """
+    Create a CoreML runtime from a Relax module.
+    """
+    compiled_funcs = []
+    for func in funcs:
+        assert isinstance(func, tvm.relax.Function)
+        model_dir = os.getcwd() + "/tmp/"
+        if not os.path.exists(model_dir):
+            os.mkdir(model_dir)
+
+        name = str(func.attrs.global_symbol)
+        builder = CodegenCoreML(name, func)
+        builder.serialize(func)
+
+        mlmodelc_path = f"{model_dir}/{name}.mlmodelc"
+
+        if os.path.exists(mlmodelc_path):
+            shutil.rmtree(mlmodelc_path)
+
+        builder.compile(model_dir)
+        dev = tvm.cpu(0)
+        compiled_funcs.append(coreml_runtime.create(name, mlmodelc_path, dev).module)
+    return compiled_funcs

python/tvm/relax/backend/pattern_registry.py

@@ -0,0 +1,119 @@
+# 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.
+
+"""Pattern registry for BYOC backends"""
+
+import atexit
+from collections.abc import Callable, Mapping
+
+from tvm.relax.dpl import DFPattern
+from tvm.relax.transform import FusionPattern
+
+from ..expr import Expr
+from . import _ffi_api
+
+_REGISTERED_PATTERN_NAMES: set[str] = set()
+
+
+def _cleanup_registered_patterns():
+    _ffi_api.RemovePatterns(list(_REGISTERED_PATTERN_NAMES))  # type: ignore # pylint: disable=no-member
+
+
+_CLEANUP_REGISTERED = False
+
+
+def _ensure_cleanup_function_registered():
+    """
+    Add a cleanup function to be called on interpreter termination, to remove all
+    patterns registered on the Python side. Without cleaning up those patterns,
+    program will segfault on termination. It's because the check functions of pattern
+    entries are referenced from the static memory of libtvm, thus they will be cleaned
+    up at the very end, making calls to Py_DecRef after Python interpreter terminates.
+    """
+    global _CLEANUP_REGISTERED  # pylint: disable=global-statement
+
+    if not _CLEANUP_REGISTERED:
+        atexit.register(_cleanup_registered_patterns)
+        _CLEANUP_REGISTERED = True
+
+
+CheckFunc = Callable[[Mapping[DFPattern, Expr], Expr], bool]
+Pattern = (
+    FusionPattern
+    | tuple[str, DFPattern]
+    | tuple[str, DFPattern, Mapping[str, DFPattern]]
+    | tuple[str, DFPattern, Mapping[str, DFPattern], CheckFunc]
+)
+
+
+def register_patterns(patterns: list[Pattern]):
+    """
+    Register patterns which will be used to partition the DataflowBlock into
+    subgraphs that are supported by external backends.
+
+    Parameters
+    ----------
+    patterns: List[Pattern]
+        Patterns to be registered. Patterns that appear later in the list have
+        higher priority when partitioning DataflowBlock.
+    """
+    _ensure_cleanup_function_registered()
+
+    entries = []
+    for item in patterns:
+        if isinstance(item, FusionPattern):
+            entries.append(item)
+        elif isinstance(item, tuple):
+            entries.append(FusionPattern(*item))
+            _REGISTERED_PATTERN_NAMES.add(item[0])
+        else:
+            raise TypeError(f"Cannot register type {type(item)} as pattern")
+    _ffi_api.RegisterPatterns(entries)
+
+
+def get_patterns_with_prefix(prefix: str) -> list[FusionPattern]:
+    """
+    Get a list of patterns whose names startwith `prefix`.
+
+    Parameters
+    ----------
+    prefix: str
+        The prefix of pattern name.
+
+    Returns
+    -------
+    patterns: FusionPattern
+        Matched patterns, ordered by priority from high to low.
+    """
+    return _ffi_api.GetPatternsWithPrefix(prefix)
+
+
+def get_pattern(name: str) -> FusionPattern | None:
+    """
+    Find the pattern with a particular name.
+
+    Parameters
+    ----------
+    name: str
+        The pattern name.
+
+    Returns
+    -------
+    pattern: Optional[FusionPattern]
+        The matched pattern. Returns None if such pattern is not found.
+    """
+    return _ffi_api.GetPattern(name)

python/tvm/relax/backend/patterns.py

@@ -0,0 +1,643 @@
+# 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
+"""Common patterns used in BYOC"""
+
+from collections.abc import Mapping
+
+from tvm.relax.dpl.pattern import (
+    DFPattern,
+    GlobalVarPattern,
+    TuplePattern,
+    is_const,
+    is_op,
+    is_tuple_get_item,
+    wildcard,
+)
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+def _with_bias_activation_pattern(
+    out: DFPattern,
+    annotations: dict[str, DFPattern],
+    with_bias: bool = False,
+    activation: str | None = None,
+    allow_reshape: bool = False,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    if with_bias:
+        annotations["bias"] = bias = wildcard()
+        if allow_reshape:
+            reshaped_bias = is_op("relax.reshape")(bias, wildcard(), varg_default_wildcard=True)
+            out = is_op("relax.add")(out, reshaped_bias, varg_default_wildcard=True)
+        else:
+            out = is_op("relax.add")(out, bias)
+
+    if activation:
+        out = is_op(activation)(out)
+
+    return out, annotations
+
+
+def make_fused_bias_activation_pattern(
+    op_name: str,
+    with_bias: bool = False,
+    activation: str | None = None,
+    allow_reshape: bool = False,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    """
+    A simple utility to create patterns for an operation fused with bias addition and activation.
+
+    Parameters
+    ----------
+    op_name: str
+        The name of a Relax op, such as "relax.nn.conv2d"
+
+    with_bias: bool
+        Whether or not to include bias addition
+
+    activation: str
+        The name of an activation Relax op, such as "relax.nn.relu"
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a fused operation
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract
+        important expressions from match result, to power the partition
+        check function and codegen.
+    """
+    lhs = wildcard()
+    rhs = wildcard()
+    out = is_op(op_name)(lhs, rhs)
+    annotations = {"lhs": lhs, "rhs": rhs, "root": out}
+
+    return _with_bias_activation_pattern(out, annotations, with_bias, activation, allow_reshape)
+
+
+def make_residual_block_pattern(
+    node_output: DFPattern | tuple[DFPattern, Mapping[str, DFPattern]],
+    binary_op="relax.add",
+    activation=None,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    """
+    Create pattern for residual block.
+
+    Parameters
+    ----------
+    node_output: Union[DFPattern, Tuple[DFPattern, Mapping[str, DFPattern]]]
+        The output of previous node.
+
+    binary_op: str
+        The op used to combine previous node output and residual input.
+
+    activation: str
+        The activation function of this residual block. It should be a name of
+        activation Relax op, such as "relax.nn.relu".
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a matrix multiplication.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract
+        important expressions from match result, to power the partition
+        check function and codegen.
+    """
+
+    if isinstance(node_output, tuple):
+        node_output, arg_patterns = node_output
+    else:
+        arg_patterns = {}
+
+    residual_input = wildcard()
+    op = is_op(binary_op)
+    output = op(node_output, residual_input) | op(residual_input, node_output)
+
+    if activation is not None:
+        output = is_op(activation)(output)
+
+    return output, {**arg_patterns, "residual": residual_input}
+
+
+def make_conv2d_pattern(
+    with_bias: bool = False,
+    activation: str | None = None,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    """
+    Create pattern for 2D convolution.
+
+    Parameters
+    ----------
+    with_bias: bool
+        Whether or not to include bias addition
+
+    activation: str
+        The name of an activation Relax op, such as "relax.nn.relu"
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a 2D convolution.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract
+        important expressions from match result, to power the partition
+        check function and codegen.
+    """
+
+    input_tensor = wildcard()
+    kernel = wildcard()
+    annotations = {"input": input_tensor, "weight": kernel}
+
+    conv2d = is_op("relax.nn.conv2d")(input_tensor, kernel)
+    annotations["root"] = conv2d
+
+    return _with_bias_activation_pattern(conv2d, annotations, with_bias, activation)
+
+
+def make_matmul_pattern(
+    with_bias: bool = False,
+    activation: str | None = None,
+    transposed_rhs: bool = False,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    """
+    Create pattern for matrix multiplication.
+
+    Parameters
+    ----------
+    with_bias: bool
+        Whether or not to include bias addition
+
+    activation: str
+        The name of an activation Relax op, such as "relax.nn.relu"
+
+    transposed_rhs: bool
+        Whether the right hand side of multiplication is transposed.
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a matrix multiplication.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract
+        important expressions from match result, to power the partition
+        check function and codegen.
+    """
+
+    lhs = wildcard()
+    rhs = wildcard()
+    annotations = {"lhs": lhs, "rhs": rhs}
+
+    if transposed_rhs:
+        rhs = is_op("relax.permute_dims")(rhs)
+
+    out = is_op("relax.matmul")(lhs, rhs)
+    annotations["root"] = out
+
+    return _with_bias_activation_pattern(out, annotations, with_bias, activation)
+
+
+def make_attention_pattern(with_bias: bool = False, var_len: bool = False):
+    """
+    Create pattern for fused multi head attention.
+
+    Parameters
+    ----------
+    with_bias: bool
+        Whether or not to include bias addition.
+
+    var_len: bool
+        Whether or not to make a pattern for batched attention with variable sequence lengths.
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a fused multi head attention.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract
+        important expressions from match result, to power the partition
+        check function and codegen.
+    """
+    query = wildcard()
+    key = wildcard()
+    value = wildcard()
+    annotations = {"query": query, "key": key, "value": value}
+    if with_bias:
+        bias = wildcard()
+        annotations["bias"] = bias
+        out = is_op("relax.nn.attention_bias")(query, key, value, bias)
+    elif var_len:
+        seqstart_q = wildcard()
+        seqstart_k = wildcard()
+        max_seqlen_q = wildcard()
+        max_seqlen_k = wildcard()
+        annotations.update(
+            {
+                "seqstart_q": seqstart_q,
+                "seqstart_k": seqstart_k,
+                "max_seqlen_q": max_seqlen_q,
+                "max_seqlen_k": max_seqlen_k,
+            }
+        )
+        out = is_op("relax.nn.attention_var_len")(
+            query, key, value, seqstart_q, seqstart_k, max_seqlen_q, max_seqlen_k
+        )
+    else:
+        out = is_op("relax.nn.attention")(query, key, value)
+
+    return out, annotations
+
+
+def make_stacked_attention_pattern(start_op: str, with_bias: bool = False, layout="BS3NH"):
+    """
+    Create pattern for fused multi head attention with stacked input.
+
+    Parameters
+    ----------
+    start_op: str
+        The starting op for pattern, i.e. `R.split` or `R.strided_slice`.
+
+    with_bias: bool
+        Whether or not to include bias addition
+
+    layout: str
+        The layout of the stacked input tensor.
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a fused multi head attention.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract
+        important expressions from match result, to power the partition
+        check function and codegen.
+    """
+    stacked_qkv = wildcard()
+    ops = {}
+    if start_op == "split":
+        ops["split"] = qkv_tuple = is_op("relax.split")(stacked_qkv)
+        query_raw = is_tuple_get_item(qkv_tuple, 0)
+        key_raw = is_tuple_get_item(qkv_tuple, 1)
+        value_raw = is_tuple_get_item(qkv_tuple, 2)
+    elif start_op == "strided_slice":
+        ops["strided_slice_query"] = query_raw = is_op("relax.strided_slice")(
+            stacked_qkv, varg_default_wildcard=True
+        )
+        ops["strided_slice_key"] = key_raw = is_op("relax.strided_slice")(
+            stacked_qkv, varg_default_wildcard=True
+        )
+        ops["strided_slice_value"] = value_raw = is_op("relax.strided_slice")(
+            stacked_qkv, varg_default_wildcard=True
+        )
+    else:
+        raise NotImplementedError()
+    query_reshape_list = wildcard()
+    key_reshape_list = wildcard()
+    value_reshape_list = wildcard()
+    if layout == "BS3NH":
+        query = is_op("relax.reshape")(query_raw, query_reshape_list)
+        key = is_op("relax.reshape")(key_raw, key_reshape_list)
+        value = is_op("relax.reshape")(value_raw, value_reshape_list)
+    elif layout == "SBN3H":
+        ops["q_transpose"] = query = is_op("relax.permute_dims")(query_raw)
+        ops["k_transpose"] = key = is_op("relax.permute_dims")(key_raw)
+        ops["v_transpose"] = value = is_op("relax.permute_dims")(value_raw)
+    annotations = {
+        "stacked_qkv": stacked_qkv,
+        "query_reshape_list": query_reshape_list,
+        "key_reshape_list": key_reshape_list,
+        "value_reshape_list": value_reshape_list,
+        **ops,
+    }
+    if with_bias:
+        bias = wildcard()
+        annotations["bias"] = bias
+        out = is_op("relax.nn.attention_bias")(query, key, value, bias)
+    else:
+        out = is_op("relax.nn.attention")(query, key, value)
+
+    if layout == "SBN3H":
+        out = is_op("relax.permute_dims")(out)
+
+    return out, annotations
+
+
+def make_layer_norm_pattern():
+    """Create a layer norm pattern."""
+    inp = wildcard()
+    gamma = wildcard()
+    beta = wildcard()
+
+    return is_op("relax.nn.layer_norm")(inp, gamma, beta), {}
+
+
+def make_rms_norm_pattern():
+    """Create a layer norm pattern."""
+    inp = wildcard()
+    weight = wildcard()
+    gv = GlobalVarPattern()
+    out = is_op("relax.call_tir")(gv, TuplePattern([inp, weight]))
+    annotations = {"gv": gv, "inp": inp, "rms_norm": out}
+    return out, annotations
+
+
+def make_matmul_dequantize_pattern(
+    transposed_rhs: bool = False,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    """
+    Create pattern for matrix multiplication and dequantize operation.
+
+    Parameters
+    ----------
+    transposed_rhs: bool
+        Whether the right hand side of multiplication is transposed.
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a matrix multiplication.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract important expressions from
+        match result, to power the partition check function and codegen.
+    """
+
+    lhs = wildcard()
+    rhs = wildcard()
+    annotations = {"lhs": lhs, "rhs": rhs}
+
+    if transposed_rhs:
+        rhs = is_op("relax.permute_dims")(rhs)
+
+    out = is_op("relax.matmul")(lhs, rhs)
+    annotations["root"] = out
+
+    scale = is_const()
+    zp = is_const()
+    annotations.update({"scale": scale, "zp": zp})
+
+    out = is_op("relax.dequantize")(out, scale, zp)
+
+    return out, annotations
+
+
+def make_matmul_multiply_pattern(
+    transposed_rhs: bool = False,
+) -> tuple[DFPattern, Mapping[str, DFPattern]]:
+    """
+    Create pattern for matrix multiplication and multiply operation.
+
+    Parameters
+    ----------
+    transposed_rhs: bool
+        Whether the right hand side of multiplication is transposed.
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing a matrix multiplication.
+
+    annotations: Mapping[str, DFPattern]
+        A mapping from name to sub pattern. It can be used to extract important expressions from
+        match result, to power the partition check function and codegen.
+    """
+
+    lhs = wildcard()
+    rhs = wildcard()
+    scaleA = wildcard()
+    scaleB = wildcard()
+    annotations = {"lhs": lhs, "rhs": rhs, "scaleA": scaleA, "scaleB": scaleB}
+
+    if transposed_rhs:
+        rhs = is_op("relax.permute_dims")(rhs)
+    out = is_op("relax.matmul")(lhs, rhs)
+    annotations["root"] = out
+    scale = is_op("relax.multiply")(scaleA.has_shape((1,)), scaleB.has_shape((1,)))
+    out = is_op("relax.multiply")(out, scale)
+    out = is_op("relax.astype")(out)
+
+    return out, annotations
+
+
+def make_attention_rewrite_pattern(
+    qkv_layout: str, out_layout: str, with_bias: bool, with_cast: bool, with_kv_repeat: bool = False
+):
+    """
+    Create pattern for implicit fused multi head attention rewriting.
+
+    Parameters
+    ----------
+    qkv_layout: str
+        The layout of the query, key and value tensor, i.e. BSNH or BSH.
+
+    out_layout: str
+        The layout of the output tensor, i.e. BSNH or BSH.
+
+    with_bias: bool
+        Whether or not to include bias addition.
+
+    with_cast: bool
+        Whether or not rewriting is intended to be applied to a module after the FP16 conversion
+        pass.
+
+    with_kv_repeat: bool
+        Whether or not to include the Relax repeat op in the pattern, which is typically used
+        in a Relax module to support multi-query attention.
+
+    Returns
+    -------
+    pattern: DFPattern
+        The resulting pattern describing an implicit fused multi head attention.
+
+    rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+        The rewriter for the pattern. It will check the matched patterns, and rewrite.
+        If the matched pattern is not able to be rewritten to `R.nn.attention`, the rewriter
+        returns the original IR.
+    """
+
+    # pylint: disable=invalid-name
+    def handle_input(tensor, layout, transpose, repeat=False):
+        if repeat:
+            tensor = is_op("relax.repeat")(tensor)
+
+        if layout == "BSNH":
+            permuted = is_op("relax.permute_dims")(tensor)
+            shape = wildcard()
+            reshaped = is_op("relax.reshape")(permuted, shape)
+            if transpose:
+                transposed = is_op("relax.permute_dims")(reshaped)
+
+            def rewriter(matchings, x):
+                if matchings[tensor].ty.ndim != 4:
+                    return None
+                if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
+                    return None
+                before_reshape = matchings[permuted].ty.shape.values
+                after_reshape = matchings[shape].ty.values
+                if not (
+                    len(before_reshape) == 4
+                    and len(after_reshape) == 3
+                    and before_reshape[-2:] == after_reshape[-2:]
+                ):
+                    return None
+                if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
+                    return None
+                return x, x.ty.shape
+
+            if transpose:
+                return transposed, rewriter
+            else:
+                return reshaped, rewriter
+        elif layout == "BSH":
+            if transpose:
+                transposed = is_op("relax.permute_dims")(tensor)
+
+            def rewriter(matchings, x):
+                if matchings[tensor].ty.ndim != 3:
+                    return None
+                if transpose and list(matchings[transposed].attrs.axes) != [0, 2, 1]:
+                    return None
+                before_reshape = x.ty.shape.values
+                after_reshape = [before_reshape[0], before_reshape[1], 1, before_reshape[2]]
+                return R.reshape(x, after_reshape), after_reshape
+
+            if transpose:
+                return transposed, rewriter
+            else:
+                return tensor, rewriter
+        else:
+            raise NotImplementedError()
+
+    def handle_output(tensor, layout):
+        if layout == "BSNH":
+            shape = wildcard()
+            reshaped = is_op("relax.reshape")(tensor, shape)
+            permuted = is_op("relax.permute_dims")(reshaped)
+
+            def rewriter(matchings, x):
+                if matchings[tensor].ty.ndim != 3:
+                    return None
+                before_reshape = matchings[tensor].ty.shape.values
+                after_reshape = matchings[shape].ty.values
+                if not (
+                    len(before_reshape) == 3
+                    and len(after_reshape) == 4
+                    and before_reshape[-2:] == after_reshape[-2:]
+                ):
+                    return None
+                if list(matchings[permuted].attrs.axes) != [0, 2, 1, 3]:
+                    return None
+                return x
+
+            return permuted, rewriter
+        elif layout == "BSH":
+
+            def rewriter(matchings, x):
+                if matchings[tensor].ty.ndim != 3:
+                    return None
+                return R.reshape(x, matchings[tensor].ty.shape.values)
+
+            return tensor, rewriter
+        else:
+            raise NotImplementedError()
+
+    q_raw, k_raw, v_raw = wildcard(), wildcard(), wildcard()
+    q, q_rewriter = handle_input(q_raw, qkv_layout, False)
+    k, k_rewriter = handle_input(k_raw, qkv_layout, True, repeat=with_kv_repeat)
+    v, v_rewriter = handle_input(v_raw, qkv_layout, False, repeat=with_kv_repeat)
+    matmul_1 = is_op("relax.matmul")(q, k)
+    scale = is_const()
+
+    if with_cast:
+        multiply = is_op("relax.multiply")(matmul_1, is_op("relax.astype")(scale))
+    else:
+        multiply = is_op("relax.multiply")(matmul_1, scale)
+
+    if with_bias:
+        bias_raw = wildcard()
+        add = is_op("relax.add")(multiply, bias_raw)
+        softmax_input = add
+    else:
+        softmax_input = multiply
+
+    if with_cast:
+        softmax_input = is_op("relax.astype")(softmax_input)
+
+    softmax = is_op("relax.nn.softmax")(softmax_input)
+
+    if with_cast:
+        softmax_output = is_op("relax.astype")(softmax)
+    else:
+        softmax_output = softmax
+
+    matmul_2 = is_op("relax.matmul")(softmax_output, v)
+
+    out, out_rewriter = handle_output(matmul_2, out_layout)
+
+    def rewriter(original, matchings):
+        query, query_shape = q_rewriter(matchings, matchings[q_raw])
+        key, key_shape = k_rewriter(matchings, matchings[k_raw])
+        value, _ = v_rewriter(matchings, matchings[v_raw])
+        if query is None or key is None or value is None:
+            return original
+        softmax_axis = matchings[softmax].attrs.axis
+        softmax_input_rank = len(matchings[softmax].ty.shape)
+        if softmax_axis == -1:
+            softmax_axis += softmax_input_rank
+        if softmax_axis != softmax_input_rank - 1:
+            return original
+        b, s, n, _ = query_shape
+        _, s_kv, _, _ = key_shape
+        if with_bias:
+            bias = matchings[bias_raw]
+            bias_shape = list(bias.ty.shape)
+            if bias_shape == [b * n, s, s_kv]:
+                bias = R.reshape(bias, [b, n, s, s_kv])
+            elif bias_shape == [b * n, 1, s_kv]:
+                bias = R.reshape(bias, [b, n, 1, s_kv])
+            elif bias_shape == [b, s, s_kv]:
+                bias = R.reshape(bias, [b, 1, s, s_kv])
+            elif bias_shape == [b, 1, s_kv]:
+                bias = R.reshape(bias, [b, 1, 1, s_kv])
+            elif bias_shape in [[1, s, s_kv], [s, s_kv]]:
+                bias = R.reshape(bias, [1, 1, s, s_kv])
+            elif bias_shape in [[1, 1, s_kv], [1, s_kv], [s_kv]]:
+                bias = R.reshape(bias, [1, 1, 1, s_kv])
+            else:
+                return original
+        else:
+            bias = None
+        out = out_rewriter(
+            matchings,
+            R.nn.attention(
+                query,
+                key,
+                value,
+                bias,
+                T.FloatImm(matchings[scale].data.dtype, float(matchings[scale].data.numpy())),
+            ),
+        )
+        return out
+
+    return out, rewriter

python/tvm/relax/backend/rocm/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+"""The Relax ROCm backend compilation pipeline and other passes."""
+
+from .pipeline import (
+    finalize_passes,
+    get_default_pipeline,
+    legalize_passes,
+    library_dispatch_passes,
+)

python/tvm/relax/backend/rocm/hipblas.py

@@ -0,0 +1,181 @@
+# 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.
+
+"""Pattern table for hipblas backend"""
+
+import operator
+from functools import reduce
+
+import tvm
+from tvm.relax import transform
+from tvm.relax.transform import PatternCheckContext
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import make_matmul_pattern
+from ..utils import has_leaking_intermediate_variables
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):  # pylint: disable=unused-argument
+    """Check if dtypes in the given workload are supported by hipblas BYOC."""
+    if lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
+        # The output cannot be 'float8_e5m2' if inputs are 'float8_e4m3fn'
+        # return out_dtype != "float8_e5m2"
+        return False
+    return (lhs_dtype == "float16" and rhs_dtype == "float16") or (
+        lhs_dtype == "int8" and rhs_dtype == "int8"
+    )
+
+
+def _check_matmul(context: PatternCheckContext) -> bool:
+    if has_leaking_intermediate_variables(context):
+        return False
+    lhs = context.annotated_expr["lhs"]
+    rhs = context.annotated_expr["rhs"]
+    matmul_call = context.annotated_expr["root"]
+
+    lhs_dtype = lhs.ty.dtype
+    rhs_dtype = rhs.ty.dtype
+    out_dtype = matmul_call.ty.dtype
+    if not _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype):
+        return False
+
+    lhs_shape = lhs.ty.shape.values
+    rhs_shape = rhs.ty.shape.values
+
+    if not isinstance(lhs_shape[-1], tvm.tirx.expr.IntImm | int):
+        # Reduction axis must be constant
+        return False
+
+    if lhs_dtype == "int8" and rhs_dtype == "int8":
+        return False
+    elif lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn":
+        return False
+
+    lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1)
+    rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1)
+
+    if "bias" in context.annotated_expr:
+        if lhs_dtype == "int8" and rhs_dtype == "int8":
+            # Non-default epilogue not supported for IGEMM
+            return False
+        bias = context.annotated_expr["bias"]
+        bias_shape = bias.ty.shape.values
+        bias_batches = reduce(operator.mul, bias_shape[:-1], 1)
+        if not isinstance(bias_batches, tvm.tirx.expr.IntImm | int) or int(bias_batches) > 1:
+            # hipblas only supports bias vector
+            return False
+
+    # hipblasLt does not seem to support batched GEMM with one of matrices having
+    # one batch (with batch_stride 0). So for batched GEMM, the two batch counts
+    # must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by
+    # flattening all batch axes into the M axis.
+    return (
+        isinstance(lhs_batches, tvm.tirx.Var)
+        or isinstance(rhs_batches, tvm.tirx.Var)
+        or (int(lhs_batches) == int(rhs_batches))
+        or (lhs_batches >= 1 and rhs_batches == 1)
+    )
+
+
+register_patterns(
+    [
+        (
+            "hipblas.matmul",
+            *make_matmul_pattern(
+                with_bias=False,
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_bias",
+            *make_matmul_pattern(
+                with_bias=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_bias_relu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_bias_gelu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.gelu",
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_transposed",
+            *make_matmul_pattern(
+                with_bias=False,
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_transposed_bias",
+            *make_matmul_pattern(
+                with_bias=True,
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_transposed_bias_relu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.relu",
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+        (
+            "hipblas.matmul_transposed_bias_gelu",
+            *make_matmul_pattern(
+                with_bias=True,
+                activation="relax.nn.gelu",
+                transposed_rhs=True,
+            ),
+            _check_matmul,
+        ),
+    ]
+)
+
+
+def partition_for_hipblas(mod):
+    """
+    Partition the input module into hipblas-supported subgraphs.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The IRModule to be partitioned.
+
+    Returns
+    -------
+    mod: tvm.IRModule
+        The resulting IRModule, containing partitioned subgraphs to be
+        offloaded to the hipblas backend.
+    """
+
+    patterns = get_patterns_with_prefix("hipblas")
+    return transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)

python/tvm/relax/backend/rocm/pipeline.py

@@ -0,0 +1,89 @@
+# 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.
+"""The Relax ROCm backend compilation pipeline and other passes."""
+
+import tvm
+from tvm import relax
+
+
+def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default library dispatch passes for ROCm backend."""
+    return [
+        relax.backend.DispatchSampling(),
+        relax.backend.DispatchSortScan(),
+    ]
+
+
+def legalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default legalization passes for ROCm backend."""
+    from tvm.s_tir import dlight as dl  # pylint: disable=import-outside-toplevel
+
+    return [
+        tvm.relax.transform.LegalizeOps(),
+        tvm.relax.transform.AnnotateTIROpPattern(),
+        tvm.relax.transform.FoldConstant(),
+        tvm.relax.transform.FuseOps(),
+        tvm.relax.transform.FuseTIR(),
+        dl.ApplyDefaultSchedule(
+            dl.gpu.Matmul(),
+            dl.gpu.GEMV(),
+            dl.gpu.Reduction(),
+            dl.gpu.GeneralReduction(),
+            dl.gpu.Fallback(),
+        ),
+    ]
+
+
+def dataflow_lower_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default dataflow lowering passes for ROCm backend."""
+    return [
+        relax.transform.RewriteDataflowReshape(),
+        relax.transform.ToNonDataflow(),
+        relax.transform.RemovePurityChecking(),
+        relax.transform.CallTIRRewrite(),
+    ]
+
+
+def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
+    """The default finalization passes for ROCm backend."""
+    return [
+        relax.transform.StaticPlanBlockMemory(),
+        relax.transform.LowerAllocTensor(),
+        relax.transform.KillAfterLastUse(),
+        relax.transform.LowerRuntimeBuiltin(),
+        relax.transform.ComputePrimValue(),
+        relax.transform.VMShapeLower(),
+        relax.transform.AttachGlobalSymbol(),
+    ]
+
+
+def get_default_pipeline(target: tvm.target.Target):
+    """Return the default compilation pipeline for ROCm."""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+        with target:
+            seq = tvm.transform.Sequential(
+                library_dispatch_passes(target)
+                + legalize_passes(target)
+                + dataflow_lower_passes(target)
+                + finalize_passes(target)
+            )
+            mod = seq(mod)
+        return mod
+
+    return _pipeline

python/tvm/relax/backend/utils.py

@@ -0,0 +1,93 @@
+# 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
+"""Utils for BYOC pattern matching"""
+
+from tvm import relax
+from tvm.relax import DataflowVar, PyExprMutator
+from tvm.relax.transform import PatternCheckContext
+from tvm.target import Target
+
+
+class BackendDispatcher(PyExprMutator):
+    """Base class for backend dispatcher"""
+
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    @staticmethod
+    def is_gpu_target(target: Target) -> bool:
+        """Check if the target is a GPU target."""
+        return "gpu" in target.keys
+
+    @staticmethod
+    def get_shape_dtype(expr: relax.Expr) -> tuple[relax.ShapeExpr, str]:
+        """Get shape and dtype from an expression.
+        If the shape and dtype is unknown, raise an error."""
+        ty = expr.ty
+        if not isinstance(expr.ty, relax.TensorType):
+            raise ValueError(f"Expecting a expr with TensorType, but got {expr} with {expr.ty}")
+
+        shape, dtype = ty.shape, ty.dtype
+        if shape is None:
+            raise ValueError(
+                f"Expecting a expr with known shape, but got {expr} with unknown shape"
+            )
+
+        return shape, dtype
+
+    def _get_target(self, ty: relax.Type) -> Target:
+        # Get target information from TensorType
+        if isinstance(ty, relax.TensorType):
+            vdevice = ty.vdevice
+            if vdevice is not None:
+                return vdevice.target
+        elif isinstance(ty, relax.TupleType):
+            for f in ty.fields:
+                tgt = self._get_target(f)
+                if tgt != Target.current():
+                    return tgt
+        # Return the target in current context
+        target = Target.current()
+        if target is None:
+            raise ValueError(
+                "Target not found. Please ensure that the target is annotated within the module, "
+                "or alternatively, execute this within a specified target context."
+            )
+        return target
+
+
+def has_leaking_intermediate_variables(context: PatternCheckContext) -> bool:
+    """
+    Check whether intermediate variables in the region to be fused are used outside
+    the fused region.
+    """
+    defined_vars = set(context.matched_bindings.keys())
+    output_var = context.value_to_bound_var[context.matched_expr]
+    intermediate_vars = {v for v in context.matched_bindings if v != output_var}
+
+    if any(not isinstance(v, DataflowVar) for v in intermediate_vars):
+        # If intermediate variable is not a DataflowVar, it can be accessed and potentially
+        # used outside the DataflowBlock.
+        return True
+
+    # Check whether all users of an intermediate variable are inside the fused region.
+    for var in intermediate_vars:
+        if any(var_user not in defined_vars for var_user in context.var_usages[var]):
+            return True
+
+    return False