chore: import upstream snapshot with attribution

python/tvm/relax/__init__.py

@@ -0,0 +1,115 @@
+# 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.
+# pylint: disable=invalid-name, wrong-import-position
+"""The Relax IR namespace containing the IR, type, operator, builder, vm, etc."""
+
+from tvm.runtime import vm
+from tvm.runtime.vm import VirtualMachine, VMInstrumentReturnKind
+from tvm.ir import Call
+
+# Expr
+from .expr import (
+    Expr,
+    Span,
+    GlobalVar,
+    Var,
+    DataflowVar,
+    Binding,
+    MatchCast,
+    VarBinding,
+    BindingBlock,
+    DataflowBlock,
+    SeqExpr,
+    ShapeExpr,
+    Tuple,
+    TupleGetItem,
+    Function,
+    ExternFunc,
+    If,
+    Constant,
+    DataTypeImm,
+    StringImm,
+    prim_value,
+)
+
+from .expr import const, extern, get_shape_of
+
+# Type
+from .type import (
+    Type,
+    AnyType,
+    ObjectType,
+    ShapeType,
+    TensorType,
+    TupleType,
+    FuncType,
+    PackedFuncType,
+)
+
+# VM
+from .exec_builder import ExecBuilder
+
+# Operator
+from .op.base import (
+    call_tir,
+    call_tir_inplace,
+    call_pure_packed,
+    call_dps_packed,
+    call_tir_with_grad,
+)
+
+# BlockBuilder
+from .block_builder import BlockBuilder
+
+# ExprFunctor
+from .expr_functor import ExprFunctor, PyExprVisitor, PyExprMutator
+
+# pipeline
+from .pipeline import get_default_pipeline
+from .pipeline import get_pipeline
+from .pipeline import register_pipeline
+
+# utils
+from .utils import convert_to_expr
+
+# BasePyModule
+from .base_py_module import BasePyModule
+
+# Import submodules in the last to avoid dependency
+from . import exec_builder
+from . import expr
+from . import ty
+from . import type
+from . import analysis
+from . import transform
+from . import block_builder
+from . import op
+from . import backend
+from . import training
+from . import distributed
+from . import frontend
+from . import utils
+
+# VM
+from .vm_build import build, VMExecutable
+
+from .binding_rewrite import DataflowBlockRewrite
+
+import tvm.script
+
+tvm.script.register_dialect("relax", "tvm.relax.script")

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

python/tvm/relax/analysis/__init__.py

@@ -0,0 +1,50 @@
+# 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 IR analysis."""
+
+from .analysis import (
+    BaseCheckResult,
+    all_global_vars,
+    all_vars,
+    bound_vars,
+    check_well_formed,
+    collect_non_negative_expressions,
+    computable_at_compile_time,
+    contains_impure_call,
+    definable_tir_vars_in_type,
+    defined_symbolic_vars,
+    derive_call_ret_type,
+    detect_recursion,
+    erase_to_well_defined,
+    free_symbolic_vars,
+    free_vars,
+    get_static_type,
+    used_vars,
+    get_var2val,
+    has_reshape_pattern,
+    name_to_binding,
+    post_order_visit,
+    remove_all_unused,
+    type_base_check,
+    type_lca,
+    suggest_layout_transforms,
+    tir_vars_in_type,
+    udchain,
+    well_formed,
+)
+from .estimate_memory_usage import estimate_memory_usage

python/tvm/relax/analysis/_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"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.analysis", __name__)

python/tvm/relax/analysis/analysis.py

@@ -0,0 +1,621 @@
+# 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=no-else-return, invalid-name
+# pylint: disable=unidiomatic-typecheck
+"""
+This file contains the set of passes for Relax, which exposes an interface for
+configuring the passes and scripting them in Python.
+"""
+
+from collections.abc import Callable
+from enum import IntEnum
+
+import tvm
+from tvm import IRModule, tirx
+from tvm.ir import Call, Type
+from tvm.relax.expr import Binding, DataflowBlock, Expr, Function, GlobalVar, Var
+from tvm.relax.type import FuncType
+from tvm.tirx import Buffer, IndexMap, PrimFunc, SBlock
+
+from . import _ffi_api
+
+
+def get_static_type(ty: Type) -> Type:
+    """Get the corresponding static type from a Type.
+
+    Parameters
+    ----------
+    ty : Type
+        The input type.
+
+    Returns
+    -------
+    ret : Type
+        The corresponding static type.
+    """
+    return _ffi_api.GetStaticType(ty)  # type: ignore
+
+
+def erase_to_well_defined(
+    ty: Type,
+    shape_var_map: dict[tirx.Var, tirx.Expr] | None = None,
+    var_map: dict[Var, Expr] | None = None,
+) -> Type:
+    """Erase ty into a well defined form.
+
+    This function removes the Type's dependencies on shape and vars that
+    are not defined in given maps.
+
+    Parameters
+    ----------
+    ty : Type
+        The input type.
+
+    shape_var_map : Dict[tirx.Var, tirx.Expr]
+        Specifies the defined shape vars and the values they should map to.
+
+    var_map : Dict[Var, Expr]
+        Specifies the defined vars and the values they should map to.
+
+    Returns
+    -------
+    ret : Type
+        The corresponding erased type.
+    """
+    shape_var_map = {} if shape_var_map is None else shape_var_map
+    var_map = {} if var_map is None else var_map
+
+    return _ffi_api.EraseToWellDefined(ty, shape_var_map, var_map)  # type: ignore
+
+
+class BaseCheckResult(IntEnum):
+    """Return result of fine-grained base check.
+
+    Note
+    ----
+    Base check comes with fine-grained fail levels.
+
+    - FAIL_L0: The lhs and rhs have no intersection at all.
+    - FAIL_L1: We get the failure by looking at static information.
+    - FAIL_L2: We get the failure due to unknown symbolic variable relations.
+    """
+
+    FAIL_L0 = 0
+    FAIL_L1 = 1
+    FAIL_L2 = 2
+    PASS = 3
+
+
+def type_base_check(base: Type, derived: Type) -> BaseCheckResult:
+    """Run a base check to see if base subsumes derived.
+
+    Parameters
+    ----------
+    base: Type
+        The base type.
+
+    derived: Type
+        The derived type.
+
+    Returns
+    -------
+    ret : Type
+        The derived return value type.
+    """
+    return _ffi_api.TypeBaseCheck(base, derived)  # type: ignore
+
+
+def derive_call_ret_type(func_ty: FuncType, call: Call, ctx: "tvm.relax.BlockBuilder") -> Type:
+    """Derive the call's ret value type from inputs.
+
+    Parameters
+    ----------
+    func_ty: FuncType
+        The call's function signature.
+
+    call: Call
+        The call expression
+
+    ctx: tvm.relax.BlockBuilder
+        The context block builder.
+
+    Returns
+    -------
+    ret : Type
+        The derived return value type.
+
+    Note
+    ----
+    This is an internal derivation function, call.op field is
+    ignored in this case and the derivation only depends on func_ty.
+    """
+    return _ffi_api.DeriveCallRetType(func_ty, call, ctx)  # type: ignore
+
+
+def type_lca(lhs: Type, rhs: Type) -> Type:
+    """Unify the two type to their least common ancestor.
+
+    Parameters
+    ----------
+    lhs: Type
+        The left operand.
+
+    rhs: Type
+        The right operand.
+
+    Returns
+    -------
+    ret : Type
+        The corresponding lca result.
+    """
+    return _ffi_api.TypeLCA(lhs, rhs)  # type: ignore
+
+
+def tir_vars_in_type(ty: Type) -> list[tirx.Var]:
+    """Get the TIR variables that appear in the input type.
+    The returned list is deduplicated - each TIR variable will appear at most once.
+
+    Parameters
+    ----------
+    ty : Type
+        The type object to be analyzed.
+
+    Returns
+    -------
+    ret : List[tirx.Var]
+        The list of TIR variables that appear in the input type.
+    """
+    return _ffi_api.TIRVarsInType(ty)  # type: ignore
+
+
+def definable_tir_vars_in_type(ty: Type) -> list[tirx.Var]:
+    """Get the TIR variables that may be defined from input type.
+    The returned list is deduplicated - each TIR variable will appear at most once.
+
+    Parameters
+    ----------
+    ty : Type
+        The type object to be analyzed.
+
+    Returns
+    -------
+    ret : List[tirx.Var]
+
+        The list of TIR variables that can be defined from the Type
+    """
+    return _ffi_api.DefinableTIRVarsInType(ty)  # type: ignore
+
+
+def collect_non_negative_expressions(ty: Type) -> list[tirx.Expr]:
+    """Collect TIR expressions used in non-negative contexts
+
+    Get TIR variables that are non-negative within the context where
+    the type is used.  For example, any expression used as a
+    tensor shape.
+
+    The returned list is deduplicated - each TIR expression will
+    appear at most once.  The order of the list is in the order of
+    occurrence within the type.
+
+    Parameters
+    ----------
+    ty : Type
+        The type object to be analyzed.
+
+    Returns
+    -------
+    ret : List[tirx.Var]
+
+        The list of TIR variables that can be defined from the Type
+
+    """
+
+    return _ffi_api.CollectNonNegativeExpressions(ty)  # type: ignore
+
+
+def defined_symbolic_vars(func: Function) -> list[Var]:
+    """Get the TIR variables that defined in the input function.
+    The returned list is deduplicated - each TIR variable will appear at most once.
+
+    Parameters
+    ----------
+    func : Function
+        The function object to be analyzed.
+
+    Returns
+    -------
+    ret : List[Var]
+        The list of symbolic variables that are defined in the input function.
+    """
+    return _ffi_api.DefinedSymbolicVars(func)  # type: ignore
+
+
+def free_symbolic_vars(func: Function) -> list[Var]:
+    """Get the TIR variables that are used but not defined in the input function.
+    The returned list is deduplicated - each TIR variable will appear at most once.
+
+    Parameters
+    ----------
+    func : Function
+        The function object to be analyzed.
+
+    Returns
+    -------
+    ret : List[Var]
+        The list of symbolic variables that are used but not defined in the input function.
+    """
+    return _ffi_api.FreeSymbolicVars(func)  # type: ignore
+
+
+def bound_vars(expr: Expr) -> list[Var]:
+    """
+    Return all bound variables from expression expr.
+    Bound variables are all variables that are declared in the expr.
+    They only have meaning inside that expr, and can only be used in it.
+    Parameters
+    ----------
+    expr: Expr
+        The expression.
+    Returns
+    -------
+    ret: List[Var]
+        List of bound vars in expr, in post-DFS order
+    """
+    return _ffi_api.bound_vars(expr)
+
+
+def free_vars(expr: Expr) -> list[Var]:
+    """
+    Return all free variables from expression expr.
+    Free variables are variables that are not bound by a
+    VarBinding or a function parameter in the expression.
+    Parameters
+    ----------
+    expr: Expr
+        The expression.
+    Returns
+    -------
+    ret: List[Var]
+        List of free vars in expr, in post-DFS order
+    """
+    return _ffi_api.free_vars(expr)
+
+
+def all_vars(expr: Expr) -> list[Var]:
+    """
+    Return all (local) variables from expression expr.
+    Parameters
+    ----------
+    expr: Expr
+        The expression.
+    Returns
+    -------
+    ret: List[Var]
+        List of vars in expr, in post-DFS order
+    """
+    return _ffi_api.all_vars(expr)
+
+
+def used_vars(expr: Expr) -> list[Var]:
+    """
+    Return all variables used in an expression.
+
+    This function collects all variable references within the given expression,
+    which is useful for analyzing variable dependencies.
+
+    Parameters
+    ----------
+    expr: Expr
+        The expression to analyze.
+
+    Returns
+    -------
+    ret: List[Var]
+        List of variables used in the expression.
+    """
+    return _ffi_api.used_vars(expr)  # type: ignore
+
+
+def all_global_vars(expr: Expr) -> list[GlobalVar]:
+    """
+    Return all global variables from expression expr.
+    Parameters
+    ----------
+    expr: Expr
+        The expression.
+    Returns
+    -------
+    ret: List[GlobalVar]
+        List of global vars in expr, in post-DFS order
+    """
+    return _ffi_api.all_global_vars(expr)
+
+
+def post_order_visit(expr, fvisit):
+    """Recursively visit the ir in post DFS order node,
+    apply fvisit. Each node is guaranteed to be visited
+    only once.
+
+    Parameters
+    ----------
+    expr : tvm.relax.Expr
+        The input expression.
+
+    fvisit : function
+        The visitor function to be applied.
+    """
+    return _ffi_api.post_order_visit(expr, fvisit)  # type: ignore
+
+
+def has_reshape_pattern(func: tirx.PrimFunc) -> bool:
+    """Check if the given PrimFunc is essentially doing a reshape operation.
+    The reshape operation also includes expand_dims, squeeze, flatten, etc.
+
+    Here the allowed reshape pattern is: for example, assume the operation is
+    `B[l_0, l_1, ..., l_b] = A[r_0, r_1, ..., r_a]`, we check if we can prove
+    that the flattened index of l_0, ..., l_b under buffer B equals to the
+    flattened index of r_0, ..., r_a under buffer A.
+
+    Parameters
+    ----------
+    func : tirx.PrimFunc
+        The function to be examined.
+
+    Returns
+    -------
+    ret : bool
+        A boolean indicating if the given PrimFunc is doing a reshape.
+
+    Notes
+    -----
+    According to the description above, the returned result can only be
+    false-negative and cannot be false-positive, since whenever we cannot
+    prove the equality, we return false. This property guarantees the safety
+    of this function.
+    """
+    return _ffi_api.has_reshape_pattern(func)  # type: ignore
+
+
+def contains_impure_call(expr: Expr, own_name: Var | GlobalVar | None = None) -> bool:
+    """
+    Check if the given expression (likely a function body) contains any impure calls.
+
+    Parameters
+    ----------
+    expr : Expr
+        The expression to be examined. If expr is a function, we check the body.
+
+    own_name : Var or GlobalVar (optional)
+        For a recursive function, the analysis can ignore the self-calls
+        for checking purity.
+
+    Returns
+    -------
+    ret : bool
+        True if there is an impure call
+        (call to a function that may have visible side effects).
+
+    Notes
+    -----
+    Relies on Type annotations, so ensure that the module has been normalized first.
+    Also, an impure call in a *nested* function does *not* mean that the outer expression contains
+    an impure call--it only does if the nested function is *later called*.
+    """
+    return _ffi_api.contains_impure_call(expr, own_name)
+
+
+def get_var2val(func: Function) -> dict[Var, Expr]:
+    """
+    Get a mapping from Var to Expr for each variable in the function.
+
+    Parameters
+    ----------
+    func : Function
+        The input function to be analyzed.
+
+    Returns
+    -------
+    Dict[Var, Expr]
+        A mapping from Var to Expr.
+    """
+    return _ffi_api.get_var2val(func)  # type: ignore
+
+
+def udchain(dfb: DataflowBlock) -> dict[Var, list[Var]]:
+    """
+    Analyze the variable use-def chain in a dataflow block.
+
+    Parameters
+    ----------
+    dfb : DataflowBlock
+        The dataflow block to analyze
+
+    Returns
+    -------
+    Dict[Var, List[Var]]
+        A mapping from variable definition to its uses.
+    """
+    return _ffi_api.udchain(dfb)  # type: ignore
+
+
+def name_to_binding(func: Function) -> dict[str, list[Binding]]:
+    """Return a map from variable name to its bindings."""
+    return _ffi_api.name_to_binding(func)  # type: ignore
+
+
+def remove_all_unused(func: Function) -> Function:
+    """It removes:
+    1. Unused local VarBindings in a DataflowBlock.
+    2. Unused DataflowBlocks in a function.
+
+    Parameters
+    ----------
+    func : Function
+        The input function to be analyzed.
+
+    Notes
+    -----
+    For IRModule-wise DCE, use py:func:`tvm.relax.transform.DeadCodeElimination`.
+
+    Returns
+    -------
+    Function
+        The function with unused variables removed.
+    """
+    return _ffi_api.remove_all_unused(func)  # type: ignore
+
+
+def well_formed(obj: IRModule | Function, check_ty: bool = True) -> None:
+    """Check if the IRModule is well formed, raising on the first violation.
+
+    Raises an error (seeded with the offending node so a pass runner can report a
+    precise access path) on the first well-formedness violation.  Use
+    :func:`check_well_formed` for a boolean answer.
+
+    Parameters
+    ----------
+    obj : Union[tvm.IRModule, Function]
+        The input IRModule or relax.Function.
+
+    check_ty : bool
+        A boolean flag indicating if the property "every Expr must
+        have defined type information" will be checked.
+
+    Note
+    ----
+    By default the type information is always checked. It is only in test cases
+    where `check_ty` might be false, so that other well-formed requirements
+    will be well tested and will not be blocked by not having type information.
+    """
+    _ffi_api.well_formed(obj, check_ty)  # type: ignore
+
+
+def check_well_formed(obj: IRModule | Function, check_ty: bool = True) -> bool:
+    """Return whether the IRModule or Function is well formed.
+
+    Wraps :func:`well_formed`, returning False instead of raising on the first violation.
+
+    Parameters
+    ----------
+    obj : Union[tvm.IRModule, Function]
+        The input IRModule or relax.Function.
+
+    check_ty : bool
+        A boolean flag indicating if the property "every Expr must
+        have defined type information" will be checked.
+
+    Returns
+    -------
+    ret: bool
+        True if the IRModule is well formed, False if not.
+    """
+    return _ffi_api.check_well_formed(obj, check_ty)  # type: ignore
+
+
+def _get_prim_func_default_dtype(func: PrimFunc):
+    """Detect default index dtype from function buffer map"""
+    for _, v in func.buffer_map.items():
+        for value in v.shape:
+            return value.ty
+    return "int64"
+
+
+def suggest_layout_transforms(
+    func: PrimFunc, write_buffer_transforms: list[IndexMap | Callable]
+) -> dict[SBlock, dict[SBlock | Buffer, IndexMap]]:
+    """Suggest Layout transformations of blocks and buffers in a PrimFunc.
+
+    Parameters
+    ----------
+    func: PrimFunc
+        PrimFunc on which analysis will be performed and transformations suggested.
+
+    write_buffer_transforms: List[Union[IndexMap, Callable]
+        List of layout transformations on the output buffers. The number of layout
+        transformations must match the number of outputs of the PrimFunc.
+
+    Returns
+    -------
+    ret: Dict[SBlock, Dict[Union[SBlock, Buffer], IndexMap]]
+         Suggested transforms per block in `func`. For each block the returned value is a map
+         from the object (block or buffer) to it's index map transformation.
+    """
+    write_buffer_index_maps = []
+    default_index_dtype = _get_prim_func_default_dtype(func)
+    for transform in write_buffer_transforms:
+        if callable(transform):
+            transform = IndexMap.from_func(transform, index_dtype=default_index_dtype)
+        assert isinstance(transform, IndexMap)
+        write_buffer_index_maps.append(transform)
+    return _ffi_api.suggest_layout_transforms(func, write_buffer_index_maps)  # type: ignore
+
+
+def detect_recursion(mod: tvm.IRModule) -> list[list[GlobalVar]]:
+    """
+    Find all sets of recursive or mutually recursive functions in the module.
+
+    Two or more functions are mutually recursive if there is some cycle of references
+    among them. For example, if there are two functions A and B, they are
+    mutually recursive if A calls B and B calls A. Another case would be with
+    three functions A, B, and C, where A calls B, B calls C, and C calls A.
+
+    (Note that functions do not have to call each other to reference each other.
+    For example, if a function returns another function, that is still a reference
+    that could potentially be recursive, even without a call.)
+
+
+    If a function is simply recursive and not mutually recursive with any other,
+    it will be reported as a group by itself.
+
+    Parameters
+    ----------
+    mod: The module
+
+    Returns
+    -------
+    ret: List[List[GlobalVar]]
+        Each member of the list is a list of global functions
+        that references each other mutually recursively.
+        If a function is simply recursive and not mutually recursive
+        with any other, it will be a singleton in this list.
+    """
+    return _ffi_api.detect_recursion(mod)  # type: ignore
+
+
+def computable_at_compile_time(func: Function) -> list[Var]:
+    """Collect variables whose value can be computed at compile-time
+
+    If a function has the `kNumInput` attribute, then the first
+    `kNumInput` parameters are provided at run-time, while all
+    remaining parameters may be known at compile-time.  This utility
+    collects all variable bindings that only depend, directly or
+    indirectly, on the parameters known at compile-time.
+
+    Parameters
+    ----------
+    func: Function
+
+        The `relax.Function` to analyze
+
+    Returns
+    -------
+    ret: List[Var]
+
+        The set of variables that can be computed at compile-time, in
+        order of their occurrence within the function.
+    """
+    return _ffi_api.computable_at_compile_time(func)  # type: ignore

python/tvm/relax/analysis/estimate_memory_usage.py

@@ -0,0 +1,182 @@
+# 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=abstract-method,unused-argument
+# pylint: disable=missing-function-docstring,missing-module-docstring
+
+import tvm
+from tvm.ir import Call, Op
+from tvm.ir.module import IRModule
+
+from ..expr import Expr, Function, ShapeExpr
+from ..expr_functor import PyExprVisitor, visitor
+
+
+def estimate_memory_usage(mod: IRModule | Function) -> str:
+    """Analysis function that estimates the memory usage of Relax functions
+    in an IRModule. The estimation includes the total memory size needed to
+    be allocated before and after memory planning.
+
+    The result might be over-estimated, as the estimation is static, which
+    does not consider control flows (such as "if" and cross-function calls).
+    It simply accumulates the size of every alloc_tensor and alloc_storage.
+
+    This analysis function is used to demonstrate the effect of memory
+    planning.
+
+    Parameters
+    ----------
+    mod : Union[IRModule, Function]
+        The input IRModule whose functions inside are to be analyzed.
+        If the input is a Function, we will wrap it with a IRModule, with
+        the function named "main".
+
+    Returns
+    -------
+    est : str
+        The estimation information, in the form of a string.
+
+    Notes
+    -----
+    We regards "relax.memory.alloc_tensor/storage" as the results produced by memory planning.
+    """
+
+    @visitor
+    class MemoryEstimator(PyExprVisitor):
+        """The IR visitor which estimates the memory usage of each Relax function.
+
+        Attributes
+        ----------
+        total_alloc_tensor_mem : int
+            The total memory size of alloc_tensor, in bytes.
+
+        total_const_size_tensor_num : int
+            The number of constant-size tensors.
+
+        total_dyn_size_tensor_num : int
+            The number of dynamic-size tensors.
+
+        planned_alloc_mem : int
+            The total memory size of memory.alloc_storage after memory planning, in bytes.
+
+        planned_mem_num : int
+            The number of memory.alloc_storages.
+        """
+
+        total_alloc_tensor_mem: int
+        total_const_size_tensor_num: int
+        total_dyn_size_tensor_num: int
+        planned_alloc_mem: int
+        planned_mem_num: int
+        builtin_alloc_tensor_op = Op.get("relax.builtin.alloc_tensor")
+        memory_alloc_tensor_op = Op.get("relax.memory.alloc_tensor")
+        memory_alloc_storage_op = Op.get("relax.memory.alloc_storage")
+
+        def estimate(self, mod: IRModule) -> str:
+            estimation: str = ""
+            for global_var, func in mod.functions_items():
+                if not isinstance(func, Function):
+                    continue
+
+                self.cleanup()
+                self.visit_expr(func)
+                estimation += self.generate_est_string(global_var.name_hint) + "\n"
+
+            if estimation != "":
+                estimation = "Memory usage estimation:\n" + estimation
+            return estimation
+
+        def cleanup(self) -> None:
+            self.total_alloc_tensor_mem = 0
+            self.total_const_size_tensor_num = 0
+            self.total_dyn_size_tensor_num = 0
+            self.planned_alloc_mem = 0
+            self.planned_mem_num = 0
+
+        def visit_call_(self, call: Call) -> None:  # pylint: disable=arguments-differ
+            if call.op == self.builtin_alloc_tensor_op:
+                self.accumulate_builtin_tensor_alloc(
+                    shape=call.args[0], dtype_str=call.args[1].value
+                )
+            elif call.op == self.memory_alloc_tensor_op:
+                self.accumulate_tensor_alloc(shape=call.args[2], dtype_str=call.args[3].value)
+            elif call.op == self.memory_alloc_storage_op:
+                self.accumulate_storage_alloc(size=call.args[0])
+
+        def calculate_size(self, shape: Expr, dtype_str: str) -> int:
+            if not isinstance(shape, ShapeExpr):
+                raise TypeError(
+                    "The shape of relax.builtin.alloc_tensor and "
+                    "relax.memory.alloc_tensor is expected to be ShapeExpr"
+                )
+            size: int = 1
+            for dim_len in shape.values:
+                if not isinstance(dim_len, tvm.tirx.IntImm):
+                    self.total_dyn_size_tensor_num += 1
+                    return -1
+                size *= dim_len.value
+            dtype = tvm.DataType(dtype_str)
+            return size * ((dtype.bits + 7) // 8) * dtype.lanes
+
+        def accumulate_builtin_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
+            size = self.calculate_size(shape, dtype_str)
+            if size == -1:
+                return
+            self.total_const_size_tensor_num += 1
+            self.total_alloc_tensor_mem += size
+            self.planned_mem_num += 1
+            self.planned_alloc_mem += size
+
+        def accumulate_tensor_alloc(self, shape: Expr, dtype_str: str) -> None:
+            size = self.calculate_size(shape, dtype_str)
+            if size == -1:
+                return
+            self.total_const_size_tensor_num += 1
+            self.total_alloc_tensor_mem += size
+
+        def accumulate_storage_alloc(self, size: Expr) -> None:
+            if not isinstance(size, ShapeExpr):
+                raise TypeError(
+                    "The size of relax.memory.alloc_storage is expected to be ShapeExpr"
+                )
+
+            self.planned_mem_num += 1
+            self.planned_alloc_mem += size.values[0].value
+
+        def generate_est_string(self, func_name: str) -> str:
+            est = (
+                f" * Without memory planning, there are {self.total_const_size_tensor_num} "
+                "constant-size memory allocation(s) with total size "
+                "{0:.4} GB".format(self.total_alloc_tensor_mem / 2**30)
+            )
+            if self.total_dyn_size_tensor_num > 0:
+                est += f", and {self.total_dyn_size_tensor_num} dynamic-size allocation(s)"
+            est += (
+                f".\n * With memory planning, there are {self.planned_mem_num} constant-size "
+                "memory allocation(s) with total size "
+                "{0:.4} GB.\n".format(self.planned_alloc_mem / 2**30)
+            )
+            if self.total_alloc_tensor_mem != 0:
+                est += (
+                    " * Memory planning reduces constant memory size to "
+                    f"{self.planned_alloc_mem / self.total_alloc_tensor_mem:.1%}."
+                )
+            return "- Function " + func_name + ":\n" + est
+
+    if isinstance(mod, Function):
+        mod = tvm.IRModule({tvm.ir.GlobalVar("foo"): mod})
+
+    return MemoryEstimator().estimate(mod)

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

python/tvm/relax/base_py_module.py

@@ -0,0 +1,614 @@
+# 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.
+# ruff: noqa: F401, F821
+"""BasePyModule: Base class for IRModules with Python function support."""
+
+import inspect
+import os
+from typing import Any, Optional, Union
+
+import numpy as np
+from tvm_ffi import Function
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir import IRModule
+from tvm.runtime import Device, Tensor
+from tvm.target import Target
+
+try:
+    from torch.utils.dlpack import to_dlpack as to_dlpack_legacy
+except ImportError:
+    to_dlpack_legacy = None
+
+try:
+    from tvm_ffi._optional_torch_c_dlpack import load_torch_c_dlpack_extension
+
+    _FASTER_DLPACK_EXTENSION = load_torch_c_dlpack_extension()
+except ImportError:
+    _FASTER_DLPACK_EXTENSION = None
+
+
+class BasePyModule:
+    """Base class that allows Python functions in IRModule with DLPack conversion.
+
+    This class provides the infrastructure for:
+    1. JIT compilation of TIR and Relax functions.
+    2. DLPack-based conversion between PyTorch tensors and TVM Tensors.
+    3. Wrapping Relax functions for easy Python calling.
+    4. Cross-function calls between Python, TIR, and Relax functions.
+
+    Only IRModules that inherit from this class are allowed to contain Python functions.
+    """
+
+    def __del__(self):
+        """Clean up registered Python functions on module destruction."""
+        try:
+            clear_func = tvm.get_global_func("vm.builtin.clear_py_func_registry")
+            clear_func()
+        except (ValueError, AttributeError):
+            pass
+
+    def __init__(
+        self,
+        ir_mod: IRModule,
+        device: Device,
+        target: Target | None = None,
+    ):
+        """Initialize BasePyModule with JIT compilation and DLPack conversion."""
+        self.device = device
+        self.ir_mod = ir_mod
+
+        # Delegate IRModule operations
+        self.functions = ir_mod.functions
+        self.attrs = ir_mod.attrs
+        self.global_infos = ir_mod.global_infos
+        self.__getitem__ = ir_mod.__getitem__
+        self.__setitem__ = ir_mod.__setitem__
+        self.functions_items = ir_mod.functions_items
+        self.with_attr = ir_mod.with_attr
+        self.get_attr = ir_mod.get_attr
+        self.update_global_info = ir_mod.update_global_info
+
+        def _getattr_python_function(name: str) -> Any:
+            """Support direct attribute access to funcs and IRModule methods."""
+            if name in self.pyfuncs:
+                return self.pyfuncs[name]
+            if name in self.compiled_tir_funcs:
+                return self.compiled_tir_funcs[name]
+            if self.relax_vm and name in self.relax_func_names:
+                try:
+                    return self.relax_vm[name]
+                except AttributeError:  # More specific exception
+                    return None
+            if hasattr(self.ir_mod, name):
+                return getattr(self.ir_mod, name)
+            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
+        self.__getattr__ = _getattr_python_function
+
+        self.compiled_tir_funcs: dict[str, Function] = {}
+        self.extern_funcs: dict[str, Function] = {}
+        self.tir_func_names: list[str] = []
+        self.relax_func_names: list[str] = []
+        self.relax_vm: relax.VirtualMachine | None = None
+        self.pyfuncs: dict[str, Any] = {}
+
+        if target is None:
+            target = Target.from_device(device)
+        elif isinstance(target, str):
+            target = Target(target)
+        self.target = target
+
+        self._collect_function_names()
+        self._compile_functions()
+        self._wrap_tir_functions()
+        self._wrap_relax_functions()
+        self._register_python_functions()
+
+    def _collect_function_names(self):
+        """Collect names of TIR and Relax functions from IRModule."""
+        for global_var, func in self.ir_mod.functions_items():
+            if isinstance(func, tirx.PrimFunc):
+                self.tir_func_names.append(global_var.name_hint)
+            elif isinstance(func, relax.Function):
+                self.relax_func_names.append(global_var.name_hint)
+
+    def _compile_functions(self):
+        """Compile TIR and Relax functions using JIT compilation."""
+        # Compile TIR functions first
+        tir_mod = tvm.IRModule(
+            {
+                gv: func
+                for gv, func in self.ir_mod.functions_items()
+                if isinstance(func, tirx.PrimFunc)
+            }
+        )
+        if tir_mod:
+            try:
+                tir_exec_mod = tvm.compile(tir_mod, target=self.target)
+                for func_name in self.tir_func_names:
+                    self.compiled_tir_funcs[func_name] = tir_exec_mod[func_name]
+            # pylint: disable=broad-exception-caught
+            except Exception as error:
+                print(f"Warning: Failed to compile one or more TIR functions: {error}")
+
+        if self.relax_func_names:
+            try:
+                exec_mod = tvm.compile(self.ir_mod, target=self.target)
+                self.relax_vm = relax.VirtualMachine(exec_mod, self.device)
+            # pylint: disable=broad-exception-caught
+            except Exception as error:
+                print(f"Warning: Failed to compile Relax VM: {error}")
+                self.relax_vm = None
+
+    def _wrap_tir_functions(self):
+        """Wrap TIR functions to make them accessible as instance attributes."""
+        for func_name, func in self.compiled_tir_funcs.items():
+            setattr(self, func_name, func)
+
+    def _wrap_relax_functions(self):
+        """Wrap Relax functions to be callable from Python with auto conversion."""
+        for func_name in self.relax_func_names:
+
+            def _create_relax_wrapper(name):
+                def wrapper(*args, **kwargs):
+                    """Wrapper for Relax function with automatic tensor conversion."""
+                    if hasattr(self.ir_mod, "pyfuncs") and name in self.ir_mod.pyfuncs:
+                        return self.ir_mod.pyfuncs[name](*args, **kwargs)
+
+                    if self.relax_vm is not None:
+                        converted_args = self._convert_pytorch_to_tvm(list(args))
+                        converted_kwargs = {
+                            k: self._convert_pytorch_to_tvm(v) for k, v in kwargs.items()
+                        }
+                        result = self.relax_vm[name](*converted_args, **converted_kwargs)
+                        return self._convert_tvm_to_pytorch(result)
+
+                    raise RuntimeError(
+                        f"Neither converted Python function nor Relax VM available for {name}"
+                    )
+
+                wrapper.__name__ = name
+                wrapper.__doc__ = f"Wrapped Relax function: {name}"
+                return wrapper
+
+            setattr(self, func_name, _create_relax_wrapper(func_name))
+
+    def _register_python_functions(self):
+        """Register Python functions with the VM runtime for call_py_func support."""
+        if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs:
+            return
+
+        try:
+            register_py_func = tvm.get_global_func("vm.builtin.register_py_func")
+        except ValueError:
+            return
+
+        for func_name, py_func in self.ir_mod.pyfuncs.items():
+
+            def create_py_func_wrapper(name, original_func):
+                def wrapper(*args, **kwargs):
+                    converted_args = [self._convert_tvm_to_pytorch(arg) for arg in args]
+                    converted_kwargs = {
+                        k: self._convert_tvm_to_pytorch(v) for k, v in kwargs.items()
+                    }
+
+                    result = original_func(self, *converted_args, **converted_kwargs)
+
+                    return self._convert_pytorch_to_tvm(result)
+
+                wrapper.__name__ = name
+                return wrapper
+
+            wrapped_func = create_py_func_wrapper(func_name, py_func)
+            register_py_func(func_name, wrapped_func)
+
+    def call_tir(self, tir_func, args, out_ty):
+        """Call a TIR function with PyTorch tensors."""
+        # Try to get function name from different sources
+        if isinstance(tir_func, str):
+            func_name = tir_func
+        elif hasattr(tir_func, "name"):
+            func_name = tir_func.name
+        elif hasattr(tir_func, "__name__"):
+            func_name = tir_func.__name__
+        else:
+            # Try to find by function object reference
+            for name, func in self.compiled_tir_funcs.items():
+                if func == tir_func:
+                    func_name = name
+                    break
+            else:
+                func_name = None
+
+        if not func_name or func_name not in self.compiled_tir_funcs:
+            available_funcs = list(self.compiled_tir_funcs.keys())
+            raise ValueError(
+                f"Could not resolve or find compiled TIR function: {tir_func}. "
+                f"Available functions: {available_funcs}"
+            )
+        func = self.compiled_tir_funcs[func_name]
+
+        out = self._create_output_tensors(out_ty, args)
+        tvm_args = self._convert_pytorch_to_tvm(args)
+        tvm_out = self._convert_pytorch_to_tvm(out)
+
+        func(*tvm_args, *tvm_out)
+
+        result = self._convert_tvm_to_pytorch(tvm_out)
+        return result[0] if len(result) == 1 else result
+
+    def call_dps_packed(self, func_name: str, args, out_ty):
+        """Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack."""
+        if hasattr(self, func_name) and callable(getattr(self, func_name)):
+            return getattr(self, func_name)(*args)
+
+        if func_name not in self.extern_funcs:
+            try:
+                self.extern_funcs[func_name] = tvm.get_global_func(func_name)
+            except ValueError as error:
+                raise ValueError(
+                    f"Function '{func_name}' not found as a global function. "
+                    f"Please implement it as a method or register it."
+                ) from error
+        func = self.extern_funcs[func_name]
+
+        out = self._create_output_tensors(out_ty, args)
+        tvm_args = self._convert_pytorch_to_tvm(args)
+        tvm_out = self._convert_pytorch_to_tvm(out)
+        func(*tvm_args, *tvm_out)
+        return out[0] if len(out) == 1 else out
+
+    def call_py_func(self, func_name: str, args):
+        """Call a Python function stored in the module's pyfuncs."""
+        if func_name not in self.pyfuncs:
+            raise ValueError(f"Python function '{func_name}' not found in module pyfuncs")
+        py_func = self.pyfuncs[func_name]
+        return py_func(self, *args)
+
+    def _create_output_tensors(self, out_ty, in_args=None):
+        # pylint: disable=import-outside-toplevel
+        import torch
+
+        ty_list = out_ty if isinstance(out_ty, list) else [out_ty]
+        out_tensors = []
+        for ty in ty_list:
+            if isinstance(ty, tuple | list) and all(isinstance(x, int | np.integer) for x in ty):
+                out_tensors.append(torch.zeros(list(map(int, ty)), dtype=torch.float32))
+                continue
+
+            if hasattr(ty, "shape") and hasattr(ty, "dtype"):
+                concrete_shape = self._infer_concrete_shape_from_args(ty.shape, in_args)
+                torch_dtype = self._convert_tvm_dtype_to_torch(ty.dtype)
+                out_tensors.append(torch.zeros(concrete_shape, dtype=torch_dtype))
+                continue
+
+            out_tensors.append(torch.zeros((1,), dtype=torch.float32))
+        return out_tensors
+
+    def _infer_concrete_shape_from_args(self, shape, in_args):
+        concrete = []
+        symbolic_positions = []
+        for idx, dim in enumerate(shape):
+            if isinstance(dim, int | np.integer):
+                concrete.append(int(dim))
+            elif isinstance(dim, tirx.IntImm):
+                concrete.append(int(dim.value))
+            else:
+                concrete.append(None)
+                symbolic_positions.append(idx)
+
+        if not symbolic_positions:
+            return concrete
+
+        candidates = []
+        if in_args is not None:
+            if not isinstance(in_args, list | tuple):
+                in_args = [in_args]
+            for obj in in_args:
+                if hasattr(obj, "shape") and isinstance(obj.shape, tuple | list):
+                    try:
+                        candidates.append(tuple(int(x) for x in obj.shape))
+                        continue
+                    except (ValueError, TypeError):
+                        # Skip objects with invalid shapes
+                        pass
+
+        target_ndim = len(shape)
+        for cand in candidates:
+            if len(cand) == target_ndim:
+                for pos in symbolic_positions:
+                    concrete[pos] = cand[pos]
+                if all(x is not None for x in concrete):
+                    return concrete
+
+        raise ValueError(
+            "Cannot infer concrete output shape from symbolic shape and inputs. "
+            "Please provide a concrete `out_ty` (e.g., a tuple/list of ints) "
+            "or ensure input tensors carry shapes that determine output extents."
+        )
+
+    def _convert_tvm_dtype_to_torch(self, tvm_dtype: str) -> "torch.dtype":
+        """Convert TVM dtype string to PyTorch dtype."""
+        # pylint: disable=import-outside-toplevel
+        import torch
+
+        dtype_mapping = {
+            "float32": torch.float32,
+            "float64": torch.float64,
+            "int32": torch.int32,
+            "int64": torch.int64,
+            "bool": torch.bool,
+        }
+        return dtype_mapping.get(str(tvm_dtype), torch.float32)
+
+    def _convert_pytorch_to_tvm(
+        self, tensors: Any | list[Any] | tuple[Any, ...]
+    ) -> Tensor | list[Tensor]:
+        """Convert PyTorch tensors to TVM Tensors using DLPack."""
+        # pylint: disable=import-outside-toplevel
+        import torch
+
+        if isinstance(tensors, list | tuple):
+            return [self._convert_single_pytorch_to_tvm(t) for t in tensors]
+        return self._convert_single_pytorch_to_tvm(tensors)
+
+    def _convert_single_pytorch_to_tvm(self, tensor: Any) -> Tensor:
+        """Convert a single PyTorch tensor to TVM Tensor with faster DLPack converter."""
+        # pylint: disable=import-outside-toplevel
+        import torch
+
+        if isinstance(tensor, Tensor):
+            return tensor
+        if isinstance(tensor, torch.Tensor):
+            # 1. Try faster C++ DLPack converter
+            if _FASTER_DLPACK_EXTENSION is not None:
+                try:
+                    dlpack = torch.to_dlpack(tensor)
+                    return tvm.runtime.from_dlpack(dlpack)
+                except (AttributeError, ValueError):
+                    pass  # Fall through to the next method
+
+            # 2. Try modern `torch.to_dlpack` (preferred for PyTorch >= 1.7)
+            try:
+                dlpack = torch.to_dlpack(tensor)
+                return tvm.runtime.from_dlpack(dlpack)
+            except (AttributeError, ValueError):
+                pass  # Fall through to the next method
+
+            # 3. Try legacy `torch.utils.dlpack.to_dlpack`
+            if to_dlpack_legacy:
+                try:
+                    dlpack = to_dlpack_legacy(tensor)
+                    return tvm.runtime.from_dlpack(dlpack)
+                except (AttributeError, ValueError) as error_legacy:
+                    print(
+                        f"Warning: Legacy DLPack conversion failed ({error_legacy}), "
+                        f"using numpy fallback."
+                    )
+
+            # 4. If all DLPack methods fail, use numpy fallback
+            numpy_array = tensor.detach().cpu().numpy()
+            return tvm.runtime.tensor(numpy_array, device=self.device)
+
+        # For other types (like scalars, lists), convert to numpy first
+        try:
+            numpy_array = np.array(tensor, dtype=np.float32)
+            return tvm.runtime.tensor(numpy_array, device=self.device)
+        except (TypeError, ValueError) as error:
+            raise TypeError(
+                f"Unsupported type for conversion to TVM Tensor: {type(tensor)}"
+            ) from error
+
+    def _convert_tvm_to_pytorch(
+        self, tvm_tensors: Any | list[Any]
+    ) -> Union["torch.Tensor", list["torch.Tensor"]]:
+        """Convert TVM Tensors to PyTorch tensors using DLPack."""
+        if isinstance(tvm_tensors, list | tuple):
+            return [self._convert_single_tvm_to_pytorch(tensor) for tensor in tvm_tensors]
+        return self._convert_single_tvm_to_pytorch(tvm_tensors)
+
+    def _convert_single_tvm_to_pytorch(self, tvm_tensor: Any) -> "torch.Tensor":
+        """Convert a single TVM Tensor to PyTorch tensor using faster DLPack converter."""
+        # pylint: disable=import-outside-toplevel
+        import torch
+
+        if isinstance(tvm_tensor, torch.Tensor):
+            return tvm_tensor
+        if not isinstance(tvm_tensor, Tensor):
+            return torch.tensor(tvm_tensor)
+
+        # 1. Try faster C++ DLPack converter
+        if _FASTER_DLPACK_EXTENSION is not None:
+            try:
+                return torch.from_dlpack(tvm_tensor)
+            except (AttributeError, ValueError):
+                pass  # Fall through to the next method
+
+        # 2. Try standard DLPack conversion
+        try:
+            return torch.from_dlpack(tvm_tensor)
+        # pylint: disable=broad-exception-caught
+        except Exception as error:
+            print(f"Warning: DLPack conversion from TVM failed ({error}), using numpy fallback")
+            numpy_array = tvm_tensor.numpy()
+            return torch.from_numpy(numpy_array)
+
+    def get_function(self, name: str) -> Function | None:
+        """Get a compiled function by name."""
+        if name in self.compiled_tir_funcs:
+            return self.compiled_tir_funcs[name]
+        if name in self.extern_funcs:
+            return self.extern_funcs[name]
+        if self.relax_vm and name in self.relax_func_names:
+            try:
+                if hasattr(self, name):
+                    return getattr(self, name)
+                return self.relax_vm[name]
+            except AttributeError as error:
+                print(f"Warning: Failed to get Relax function '{name}': {error}")
+        return None
+
+    def list_functions(self) -> dict[str, list[str]]:
+        """List all available functions."""
+        return {
+            "tirx": self.tir_func_names,
+            "relax": self.relax_func_names,
+            "extern": list(self.extern_funcs.keys()),
+        }
+
+    def add_python_function(self, name: str, func: callable):
+        """Add a Python function to the module."""
+        self.pyfuncs[name] = func
+
+        # Create a wrapper that handles both instance methods and static functions
+        # pylint: disable=import-outside-toplevel
+        import functools
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            sig = inspect.signature(func)
+            params = list(sig.parameters.keys())
+
+            if params and params[0] == "self":
+                return func(self, *args, **kwargs)
+            else:
+                return func(*args, **kwargs)
+
+        # Set the wrapper as an instance attribute
+        setattr(self, name, wrapper)
+
+    def script(
+        self,
+        *,
+        name: str | None = None,
+        show_meta: bool = False,
+        ir_prefix: str = "I",
+        module_alias: str = "cls",
+        int_dtype: str = "int32",
+        float_dtype: str = "void",
+        verbose_expr: bool = False,
+        indent_spaces: int = 4,
+        print_line_numbers: bool = False,
+        num_context_lines: int = -1,
+        syntax_sugar: bool = True,
+        show_object_address: bool = False,
+        show_all_ty: bool = True,
+        extra_config: dict | None = None,
+    ) -> str:
+        """Print TVM IR into TVMScript text format with Python function support.
+
+        This method extends the standard IRModule script() method to handle
+        Python functions stored in the IRModule's pyfuncs attribute.
+        """
+        # First get the standard IRModule script
+        base_script = self.ir_mod.script(
+            name=name,
+            show_meta=show_meta,
+            ir_prefix=ir_prefix,
+            module_alias=module_alias,
+            int_dtype=int_dtype,
+            float_dtype=float_dtype,
+            verbose_expr=verbose_expr,
+            indent_spaces=indent_spaces,
+            print_line_numbers=print_line_numbers,
+            num_context_lines=num_context_lines,
+            syntax_sugar=syntax_sugar,
+            show_object_address=show_object_address,
+            show_all_ty=show_all_ty,
+            extra_config=extra_config,
+        )
+
+        # If there are no Python functions, return the base script
+        if not hasattr(self.ir_mod, "pyfuncs") or not self.ir_mod.pyfuncs:
+            return base_script
+
+        # Insert Python functions into the script
+        return self._insert_python_functions(base_script, indent_spaces)
+
+    def _insert_python_functions(self, base_script: str, indent_spaces: int) -> str:
+        """Insert Python functions into the TVMScript output."""
+        lines = base_script.split("\n")
+        result_lines = []
+
+        # Find the class definition line and insert Python functions after it
+        class_found = False
+        class_indent = 0
+
+        for line in lines:
+            result_lines.append(line)
+
+            # Look for class definition
+            if not class_found and line.strip().startswith("class "):
+                class_found = True
+                class_indent = len(line) - len(line.lstrip())
+
+                # Insert Python functions after the class definition
+                if hasattr(self.ir_mod, "pyfuncs") and self.ir_mod.pyfuncs:
+                    for func_name, func in self.ir_mod.pyfuncs.items():
+                        # Get the function source code
+                        func_source = self._get_function_source(func)
+                        if func_source:
+                            # Format the function with proper indentation
+                            formatted_func = self._format_python_function(
+                                func_name, func_source, class_indent + indent_spaces
+                            )
+                            result_lines.append(formatted_func)
+                            result_lines.append("")  # Add empty line for separation
+
+        return "\n".join(result_lines)
+
+    def _get_function_source(self, func: callable) -> str | None:
+        """Get the source code of a Python function."""
+        try:
+            source = inspect.getsource(func)
+            return source
+        except (OSError, TypeError):
+            # If we can't get the source, return None
+            return None
+
+    def _format_python_function(self, _func_name: str, func_source: str, indent: int) -> str:
+        """Format a Python function with proper indentation for TVMScript."""
+        lines = func_source.split("\n")
+        formatted_lines = []
+
+        for line in lines:
+            # Skip the function definition line if it's already properly indented
+            if line.strip().startswith("def ") or line.strip().startswith("@"):
+                # Keep decorators and function definition as is
+                formatted_lines.append(" " * indent + line.strip())
+            else:
+                # Add proper indentation for the function body
+                formatted_lines.append(" " * indent + line.strip())
+
+        return "\n".join(formatted_lines)
+
+    def show(self, style: str | None = None, black_format: bool | None = None, **kwargs) -> None:
+        """A sugar for print highlighted TVM script with Python function support.
+
+        This method extends the standard IRModule show() method to handle
+        Python functions stored in the IRModule's pyfuncs attribute.
+        """
+        from tvm.script.highlight import cprint  # pylint: disable=import-outside-toplevel
+
+        if black_format is None:
+            env = os.environ.get("TVM_BLACK_FORMAT")
+            black_format = env and int(env)
+
+        script_content = self.script(**kwargs)
+        cprint(script_content, style=style, black_format=black_format)

python/tvm/relax/binding_rewrite.py

@@ -0,0 +1,158 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=no-else-return, invalid-name
+"""Developer API of add/remove/replace bindings in Relax."""
+
+import tvm_ffi
+
+import tvm
+from tvm.runtime import Object
+
+from . import Binding, DataflowBlock, Expr, Function, Var, _ffi_api
+
+
+@tvm_ffi.register_object("relax.DataflowBlockRewrite")
+class DataflowBlockRewrite(Object):
+    """
+    A binding/statement-level dataflow block rewriter.
+
+    Notes
+    -----
+    Due to the immutable and copy-on-write nature of TVM AST nodes, the rewriting is not done in
+    place. Instead, a new DataflowBlock is created and returned with mutated_dfb. Similarly, its new
+    root Function is created and returned by mutated_root_fn. To apply this change for an IRModule,
+    use mutate_irmodule which rewrites the old function that registered in the constructor.
+    """
+
+    __slots__ = ("__dict__",)
+
+    def __init__(self, dfb: DataflowBlock, root_fn: Function):
+        """
+        Construct a rewriter with the DataflowBlock to rewrite and its root function.
+
+        Parameters
+        ----------
+        dfb : DataflowBlock
+            The DataflowBlock to rewrite.
+        root_fn : Function
+            The root function of the DataflowBlock.
+        """
+        self.func_name = root_fn.__name__ if hasattr(root_fn, "__name__") else None
+        self.__init_handle_by_constructor__(
+            _ffi_api.DataflowBlockRewrite,
+            dfb,
+            root_fn,  # type: ignore
+        )
+
+    def replace_all_uses(self, old_var: Var, new_var: Var) -> None:
+        """
+        Replace all uses of old_var with new_var.
+
+        Parameters
+        ----------
+        old_var : Var
+            The old variable to replace.
+        new_var : Var
+            The new variable to replace with.
+        """
+        _ffi_api.dfb_rewrite_replace_all_uses(self, old_var, new_var)  # type: ignore
+
+    def add_binding(self, binding: Binding) -> None:
+        return _ffi_api.dfb_rewrite_add_binding(self, binding)  # type: ignore
+
+    def add(self, expr: Expr, name: str | None = None, is_dfvar: bool = False) -> None:
+        """
+        Add a new statement to the DataflowBlock with an automatically generated variable name.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression to add.
+        name : Optional[str], optional
+            Variable name, by default None
+        is_dfvar : bool, optional
+            The variable type, by default False
+
+        Notes
+        -----
+        If the variable name is not given, it will be automatically generated in a form of
+        "tmp${COUNTER}". The variable type will be DataflowVar if is_dfvar is True, otherwise
+        it will be Var. Being Var means the variables are output variables of the DataflowBlock.
+        While being DataflowVar means the variables are internal variables of the DataflowBlock.
+        """
+        _ffi_api.dfb_rewrite_add(self, expr, name, is_dfvar)  # type: ignore
+
+    def remove_unused(self, var: Var, allow_undef=False) -> None:
+        """
+        Remove a statement by its variable definition if and only if it is unused.
+
+        Parameters
+        ----------
+        var : Var
+            The unused variable definition.
+        allow_undef : bool, optional
+            Whether to allow var being undefined variable, by default False
+
+        Raises
+        ------
+        RuntimeError if the variable is used or undefined (allow_undef=False).
+        """
+        _ffi_api.dfb_rewrite_remove_unused(self, var, allow_undef)  # type: ignore
+
+    def remove_all_unused(self) -> None:
+        """
+        Remove all unused variables.
+
+        Notes
+        -----
+        This could remove unused variables in other DataflowBlocks as well.
+        """
+        _ffi_api.dfb_rewrite_remove_all_unused(self)  # type: ignore
+
+    def mutated_dfb(self) -> DataflowBlock:
+        """
+        Returns the mutated DataflowBlock.
+        """
+        return self.dfb
+
+    def mutated_root_fn(self) -> Function:
+        """
+        Returns the mutated root function.
+        """
+        ret = self.root_fn
+        if self.func_name:
+            ret.__name__ = self.func_name
+        return ret
+
+    def mutate_irmodule(self, irmodule: tvm.IRModule) -> tvm.IRModule:
+        """
+        Return an updated IRModule by replacing the old function with the mutated root function.
+
+        Parameters
+        ----------
+        irmodule : tvm.IRModule
+            The base IRModule to update.
+
+        Returns
+        -------
+        tvm.IRModule
+            The updated IRModule.
+        """
+        ret = _ffi_api.dfb_rewrite_mutate_irmodule(self, irmodule)  # type: ignore
+        if hasattr(irmodule, "__name__"):
+            ret.__name__ = irmodule.__name__
+        return ret

python/tvm/relax/block_builder.py

@@ -0,0 +1,807 @@
+# 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=no-else-return, invalid-name, unused-argument, import-outside-toplevel
+# ruff: noqa: RUF012
+"""Developer API of constructing Relax AST."""
+
+from collections.abc import Callable, Sequence
+from typing import Any, Optional
+
+import tvm_ffi
+
+import tvm
+from tvm import relax as rx
+from tvm import tirx
+from tvm.ir.module import IRModule
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .expr import BaseFunc, Binding, BindingBlock, Expr, GlobalVar, Tuple, Var
+from .op.base import call_tir, call_tir_with_grad
+from .type import Type
+from .utils import gen_call_tir_inputs
+
+
+class FunctionScope:
+    """Auxiliary scope for function"""
+
+    def __init__(self, block_builder, name, params, attrs, is_pure):
+        self._bb = block_builder
+        self._name = name
+        self._params = params
+        self._attrs = attrs
+        self._is_pure = is_pure
+
+        # Blocks that have been collected within the function
+        self._blocks = []
+        # a boolean flag that tracks if emit_func_output has been called
+        self._is_emit_func_output_called = False
+
+    def __enter__(self):
+        self._bb._enter_function_scope(self)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        # __exit__ should properly handle the case where the with block exits with an exception
+        # when handling error case in exit, always check if there is already an exception
+        # been thrown in the with block
+        self._bb._exit_function_scope(exc_type, exc_val, exc_tb)
+
+
+class DataflowScope:
+    """Auxiliary scope for Dataflow block"""
+
+    def __init__(self, block_builder):
+        self._bb = block_builder
+
+    def __enter__(self):
+        block = self._bb._end_block()
+        if len(block.bindings) > 0:
+            self._bb._func._blocks.append(block)
+        self._bb._begin_dataflow_block()
+
+    def __exit__(self, ptype, value, trace):
+        block = self._bb._end_block()
+        if len(block.bindings) > 0:
+            self._bb._func._blocks.append(block)
+        self._bb._begin_binding_block()
+
+
+class TestingScope:
+    """Auxiliary scope for testing purposes"""
+
+    def __init__(self, block_builder, def_vars):
+        self._bb = block_builder
+        shape_vars = []
+        for var in def_vars:
+            if isinstance(var, tvm.tirx.Var):
+                shape_vars.append(var)
+            else:
+                raise ValueError("def_vars only can take tirx.Var")
+        # setup a dummy var so shape is in scope.
+        sparam = rx.Var("sparam", rx.ShapeType(shape_vars))
+        self._scope_params = [sparam]
+
+    def __enter__(self):
+        self._bb.begin_scope(self._scope_params)
+        self._bb._begin_dataflow_block()
+
+    def __exit__(self, ptype, value, trace):
+        self._bb._end_block()
+        self._bb.end_scope()
+
+
+@tvm_ffi.register_object("relax.BlockBuilder")
+class BlockBuilder(Object):
+    """A builder to build Relax IR for testing and dev.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        m = tirx.Var("m", "int32")
+        n = tirx.Var("n", "int32")
+        x = rx.Var("x", rx.TensorType([m, n], "float16"))
+        y = rx.Var("y", rx.TensorType([n], "float16"))
+        bb = rx.BlockBuilder()
+        with bb.function([x, y], "func"):
+            with bb.dataflow() as df:
+                lv0 = bb.emit(rx.add(x, y))
+                lv1 = bb.emit(rx.multiply(lv0, y))
+                gv0 = bb.emit_output(lv1)
+            bb.emit_func_output(gv0)
+        mod = bb.get()
+
+    BlockBuilder can also be used to construct neural networks with nn.Module API
+
+    .. code-block:: python
+
+        from tvm.relax.testing import nn
+
+        n = tirx.Var("n", "int64")
+        input_size = 784
+        hidden_sizes = [128, 32]
+        output_size = 10
+        bb = rx.BlockBuilder()
+
+        with bb.function("main"):
+            model = nn.Sequential(
+                nn.Linear(input_size, hidden_sizes[0]),
+                nn.ReLU(),
+                nn.Linear(hidden_sizes[0], hidden_sizes[1]),
+                nn.ReLU(),
+                nn.Linear(hidden_sizes[1], output_size),
+                nn.LogSoftmax(),
+            )
+            data = nn.Placeholder((n, input_size), name="data")
+            output = model(data)
+            params = [data] + model.parameters()
+            builder.emit_func_output(output, params=params)
+        mod = bb.get()
+    """
+
+    __slots__ = ("__dict__",)
+
+    _stack = []
+
+    @staticmethod
+    def current() -> Optional["BlockBuilder"]:
+        """Returns the current BlockBuilder."""
+        if BlockBuilder._stack:
+            return BlockBuilder._stack[-1]
+        else:
+            return None
+
+    def __init__(self, mod: IRModule = None):
+        # Which functions are currently being defined
+        self._func_stack: list[FunctionScope] = []
+        self.__init_handle_by_constructor__(_ffi_api.BlockBuilderCreate, mod)  # type: ignore
+
+    def _begin_dataflow_block(self) -> None:
+        _ffi_api.BlockBuilderBeginDataflowBlock(self)  # type: ignore
+
+    def _begin_binding_block(self) -> None:
+        _ffi_api.BlockBuilderBeginBindingBlock(self)  # type: ignore
+
+    def _end_block(self) -> BindingBlock:
+        return _ffi_api.BlockBuilderEndBlock(self)  # type: ignore
+
+    @property
+    def _func(self):
+        if self._func_stack:
+            return self._func_stack[-1]
+        else:
+            raise RuntimeError(
+                "Cannot access BlockBuilder._func when outside a bb._function() block"
+            )
+
+    def _enter_function_scope(self, func_scope):
+        BlockBuilder._stack.append(self)
+        self._func_stack.append(func_scope)
+        self.begin_scope(func_scope._params)
+        self._begin_binding_block()
+
+    def _exit_function_scope(self, exc_type, exc_val, exc_tb):
+        # record
+        is_emit_func_output_called = self._func._is_emit_func_output_called
+        # recover to default state
+        self._func_stack.pop()
+
+        assert BlockBuilder._stack
+        assert BlockBuilder._stack[-1] is self
+        BlockBuilder._stack.pop()
+
+        # NOTE: we must raise after we recover the state so future
+        # block builder scoping functions correctly
+        if exc_type is None:
+            if not is_emit_func_output_called:
+                raise RuntimeError("emit_func_output must be called in a relax function.")
+
+    def function(
+        self,
+        name: str,
+        params: Var | Tuple | list[Var] | None = None,
+        attrs: dict[str, Object] | None = None,
+        pure: bool = True,
+        private: bool = False,
+    ) -> FunctionScope:
+        """Annotate a Relax function.
+
+        Parameters
+        ----------
+        name : str, optional
+            The name of the function
+
+        params : tvm.relax.Var | Tuple | List[tvm.relax.Var], optional
+            The parameters of the function.
+            If params is None, it means deferring initialization of function parameters
+            until emit_func_output.
+
+        attrs : Dict[str, Object], optional
+            The function attrs
+
+        pure : bool, optional
+            Whether the function is annotated as pure.
+
+        private : bool, optional
+            Whether the function is annotated as private.
+            If the function is private, it will not have a global symbol attribute.
+            If it is not private and not an inner function, then it will have
+            a global symbol attribute (mapped to the function's name)
+
+        Returns
+        -------
+        ret: FunctionScope
+            A FunctionScope for building a Relax function node.
+        """
+        if isinstance(params, rx.Var):
+            params = [params]
+        elif isinstance(params, list | tuple):
+            for param in params:
+                if not isinstance(param, rx.Var):
+                    raise TypeError(
+                        f"each element of function parameters must be of type tvm.relax.Var,\
+                                    but got: {type(param)}"
+                    )
+        if attrs is None:
+            attrs = {}
+        # The block builder does not permit nesting functions, per above comment,
+        # so no further check should be needed
+        if not private:
+            attrs["global_symbol"] = name
+
+        return FunctionScope(self, name, params, attrs, is_pure=pure)
+
+    def testing_scope(self, def_vars: list[tirx.Var]) -> TestingScope:
+        """Start a scope for unit-testing purposes.
+
+        Parameters
+        ----------
+        def_vars: List[tirx.Var]
+            List of symbolic variables that are marked as defined in scope.
+
+        Returns
+        -------
+        ret: TestingScope
+            A TestingScope to setup builder for emit and other purposes.
+        """
+        return TestingScope(self, def_vars)
+
+    def dataflow(self) -> DataflowScope:
+        """Annotate a Relax dataflow block.
+
+        Returns
+        -------
+        ret: DataflowScope
+            A DataflowScope for building a Relax dataflow block.
+        """
+        return DataflowScope(self)
+
+    def _normalize_python_tuple(self, expr: Expr | Sequence[Expr]):
+        """Internal utility function to convert to relax.Tuple
+
+        The `emit`, `emit_output`, and `emit_func_output` can be
+        called with python `list` or `tuple` objects.  These objects
+        should be converted to `relax.Tuple` prior to calling an FFI
+        function, as they would otherwise be converted to
+        `tvm_ffi.Array`.  In addition, any nested tuple objects
+        should be converted.
+        """
+        if isinstance(expr, list | tuple):
+            return Tuple([self._normalize_python_tuple(element) for element in expr])
+        elif expr is None:
+            from . import op
+
+            return op.null_value()
+        else:
+            return expr
+
+    def emit(self, expr: Expr, name_hint: str = "") -> Var:
+        """Emit an expr.
+        This infers the shape and type of the expr, create a variable,
+        and bind the expr to the variable.
+
+        Parameters
+        ----------
+        expr : tvm.relax.Expr
+            The Expr to be emitted.
+
+        name_hint : str
+            Name hint for the bound variable.
+
+        Returns
+        -------
+        ret : tvm.relax.Var
+            A newly created variable that gets bound to the input expr.
+        """
+        expr = self._normalize_python_tuple(expr)
+        return _ffi_api.BlockBuilderEmit(self, expr, name_hint)  # type: ignore
+
+    def call_te(self, func: Callable, *args: Any, **kwargs: Any) -> Expr:
+        """Generate a call node according to the te function.
+        This function converts arguments from relax expression to te tensor,
+        The callback func should return a te tensor or a list of te tensors.
+        Please see detailed example in emit_te
+
+        Parameters
+        ----------
+        func : Callable
+            A function that returns a te tensor or a list of te tensors.
+
+        args : Any, optional
+            arguments passed to the function.
+
+        kwargs : Any, optional
+            The keyword arguments passed to the function.
+            Note that the following keyword args are reserved:
+
+                - 'primfunc_name_hint' for passing name hint to the PrimFunc
+                  that gets generated.
+                - 'primfunc_attrs' is reserved for passing func attributes to
+                  be added to the PrimFunc that gets created.
+
+
+        Returns
+        -------
+        ret : tvm.relax.Call
+            A newly created call node
+        """
+
+        primfunc_name = kwargs.pop("primfunc_name_hint", None)
+        tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
+
+        if not primfunc_name:
+            primfunc_name = func.__name__
+        gvar = self.add_func(tir_func, primfunc_name)
+
+        return call_tir(gvar, call_args, output_ty, tir_vars)
+
+    def call_te_with_grad(
+        self,
+        func: Callable,
+        *args: Any,
+        te_grad_name: str,
+        te_grad_kwargs: dict[str, Object] | None = None,
+        **kwargs: Any,
+    ) -> Expr:
+        """Generate a call node according to the te function.
+        This method will generate a call_tir_with_grad node, i.e. a call_tir node bound with a
+        te gradient function (refered by te_grad_name).
+
+        Parameters
+        ----------
+        func : Callable
+            A function that returns a te tensor or a list of te tensors.
+
+        args : Any, optional
+            arguments passed to the function.
+
+        te_grad_name : str
+            The registered name of the te gradient function associated with the call_tir_with_grad
+            node. Must be provided as a keyword argument.
+
+        te_grad_kwargs : Dict[str, Object], optional
+            The keyword arguments passed to the te gradient function.
+            Optionally provided as a keyword argument. Default: {}.
+
+        kwargs : Any, optional
+            The keyword arguments passed to the function.
+            Note that the following keyword args are reserved:
+
+                - 'primfunc_name_hint' for passing name hint to the PrimFunc
+                  that gets generated.
+                - 'primfunc_attrs' is reserved for passing func attributes to
+                  be added to the PrimFunc that gets created.
+
+        Returns
+        -------
+        ret : tvm.relax.Call
+            A newly created call node
+        """
+
+        primfunc_name = kwargs.pop("primfunc_name_hint", None)
+        tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
+
+        if te_grad_kwargs is None:
+            te_grad_kwargs = {}
+
+        if not primfunc_name:
+            primfunc_name = func.__name__
+        gvar = self.add_func(tir_func, primfunc_name)
+
+        return call_tir_with_grad(
+            gvar, call_args, output_ty, te_grad_name, te_grad_kwargs, tir_vars
+        )
+
+    def emit_te(self, func: Callable, *args: Any, **kwargs: Any) -> Var:
+        """Emit a call node according to the te function.
+        This function converts arguments from relax expression to te tensor,
+        The callback func should return a te tensor or a list of te tensors.
+
+        Parameters
+        ----------
+        func : Callable
+            A function that returns a te tensor or a list of te tensors.
+
+        args : Any, optional
+            arguments passed to the function.
+
+        kwargs : Any, optional
+            The keyword arguments passed to the function.
+            Note that the key "primfunc_name_hint" is reserved for passing name hint
+            to the PrimFunc that gets generated.
+
+        Returns
+        -------
+        ret : tvm.relax.Var
+            A newly created variable that gets bound to the call code.
+
+        Example
+        -------
+
+        .. code-block:: python
+
+            bb = rx.BlockBuilder()
+            n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
+            x = rx.Var("x", rx.TensorType([n, m], "float32"))
+            y = rx.Var("y", rx.TensorType([n, m], "float32"))
+
+            def te_func(args, args_dict, msg):
+                A = args[0]
+                B = args_dict["B"]
+                return te.compute((128, 128), lambda i, j: A[i, j] + B[i, j])
+
+            with bb.function([x, y], "rx_func"):
+                out = bb.emit_te(te_func, [x], {"B": y}, msg="hello")
+                bb.emit_func_output(out)
+
+        will result in TVMScript
+
+        .. code-block:: python
+
+            @tvm.script.ir_module
+            class Module:
+                @T.prim_func(s_tir=True)
+                def te_func(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle,
+                            var_compute: T.handle) -> None:
+                    # function attr dict
+                    T.func_attr({"tirx.noalias": True})
+                    m = T.int64()
+                    n = T.int64()
+                    rxplaceholder = T.match_buffer(var_rxplaceholder, [n, m], dtype="float32")
+                    rxplaceholder_1 = T.match_buffer(var_rxplaceholder_1, [n, m], dtype="float32")
+                    compute = T.match_buffer(var_compute, [128, 128], dtype="float32")
+                    # body
+                    # with T.sblock("root")
+                    for i0, i1 in T.grid(128, 128):
+                        with T.sblock("compute"):
+                            i, j = T.axis.remap("SS", [i0, i1])
+                            T.reads([rxplaceholder[i, j], rxplaceholder_1[i, j]])
+                            T.writes([compute[i, j]])
+                            compute[i, j] = rxplaceholder[i, j] + rxplaceholder_1[i, j]
+
+                @R.function
+                def rx_func(x: Tensor((n, m), "float32"), y: Tensor((n, m), "float32")) -> Tensor:
+                    # block 0
+                    gv = relax.call_tir("te_func", (x, y), R.Tensor((128, 128), "float32"))
+                    return gv
+
+        Example
+        -------
+
+        .. code-block:: python
+
+            bb = relax.BlockBuilder()
+            n = tirx.Var("n", "int64")
+            x = relax.Var("x", relax.TensorType([n], "float32"))
+            y = relax.Var("y", relax.TensorType([n + 1], "float32"))
+
+            def te_func(A):
+                C = te.compute((n + 1), lambda i: A[i])
+                return C
+
+            with bb.function("rx_func", [x, y]):
+                x1 = bb.emit_te(te_func, y)
+                bb.emit_func_output(x1)
+
+        will result in TVMScript
+
+        .. code-block:: python
+
+            @tvm.script.ir_module
+            class Module:
+                @T.prim_func(s_tir=True)
+                def te_func(var_rxplaceholder: T.handle, var_compute: T.handle, n: T.int64) -> None:
+                    rxplaceholder = T.match_buffer(var_rxplaceholder, [n + T.int64(1)],
+                                                   dtype="float32")
+                    compute = T.match_buffer(var_compute, [n + T.int64(1)], dtype="float32")
+                    # body
+                    # with T.sblock("root")
+                    for i0 in T.serial(0, n + T.int64(1)):
+                        with T.sblock("compute"):
+                            i = T.axis.spatial(n + T.int64(1), i0)
+                            T.reads([rxplaceholder[i]])
+                            T.writes([compute[i]])
+                            compute[i] = rxplaceholder[i]
+
+                @R.function
+                def rx_func(x: Tensor((n,), "float32"), y: Tensor(((n + 1),), "float32"))
+                    -> Tensor(None, "float32", ndim=-1):
+                    # block 0
+                    gv = relax.call_tir(te_func, (y,), R.Tensor((n + 1,), "float32"), (n,))
+                    return gv
+        """
+        name_hint = kwargs.pop("name_hint", "")
+        return self.emit(self.call_te(func, *args, **kwargs), name_hint=name_hint)
+
+    def match_cast(self, value: Expr, ty: Type, name_hint: str = "") -> Var:
+        """Emit a MatchCast.
+
+        Parameters
+        ----------
+        value : tvm.relax.Expr
+            The value of the MatchCast to be emitted.
+
+        ty : Type
+            The type to be matched.
+
+        name_hint : str
+            The name of the match cast
+
+        Returns
+        -------
+        ret : tvm.relax.Var
+            A newly created variable that get bounds to be the casted result.
+        """
+        return _ffi_api.BlockBuilderEmitMatchCast(
+            self,
+            value,
+            ty,
+            name_hint,
+        )  # type: ignore
+
+    def emit_output(self, output: Expr | Tuple | list[Expr], name_hint: str = "") -> Var:
+        """Emit output for the current dataflow block or function.
+
+        Parameters
+        ----------
+        output : Expr | Tuple | List[Expr]
+            The output of the current block/function.
+
+        name_hint : str
+            Name hint for the bound variable.
+
+        Returns
+        -------
+        ret : tvm.relax.Var
+            The return variable which gets bound to the output.
+        """
+        output = self._normalize_python_tuple(output)
+        return _ffi_api.BlockBuilderEmitOutput(self, output, name_hint)  # type: ignore
+
+    def emit_func_output(
+        self,
+        output: Expr | Tuple | list[Expr],
+        params: Var | Tuple | list[Var] | None = None,
+    ) -> GlobalVar:
+        """Emit output for the function.
+
+        Parameters
+        ----------
+        output : Expr | Tuple | List[Expr]
+            The output of the current block/function.
+
+        params : tvm.relax.Var | Tuple | List[tvm.relax.Var], optional
+            The parameters of the function to be built.
+            If params is None, it means the params have been initialized in the function with scope.
+
+        Returns
+        -------
+        gvar: tvm.ir.GlobalVar
+
+            A GlobalVar representing the function
+        """
+        if self._func._is_emit_func_output_called:
+            raise RuntimeError("emit_func_output must be called exactly once in a relax function.")
+        self._func._is_emit_func_output_called = True
+
+        if self._func._params is not None and params is not None:
+            raise RuntimeError(
+                "function parameters have been initialized in the function with scope."
+            )
+
+        if self._func._params is None and params is None:
+            raise RuntimeError("Relax function must have parameter.")
+
+        if self._func._params is None:
+            self._func._params = params
+
+        if BlockBuilder.current() is not self:
+            raise RuntimeError("BlockBuilder.current() must be self.")
+
+        output = self._normalize_python_tuple(output)
+
+        block = self._end_block()
+        if len(block.bindings) > 0:
+            self._func._blocks.append(block)
+
+        seqe = rx.SeqExpr(self._func._blocks, output)
+
+        # If the parameters were not provided as part of
+        # `bb.function()`, then any variables provided from the params
+        # are not in scope.  Otherwise, TIR variables used in dynamic
+        # inputs are removed as undefined (e.g. Replacing
+        # `R.Tensor(["batch_size"])` with `R.Tensor(ndims=1)`).
+        self.begin_scope(self._func._params)
+        try:
+            seqe = self.normalize(seqe)
+        finally:
+            self.end_scope()
+
+        # do not specify ret_ty and let constructor deduce
+        # from seqe.ty
+        func = rx.Function(self._func._params, seqe, is_pure=self._func._is_pure)
+        for key, value in self._func._attrs.items():
+            func = func.with_attr(key, value)
+        self.end_scope()
+        return self.add_func(func, self._func._name)
+
+    def normalize(self, expr: Expr) -> Expr:
+        """Normalize an Expr to complete its shape and type.
+
+        Parameters
+        ----------
+        expr : Expr
+            The input expr.
+
+        Returns
+        -------
+        ret : Expr
+            The expr with normalized shape and type.
+        """
+        return _ffi_api.BlockBuilderNormalize(self, expr)  # type: ignore
+
+    def get(self) -> tvm.IRModule:
+        """Return intermediate IRModule. For the situation where the IRModule is needed in the
+        middle of a building process.
+
+        Returns
+        -------
+        ret : tvm.IRModule
+            An IRModule with Relax and TIR functions being built.
+        """
+        return _ffi_api.BlockBuilderGetContextIRModule(self)  # type: ignore
+
+    def finalize(self) -> tvm.IRModule:
+        """Finalize the building process and return the result IRModule.
+
+        Possibly rename GlobalVars in the IRModule to ensure name uniqueness and the invariant:
+        every public function has the same name as its "global_symbol" attribute.
+
+        Note this method should be called only once at the end of the building process, since it may
+        invalidate global vars previously returned by this builder.
+        See also tvm.relax.transform.NormalizeGlobalVar.
+
+        Returns
+        -------
+        ret : tvm.IRModule
+            An IRModule with Relax and TIR functions being built.
+        """
+        return _ffi_api.BlockBuilderFinalize(self)  # type: ignore
+
+    def get_unique_name(self, name_prefix: str) -> str:
+        """Generate a unique name with a specified prefix.
+
+        Parameters
+        ----------
+        name_hint : str
+            The name prefix.
+
+        Returns
+        -------
+        ret : str
+            The generated name.
+        """
+        return _ffi_api.BlockBuilderGetUniqueName(self, name_prefix)  # type: ignore
+
+    def add_func(self, func: BaseFunc, func_name: str) -> GlobalVar:
+        """Add a Relax function or a TIR PrimFunc to the IRModule being built.
+
+        Parameters
+        ----------
+        func : BaseFunc
+            The function to be added.
+
+        func_name : str
+            The name of the function to be added.
+
+        Returns
+        -------
+        gvar : GlobalVar
+            The global var bound to the added function.
+        """
+        return _ffi_api.BlockBuilderAddFunction(self, func, func_name)  # type: ignore
+
+    def update_func(self, gv: GlobalVar, updated_func: BaseFunc) -> None:
+        """Add a Relax function or a TIR PrimFunc to the IRModule being built.
+
+        Parameters
+        ----------
+        gv : GlobalVar
+            The global var referring the function to be updated.
+
+        updated_func : BaseFunc
+            The updated function.
+        """
+        return _ffi_api.BlockBuilderUpdateFunction(self, gv, updated_func)  # type: ignore
+
+    def current_block_is_dataflow(self) -> bool:
+        """Check if the block being built is DataflowBlock or not.
+
+        Returns
+        -------
+        ret : bool
+            A boolean that indicates if the block being built is DataflowBlock or not.
+        """
+        return _ffi_api.BlockBuilderCurrentBlockIsDataFlow(self)  # type: ignore
+
+    def emit_normalized(self, binding: Binding) -> None:
+        """Emit an already normalized binding.
+
+        Parameters
+        ----------
+        binding: Binding
+            The binding to be emitted.
+        """
+        _ffi_api.BlockBuilderEmitNormalized(self, binding)  # type: ignore
+
+    def lookup_binding(self, var: Var) -> Expr | None:
+        """Lookup a var in the binding table.
+
+        Parameters
+        ----------
+        var: Var
+            The input var.
+
+        Returns
+        -------
+        expr: Expr
+            The Expr bound to the input var.
+        """
+        return _ffi_api.BlockBuilderLookupBinding(self, var)  # type: ignore
+
+    def begin_scope(self, params: list[Var] | None = None) -> None:
+        """Begin a new scope, with optional parameters that
+        are visible within the scope.
+
+        Parameters
+        ----------
+        params: Optional[List[Var]]
+            Parameters that are visible within the scope.
+
+        Note
+        ----
+        This function should be called when new scope is introduced
+        (function, seq) to properly track the variable availability
+        and help the best effort deduction.
+        """
+
+        return _ffi_api.BlockBuilderBeginScope(self, params)  # type: ignore
+
+    def end_scope(self) -> None:
+        """End the current scope. Please see `begin_scope` for details"""
+
+        return _ffi_api.BlockBuilderEndScope(self)  # type: ignore

python/tvm/relax/distributed/__init__.py

@@ -0,0 +1,24 @@
+# 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 infrastructure for distributed inference on Relax."""
+
+from .global_info import DeviceMesh, device_mesh
+from .type import Placement, DTensorType, PlacementSpec
+
+from . import transform

python/tvm/relax/distributed/_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 APIs for tvm.relax.distributed"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.distributed", __name__)

python/tvm/relax/distributed/global_info.py

@@ -0,0 +1,66 @@
+# 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=redefined-builtin, invalid-name
+"""Global Info Data structures for distributed tensor."""
+
+import tvm_ffi
+from tvm_ffi import Shape
+
+from tvm.ir import Range
+from tvm.ir.global_info import GlobalInfo
+
+from . import _ffi_api as ffi
+
+
+@tvm_ffi.register_object("relax.distributed.DeviceMesh")
+class DeviceMesh(GlobalInfo):
+    """Device mesh express a view of topology of devices,
+       represented by an n-d matrix of device ids.
+
+    Parameters
+    ----------
+    shape: Union[Shape, List[int], Tuple[int]]
+        Logical shape of device mesh
+    device_ids: Union[List[int], Range]
+        Represents the device id in the mesh
+    """
+
+    def __init__(self, shape: Shape | list[int] | tuple[int], device_ids: list[int] | Range):
+        if not isinstance(shape, Shape):
+            shape = Shape(shape)
+        device_range = None
+        if isinstance(device_ids, Range):
+            device_range = device_ids
+            device_ids = []
+        self.__init_handle_by_constructor__(ffi.DeviceMesh, shape, device_ids, device_range)  # type: ignore
+
+
+def device_mesh(shape: Shape, device_ids: list[int] | Range) -> DeviceMesh:
+    """Create a device mesh expression.
+    Parameters
+    ----------
+    shape : Shape
+        The shape of the device mesh.
+    device_ids: Union[List[int], Range]
+        Represents the device id in the mesh
+
+    Returns
+    -------
+    res : DeviceMesh
+        The device mesh.
+    """
+    return DeviceMesh(shape, device_ids)  # pylint: disable=no-member # type: ignore

python/tvm/relax/distributed/transform/__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.
+"""Relax distributed-related transformations."""
+
+from .transform import (
+    PropagateSharding,
+    LowerGlobalViewToLocalView,
+    LegalizeRedistribute,
+    LowerDistIR,
+)

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

@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+"""FFI APIs for tvm.relax.distributed.transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.distributed.transform", __name__)

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

@@ -0,0 +1,68 @@
+# 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
+"""Relax distributed-related transformation passes."""
+
+import tvm.ir
+
+from . import _ffi_api
+
+
+def PropagateSharding() -> tvm.ir.transform.Pass:
+    """Propagate sharding information.
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.PropagateSharding()  # type: ignore
+
+
+def LowerGlobalViewToLocalView() -> tvm.ir.transform.Pass:
+    """Lower global view TIR to local view
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.LowerGlobalViewToLocalView()  # type: ignore
+
+
+def LegalizeRedistribute() -> tvm.ir.transform.Pass:
+    """Legalize redistribute op to ccl op.
+    S->R: R.ccl.allgather
+    R->S: R.dist.redistribute_replica_to_shard
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.LegalizeRedistribute()  # type: ignore
+
+
+def LowerDistIR() -> tvm.ir.transform.Pass:
+    """Lower DistIR to Relax
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The registered pass
+    """
+    return _ffi_api.LowerDistIR()  # type: ignore

python/tvm/relax/distributed/type.py

@@ -0,0 +1,146 @@
+# 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=redefined-builtin, invalid-name
+"""Types for distributed tensor."""
+
+import enum
+
+import tvm_ffi
+
+from tvm.ir import Span
+from tvm.relax.type import TensorType, Type
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .global_info import DeviceMesh
+
+
+class PlacementSpecKind(enum.IntEnum):
+    kSharding = 0
+    kReplica = 1
+
+
+@tvm_ffi.register_object("relax.distributed.PlacementSpec")
+class PlacementSpec(Object):
+    """Describes how data is distributed in one dimension of the device mesh
+
+    Parameters
+    ----------
+    axis: int
+        If the kind is sharding, this value represents the tensor dimension to shard.
+        otherwise, axis is -1
+    kind: PlacementSpecKind
+        The kind of placement spec. Possible values: kSharding and kReplica.
+    """
+
+    axis: int
+    kind: PlacementSpecKind
+
+    def __init__(self, *args, **kwargs):
+        raise RuntimeError("PlacementSpec is not intended to be constructed directly, ")
+
+    @staticmethod
+    def sharding(axis: int) -> "PlacementSpec":
+        """Create a sharding placement spec
+
+        Parameters
+        ----------
+        axis: int
+            The tensor dimension to shard.
+
+        Returns
+        -------
+        placement_spec: PlacementSpec
+            The placement spec.
+        """
+        return _ffi_api.Sharding(axis)
+
+    @staticmethod
+    def replica() -> "PlacementSpec":
+        """Create a replica placement spec
+
+        Returns
+        -------
+        placement_spec: PlacementSpec
+            The placement spec.
+        """
+        return _ffi_api.Replica()
+
+
+@tvm_ffi.register_object("relax.distributed.Placement")
+class Placement(Object):
+    """Describes how data is distributed in each dimension of the device mesh
+
+    Parameters
+    ----------
+    dim_specs: List[PlacementSpec]
+        The placement spec for each dimension of the device mesh.
+    """
+
+    def __init__(self, dim_specs: list[PlacementSpec]):
+        self.__init_handle_by_constructor__(_ffi_api.Placement, dim_specs)  # type: ignore
+
+    @staticmethod
+    def from_text(text: str) -> "Placement":
+        """Create a placement from a text string.
+
+        Parameters
+        ----------
+        text: str
+            The text string.
+
+        Returns
+        -------
+        placement: Placement
+            The placement.
+        """
+        return _ffi_api.PlacementFromText(text)
+
+
+@tvm_ffi.register_object("relax.DTensorType")
+class DTensorType(Type):
+    """Type of a Distributed Tensor value.
+
+    Parameters
+    ----------
+    tensor_ty: TensorType
+        The tensor type carried by the distributed tensor.
+    device_mesh: DeviceMesh
+        The device mesh of the tensor.
+    placement: Placement
+        The placement of the tensor among the device mesh
+
+    """
+
+    tensor_ty: TensorType
+    device_mesh: DeviceMesh
+    placement: Placement
+
+    def __init__(
+        self,
+        tensor_ty: TensorType,
+        device_mesh: DeviceMesh,
+        placement: Placement,
+        span: Span = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DTensorType,
+            tensor_ty,
+            device_mesh,
+            placement,
+            span,  # type: ignore
+        )

python/tvm/relax/dpl/__init__.py

@@ -0,0 +1,29 @@
+# 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 Dataflow Pattern Language."""
+
+from .pattern import *
+from .context import *
+from .rewrite import (
+    rewrite_call,
+    rewrite_bindings,
+    PatternMatchingRewriter,
+    ExprPatternRewriter,
+    OrRewriter,
+)

python/tvm/relax/dpl/_ffi.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.
+"""DataFlow Pattern Language FFI bindings."""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.dpl", __name__)

python/tvm/relax/dpl/context.py

@@ -0,0 +1,79 @@
+# 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 Graph Matching Context Manager for Dataflow Pattern Language."""
+
+import tvm
+
+from ..expr import DataflowBlock, Var
+from . import _ffi as ffi
+from .pattern import DFPattern
+
+
+class PatternContext(tvm.runtime.Object):
+    """A context object for doing graph (topogical) pattern matching."""
+
+    def __init__(self, incremental=False):
+        """
+        Initialize the PatternContext
+
+        Parameters
+        ----------
+        incremental : bool, optional
+            perform incremental matching based on the recent context, by default False
+        """
+        self.__init_handle_by_constructor__(ffi.PatternContext, incremental)  # type: ignore
+
+    def __enter__(self):
+        """Enter the context"""
+        ffi.enter_context(self)  # type: ignore
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit the context"""
+        ffi.exit_context(self)  # type: ignore
+
+    @staticmethod
+    def current() -> "PatternContext":
+        """
+        Get the current context
+
+        Returns
+        -------
+        PatternContext
+            The current context
+        """
+        return ffi.current_context()  # type: ignore
+
+    def match_dfb(
+        self,
+        dfb: DataflowBlock,
+    ) -> dict[DFPattern, Var]:
+        """
+        Match a DataflowBlock via a graph of DFPattern and corresponding constraints
+
+        Parameters
+        ----------
+        dfb : DataflowBlock
+            The DataflowBlock to match
+
+        Returns
+        -------
+        Dict[DFPattern, Var]
+            The mapping from DFPattern to matched expression
+        """
+        return ffi.match_dfb(self, dfb)  # type: ignore

python/tvm/relax/dpl/rewrite.py

@@ -0,0 +1,299 @@
+# 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.
+"""APIs for pattern-based rewriting."""
+
+from collections.abc import Callable
+
+from tvm_ffi import register_object
+
+from tvm.ir import IRModule
+from tvm.runtime import Object
+
+from ..expr import Expr, Function, Var
+from . import _ffi as ffi
+from .context import PatternContext
+from .pattern import DFPattern
+
+
+@register_object("relax.dpl.PatternMatchingRewriter")
+class PatternMatchingRewriter(Object):
+    """A pattern-matching rewriter for Relax"""
+
+    @staticmethod
+    def from_pattern(
+        pattern: DFPattern,
+        func: Callable[[Expr, dict[DFPattern, Expr]], Expr],
+    ) -> "PatternMatchingRewriter":
+        """Construct from a pattern and rewriter-function
+
+        The replacements performed by the rewriter will be equivalent
+        to using the `pattern` and `func` as arguments to
+        `rewrite_call`.
+
+        Parameters
+        ----------
+        pattern: DFPattern
+
+            The pattern to be matched against.
+
+        func: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+
+            A function that returns the rewritten expression.  See
+            `rewrite_call` for details and examples.
+
+
+        Returns
+        -------
+        rewriter_obj: PatternMatchingRewriter
+
+            The rewriter object
+
+        """
+        return ffi.PatternMatchingRewriterFromPattern(
+            pattern,
+            func,
+        )  # type: ignore
+
+    @staticmethod
+    def from_module(mod: IRModule) -> "PatternMatchingRewriter":
+        """Construct a rewriter from an IRModule
+
+        The IRModule must have two publicly-exposed functions,
+        `pattern` and `replacement`, where `pattern` and `replacement`
+        have the same function signature, as shown in the example
+        below.
+
+        .. code-block:: python
+
+            @I.ir_module
+            class RewriteAddIntoMultiply:
+                @R.function
+                def pattern(A: R.Tensor):
+                    B = A + A
+                    return B
+
+                @R.function
+                def replacement(A: R.Tensor):
+                    B = A * 2
+                    return B
+
+            rewriter = PatternMatchingRewriter.from_module(RewriteAddIntoMultiply)
+            rewritten_ir_module = rewriter(ir_module)
+
+        To support the common case of defining an IRModule with
+        TVMScript, then immediately turning it into a rewriter, the
+        `@R.rewriter` annotation can be used.
+
+        .. code-block:: python
+
+            @R.rewriter
+            class RewriteAddIntoMultiply:
+                @R.function
+                def pattern(A: R.Tensor):
+                    B = A + A
+                    return B
+
+                @R.function
+                def replacement(A: R.Tensor):
+                    B = A * 2
+                    return B
+
+            rewritten_ir_module = RewriteAddIntoMultiply(ir_module)
+
+        Parameters
+        ----------
+        mod: IRModule
+
+            A module with `pattern` and `replacement` functions,
+            defining a rewrite rule.
+
+
+        Returns
+        -------
+        rewriter_obj: PatternMatchingRewriter
+
+            The rewriter object
+
+        """
+        return ffi.PatternMatchingRewriterFromModule(mod)  # type: ignore
+
+    def __call__(self, obj: Expr | IRModule) -> Expr | IRModule:
+        """Apply the rewriter
+
+        Parameters
+        ----------
+        obj: Union[Expr, IRModule])
+
+            The object to be rewritten.  May be applied to either a
+            relax expression, or an IRModule.
+
+        Returns
+        -------
+        updated: Union[Expr, IRModule]
+
+            The rewritten object
+
+        """
+        return ffi.PatternMatchingRewriterApply(self, obj)
+
+    def __or__(self, other: "PatternMatchingRewriter") -> "PatternMatchingRewriter":
+        """Compose two rewriters
+
+        Composing two rewrite rules together allows them to be applied
+        in a single Relax-level transformation.
+
+        Parameters
+        ----------
+        other: PatternMatchingRewriter
+
+            Another rewrite rule
+
+        Returns
+        -------
+        PatternMatchingRewriter
+
+            A rewriter that will apply either rewrite pattern
+
+        """
+        return OrRewriter(self, other)
+
+
+@register_object("relax.dpl.ExprPatternRewriter")
+class ExprPatternRewriter(PatternMatchingRewriter):
+    def __init__(self, pattern, func):
+        self.__init_handle_by_constructor__(
+            ffi.PatternRewriter,
+            pattern,
+            func,
+        )  # type: ignore
+
+
+@register_object("relax.dpl.OrRewriter")
+class OrRewriter(PatternMatchingRewriter):
+    def __init__(self, lhs, rhs):
+        self.__init_handle_by_constructor__(
+            ffi.OrRewriter,
+            lhs,
+            rhs,
+        )  # type: ignore
+
+
+@register_object("relax.dpl.TupleRewriter")
+class TupleRewriter(PatternMatchingRewriter):
+    def __init__(self, patterns, func):
+        self.__init_handle_by_constructor__(
+            ffi.TupleRewriter,
+            patterns,
+            func,
+        )  # type: ignore
+
+
+def rewrite_call(
+    pattern: DFPattern,
+    rewriter: Callable[[Expr, dict[DFPattern, Expr]], Expr],
+    func: Function,
+) -> Function:
+    """
+    Rewrite a function with the given pattern and the rewriter function.
+
+    Parameters
+    ----------
+    pattern: DFPattern
+        The pattern to match.
+
+    rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+        The function to be called on a successful matching for rewriting. Given the matched
+        call node and the map of patterns and matched expressions, it should return a new call node
+        to replace the original one or the original matched call node as is.
+
+        For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+        ```
+        x = wildcard()
+        pattern = is_op("relax.add")(x, x)
+
+        def rewriter(orig, matchings):
+            return R.multiply(matchings[x], R.const(2, "float32"))
+        ```
+
+    func: Function
+        The function to rewrite.
+
+    Returns
+    -------
+    rewritten_func: Function
+        The rewritten or the input function, depending on the pattern matching result.
+    """
+    return ffi.rewrite_call(pattern, rewriter, func)
+
+
+def rewrite_bindings(
+    ctx: PatternContext,
+    rewriter: Callable[[dict[DFPattern, Var], dict[Var, Expr]], dict[Var, Expr]],
+    func: Function,
+) -> Function:
+    """
+    Rewrite a function with the given pattern and the rewriter function.
+
+    Parameters
+    ----------
+    ctx: PatternContext
+        The pattern constraint context under which rewriting takes place.
+
+    rewriter: Callable[[Dict[DFPattern, Var], Dict[Var, Expr]], Dict[Var, Expr]]
+        The function to be called on a successful matching for rewriting. Given the map of patterns
+        and corresponding variables (bound variables or parameters), it should return a map that
+        specifies new values for matched bound variables. It can refer to the passed bindings to
+        create the replacement expressions.
+
+        For example, to rewrite three matmuls for QKV projection in transformer models into one
+        matmul followed by slicing, one can use the follwoing rewriter:
+        ```
+        inp_pat = wildcard()
+        Q_weight_pat, K_weight_pat, V_weight_pat = wildcard(), wildcard(), wildcard()
+
+        matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat)
+        matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat)
+        matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat)
+
+        def rewriter(matchings):
+            inp = matchings[inp_pat]
+            Q_weight = matchings[Q_weight_pat]
+            K_weight = matchings[K_weight_pat]
+            V_weight = matchings[V_weight_pat]
+            width = Q_weight.ty.shape[1]
+
+            concat = R.concat([Q_weight, K_weight, V_weight], axis=1)
+            matmul = R.matmul(inp, concat)
+            Q = R.strided_slice(matmul, axes=[2], begin=[0], end=[width])
+            K = R.strided_slice(matmul, axes=[2], begin=[width], end=[width * 2])
+            V = R.strided_slice(matmul, axes=[2], begin=[width * 2], end=[width * 3])
+
+            # matchings[matmul1] gives the bound variable in the binding whose RHS matches with
+            # the matmul1 pattern. For example, lv0 in lv0 = R.matmul(x1, w0).
+            # We want to replace the RHS of this binding with Q.
+            return {matchings[matmul1]: Q, matchings[matmul2]: K, matchings[matmul3]: V}
+        ```
+
+    func: Function
+        The function to rewrite.
+
+    Returns
+    -------
+    rewritten_func: Function
+        The rewritten or the input function, depending on the pattern matching result.
+    """
+    return ffi.rewrite_bindings(ctx, rewriter, func)

python/tvm/relax/exec_builder.py

@@ -0,0 +1,151 @@
+# 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: RUF012
+"""A builder to build Relax VM executable."""
+
+from enum import IntEnum
+
+import tvm_ffi
+from tvm_ffi import Shape
+
+import tvm
+
+from . import _ffi_api
+from .vm_build import VMExecutable
+
+
+class SpecialReg(IntEnum):
+    """Magic numbers that represent special registers in vm."""
+
+    VOID_ARG = (1 << 54) + 0
+    VM_STATE = (1 << 54) + 1
+
+
+class VMFuncKind(IntEnum):
+    """VM function kind code."""
+
+    PACKED_FUNC = 0
+    VM_FUNC = 1
+
+
+class VMFuncScope:
+    """An object corresponds to each VM function, working as a context manager."""
+
+    stack: list["VMFuncScope"] = []
+
+    def __init__(self, exit_callback):
+        self.exit_callback = exit_callback
+
+    def __enter__(self):
+        VMFuncScope.stack.append(self)
+        return self
+
+    def __exit__(self, ptype, value, trace):
+        VMFuncScope.stack.pop()
+        self.exit_callback()
+
+
+@tvm_ffi.register_object("relax.ExecBuilder")
+class ExecBuilder(tvm_ffi.core.Object):
+    """A builder to emit instructions and build executable for the virtual machine."""
+
+    def __init__(self) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.ExecBuilderCreate)  # type: ignore
+
+    def r(self, idx: int) -> int:
+        """set instruction's argument as a register."""
+        return _ffi_api.ExecBuilderR(self, idx)  # type: ignore
+
+    def imm(self, value: int) -> int:
+        """set instruction's argument as an immediate."""
+        return _ffi_api.ExecBuilderImm(self, value)  # type: ignore
+
+    def c(self, idx: int) -> int:
+        """set instruction's argument as a constant."""
+        return _ffi_api.ExecBuilderC(self, idx)  # type: ignore
+
+    def f(self, name: str) -> int:
+        """set instruction's argument as a function."""
+        return _ffi_api.ExecBuilderF(self, name)  # type: ignore
+
+    def void_arg(self) -> int:
+        return self.r(SpecialReg.VOID_ARG)
+
+    def vm_state(self) -> int:
+        return self.r(SpecialReg.VM_STATE)
+
+    def declare_function(self, func_name: str, kind: VMFuncKind = VMFuncKind.PACKED_FUNC) -> None:
+        """Declare a function"""
+        _ffi_api.ExecBuilderDeclareFunction(self, func_name, kind)  # type: ignore
+
+    def function(
+        self, func_name: str, num_inputs: int | None = 0, param_names: list[str] | None = None
+    ) -> VMFuncScope:
+        """annotate a VM function."""
+        _ffi_api.ExecBuilderEmitFunction(self, func_name, num_inputs, param_names)  # type: ignore
+        return VMFuncScope(lambda: _ffi_api.ExecBuilderEndFunction(self, func_name))  # type: ignore
+
+    def _check_scope(self) -> None:
+        if len(VMFuncScope.stack) == 0:
+            raise ValueError("emit should happen in a function scope")
+
+    def convert_constant(self, const: object) -> int:
+        return _ffi_api.ExecBuilderConvertConstant(self, const)  # type: ignore
+
+    def emit_call(
+        self,
+        name: str,
+        args: list[tvm.runtime.Tensor | tvm.DataType] | None = None,
+        dst: int | None = None,
+    ) -> None:
+        """emit a call instruction which calls a packed function."""
+        self._check_scope()
+        if dst is None:
+            dst = SpecialReg.VOID_ARG
+        args_ = []
+        if args is not None:
+            for arg in args:
+                if isinstance(arg, tuple):
+                    shape_tuple = Shape(arg)
+                    new_arg = self.convert_constant(shape_tuple)
+                    args_.append(new_arg)
+                elif isinstance(arg, tvm.runtime.Tensor | tvm.DataType | Shape):
+                    new_arg = self.convert_constant(arg)
+                    args_.append(new_arg)
+                else:
+                    args_.append(arg)
+        _ffi_api.ExecBuilderEmitCall(self, name, args_, dst)  # type: ignore
+
+    def emit_ret(self, result: int) -> None:
+        """emit a return instruction"""
+        self._check_scope()
+        _ffi_api.ExecBuilderEmitRet(self, result)  # type: ignore
+
+    def emit_goto(self, pc_offset):
+        """emit a goto instruction"""
+        self._check_scope()
+        _ffi_api.ExecBuilderEmitGoto(self, pc_offset)  # type: ignore
+
+    def emit_if(self, cond, false_offset):
+        """emit an if instruction"""
+        self._check_scope()
+        _ffi_api.ExecBuilderEmitIf(self, cond, false_offset)  # type: ignore
+
+    def get(self) -> VMExecutable:
+        """return the executable"""
+        return VMExecutable(_ffi_api.ExecBuilderGet(self))  # type: ignore

python/tvm/relax/expr.py

@@ -0,0 +1,884 @@
+# 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.
+# ruff: noqa: F401
+"""The expression nodes of Relax."""
+
+import typing
+from collections.abc import Callable, Mapping
+from numbers import Integral, Number, Real
+from typing import Any, Optional, Union
+
+import numpy as _np  # type: ignore
+import tvm_ffi
+from tvm_ffi.core import String
+
+import tvm.ir
+import tvm.relax
+import tvm.runtime
+from tvm import DataType
+
+from ..ir import BaseFunc, Node, Span
+from ..runtime import Scriptable
+from . import _ffi_api
+
+# It is a workaround for mypy: https://github.com/python/mypy/issues/7866#issuecomment-549454370
+# This feature is not supported until python 3.10:
+# https://docs.python.org/3.10/whatsnew/3.10.html#pep-613-typealias
+Expr = tvm.ir.Expr
+Type = tvm.ir.Type  # pylint: disable=invalid-name
+GlobalVar = tvm.ir.GlobalVar
+
+
+def prim_value(value: Expr | int | float, dtype: str | None = None) -> Expr:
+    """Convert a Python scalar or primitive expression to ``Expr``.
+
+    Parameters
+    ----------
+    value : Expr | int | float
+        The value to convert.
+
+    dtype : Optional[str]
+        The dtype to use when converting Python numeric values.
+
+    Returns
+    -------
+    result : Expr
+        The converted primitive expression.  Existing ``Expr`` inputs are
+        returned unchanged.
+    """
+    if tvm.ir.is_prim_expr(value):
+        return value
+    if isinstance(value, bool | _np.bool_):
+        return tvm.tirx.IntImm(dtype or "bool", int(value))
+    if isinstance(value, Integral):
+        return tvm.tirx.IntImm(dtype or "int64", int(value))
+    if isinstance(value, Real):
+        return tvm.tirx.FloatImm(dtype or "float64", float(value))
+    tvm_value = tvm_ffi.convert(value)
+    if tvm.ir.is_prim_expr(tvm_value):
+        return tvm_value
+    raise TypeError(f"Cannot convert {value} with type {type(value)} to `Expr`")
+
+
+def _relax_type_is_base_of(self: Type, derived: Type) -> bool:
+    """Check if this Relax type is a base of another Relax type."""
+
+    return _ffi_api.TypeIsBaseOf(self, derived)  # type: ignore
+
+
+Type.is_base_of = _relax_type_is_base_of  # type: ignore[attr-defined]
+
+
+# will be registered afterwards in python/tvm/relax/op/init.py
+_op_ffi_api = None  # pylint: disable=invalid-name
+
+
+def _binary_op_helper(lhs: "ExprWithOp", rhs: "ExprWithOp", op: Callable) -> "ExprWithOp":
+    if not isinstance(lhs, Expr):  # type: ignore
+        raise ValueError("lhs must be Expr")
+    if isinstance(rhs, Expr):  # type: ignore
+        return op(lhs, rhs)
+    elif isinstance(rhs, Number):
+        raise TypeError(f"Please convert {rhs} with `const` first")
+    else:
+        raise TypeError(f"type {type(rhs)} not supported")
+
+
+def _binary_rhs_helper(rhs: "ExprWithOp") -> "ExprWithOp":
+    if isinstance(rhs, Number):
+        raise TypeError(f"Please convert {rhs} with `const` first")
+    raise TypeError(f"type {type(rhs)} not supported")
+
+
+class ExprWithOp(Expr, Scriptable):
+    """Basetype of all relax expressions that defines op overloading."""
+
+    def astype(self, dtype: str | DataType) -> "ExprWithOp":
+        """Cast the content type of the current data to dtype.
+
+        Parameters
+        ----------
+        dtype : str
+            The target data type.
+
+        Note
+        ----
+        This function only works for TensorType Exprs.
+
+        Returns
+        -------
+        result : ExprWithOp
+            The result expression.
+        """
+        return _op_ffi_api.astype(self, dtype)  # type: ignore
+
+    def __neg__(self) -> "ExprWithOp":
+        return _op_ffi_api.negative(self)  # type: ignore
+
+    def __lt__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.less)  # type: ignore
+
+    def __gt__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.greater)  # type: ignore
+
+    def __ge__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.greater_equal)  # type: ignore
+
+    def __le__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.less_equal)  # type: ignore
+
+    # NOTE: Cannot override __eq__ and __ne__, which will influence object equal
+
+    def __add__(self, other: Expr) -> "ExprWithOp":
+        if isinstance(self.ty, tvm.relax.TupleType) and isinstance(other, tuple):
+            return tuple([*self, *other])
+
+        return _binary_op_helper(self, other, _op_ffi_api.add)  # type: ignore
+
+    def __radd__(self, other: Expr) -> "ExprWithOp":
+        return self.__add__(other)
+
+    def __sub__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.subtract)  # type: ignore
+
+    def __rsub__(self, other: Expr) -> "ExprWithOp":
+        return _binary_rhs_helper(other)
+
+    def __mul__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.multiply)  # type: ignore
+
+    def __rmul__(self, other: Expr) -> "ExprWithOp":
+        return self.__mul__(other)
+
+    def __truediv__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.divide)  # type: ignore
+
+    def __rtruediv__(self, other: Expr) -> "ExprWithOp":
+        return _binary_rhs_helper(other)
+
+    def __floordiv__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.floor_divide)  # type: ignore
+
+    def __rfloordiv__(self, other: Expr) -> "ExprWithOp":
+        return _binary_rhs_helper(other)
+
+    def __mod__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.mod)  # type: ignore
+
+    def __rmod__(self, other: Expr) -> "ExprWithOp":
+        return _binary_rhs_helper(other)
+
+    def __pow__(self, other: Expr) -> "ExprWithOp":
+        return _binary_op_helper(self, other, _op_ffi_api.power)  # type: ignore
+
+    def __rpow__(self, other: Expr) -> "ExprWithOp":
+        return _binary_rhs_helper(other)
+
+    def __call__(self, *args: list[Expr], attrs: dict[str, Any] | None = None) -> "ExprWithOp":
+        """Call the variable (if it represents a function).
+
+        Parameters
+        ----------
+        args: List[Expr]
+            The arguments to the call.
+
+        attr: Optional[Dict[str, object]]
+            The additional attributes to the call.
+
+        Returns
+        -------
+        call: ExprWithOp
+            A call taking the variable as a function.
+        """
+        return tvm.ir.Call(self, args, attrs=attrs)
+
+    def __getitem__(self, index: int) -> "ExprWithOp":
+        """Get the i-th element of the tuple or Expr with TupleType.
+
+        Parameters
+        ----------
+        index: int
+            The index of the element to be retrieved.
+
+        Note
+        ----
+        This function will be overridden by Tuple and ShapeExpr
+
+        Returns
+        -------
+        result: ExprWithOp
+            The result expression.
+        """
+        try:
+            return TupleGetItem(self, index)
+        except RuntimeError as err:
+            # For Python objects with __getitem__, but without
+            # __len__, tuple unpacking is done by iterating over
+            # sequential indices until IndexError is raised.
+            # Therefore, convert from RuntimeError to IndexError for
+            # compatibility.
+            if "Index out of bounds" in err.args[0]:
+                raise IndexError from err
+            raise
+
+
+@tvm_ffi.register_object("relax.expr.If")
+class If(ExprWithOp):
+    """A conditional expression in Relax.
+
+    Parameters
+    ----------
+    cond: Expr
+        The condition.
+
+    true_branch: Expr
+        The expression evaluated when condition is true.
+
+    false_branch: Expr
+        The expression evaluated when condition is false.
+
+    span: Optional[Span]
+        Span that points to original source code
+    """
+
+    cond: Expr
+    true_branch: Expr
+    false_branch: Expr
+    span: Span | None
+
+    def __init__(self, cond: Expr, true_branch: Expr, false_branch: Expr, span: Span | None = None):
+        self.__init_handle_by_constructor__(
+            _ffi_api.If,
+            cond,
+            true_branch,
+            false_branch,
+            span,  # type: ignore
+        )
+
+
+@tvm_ffi.register_object("relax.expr.Tuple")
+class Tuple(ExprWithOp):
+    """Tuple expression that groups several fields together.
+
+    Parameters
+    ----------
+    fields : Union[List[Expr], typing.Tuple[Expr, ...]]
+        The fields in the tuple.
+
+    span: Optional[Span]
+        Span that points to original source code
+    """
+
+    fields: list[Expr]
+    span: Span | None
+
+    def __init__(self, fields: list[Expr] | tuple[Expr, ...], span: Span | None = None):
+        if isinstance(fields, tvm.relax.Tuple):
+            fields = fields.fields
+        elif isinstance(getattr(fields, "ty", None), tvm.relax.TupleType):
+            fields = [*fields]
+
+        self.__init_handle_by_constructor__(_ffi_api.Tuple, fields, span)  # type: ignore
+
+    def __getitem__(self, index: int) -> Expr:
+        if index >= len(self) or index < -len(self):
+            raise IndexError("Tuple index out of range")
+        return self.fields[index]
+
+    def __len__(self) -> int:
+        return len(self.fields)
+
+
+@tvm_ffi.register_object("relax.expr.TupleGetItem")
+class TupleGetItem(ExprWithOp):
+    """Get index-th item from a tuple.
+
+    Parameters
+    ----------
+    tuple_value: Expr
+        The input tuple expression.
+
+    index: int
+        The index.
+
+    span: Optional[Span]
+        Span that points to original source code
+    """
+
+    tuple_value: Expr
+    index: int
+    span: Span | None
+
+    def __init__(self, tuple_value: Expr, index: int, span: Span | None = None):
+        self.__init_handle_by_constructor__(
+            _ffi_api.TupleGetItem,
+            tuple_value,
+            index,
+            span,  # type: ignore
+        )
+
+
+@tvm_ffi.register_object("relax.expr.ShapeExpr")
+class ShapeExpr(ExprWithOp):
+    """A shape expression which allows users to construct a shape containing Expr.
+
+    Parameters
+    ----------
+    values: Union[List[Expr], typing.Tuple[Expr, ...], tvm_ffi.Array]
+        The values of the shape expression.
+
+    span: Optional[Span]
+        Span that points to original source code
+    """
+
+    values: list[Expr]
+    span: Span | None
+
+    def __init__(
+        self,
+        values: list[Expr] | tuple[Expr, ...] | tvm_ffi.Array,
+        span: Span | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.ShapeExpr, values, span)  # type: ignore
+
+    def __getitem__(self, index):
+        if index >= len(self) or index < -len(self):
+            raise IndexError("ShapeExpr index out of range")
+        return self.values[index]
+
+    def __len__(self):
+        return len(self.values)
+
+
+def make_shape(shape: list[Any] | tuple[Any, ...]) -> ShapeExpr:
+    if isinstance(shape, list | tuple):
+        return ShapeExpr(shape)
+    raise TypeError(
+        "make_shape expects a list or tuple of shape values, "
+        f"but received type {type(shape).__name__}"
+    )
+
+
+@tvm_ffi.register_object("relax.expr.Constant")
+class Constant(ExprWithOp):
+    """Constant Tensor
+
+    Parameters
+    ----------
+    data: tvm.runtime.Tensor
+        The data of the constant tensor.
+
+    ty: Optional[Type]
+        The type of the constant tensor. If not specified, infer it from data.
+
+    span: Optional[Span]
+        Span that points to original source code
+
+    Note
+    ----
+    Scalar constants are represented by ndim-0 constant tensors.
+    """
+
+    data: tvm.runtime.Tensor
+    span: Span | None
+
+    def __init__(
+        self,
+        data: tvm.runtime.Tensor,
+        ty: Type | None = None,
+        span: Span | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.Constant,
+            data,
+            ty,
+            span,  # type: ignore
+        )
+
+
+@tvm_ffi.register_object("relax.expr.Var")
+class Var(ExprWithOp):
+    """The variable class for all Relax bindings.
+
+    Parameters
+    ----------
+    name_hint: str
+        The name hint of the variable.
+
+    ty: Optional[Type]
+        The type annotation of the variable.
+
+    span: Optional[Span]
+        Span that points to original source code
+    """
+
+    name_hint: str
+    span: Span | None
+
+    def __init__(
+        self,
+        name_hint: str,
+        ty: Type | None = None,
+        span: Span | None = None,
+    ) -> None:
+        if ty is not None:
+            ty = tvm.runtime.convert(ty)
+            if not isinstance(ty, Type):
+                raise TypeError(
+                    "ty needs to be an instance of Type. "
+                    "If you attempt to pass in shape, "
+                    "use relax.TensorType(shape, dtype)."
+                )
+        self.__init_handle_by_constructor__(
+            _ffi_api.Var,  # type: ignore
+            name_hint,
+            ty,
+            span,
+        )
+
+
+@tvm_ffi.register_object("relax.expr.DataflowVar")
+class DataflowVar(Var):
+    """A sub-type of the variable node used to mark dataflow variables from
+    normal visible "function local" bindings.
+
+
+    Parameters
+    ----------
+    name_hint: str
+        The name hint of the variable.
+
+    ty: Optional[Type]
+        The type annotation of the variable.
+
+    span: Optional[Span]
+        Span that points to original source code
+    """
+
+    name_hint: str
+    span: Span | None
+
+    def __init__(
+        self,
+        name_hint: str,
+        ty: Type | None = None,
+        span: Span | None = None,
+    ) -> None:
+        # pylint: disable=super-init-not-called
+        if ty is not None:
+            ty = tvm.runtime.convert(ty)
+            if not isinstance(ty, Type):
+                raise TypeError(
+                    "ty needs to be an instance of Type. "
+                    "If you attempt to pass in shape, "
+                    "use relax.TensorType(shape, dtype)."
+                )
+
+        self.__init_handle_by_constructor__(_ffi_api.DataflowVar, name_hint, ty, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.StringImm")
+class StringImm(Expr, Scriptable):
+    """Represent a string literal constant."""
+
+    value: str
+    span: Span | None
+
+    def __init__(self, value: str, span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.StringImm, value, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.DataTypeImm")
+class DataTypeImm(Expr, Scriptable):
+    """Represent a data type constant."""
+
+    value: DataType
+    span: Span | None
+
+    def __init__(self, value: DataType | str, span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.DataTypeImm, value, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.Binding")
+class Binding(Node, Scriptable):
+    """The base class of a binding in Relax."""
+
+    var: Var
+    span: Span | None
+
+
+@tvm_ffi.register_object("relax.expr.MatchCast")
+class MatchCast(Binding):
+    """Runtime-match the value to the type.
+
+    This operation does runtime check, populates the un-defined symbolic shape vars
+    and vars in ty in the first occurrence, and insert equality assertions in
+    other cases.
+
+    Parameters
+    ----------
+    var: Var
+        The return variable that the match cast bind to.
+
+    value: Expr
+        The input value expression.
+
+    ty: tvm.relax.Type
+        The type to match cast to.
+    """
+
+    ty: Type
+    value: Expr
+    span: Span | None
+
+    def __init__(self, var: Var, value: Expr, ty: Type, span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.MatchCast,
+            var,
+            value,
+            ty,
+            span,  # type: ignore
+        )
+
+
+@tvm_ffi.register_object("relax.expr.VarBinding")
+class VarBinding(Binding):
+    """Variable binding, bind he variable of the lhs with the rhs.
+
+    Parameters
+    ----------
+    var: Var
+        The return variable that the match cast bind to.
+
+    value: Expr
+        The input value expression.
+
+    """
+
+    var: Var
+    value: Expr
+    span: Span | None
+
+    def __init__(self, var: Var, value: Expr, span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.VarBinding, var, value, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.BindingBlock")
+class BindingBlock(Node, Scriptable):
+    """base class of binding block, bindings inside can be impure
+    (with side effect or control flow)"""
+
+    bindings: list[Binding]
+    span: Span | None
+
+    def __init__(self, bindings: list[Binding], span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.BindingBlock, bindings, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.DataflowBlock")
+class DataflowBlock(BindingBlock):
+    """dataflow block, bindings inside are pure (no side effect and no control flow)"""
+
+    bindings: list[Binding]
+    span: Span | None
+
+    def __init__(self, bindings: list[Binding], span: Span | None = None) -> None:
+        # pylint: disable=super-init-not-called
+        self.__init_handle_by_constructor__(_ffi_api.DataflowBlock, bindings, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.SeqExpr")
+class SeqExpr(ExprWithOp):
+    """A sequence of binding blocks followed by an expression."""
+
+    blocks: list[BindingBlock]
+    body: Expr
+    span: Span | None
+
+    def __init__(self, blocks: list[BindingBlock], body: Expr, span: Span | None = None) -> None:
+        self.__init_handle_by_constructor__(_ffi_api.SeqExpr, blocks, body, span)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.Function")
+class Function(BaseFunc, Scriptable):
+    """A Relax function."""
+
+    params: list[Var]
+    body: Expr
+    ret_ty: Type
+    is_pure: bool
+    attrs: tvm.ir.DictAttrs
+    span: Span | None
+
+    def __init__(
+        self,
+        params: list[Var],
+        body: Expr,
+        ret_ty: Type | None = None,
+        is_pure: bool | None = True,
+        attrs: tvm.ir.DictAttrs | None = None,
+        span: Span | None = None,
+    ) -> None:
+        if attrs is None:
+            attrs = tvm.ir.DictAttrs({})
+        self.__init_handle_by_constructor__(
+            _ffi_api.Function,
+            params,
+            body,
+            ret_ty,
+            is_pure,
+            attrs,
+            span,
+        )  # type: ignore
+
+    @staticmethod
+    def create_empty(
+        params: list[Var],
+        ret_ty: Type,
+        is_pure: bool | None = True,
+        attrs: tvm.ir.DictAttrs | None = None,
+        span: Span | None = None,
+    ):
+        """Construct a relax.Function but without body"""
+        if attrs is None:
+            attrs = tvm.ir.DictAttrs({})
+        return _ffi_api.FunctionCreateEmpty(params, ret_ty, is_pure, attrs, span)  # type: ignore
+
+    def __call__(self, *args):
+        """Invoke the global function.
+
+        Parameters
+        ----------
+        args: List[relax.Expr]
+            Arguments.
+        """
+        return tvm.ir.Call(self, args, None, None)
+
+    def bind_symbolic_vars(self, binding_map: Mapping[str | tvm.tirx.Var, Expr]) -> "Function":
+        """Return a new function with updated symbolic variable
+
+        Parameters
+        ----------
+        binding_map: Mapping[str | tvm.tirx.Var, Expr]
+
+            The mapping of values to be replaced.  Keys may be either
+            a `tirx.Var` or a string name of the variable.  If the
+            variables are referred to by name, the name must uniquely
+            identify a symbolic variable in the function.
+
+        Returns
+        -------
+        func: Function
+
+            The updated function
+        """
+
+        # Relax uses int64 for symbolic variables, but the FFI
+        # converts python integers into int32.
+        binding_map = {
+            key: tvm.tirx.const(value, "int64") if isinstance(value, int) else value
+            for key, value in binding_map.items()
+        }
+
+        return _ffi_api.FunctionBindSymbolicVars(self, binding_map)  # type: ignore
+
+    def bind_params(
+        self,
+        binding_map: Mapping[
+            str | Var,
+            int | float | Expr | tvm.runtime.Tensor | _np.ndarray,
+        ],
+    ) -> "Function":
+        """Return a new function with updated symbolic variable
+
+        Parameters
+        ----------
+        binding_map: Mapping[
+                str | Var,
+                int | float | Expr | tvm.runtime.Tensor | _np.ndarray,
+        ]
+
+            The mapping of values to be replaced.
+
+            Keys may be either a `relax.Var` or a string name of the
+            Relax variable.  If the variables are referred to by name,
+            the name must uniquely identify a parameter in the
+            function.
+
+            Values must be a relax expression, or a value that is
+            convertible into a relax expression.  The value must be
+            compatible with the variable being replaced.
+
+        Returns
+        -------
+        func: Function
+
+            The updated function
+        """
+
+        def _normalize_value(value):
+            # Conversions that must occur prior to the FFI
+            # conversions.
+            if isinstance(value, int):
+                # Relax uses int64 for symbolic variables, but the FFI
+                # converts python integers into int32.
+                return tvm.tirx.const(value, "int64")
+            elif isinstance(value, _np.ndarray | tvm.runtime.Tensor):
+                return tvm.relax.const(value)
+            else:
+                return value
+
+        binding_map = {key: _normalize_value(value) for key, value in binding_map.items()}
+
+        return _ffi_api.FunctionBindParams(self, binding_map)  # type: ignore
+
+    def inline_functions(
+        self, function_map: Mapping[str | tvm.ir.GlobalVar, "Function"]
+    ) -> "Function":
+        return _ffi_api.FunctionInlineFunctions(self, function_map)  # type: ignore
+
+
+@tvm_ffi.register_object("relax.expr.ExternFunc")
+class ExternFunc(BaseFunc, ExprWithOp):
+    """extern function, which represents a PackedFunc."""
+
+    global_symbol: String
+    span: Span | None
+
+    def __init__(
+        self,
+        global_symbol: String,
+        ty: Type | None = None,
+        span: Span | None = None,
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.ExternFunc,
+            global_symbol,
+            ty,
+            span,  # type: ignore
+        )
+
+
+def extern(name: str, ty: Type | None = None, span: Span | None = None):
+    """Create extern function."""
+    return ExternFunc(name, ty, span)
+
+
+def const(
+    value: bool | int | float | _np.ndarray | tvm.runtime.Tensor, dtype: str | None = None
+) -> Constant:
+    """Create a constant value.
+
+    Parameters
+    ----------
+    value: bool | int | float | numpy.ndarray | tvm.runtime.Tensor
+        The constant value.
+
+    dtype: Optional[str]
+        The data type of the resulting constant.
+
+    Note
+    ----
+    When dtype is None, we use the following rule:
+
+    - int maps to "int32"
+    - float maps to "float32"
+    - bool maps to "bool"
+    - other using the same default rule as numpy.
+    """
+    # Needed for bf16 and fp8 support (does not come with numpy)
+    import ml_dtypes  # pylint: disable=unused-import,import-outside-toplevel
+
+    if isinstance(dtype, tvm.ir.PrimType):
+        dtype = dtype.dtype
+
+    if isinstance(value, Number | (bool | list)):
+        value = _np.array(value, dtype=dtype)
+
+    if not dtype:
+        # when dtype is None: int maps to "int32", float maps to "float32"
+        dtype = {  # type: ignore
+            _np.dtype("int64"): _np.int32,  # type: ignore
+            _np.dtype("float64"): _np.float32,  # type: ignore
+        }.get(
+            value.dtype,
+            None,  # type: ignore
+        )
+
+    if isinstance(value, _np.ndarray | _np.generic):
+        if dtype is not None:
+            value = value.astype(dtype)
+        value = tvm.runtime.tensor(value)
+
+    if not isinstance(value, tvm.runtime.Tensor):
+        raise ValueError("value has to be scalar or Tensor")
+
+    return Constant(value)
+
+
+@tvm_ffi.register_object("relax.TEPlaceholderOp")
+class TEPlaceholderOp(tvm.te.tensor.Operation):
+    """The placeholder op that represents a relax expression."""
+
+
+def te_tensor(
+    value: Expr, tir_var_map: dict[tvm.tirx.Var, tvm.tirx.Expr], name: str = "rxplaceholder"
+):
+    """Create a TE tensor from relax expression, with TIR variables in the
+    tensor shape substituted by the given mapping
+
+    Parameters
+    ----------
+    value : Expr
+        The relax expression, which is required to have TensorType.
+
+    tir_var_map : Dict[tvm.tirx.Var, tvm.tirx.Expr]
+        The mapping to substitute the TIR variables appeared in the
+        shape of the input Expr.
+
+    name : str
+        The name of the created tensor.
+    """
+    return _ffi_api.TETensor(value, tir_var_map, name)  # type: ignore
+
+
+def get_shape_of(expr: Expr) -> Expr:
+    """Get shape of expr.
+
+    Parameters
+    ----------
+    expr: Expr
+        The input expr.
+
+    Returns
+    -------
+    shape: Expr
+        The shape expression
+
+    Note
+    ----
+    This function requires expr to be normalized.
+    The function will report an error if expr's Type is not TensorType.
+    It will try to return symbolic function when possible. If the tensor do not
+    have a compile-time symbolic shape, the function will then choose to return
+    `Call(relax.op.shape_of, [expr])`.
+    """
+    return _ffi_api.GetShapeOf(expr)  # type: ignore
+
+
+def _update_type(expr: Expr, ty: Type | None) -> None:
+    _ffi_api.UpdateType(expr, ty)  # type: ignore

python/tvm/relax/frontend/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+"""Frontends for constructing Relax programs, with the model importers"""
+
+from . import nn
+from .common import detach_params

python/tvm/relax/frontend/common.py

@@ -0,0 +1,127 @@
+# 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
+"""Commons for Relax frontend."""
+
+import numpy as _np
+
+import tvm
+from tvm import topi
+
+
+def detach_params(mod: tvm.IRModule) -> tuple[tvm.IRModule, dict[str, list[tvm.runtime.Tensor]]]:
+    """Detach the attribute "params" in the functions of the input IRModule as
+    separate dictionary of params.
+
+    Parameters
+    ----------
+    mod : tvm.IRModule
+        The IRModule whose functions' "param" attribute is going to be detached.
+
+    Returns
+    -------
+    detached_mod : tvm.IRModule
+        The IRModule after the detachment.
+
+    params_dict : Dict[str, List[tvm.runtime.Tensor]]
+        The detached params. The dict keys corresponds to the names of the
+        functions in the input IRModule that have attribute "params".
+    """
+    detached_mod = tvm.IRModule()
+    params_dict = dict()
+    for gv, func in mod.functions_items():
+        if "params" in func.attrs:
+            params = list(func.attrs["params"])
+            if not all([isinstance(param, tvm.runtime.Tensor) for param in params]):
+                raise ValueError('The value "params" attribute is expected to be a list of Tensor.')
+            params_dict[gv.name_hint] = params
+            detached_mod[gv] = func.without_attr("params")
+        else:
+            detached_mod[gv] = func
+    return detached_mod, params_dict
+
+
+def autopad(
+    bb,
+    data,
+    strides,
+    kernel_shape,
+    dilations=(1, 1),
+    pad_type="constant",
+    deconv=False,
+    mode="SAME_UPPER",
+    pad_value=0.0,
+):
+    """
+    Perform autopadding with dynamic input shapes
+    """
+    # get attributes as constants
+    strides = _np.array(strides)
+    dilated_kernel_shape = _np.array(
+        [(kernel - 1) * dilation + 1 for kernel, dilation in zip(kernel_shape, dilations)]
+    )
+    # get input shape
+    ndim = data.ty.ndim
+    data_shape = list(data.ty.shape)
+    shape = data_shape[2:ndim]
+
+    # set up integer constants
+    zero = 0
+    one = 1
+    two = 2
+
+    # Calculate total padding
+    mod = shape % strides
+
+    left = _np.maximum(dilated_kernel_shape - strides, zero)
+    right = _np.maximum(dilated_kernel_shape - mod, zero)
+
+    total_pad = _np.where(_np.equal(mod, zero), left, right)
+    if deconv:
+        total_pad = _np.array(kernel_shape) - one - total_pad
+
+    # split total padding into before and after
+    pad_before = _np.floor_divide(total_pad, two)
+    pad_after = total_pad - pad_before
+
+    # combine
+    if "LOWER" in mode:
+        pad = _np.concatenate(
+            [_np.reshape(pad_after, [-1, 1]), _np.reshape(pad_before, [-1, 1])], axis=1
+        )
+    else:
+        pad = _np.concatenate(
+            [_np.reshape(pad_before, [-1, 1]), _np.reshape(pad_after, [-1, 1])], axis=1
+        )
+
+    # pad N and C with zeros
+    pad = _np.concatenate([_np.zeros([2, 2], dtype="int64"), pad], axis=0)
+
+    if pad_type not in ["constant", "edge", "reflect"]:
+        raise tvm.error.OpAttributeInvalid(
+            "Value " + pad_type + ' in attribute "mode" is invalid for operator Pad.'
+        )
+
+    if pad_type == "constant":
+        return bb.emit_te(topi.nn.pad, data, pad[:, 0].tolist(), pad[:, 1].tolist(), pad_value)
+    elif pad_type == "reflect":
+        return bb.emit_te(
+            topi.nn.mirror_pad, data, pad[:, 0].tolist(), pad[:, 1].tolist(), "REFLECT"
+        )
+    else:
+        # edge mode - replicate border values
+        return bb.emit_te(topi.nn.replicate_pad, data, pad[:, 0].tolist(), pad[:, 1].tolist())

python/tvm/relax/frontend/nn/__init__.py

@@ -0,0 +1,53 @@
+# 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.
+"""A PyTorch-like API to build IRModules."""
+
+# pylint: disable=redefined-builtin
+from . import op, spec
+from .core import (
+    Effect,
+    Module,
+    ModuleDict,
+    ModuleList,
+    Object,
+    Parameter,
+    ParameterDict,
+    ParameterList,
+    Tensor,
+)
+from .exporter import add_extern
+from .extern import ExternModule, ObjectModule, SourceModule
+from .modules import (
+    GELU,
+    Conv1D,
+    Conv2D,
+    Conv3D,
+    ConvTranspose1D,
+    Embedding,
+    GroupNorm,
+    IOEffect,
+    KVCache,
+    LayerNorm,
+    Linear,
+    ReLU,
+    RMSNorm,
+    SiLU,
+)
+from .op import *
+from .subroutine import SubroutineMixin
+from .visitor import Mutator

python/tvm/relax/frontend/nn/_tensor_op.py

@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F821
+"""Adding member operators to nn.Tensor."""
+
+from tvm import tirx
+
+
+def _op():
+    from tvm.relax.frontend.nn import op  # pylint: disable=import-outside-toplevel
+
+    return op
+
+
+def _convert_scalar(scalar, ref) -> "Tensor":
+    from .core import Tensor  # pylint: disable=import-outside-toplevel
+
+    if isinstance(scalar, Tensor):
+        return scalar
+    if isinstance(scalar, tirx.FloatImm | tirx.IntImm):
+        return Tensor.from_scalar(scalar.value, dtype=ref.dtype)
+    if isinstance(scalar, int | float):
+        return Tensor.from_scalar(scalar, dtype=ref.dtype)
+    return scalar
+
+
+class _TensorOp:
+    def __add__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().add(self, other)
+
+    def __radd__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().add(self, other)
+
+    def __sub__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().subtract(self, other)
+
+    def __rsub__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().subtract(other, self)
+
+    def __mul__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().multiply(self, other)
+
+    def __rmul__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().multiply(self, other)
+
+    def __truediv__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().divide(self, other)
+
+    def __lt__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().less(self, other)
+
+    def __le__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().less_equal(self, other)
+
+    def __gt__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().greater(self, other)
+
+    def __ge__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().greater_equal(self, other)
+
+    def astype(self, dtype):
+        return _op().astype(self, dtype)
+
+    def maximum(self, other):
+        other = _convert_scalar(other, self)
+        return _op().maximum(self, other)
+
+    def minimum(self, other):
+        other = _convert_scalar(other, self)
+        return _op().minimum(self, other)
+
+    def reshape(self, *shape):
+        return _op().reshape(self, shape)
+
+    def permute_dims(self, *axes):
+        return _op().permute_dims(self, axes)
+
+    def repeat(self, repeats: int, axis: int | None = None):
+        return _op().repeat(self, repeats, axis)

python/tvm/relax/frontend/nn/core.py

@@ -0,0 +1,879 @@
+# 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 core infra for nn.Module, which includes the following pieces:
+- Tensor, a wrapper on top of relax.Expr whose ty is a TensorType,
+  providing more convenient access shape and dtype information.
+  Tensor is always symbolic and not bound to any concrete values.
+- Parameter, a special tensor which could be bound or not bound to concrete values.
+- Module, a container of nn.Parameters and sub nn.Modules.
+- Effect, a non-user-facing class that encloses potential side effects, for example, IO,
+  impure external function callings, inplace mutation, etc.
+"""
+
+from collections import OrderedDict
+from collections.abc import Callable, Iterator, Sequence
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Union,
+)
+
+import numpy as np  # type: ignore
+
+import tvm.runtime
+from tvm import tirx
+from tvm.ir import IRModule
+from tvm.ir.transform import Pass
+from tvm.runtime import Device
+from tvm.runtime import device as as_device
+from tvm.runtime.vm import VirtualMachine
+from tvm.target import Target
+
+from .... import relax as rx
+from ...block_builder import BlockBuilder
+from ...type import (
+    AnyType,
+    ShapeType,
+    TensorType,
+    TupleType,
+)
+from ._tensor_op import _TensorOp
+from .subroutine import SubroutineMixin
+
+if TYPE_CHECKING:
+    import torch  # type: ignore
+
+    from . import spec as _spec
+    from .extern import ExternModule
+
+
+_DEFAULT_DTYPE = "float32"
+
+
+def get_default_dtype() -> str:
+    """Get the default parameter dtype if not specified. By default it is float32.
+
+    Returns
+    -------
+    dtype : str
+        The default dtype
+    """
+    return _DEFAULT_DTYPE
+
+
+def set_default_dtype(dtype: str) -> None:
+    """Set the default parameter dtype.
+
+    Parameters
+    ----------
+    dtype : str
+        The default dtype to be set
+    """
+    global _DEFAULT_DTYPE  # pylint: disable=global-statement
+    _DEFAULT_DTYPE = dtype
+
+
+class Tensor(_TensorOp):
+    """A wrapper on top of relax.Expr whose ty is a TensorType, providing more
+    convenient access shape and dtype information. Tensor is always symbolc and not bound to any
+    concrete values. Shape and dtype inference is done eagerly upon tensor creation, i.e. when
+    operators are applied on tensors, the shape and dtype information is already available.
+    """
+
+    _expr: rx.Expr
+
+    def __init__(self, *, _expr: rx.Expr) -> None:
+        """Private constructor. Tensor is never supposed to be constructed directly by users."""
+
+        def _check_tensor(expr: rx.Expr) -> None:
+            assert expr.ty is not None
+            assert isinstance(expr.ty, TensorType)
+            assert expr.ty.ndim != -1
+            assert expr.ty.shape is not None
+            assert expr.ty.shape.ty is not None
+            assert isinstance(expr.ty.shape.ty, ShapeType)
+            assert expr.ty.shape.ty.values is not None
+
+        _check_tensor(_expr)
+        self._expr = _expr
+
+    @staticmethod
+    def from_const(data) -> "Tensor":
+        """Construct a tensor from numpy constants."""
+        return Tensor(_expr=rx.const(data))
+
+    @staticmethod
+    def from_scalar(data: int | float, dtype: str) -> "Tensor":
+        """Construct a tensor from a scalar with dtype specified."""
+        return Tensor(_expr=rx.const(data, dtype=dtype))
+
+    @staticmethod
+    def from_ty(ty: rx.TensorType, name: str = "tensor") -> "Tensor":
+        """Construct a nn.Tensor from a Relax TensorType.
+
+        TensorType is the Relax type-level description of a tensor, carrying its shape
+        and dtype without holding actual data. This factory creates an unbound placeholder
+        ``nn.Tensor`` that can be used as a symbolic input when tracing an ``nn.Module``.
+
+        Parameters
+        ----------
+        ty : rx.TensorType
+            The type describing the tensor's shape and dtype.
+
+        name : str
+            Name hint for the underlying Relax variable.
+
+        Returns
+        -------
+        tensor : Tensor
+            A symbolic ``nn.Tensor`` backed by a ``relax.Var`` with the given type.
+        """
+        return Tensor(
+            _expr=rx.Var(
+                name_hint=name,
+                ty=ty,
+            )
+        )
+
+    @staticmethod
+    def placeholder(
+        shape: Sequence[int | str | tirx.Expr],
+        dtype: str,
+        name: str = "tensor",
+    ) -> "Tensor":
+        """Create a placeholder tensor with given shape and dtype. A placeholder tensor should
+        never be created directly by users in usual cases, and the only exception is to indicate
+        the shape/dtype of return values of an external function.
+
+        If shape is a string `name`, we create a symbolic shape `tvm.tirx.Var(name, "int64")`.
+        """
+        new_shape = []
+        for expr in shape:
+            if isinstance(expr, int | tirx.IntImm):
+                expr = int(expr)
+                assert expr >= 0
+                new_shape.append(expr)
+                continue
+            if isinstance(expr, str):
+                expr = tirx.Var(expr, "int64")
+                new_shape.append(expr)
+                continue
+            if not tvm.ir.is_prim_expr(expr):
+                raise TypeError(f"Invalid shape: {shape}")
+            assert expr.ty == tvm.ir.PrimType("int64")
+            new_shape.append(expr)
+        return Tensor(
+            _expr=rx.Var(
+                name_hint=name,
+                ty=TensorType(
+                    shape=new_shape,  # type: ignore[arg-type]
+                    dtype=dtype,
+                ),
+            )
+        )
+
+    @property
+    def shape(self) -> list[int | tirx.Expr]:
+        """Returns the shape of the tensor as a list of integers.
+
+        An integer can be a python int or tvm.tirx.Expr, depending on whether the shape is
+        fully static, for example, [1, 2, tvm.tirx.Var("n")] is a valid shape where the last
+        dimension is dynamic while the first two dimensions are always static constants.
+
+        Returns
+        -------
+        shape : List[Union[int, tirx.Expr]]
+            The shape of the tensor
+        """
+
+        def _simplify(expr: tirx.Expr):
+            return expr.value if isinstance(expr, tirx.IntImm) else expr
+
+        shape_ty: ShapeType = self._expr.ty.shape.ty
+        return [_simplify(x) for x in shape_ty.values]
+
+    @property
+    def ndim(self) -> int:
+        """Returns the number of dimensions of the tensor.
+
+        Returns
+        -------
+        ndim : int
+            The number of dimensions of the tensor
+        """
+        return self._expr.ty.ndim
+
+    @property
+    def dtype(self) -> str:
+        """Returns the data type of the tensor.
+
+        Returns
+        -------
+        dtype : str
+            The data type of the tensor
+        """
+        return self._expr.ty.dtype
+
+    def __repr__(self) -> str:
+        return f'Tensor({self.shape}, "{self.dtype}")'
+
+
+class Parameter(Tensor):
+    """A parameter represents the weight of a neural network layer. It is a special tensor which
+    could be bound or not bound to concrete values. If a parameter is bound to a concrete value,
+    it is called a bound parameter, otherwise it is called an unbound parameter.
+    """
+
+    _data: Tensor | None
+    attrs: dict[str, Any]
+
+    def __init__(
+        self,
+        shape: Sequence[int | str | tirx.Expr],
+        dtype: str | None = None,
+    ) -> None:
+        """Create a parameter with given shape and dtype. The parameter is not bound to any
+        concrete values.
+
+        Parameters
+        ----------
+        shape : Sequence[Union[int, str, tirx.Expr]]
+            The shape of the parameter. If it is a string `name`, we create a symbolic shape
+            `tvm.tirx.Var(name, "int64")`.
+        dtype : Optional[str]
+            The data type of the parameter. If not specified, the default dtype will be used.
+        """
+        if dtype is None:
+            dtype = get_default_dtype()
+        super().__init__(_expr=Tensor.placeholder(shape, dtype=dtype, name="param")._expr)
+        self._data = None
+        self.attrs = OrderedDict()
+
+    @property
+    def data(self) -> Tensor | None:
+        """Returns the concrete value of the parameter if it is bound to a concrete value,
+        otherwise returns None. The returned value is a tvm.runtime.Tensor."""
+        return self._data
+
+    @data.setter
+    def data(self, data: Union[None, tvm.runtime.Tensor, np.ndarray, "torch.Tensor"]) -> None:
+        """Set the concrete value of the parameter. The data should be one of the following:
+        - None: unbind the parameter to concrete values
+        - tvm.runtime.Tensor
+        - numpy.ndarray
+        - torch.Tensor and any other DLPack-compliant tensors
+        """
+        if data is None:
+            self._data = data
+            return
+        # Try to do zero-copy if possible
+        if isinstance(data, tvm.runtime.Tensor):
+            pass
+        elif isinstance(data, np.ndarray):
+            data = tvm.runtime.tensor(data)
+        elif hasattr(data, "__dlpack__"):
+            data = _from_dlpack(data)
+        else:
+            raise TypeError(f"Unsupported data type: {type(data)}")
+        if data.shape != tuple(self.shape):
+            raise ValueError(f"Shape mismatch: expected {tuple(self.shape)}, got {data.shape}")
+        if data.dtype != self.dtype:
+            raise ValueError(f"Dtype mismatch: expected {self.dtype}, got {data.dtype}")
+        self._data = data
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        """Change the dtype of the parameter if it is not bound to any concrete data"""
+        if dtype is not None:
+            if self._data is not None:
+                raise ValueError(
+                    "Changing the dtype of a Parameter that has been bound to concrete "
+                    "data is not recommended. It might lead to potential precision loss "
+                    "or other unexpected behaviors"
+                )
+            self._expr = Tensor.placeholder(  # pylint: disable=protected-access
+                self.shape, dtype=dtype, name="param"
+            )._expr
+
+
+class Object:
+    """A wrapper on top of relax.Expr whose ty is the base
+    AnyType, rather than a more specific subtype. Object effectively
+    represents non-tensor frontend components such as KV caches.
+    """
+
+    _expr: rx.Var
+
+    def __init__(self, *, _expr: rx.Expr, _name: str) -> None:
+        """Private constructor. Object is never supposed to be constructed directly by users."""
+        if not isinstance(_expr, rx.Var):
+            _expr = BlockBuilder.current().emit(_expr, _name)
+        self._expr = _expr
+        assert isinstance(self._expr.ty, AnyType)
+
+
+class Effect:
+    """Effect is a special non-user facing type that is used to represent operations with side
+    effects, for example, print. It is used to represent the output of a computation.
+    """
+
+    def emit_init(self, name_hint: str, builder: BlockBuilder) -> list[rx.DataflowVar]:
+        """Emit the initialization of the effect. This method is called by the compiler to
+        initialize the effect."""
+        raise NotImplementedError
+
+    def create(self, name_hint: str) -> list[rx.Var]:
+        """Create the implicit inputs to a relax.Function that represents the side effect"""
+        raise NotImplementedError
+
+    def set_state(self, state_vars: list[rx.Var]) -> None:
+        """Set the variables that represents the effect"""
+        raise NotImplementedError
+
+    def finalize(self) -> list[rx.Var]:
+        """finalize the effect as the implicit return value of a relax.Function"""
+        raise NotImplementedError
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        """Convert the effect to specific dtype. Usually it is no-op for most of the effects"""
+
+
+class Module(SubroutineMixin):
+    """Base class for neural network components. Subclass it to build your models.
+    Modules can nest within each other in a tree structure using regular attribute assignment."""
+
+    def named_parameters(self, prefix: str = "") -> Iterator[tuple[str, Parameter]]:
+        """This method provides an iterator over module parameters,
+        yielding both the parameter name and its corresponding value.
+
+        Parameters
+        ----------
+        prefix : str
+            Prefix to prepend to all parameter names.
+
+        Yields
+        ------
+        (str, Parameter) - Tuple containing the name and parameter
+        """
+        yield from _attribute_finder(
+            self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
+        )
+
+    def parameters(self) -> Iterator[Parameter]:
+        """This method provides an iterator over module parameters,
+        yielding only the Parameter value.
+
+        Yields
+        ------
+        Parameter - The module's parameter
+        """
+        for _, param in self.named_parameters():
+            yield param
+
+    def state_dict(
+        self, *, prefix: str = "", destination: dict[str, Parameter] | None = None
+    ) -> dict[str, Parameter]:
+        """Returns a dictionary containing references to the whole state of the module.
+
+        Parameters
+        ----------
+        prefix : str
+            Prefix to prepend to all parameter names.
+        destination : Optional[Dict[str, Parameter]]
+            Dictionary to which state will be saved. If None, a new dictionary is created.
+
+        Returns
+        -------
+        dict : Dict[str, Parameter]
+            a dictionary containing a whole state of the module
+        """
+        if destination is None:
+            destination = OrderedDict()
+        for name, param in _attribute_finder(
+            self, prefix, condition_yield=lambda x: isinstance(x, Parameter)
+        ):
+            destination[name] = param
+        return destination
+
+    def load_state_dict(
+        self, state_dict: dict[str, Parameter], strict: bool = True
+    ) -> tuple[list[str], list[str]]:
+        """This function copies parameters and buffers from the state_dict into the current module
+        and its descendants. If `strict` is set to True, the keys in the `state_dict` must exactly
+        match the keys returned by the `state_dict()` function of this module.
+
+        Parameters
+        ----------
+        state_dict : Dict[str, Parameter]
+            A dictionary containing a whole state of the module
+        strict : bool = True
+            Whether to strictly enforce that the keys in `state_dict` match the keys returned by
+            this module's `state_dict()` function.
+
+        Returns
+        -------
+        (missing_keys, unexpected_keys) : Tuple[List[str], List[str]]
+            A tuple of two lists: the missing keys and the unexpected keys.
+        """
+        self_state_dict = self.state_dict()
+        missing_keys: list[str] = []
+        unexpected_keys: list[str] = []
+        for key, value in state_dict.items():
+            if key not in self_state_dict:
+                unexpected_keys.append(key)
+                continue
+            if value.data is None:
+                raise ValueError(f"Parameter {key} is not set to any concrete tensor")
+            self_state_dict.pop(key).data = value.data
+        missing_keys = list(self_state_dict.keys())
+        if strict and (missing_keys or unexpected_keys):
+            raise KeyError(f"Missing keys: {missing_keys}, Unexpected keys: {unexpected_keys}")
+        return missing_keys, unexpected_keys
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Call the module with the given inputs and returns the output."""
+        if not hasattr(self, "forward"):
+            raise NotImplementedError(f"Module {type(self)} does not have a `forward` method")
+        return self.forward(*args, **kwargs)  # pylint: disable=no-member
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        """Convert the module to specific dtype recursively"""
+        for _, item in self.__dict__.items():
+            if hasattr(item, "to") and callable(item.to):
+                item.to(dtype=dtype)
+        if dtype is not None and isinstance(getattr(self, "dtype", None), str):
+            self.dtype = dtype  # pylint: disable=attribute-defined-outside-init
+
+    def export_tvm(
+        self,
+        spec: "_spec.ModuleSpecType",
+        debug: bool = False,
+        allow_extern: bool = False,
+    ) -> (
+        tuple[IRModule, list[tuple[str, Parameter]]]
+        | tuple[IRModule, list[tuple[str, Parameter]], list["ExternModule"]]
+    ):
+        """Export the module to TVM IRModule and parameters
+
+        Parameters
+        ----------
+        spec : _spec.ModuleSpecType
+            A dictionary mapping each input name to a specification
+            that defines the inputs shape and dtype.
+        debug : bool
+            If set to True, then the exported module will support
+            effects. This enables things like printing in the graph.
+
+        Returns
+        -------
+        irmodule : tvm.ir.IRModule
+            The converted tvm IR representation of the model.
+        params : List[Tuple[str, Parameter]]
+            A list of Parameters corresponding to the weights of the model.
+        ext_mods : List[nn.ExternModule]
+            A list of ExternModules that are used in the model.
+        """
+        # pylint: disable=import-outside-toplevel
+        from . import spec as _spec
+        from .exporter import Exporter
+
+        # pylint: enable=import-outside-toplevel
+        spec = _spec.ModuleSpec.from_raw(spec, self)
+        mod, params, ext_mods = Exporter(debug=debug).build(spec)
+        if allow_extern:
+            return mod, params, ext_mods
+        if ext_mods:
+            raise ValueError(
+                "`ExternModule`(s) exist when they are not allowed. "
+                "Turn on flag `allow_extern` to allow."
+            )
+        return mod, params
+
+    def jit(  # pylint: disable=too-many-arguments
+        self,
+        spec: "_spec.ModuleSpec",
+        device: str | Device = "cpu",
+        pipeline: None | str | Pass = "default_build",
+        out_format: str = "torch",
+        debug: bool = False,
+    ) -> Any:
+        """Just-in-time compile an ``nn.Module`` into a callable executable.
+
+        The method exports the module to a Relax IRModule, applies the given compilation
+        pipeline, builds a Relax VM executable, and wraps the result so it can be called
+        directly (e.g. with PyTorch tensors when ``out_format="torch"``).
+
+        Parameters
+        ----------
+        spec : _spec.ModuleSpec
+            A specification mapping each module input to its shape and dtype.
+
+        device : Union[str, Device]
+            The device to compile and run on (e.g. ``"cpu"``, ``"cuda"``).
+
+        pipeline : Union[None, str, Pass]
+            The Relax compilation pipeline to apply. ``"default_build"`` uses the standard
+            optimization pipeline; ``None`` skips pipeline passes.
+
+        out_format : str
+            Output wrapper format. ``"torch"`` returns a ``TorchModule`` whose ``forward``
+            accepts and returns PyTorch tensors.
+
+        debug : bool
+            If ``True``, enable effect-based debugging (e.g. printing) in the compiled graph.
+
+        Returns
+        -------
+        module : Any
+            A callable wrapper (type depends on *out_format*) around the compiled VM.
+        """
+
+        def _compile(spec, device, pipeline, debug):
+            # pylint: disable=import-outside-toplevel
+            from ...transform import AttachExternModules
+            from ...vm_build import build as relax_build
+            from . import spec as _spec
+            from .exporter import Exporter
+
+            # pylint: enable=import-outside-toplevel
+
+            spec = _spec.ModuleSpec.from_raw(spec, self)
+            mod, params, ext_mods = Exporter(debug=debug).build(spec)
+            mod = AttachExternModules(ext_mods)(mod)  # pylint: disable=not-callable
+            vm = VirtualMachine(  # pylint: disable=invalid-name
+                relax_build(
+                    mod,
+                    target=Target.from_device(device),
+                    relax_pipeline=pipeline,
+                ),
+                device,
+            )
+            params = _param_to_tensor(params, device)
+            return spec, vm, params
+
+        device = as_device(device)
+        spec, vm, params = _compile(spec, device, pipeline, debug)  # pylint: disable=invalid-name
+
+        if out_format == "torch":
+            from . import torch  # pylint: disable=import-outside-toplevel
+
+            return torch.TorchModule(spec=spec, params=params, vm=vm)
+
+        raise ValueError(f"Unknown out_format: {out_format}")
+
+
+class ModuleDict(Module):
+    """Holds submodules in a dict."""
+
+    def __init__(self, modules: OrderedDict[str, Module] | None = None):
+        if modules is None:
+            self.modules = OrderedDict()
+        else:
+            self.modules = OrderedDict(modules)
+
+    def __iter__(self):
+        return iter(self.modules.values())
+
+    def __getitem__(self, key: str) -> Module:
+        return self.modules[key]
+
+    def __setitem__(self, key: str, module: Module) -> None:
+        self.modules[key] = module
+
+    def __len__(self) -> int:
+        return len(self.modules)
+
+    def keys(self) -> Iterator[str]:
+        return self.modules.keys()
+
+    def values(self) -> Iterator[Module]:
+        return self.modules.values()
+
+    def items(self) -> Iterator[tuple[str, Module]]:
+        return self.modules.items()
+
+    def get(self, key: str, default: Module | None = None) -> Module | None:
+        return self.modules.get(key, default)
+
+    def update(self, modules: dict[str, Module]) -> None:
+        self.modules.update(modules)
+
+    def clear(self) -> None:
+        self.modules.clear()
+
+    def pop(self, key: str) -> Module:
+        return self.modules.pop(key)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.modules
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for module in self.modules.values():
+            module.to(dtype=dtype)
+
+
+class ParameterDict(Module):
+    """Holds parameters in a dict."""
+
+    def __init__(
+        self,
+        params: OrderedDict[str, Parameter] | dict[str, Parameter] | None = None,
+    ):
+        self.params: OrderedDict[str, Parameter] = OrderedDict()
+        if params is not None:
+            self.update(params)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self.params)
+
+    def __getitem__(self, key: str) -> Parameter:
+        return self.params[key]
+
+    def __setitem__(self, key: str, param: Parameter) -> None:
+        if not isinstance(key, str):
+            raise TypeError(f"ParameterDict keys must be strings, but got {type(key).__name__}")
+        if not isinstance(param, Parameter):
+            raise TypeError(
+                f"ParameterDict values must be nn.Parameter, but got {type(param).__name__}"
+            )
+        self.params[key] = param
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def keys(self) -> Iterator[str]:
+        return self.params.keys()
+
+    def values(self) -> Iterator[Parameter]:
+        return self.params.values()
+
+    def items(self) -> Iterator[tuple[str, Parameter]]:
+        return self.params.items()
+
+    def get(self, key: str, default: Parameter | None = None) -> Parameter | None:
+        return self.params.get(key, default)
+
+    def update(self, params: dict[str, Parameter]) -> None:
+        for key, param in params.items():
+            self[key] = param
+
+    def clear(self) -> None:
+        self.params.clear()
+
+    def pop(self, key: str) -> Parameter:
+        return self.params.pop(key)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.params
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for param in self.params.values():
+            param.to(dtype=dtype)
+
+
+class ModuleList(Module):
+    """Holds submodules in a list."""
+
+    def __init__(self, modules: list[Module]):
+        self.modules = modules
+
+    def __iter__(self):
+        return iter(self.modules)
+
+    def __getitem__(self, idx: int) -> Module:
+        return self.modules[idx]
+
+    def __setitem__(self, idx: int, module: Module) -> None:
+        self.modules[idx] = module
+
+    def __len__(self):
+        return len(self.modules)
+
+    def append(self, module: Module):
+        """Add a module to the end of the ModuleList"""
+        self.modules.append(module)
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for module in self.modules:
+            module.to(dtype=dtype)
+
+    def forward(self, x):  # pylint: disable=invalid-name
+        """Feed-forward pass of the module"""
+        for module in self.modules:
+            x = module(x)
+        return x
+
+
+class ParameterList(Module):
+    """Holds parameters in a list."""
+
+    def __init__(self, params: list[Parameter] | None = None):
+        self.params: list[Parameter] = []
+        if params is not None:
+            self.extend(params)
+
+    def __iter__(self) -> Iterator[Parameter]:
+        return iter(self.params)
+
+    def __getitem__(self, idx: int) -> Parameter:
+        return self.params[idx]
+
+    def __setitem__(self, idx: int, param: Parameter) -> None:
+        if not isinstance(param, Parameter):
+            raise TypeError(
+                f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}"
+            )
+        self.params[idx] = param
+
+    def __len__(self) -> int:
+        return len(self.params)
+
+    def append(self, param: Parameter) -> None:
+        """Add a parameter to the end of the ParameterList"""
+        if not isinstance(param, Parameter):
+            raise TypeError(
+                f"ParameterList elements must be nn.Parameter, but got {type(param).__name__}"
+            )
+        self.params.append(param)
+
+    def extend(self, params: list[Parameter]) -> None:
+        """Add parameters to the end of the ParameterList"""
+        for param in params:
+            self.append(param)
+
+    def to(self, dtype: str | None = None) -> None:  # pylint: disable=invalid-name
+        for param in self.params:
+            param.to(dtype=dtype)
+
+
+def wrap_nested(expr: rx.Expr, name: str) -> Tensor | Sequence[Tensor]:
+    """Wrap the given relax.Expr, emit it using the current BlockBuilder,
+    and automatically handle nested cases if the expr represents a Tuple.
+
+    Parameters
+    ----------
+    expr : relax.Expr
+        The Expr to be wrapped.
+
+    name : str
+        Name hint.
+
+    Returns
+    -------
+    result : Union[Tensor, Tuple[Tensor]]
+        The computed result.
+    """
+    if not isinstance(expr, rx.DataflowVar):
+        expr = BlockBuilder.current().emit(expr, name)
+    if isinstance(expr.ty, TensorType):
+        return Tensor(_expr=expr)
+    if isinstance(expr.ty, TupleType):
+        return tuple(
+            wrap_nested(  # type: ignore
+                rx.TupleGetItem(expr, i),
+                name=f"{name}.{i}",
+            )
+            for i in range(len(expr.ty.fields))
+        )
+    raise TypeError(f"Unsupported return type: {expr.ty}")
+
+
+def _attribute_finder(root: Module, prefix: str, condition_yield: Callable[[Any], bool]):
+    """Find attributes that satisfy the condition recursively"""
+    if isinstance(root, ParameterList):
+        for i, param in enumerate(root):
+            if condition_yield(param):
+                yield prefix + f"{i}", param
+        return
+    elif isinstance(root, ParameterDict):
+        for name, param in root.items():
+            if condition_yield(param):
+                yield prefix + name, param
+        return
+    elif isinstance(root, ModuleList):
+        for i, subitem in enumerate(root):
+            yield from _attribute_finder(subitem, prefix + f"{i}.", condition_yield)
+        return
+    elif isinstance(root, ModuleDict):
+        for name, subitem in root.items():
+            yield from _attribute_finder(subitem, prefix + f"{name}.", condition_yield)
+        return
+    for name, item in root.__dict__.items():
+        if condition_yield(item):
+            yield prefix + name, item
+        elif isinstance(item, ParameterList):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ParameterDict):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ModuleList):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+        elif isinstance(item, ModuleDict):
+            for sub_name, sub_item in item.items():
+                yield from _attribute_finder(
+                    sub_item,
+                    prefix + name + f".{sub_name}.",
+                    condition_yield,
+                )
+        elif isinstance(item, Module):
+            yield from _attribute_finder(
+                item,
+                prefix + name + ".",
+                condition_yield,
+            )
+
+
+def _from_dlpack(tensor) -> tvm.runtime.Tensor:
+    try:
+        return tvm.runtime.from_dlpack(tensor)
+    except RuntimeError:
+        pass
+    # special logic for PyTorch
+    device_type = tensor.device.type
+    device_id = tensor.device.index or 0
+    return tvm.runtime.tensor(
+        tensor.numpy(),
+        device=Device(
+            Device._DEVICE_NAME_TO_TYPE[device_type],
+            device_id,
+        ),
+    )
+
+
+def _param_to_tensor(
+    params: list[tuple[str, Parameter]], device: Device
+) -> list[tvm.runtime.Tensor]:
+    results = []
+    missing = []
+    for name, param in params:
+        if param.data is None:
+            missing.append(name)
+        else:
+            results.append(param.data.copyto(target=device))
+    if missing:
+        raise ValueError(f"Parameters are not set to any concrete values: {', '.join(missing)}")
+    return results

python/tvm/relax/frontend/nn/exporter.py

@@ -0,0 +1,334 @@
+# 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.
+"""Export `nn.Module` to TVM's IRModule."""
+
+import functools
+import operator
+import threading
+import typing
+
+from tvm import tirx
+from tvm.ir import IRModule
+
+from .... import relax as rx
+from ...block_builder import BlockBuilder
+from ...type import AnyType, ShapeType, TupleType
+from . import core, extern
+from . import spec as _spec
+from .modules import IOEffect
+
+
+def add_extern(mod: extern.ExternModule) -> None:
+    """Add an external module to the exporter."""
+    try:
+        exporter = Exporter.current()
+    except Exception as exception:
+        raise RuntimeError(
+            "`nn.add_extern` should only be invoked when exporting a module."
+        ) from exception
+    exporter.add_external_module(mod)
+
+
+class Exporter:
+    """Builder of ModuleSpec, which exports an nn.Module to TVM IRModule."""
+
+    _tls = threading.local()
+
+    builder: BlockBuilder
+    io_effect: core.Effect
+    extern_mods: list[extern.ExternModule]
+
+    def __init__(self, debug: bool) -> None:
+        self.builder = BlockBuilder()
+        self.io_effect = IOEffect() if debug else None
+        self.extern_mods = []
+
+    @staticmethod
+    def current() -> "Exporter":
+        """Get the current Exporter under the with scope."""
+        assert hasattr(Exporter._tls, "current")
+        return Exporter._tls.current
+
+    def __enter__(self) -> "Exporter":
+        assert not hasattr(Exporter._tls, "current")
+        Exporter._tls.current = self
+        return self
+
+    def __exit__(self, exc_type, exc, traceback) -> None:
+        assert hasattr(Exporter._tls, "current")
+        delattr(Exporter._tls, "current")
+
+    def add_external_module(self, mod: extern.ExternModule) -> None:
+        """Add an external module to the exporter."""
+        # pylint: disable=protected-access
+        all_symbols: list[str] = []
+        for extern_mod in self.extern_mods:
+            all_symbols.extend(extern_mod._symbols.keys())
+        duplicated_symbols = list(set(mod._symbols.keys()) & set(all_symbols))
+        # pylint: enable=protected-access
+        if duplicated_symbols:
+            raise ValueError(f"Duplicate symbols: {duplicated_symbols}")
+        self.extern_mods.append(mod)
+
+    def build(  # pylint: disable=too-many-locals
+        self,
+        spec: _spec.ModuleSpec,
+    ) -> tuple[
+        IRModule,
+        list[tuple[str, core.Parameter]],
+        list[extern.ExternModule],
+    ]:
+        """Build the ModuleSpec to TVM IRModule. Returns the IRModule and the parameters."""
+
+        # pylint: disable=protected-access
+        def _params() -> list[tuple[str, core.Parameter]]:
+            params = []
+            for name, param in core._attribute_finder(
+                spec.module, prefix="", condition_yield=lambda x: isinstance(x, core.Parameter)
+            ):
+                params.append((name, param))
+            return params
+
+        def _effects() -> list[tuple[str, core.Effect]]:
+            result = []
+            if self.io_effect is not None:
+                result.append(("", self.io_effect))
+            for name, effect in core._attribute_finder(
+                spec.module, "", condition_yield=lambda x: isinstance(x, core.Effect)
+            ):
+                result.append((name, effect))
+            return result
+
+        # pylint: enable=protected-access
+        params = None
+        effects = _effects()
+        ext_mods = self.extern_mods
+        with self:
+            if effects:
+                with self.builder.function("_initialize_effect"):
+                    with self.builder.dataflow():
+                        outputs = _emit_effect_init(self.builder, effects)
+                    self.builder.emit_func_output(outputs, params=[])
+            for method_name, method_spec in zip(spec.method_names, spec.method_specs):
+                params = _params()  # Re-initialize so symbolic shapes not shared across methods
+                len_args = len(method_spec.arg_specs)
+                len_effects = {
+                    "packed": 1,
+                    "none": 0,
+                    "plain": len(effects),
+                }[method_spec.effect_mode]
+                with self.builder.function(
+                    method_name,
+                    attrs={"num_input": len_args + len_effects},  # type: ignore
+                ):
+                    with self.builder.dataflow():
+                        outputs, inputs = _emit_method(self.builder, method_spec, params, effects)
+                    self.builder.emit_func_output(outputs, inputs)
+        mod = self.builder.finalize()
+        rx.analysis.well_formed(mod)
+
+        return mod, params, ext_mods
+
+
+def _emit_effect_init(
+    builder: BlockBuilder,
+    effects: list[tuple[str, core.Effect]],
+):
+    outputs = []
+    for prefix, effect in effects:
+        inits = effect.emit_init(prefix, builder)
+        assert isinstance(inits, list)
+        outputs.extend(inits)
+    outputs = builder.emit_output(builder.emit(rx.Tuple(outputs)))
+    return outputs
+
+
+def _emit_method(  # pylint: disable=too-many-locals,too-many-branches,too-many-statements
+    builder: BlockBuilder,
+    spec: _spec.MethodSpec,
+    params: list[tuple[str, core.Parameter]],
+    effects: list[tuple[str, core.Effect]] | None,
+):
+    # pylint: disable=protected-access
+    # symbolic shape's name mapping to its tirx.Var for reuse
+    str2var_params: dict[str, tirx.Var] = {}
+
+    def _unwrap_ret(expr: typing.Any) -> typing.Any:
+        if isinstance(expr, core.Tensor | core.Object):
+            return expr._expr
+        if isinstance(expr, tuple):
+            return rx.Tuple([_unwrap_ret(x) for x in expr])
+        if isinstance(expr, list):
+            return rx.Tuple([_unwrap_ret(x) for x in expr])
+        raise TypeError(f"Unsupported return type: {type(expr)}")
+
+    def _convert_input(arg):
+        if isinstance(arg, tirx.Var):
+            return rx.Var(arg.name, ty=ShapeType(values=[arg]))
+        if isinstance(arg, core.Tensor | core.Object):
+            return arg._expr  # pylint: disable=protected-access
+        if isinstance(arg, _spec.Tuple):
+            return rx.Var(
+                arg.name,
+                ty=TupleType([_convert_input(arg_i).ty for arg_i in arg.elements]),
+            )
+        raise TypeError(f"Unsupported input type: {type(arg)}")
+
+    def _params(mode: str) -> list[rx.Var]:
+        inputs: list[rx.Var] = []
+
+        def _get_var(shape_var: tirx.Var) -> tirx.Var:
+            name = shape_var.name
+            if name in str2var_params:
+                return str2var_params[name]
+            var = tirx.Var(name, "int64")
+            str2var_params[name] = var
+            return var
+
+        for name, param in params:
+            # Make sure the a symbolic shape is not re-registered (same as _method_spec_to_inputs)
+            # e.g. we do not see `vocab_size` for `lm_head` and `vocab_size_1` for `embed_tokens`
+            new_shape = [_get_var(x) if isinstance(x, tirx.Var) else x for x in param.shape]
+            var = core.Tensor.placeholder(new_shape, param.dtype, name)._expr
+            inputs.append(var)
+            param._expr = var
+        if mode == "none":
+            return []
+        if mode == "plain":
+            return inputs
+        if mode == "packed":
+            input_var = rx.Var(
+                "packed_params",
+                TupleType(fields=[x.ty for x in inputs]),
+            )
+            for i, (name, param) in enumerate(params):
+                param._expr = builder.emit(rx.TupleGetItem(input_var, i), name_hint=name)
+            return [input_var]
+        raise ValueError(f"Invalid param_mode: {mode}")
+
+    def _effects(mode: str) -> list[rx.Var]:
+        unflat_inputs: list[list[rx.Var]] = []
+        for name, effect in effects:
+            effect_input = effect.create(name)
+            effect.set_state(effect_input)
+            unflat_inputs.append(effect_input)
+        inputs: list[rx.Var] = functools.reduce(operator.iadd, unflat_inputs, [])
+        if mode == "none":
+            return []
+        if mode == "plain":
+            return inputs
+        if mode == "packed":
+            input_var = rx.Var(
+                "packed_effects",
+                TupleType(fields=[x.ty for x in inputs]),
+            )
+            i = 0
+            for effect_input, (_, effect) in zip(unflat_inputs, effects):
+                updated_effect_input = []
+                for effect_input_i in effect_input:
+                    updated_effect_input.append(
+                        builder.emit(
+                            rx.TupleGetItem(input_var, i),
+                            name_hint=effect_input_i.name_hint,
+                        )
+                    )
+                    i += 1
+                effect.set_state(updated_effect_input)
+            return [input_var]
+
+        raise ValueError(f"Invalid effect_mode: {mode}")
+
+    # pylint: enable=protected-access
+
+    def _detuple(arg, var: rx.Var, builder: BlockBuilder):
+        if isinstance(arg, _spec.Tuple):
+            ret = []
+            for i, elem in enumerate(arg.elements):
+                field = builder.emit(rx.TupleGetItem(var, i), name_hint=f"{arg.name}_{i}")
+                ret.append(_detuple(elem, field, builder))
+            return type(arg.elements)(ret)
+        if isinstance(arg, core.Tensor):
+            return core.Tensor(_expr=var)
+        if isinstance(arg, tirx.Var):
+            return arg
+        raise TypeError(f"Unsupported input type: {type(arg)}")
+
+    # TODO(@junrushao): Warn if params/effects are used when their mode is "none"
+    explicit_inputs = _method_spec_to_inputs(spec)
+    inputs = [_convert_input(x) for x in explicit_inputs]
+    inputs = inputs + _effects(spec.effect_mode)
+    inputs = inputs + _params(spec.param_mode)
+
+    for arg_idx, (arg, var) in enumerate(zip(explicit_inputs, inputs)):
+        if isinstance(arg, _spec.Tuple):
+            explicit_inputs[arg_idx] = _detuple(arg, var, builder)
+
+    outputs = spec.method(*explicit_inputs)
+    effect_outputs = []
+    for _, effect in effects:
+        effect_outputs.extend(effect.finalize())
+    if effect_outputs and spec.effect_mode != "none":
+        outputs = builder.emit_output(rx.Tuple([_unwrap_ret(outputs), rx.Tuple(effect_outputs)]))
+    else:
+        outputs = builder.emit_output(_unwrap_ret(outputs))
+    return outputs, inputs
+
+
+def _method_spec_to_inputs(
+    spec: _spec.MethodSpec,
+) -> list[tirx.Var | core.Tensor]:
+    """Convert the MethodSpec to a list of inputs to Module's method."""
+    str2var: dict[str, tirx.Var] = {}
+
+    def _get_var(name: str) -> tirx.Var:
+        if name in str2var:
+            return str2var[name]
+        var = tirx.Var(name, "int64")
+        str2var[name] = var
+        return var
+
+    def _convert_input(arg_name, arg_spec):
+        if isinstance(arg_spec, _spec.Int):
+            arg = _get_var(arg_name)
+        elif isinstance(arg_spec, _spec.Tensor):
+            arg = core.Tensor.placeholder(  # pylint: disable=protected-access
+                shape=[_get_var(x) if isinstance(x, str) else x for x in arg_spec.shape],
+                dtype=arg_spec.dtype,
+                name=arg_name,
+            )
+        elif isinstance(arg_spec, _spec.Object):
+            arg = arg_spec.object_type(_expr=rx.Var(arg_name, AnyType()), _name=arg_name)
+        elif isinstance(arg_spec, _spec.Tuple):
+            elements = type(arg_spec.elements)(
+                [
+                    _convert_input(arg_name=arg_name + f"_{i}", arg_spec=arg_spec.elements[i])
+                    for i in range(len(arg_spec.elements))
+                ]
+            )
+            arg = _spec.Tuple(
+                name=arg_name,
+                elements=elements,
+            )
+        else:
+            raise TypeError(f"Invalid spec for argument {arg_name}: {arg_spec}")
+        return arg
+
+    args = []
+    for arg_name, arg_spec in zip(spec.arg_names, spec.arg_specs):
+        arg = _convert_input(arg_name=arg_name, arg_spec=arg_spec)
+        args.append(arg)
+    return args

python/tvm/relax/frontend/nn/extern.py

@@ -0,0 +1,402 @@
+# 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.
+# ruff: noqa: E722
+"""External modules to be linked into the exported IRModule."""
+
+import os
+import shutil
+import sys
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+
+import tvm_ffi
+
+import tvm
+from tvm import libinfo, tirx
+from tvm.runtime import Module, load_static_library
+from tvm.support import cc as _cc
+
+from ...op import call_dps_packed
+from . import core
+from .core import wrap_nested
+
+
+class ExternModule:
+    """The abstract base class for external modules. External modules are designed to help
+    incorporate user-provided handcrafted kernels into the exported TVM IRModule.
+    """
+
+    _symbols: dict[str, Callable]
+
+    def __init__(self, symbols: dict[str, Callable]) -> None:
+        self._symbols = symbols
+
+    def __getitem__(self, func_name: str) -> Callable:
+        _inference_function = self._symbols[func_name]
+
+        def _call(*input_args):
+            def _convert(arg, name: str):
+                from tvm import relax as rx  # pylint: disable=import-outside-toplevel
+
+                if isinstance(arg, core.Tensor):
+                    return arg._expr  # pylint: disable=protected-access
+                if isinstance(arg, int):
+                    return rx.prim_value(tirx.IntImm("int64", arg))
+                if isinstance(arg, float):
+                    return rx.prim_value(tirx.FloatImm("float64", arg))
+                if isinstance(arg, str):
+                    return rx.StringImm(arg)
+                if tvm.ir.is_prim_expr(arg):
+                    return rx.prim_value(arg)
+                if isinstance(arg, tuple | list):
+                    return rx.Tuple([_convert(e, f"{name}_{i}") for i, e in enumerate(arg)])
+                raise TypeError(f"Unsupported input type: {type(arg)}")
+
+            rx_inputs = _convert(input_args, "input")
+            rx_outputs_ty = _convert(_inference_function(*input_args), "dummy").ty
+            return wrap_nested(call_dps_packed(func_name, rx_inputs, rx_outputs_ty), func_name)
+
+        return _call
+
+    def _load(self, path: Path) -> Module:
+        return load_static_library(str(path), func_names=list(self._symbols.keys()))
+
+    def load(self) -> Module:
+        """Loads the external module into a TVM runtime module."""
+        raise NotImplementedError
+
+
+class ObjectModule(ExternModule):  # pylint: disable=too-few-public-methods
+    """A subclass of `nn.ExternModule`, which allows
+    users to provide an object `.o` file to be linked into compiled
+    artifact;
+    """
+
+    def __init__(
+        self,
+        symbols: dict[str, Callable],
+        filepath: Path,
+    ) -> None:
+        if not isinstance(filepath, Path):
+            filepath = Path(filepath)
+        if not filepath.is_file():
+            raise ValueError(f"Not a file: {filepath!s}")
+        self.filepath = filepath
+        super().__init__(symbols)
+
+    def load(self) -> Module:
+        return self._load(self.filepath)
+
+
+class SourceModule(ExternModule):  # pylint: disable=too-few-public-methods
+    """A subclass of `nn.ExternModule`. It compiles C++/CUDA source code and link them into the
+    eventual IRModule.
+
+    **Shape/dtype inference.** The `nn.ExternModule` system requires users to provide additional
+    information to work, namely, `symbols`. It is a dictionary that maps each symbol in the
+    external object file to its shape/dtype inference function. Consider a case where function
+    `my_func` accepts two tensors, `a` of shape `(x, y, 1)`, and `b` of shape `(y, z, 5)`, and
+    produces a tensor `c` of shape `(x, y, z, 9)`, the shape/dtype inference function should look
+    like:
+
+    .. code-block:: python
+
+        def shape_dtype_inference(a, b):
+            x, y, _ = a.shape
+            _, z, _ = b.shape
+            return nn.Tensor.placeholder((x, y, z, 9), dtype="float32")
+
+
+    and the `symbols` dictionary should be provided as:
+
+    .. code-block:: python
+
+        symbols={
+            "my_func": shape_dtype_inference,
+        }
+
+
+    **Calling convention.** All external modules now follows "destination-passing-style" (DPS)
+    calling convention, which means the returned tensors are pre-allocated by the system already
+    and passed in as an argument of the external function.
+
+    Reuse the example above, the implementation of `my_func` should include three parameters in
+    its signature, where tensors are represented using DLTensor from DLPack, the de facto standard
+    of in-memory representation of tensors. More details:
+    https://github.com/dmlc/dlpack/blob/v0.8/include/dlpack/dlpack.h#L163-L206.
+
+    To expose the symbol, `TVM_FFI_DLL_EXPORT_TYPED_FUNC(symbol, function)` is guaranteed available:
+
+    .. code-block:: C++
+
+        // those headers are guaranteed to be available
+        #include <dlpack/dlpack.h>
+        #include <tvm/ffi/dtype.h>
+        #include <tvm/ffi/function.h>
+
+        namespace {
+        // anonymous namespace hides the symbol `_my_func_impl` from other translation units
+        int _my_func_impl(DLTensor* a, DLTensor* b, DLTensor* c) {
+            // `a` and `b` are inputs, and `c` is the output
+        }
+        }
+        // expose symbol `my_func` instead of `_my_func_impl`
+        TVM_FFI_DLL_EXPORT_TYPED_FUNC(my_func, _my_func_impl);
+
+    **A compiler pass `AttachExternModules`.** It is introduced to attach a list of
+    `nn.ExternModule`s into an IRModule at any stage of the compilation pipeline,
+    and attach the compiled external modules as `runtime.Module`s into IRModule's `external_mods`
+    attribute. It is required by linking in `tvm.compile`, but with the existence of this pass,
+    source compilation can be deferred to arbitrary stage of TVM compilation.
+
+    **Caveats.** It is required to call `nn.add_extern` to register external modules exactly once
+    during `export_tvm`. Each symbol should be registered exactly once to avoid potential conflicts,
+    and otherwise an error will be raised.
+    """
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        symbols: dict[str, Callable],
+        source_code: str | Path,
+        source_format: str,  # "cpp", "cu"
+        compile_options: list[str] | None = None,
+        compiler: str | None = None,
+        output_format: str = "obj",  # "obj", "wasm"
+    ):
+        """Constructs a `nn.SourceModule` from source code.
+
+        Parameters
+        ----------
+        symbols : Dict[str, Callable]
+            The dictionary that maps each symbol in the external object file to its shape/dtype
+            inference function.
+
+        source_code : Union[str, Path]
+            Source code or path to the source code to be compiled.
+
+        source_format : str
+            The source code format. It can be either "cpp" or "cu".
+
+        compile_options : Optional[List[str]]
+            The compile options. If not provided, the default compile options will be used.
+
+        compiler : Optional[str]
+            The compiler. If not provided, the default compiler will be used. On Windows,
+            compilation requires `clang` by default.
+
+        output_format : str
+            The output format. It can be either "obj" or "wasm". "obj" is the default format,
+            which is a shared object file. "wasm" is the WebAssembly format, which is a binary
+            file.
+        """
+
+        def _detect_input_suffix(source_format: str) -> str:
+            if source_format == "cpp":
+                return ".cpp"
+            if source_format == "cu":
+                return ".cu"
+            raise ValueError(f"Invalid source format: {source_format}")
+
+        def _detect_output_suffix(output_format: str) -> str:
+            if output_format == "obj":
+                if _cc._is_linux_like():  # pylint: disable=protected-access
+                    return ".o"
+                if _cc._is_windows_like():  # pylint: disable=protected-access
+                    return ".obj"
+                raise ValueError(f"Unsupported platform: {sys.platform}")
+            if output_format == "wasm":
+                return ".wasm"
+            raise ValueError(f"Invalid output format: {output_format}")
+
+        def _detect_source_code(source_code) -> str:
+            if isinstance(source_code, Path):
+                path = source_code
+                if not path.is_file():
+                    raise ValueError(f"Not a file: {path!s}")
+            else:
+                try:
+                    path = Path(source_code)
+                except:  # pylint: disable=bare-except
+                    return source_code
+                try:
+                    if not path.is_file():
+                        return source_code
+                except:  # pylint: disable=bare-except
+                    return source_code
+            with path.open("r", encoding="utf-8") as file:
+                return file.read()
+
+        self.source_code = _detect_source_code(source_code)
+        if compile_options is None:
+            self.compile_options = SourceModule.get_compile_options(source_format=source_format)
+        else:
+            self.compile_options = list(compile_options)
+        self.compiler = compiler
+        self.source_suffix = _detect_input_suffix(source_format)
+        self.output_suffix = _detect_output_suffix(output_format)
+        super().__init__(symbols)
+
+    @staticmethod
+    def tvm_home() -> Path:
+        """Find TVM's home directory. If `TVM_HOME` environment variable is set, use it.
+        Otherwise, use the directory where the `tvm` Python package is installed.
+        As a sanity check, it is required to have `include` and `3rdparty` as direct subdirectories.
+
+        Returns
+        -------
+        tvm_home : pathlib.Path
+            The TVM home directory, and it is guaranteed to have `include` and `3rdparty` as
+            direct subdirectories.
+        """
+        if os.environ.get("TVM_HOME", None):
+            tvm_path = Path(os.environ["TVM_HOME"])
+            assert tvm_path.exists(), (
+                f"Using environment variable `TVM_HOME`, but directory not found: {tvm_path!s}"
+            )
+            assert tvm_path.is_dir(), (
+                f"Using environment variable `TVM_HOME`, but it is not a directory: {tvm_path!s}"
+            )
+        else:
+            import tvm  # pylint: disable=import-outside-toplevel
+
+            tvm_path = Path(tvm.__file__).parent
+            assert tvm_path.is_dir()
+        tvm_path = tvm_path.resolve()
+        while True:
+            exists_include = (tvm_path / "include").is_dir()
+            exists_3rdparty = (tvm_path / "3rdparty").is_dir()
+            if exists_include and exists_3rdparty:
+                return tvm_path.resolve()
+            parent = tvm_path.parent
+            if parent == tvm_path:
+                raise ValueError(
+                    "Cannot detect TVM directory. "
+                    "Please explicitly specify it by setting `TVM_HOME` environment variable, "
+                    "and make sure it contains `include` and `3rdparty` as direct sub-directories."
+                )
+            tvm_path = parent
+        return tvm_path.resolve()
+
+    @staticmethod
+    def get_includes(tvm_pkg: list[str] | None = None) -> list[Path]:
+        """Returns the default include paths according to `tvm_home()`.
+        By default, it includes TVM, DLPack. With `tvm_pkg` provided, it also
+        includes the specified package under `tvm_home/3rdparty`.
+
+        Parameters
+        ----------
+        tvm_pkg : Optional[List[str]]
+            The list of packages to be included under `tvm_home/3rdparty`. Each element should be
+            a relative path to `tvm_home/3rdparty`.
+
+        Returns
+        -------
+        includes : List[pathlib.Path]
+            The list of include paths.
+        """
+        results = [
+            Path(libinfo.find_include_path()),
+            Path(tvm_ffi.libinfo.find_include_path()),
+            Path(tvm_ffi.libinfo.find_dlpack_include_path()),
+        ]
+        if tvm_pkg:
+            tvm_home = SourceModule.tvm_home()
+            for relative in tvm_pkg:
+                results.append(tvm_home / "3rdparty" / relative)
+        results = list(dict.fromkeys(results))
+        for path in results:
+            assert path.exists(), f"Not found: {path!s}"
+            assert path.is_dir(), f"Not a directory: {path!s}"
+        return results
+
+    @staticmethod
+    def get_compile_options(
+        source_format: str,
+        tvm_pkg: list[str] | None = None,
+    ) -> list[str]:
+        """Returns the default compile options depending on `source_format`, including the default
+        inlcude paths w.r.t. `tvm_home()`, and by default,
+        it uses "-O3" and "-std=c++17".
+
+        Parameters
+        ----------
+        source_format : str
+            The source code format. It can be either "cpp" or "cu".
+
+        tvm_pkg : Optional[List[str]]
+            The list of packages to be included under `tvm_home/3rdparty`. Each element should be
+            a relative path to `tvm_home/3rdparty`.
+
+        Returns
+        -------
+        compile_options : List[str]
+            The list of compilation flags.
+        """
+        include_flags = []
+        for include_path in SourceModule.get_includes(tvm_pkg=tvm_pkg):
+            include_flags += ["-I", str(include_path)]
+        if source_format == "cpp":
+            host_flags = [
+                "-c",  # generate object file
+                "-O3",
+                "-std=c++17",
+            ]
+        elif source_format == "cu":
+            host_flags = [
+                "-c",  # generate object file
+                "-O3",
+                "-std=c++17",
+                # Enable `-fPIC` for the host compiler
+                "-Xcompiler=-fPIC",
+            ]
+        else:
+            raise ValueError(f"Invalid source format: {source_format}")
+        return include_flags + host_flags
+
+    def compile(self, output_path: Path) -> None:
+        """Compiles the source code in a provided directory and returns the compiled artifact."""
+        with tempfile.TemporaryDirectory() as temp_dir_str:
+            temp_dir = Path(temp_dir_str)
+            source_filename = f"main{self.source_suffix}"
+            object_filename = f"main{self.output_suffix}"
+            source_path = temp_dir / source_filename
+            object_path = temp_dir / object_filename
+            with source_path.open("w", encoding="utf-8") as file:
+                file.write(self.source_code)
+            _cc.create_shared(
+                output=object_filename,
+                objects=[source_filename],
+                options=self.compile_options,
+                cc=self.compiler,
+                cwd=temp_dir,
+                ccache_env=(
+                    {
+                        "CCACHE_COMPILERCHECK": "content",
+                        "CCACHE_NOHASHDIR": "1",
+                    }
+                    if shutil.which("ccache")
+                    else None
+                ),
+            )
+            shutil.move(str(object_path), str(output_path))
+
+    def load(self) -> Module:
+        with tempfile.TemporaryDirectory() as temp_dir_str:
+            output_path = Path(temp_dir_str) / f"main{self.output_suffix}"
+            self.compile(output_path)
+            return self._load(output_path)

python/tvm/relax/frontend/nn/llm/__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.
+"""LLM support for PyTorch-like API to build IRModules."""
+
+from . import kv_cache, position_embedding
+from .position_embedding import llama_rope
+from .tree_attn import tree_attn
+from .kv_cache import PagedKVCache

python/tvm/relax/frontend/nn/llm/_decode_kernels.py

@@ -0,0 +1,526 @@
+# 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.
+# ruff: noqa: E501
+# fmt: off
+
+"""Single-token decode attention kernels and attention-state merge helpers.
+
+Contents:
+- ``_attention_decode_cpu`` / ``_attention_decode`` — paged-KV decode (one Q token
+  per sequence), CPU scalar and GPU allreduce variants.
+- ``_merge_state_inplace_cpu`` / ``_merge_state_inplace`` — combine two
+  log-sum-exp attention outputs in place. Used by multi-stage decoding and by
+  the distributed KV-transfer path.
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import math
+from typing import Any
+
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from ._kernel_common import (
+    _declare_length_info,
+    _get_kv_chunk_len,
+    _get_seq_offset,
+    _rope,
+    _var,
+    _var_cpu,
+    check_thread_limits,
+    get_max_num_threads_per_block,
+)
+
+
+def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], page_size: int = 16):
+    H_qo = num_qo_heads
+    H_kv = num_kv_heads
+    D = head_dim
+    group_size = num_qo_heads // num_kv_heads
+
+    global_symbol = "batch_decode_paged_kv_cpu"
+    if sliding_window:
+        global_symbol += "_sliding_window"
+
+    @T.prim_func(s_tir=True)
+    def batch_decode_paged_kv(
+        Q_handle: T.handle,
+        pages_handle: T.handle,
+        page_table_indptr_handle: T.handle,
+        page_table_values_handle: T.handle,
+        var_length_info: T.handle,  # [b] when sliding window = False, or otherwise [3, b]
+        k_rope_pos_offset_handle: T.handle,
+        q_rope_position_handle: T.handle,
+        output_handle: T.handle,
+        lse_handle: T.handle,
+        rotary_mode: T.int32,
+        rope_scale: T.float32,
+        rope_theta: T.float32,
+        sm_scale: T.float32,
+    ):
+        T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
+        B = T.int32()
+        nnz_pages = T.int32()
+        max_num_pages = T.int32()
+        page_indptr_elem_offset = T.int32()
+        page_values_elem_offset = T.int32()
+        k_rope_pos_offset_elem_offset = T.int32()
+        q_rope_position_elem_offset = T.int32()
+        length_info_elem_offset = T.int32()
+
+        Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
+        pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype)
+        page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
+        page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
+        k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
+        q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
+        output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
+        lse = T.match_buffer(lse_handle, (B, H_qo), "float32")  # pylint: disable=unused-variable
+        # The length information of the sequences.
+        # - It is in shape `(3, batch_size)` when sliding window is enabled.
+        #   For a sequence "i", location
+        #   - "(0, i)" is the number of KV slots used in the last page of the seq ("last_page_len"),
+        #   - "(1, i)" is the starting offset of the sliding window in the seq,
+        #   - "(2, i)" is the attn sink length of the sequence.
+        # - It is in shape `(batch_size,)` when sliding window is disabled,
+        #   denoting the "last_page_len".
+        length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
+
+        for b in T.serial(B):
+            with T.sblock("attn"):
+                O_local = T.sblock_alloc_buffer((D,), "float32")
+                Q_local = T.sblock_alloc_buffer((D,), "float32")
+                K_local = T.sblock_alloc_buffer((D,), "float32")
+                V_local = T.sblock_alloc_buffer((D,), "float32")
+
+                kv_chunk_len = T.sblock_alloc_buffer((1,), "int32")
+
+                m_val = T.sblock_alloc_buffer((1,), "float32")
+                new_m = T.sblock_alloc_buffer((1,), "float32")
+                d_val = T.sblock_alloc_buffer((1,), "float32")
+                S_val = T.sblock_alloc_buffer((1,), "float32")
+                scale_O = T.sblock_alloc_buffer((1,), "float32")
+                factor = T.sblock_alloc_buffer((1,), "float32")
+
+                cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[b]
+                cur_page_indptr_end: T.let[T.int32] = page_table_indptr[b + 1]
+
+                kv_chunk_len[0] = T.if_then_else(
+                    cur_page_indptr_begin != cur_page_indptr_end,
+                    _get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b, length_info, sliding_window),
+                    0,
+                )
+
+                for h_qo in T.serial(H_qo):
+                    m_val[0] = -5e4
+                    d_val[0] = 1.0
+
+                    for d in T.serial(D):
+                        O_local[d] = 0.0
+
+                    for d in T.serial(D):
+                        Q_local[d] = T.if_then_else(
+                            rotary_mode == 1,
+                            _rope(Q, q_rope_position[b], head_dim, rope_theta, rope_scale, (b, h_qo, d), qkv_dtype, rope_scaling),
+                            Q[b, h_qo, d],
+                        )
+
+                    for row_idx in T.serial(kv_chunk_len[0]):
+                        seq_offset: T.let[T.int32()] = _get_seq_offset(row_idx, b, length_info, sliding_window)
+                        page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + (seq_offset // page_size)]
+                        page_offset: T.let[T.int32()] = seq_offset % page_size
+
+                        for d in T.serial(D):
+                            K_local[d] = T.if_then_else(
+                                rotary_mode == 1,
+                                _rope(pages, k_rope_pos_offset[b] + row_idx, head_dim, rope_theta, rope_scale, (page_no, 0, h_qo // group_size, page_offset, d), qkv_dtype, rope_scaling),
+                                pages[page_no, 0, h_qo // group_size, page_offset, d],
+                            )
+                        S_val[0] = 0.0
+                        for d in T.serial(D):
+                            S_val[0] += Q_local[d] * K_local[d]
+                        S_val[0] *= sm_scale * math.log2(math.exp(1))
+
+                        new_m[0] = T.max(m_val[0], S_val[0])
+                        d_val[0] = (d_val[0] * T.exp2(m_val[0] - new_m[0])) + T.exp2(S_val[0] - new_m[0])
+
+                        scale_O[0] = T.exp2(m_val[0] - new_m[0])
+
+                        for d in T.serial(D):
+                            O_local[d] = O_local[d] * scale_O[0]
+
+                        m_val[0] = new_m[0]
+                        for d in T.serial(D):
+                            V_local[d] = pages[page_no, 1, h_qo // group_size, page_offset, d]
+
+                        factor[0] = T.exp2(S_val[0] - m_val[0])
+                        for d in T.serial(D):
+                            O_local[d] = O_local[d] + V_local[d] * factor[0]
+                    for d in T.serial(D):
+                        O_local[d] = O_local[d] / d_val[0]
+                        output[b, h_qo, d] = O_local[d]
+                    lse[b, h_qo] = m_val[0] + T.log2(d_val[0])
+
+    return batch_decode_paged_kv
+
+
+def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_window: bool, rope_scaling: dict[str, Any], target: Target, page_size: int = 16):
+    qkv_dtype_bytes = 2
+    H_qo = num_qo_heads
+    H_kv = num_kv_heads
+    D = head_dim
+
+    THREAD_LIMIT = 512
+    TILE_SIZE_PER_BDX = 2
+    if target.kind.name == "opencl" and (("android" in str(target.host)) or ("adreno" in str(target.attrs))):
+        # Keeping lower thread limit for this kernel on adreno target
+        # to avoid register spill
+        THREAD_LIMIT = 256
+        TILE_SIZE_PER_BDX = 1
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    thread_limit = min(max_num_threads_per_block, THREAD_LIMIT)
+
+    GROUP_SIZE = H_qo // H_kv
+    VEC_SIZE = min(max(8 // qkv_dtype_bytes, D // 32), 4)
+    bdx = D // VEC_SIZE
+    bdy = GROUP_SIZE
+    while bdx * bdy > thread_limit and bdy > 1:
+        bdy //= 2
+    gdz = GROUP_SIZE // bdy
+    threads_per_CTA = max(thread_limit, bdx * bdy)
+    bdz = threads_per_CTA // (bdx * bdy)
+    tile_size_per_bdx = TILE_SIZE_PER_BDX if GROUP_SIZE == 1 else 1
+    check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=bdz, gdz=1)
+
+    global_symbol = "batch_decode_paged_kv"
+    if sliding_window:
+        global_symbol += "_sliding_window"
+
+    # pylint: disable=too-many-branches
+    @T.prim_func(s_tir=True)
+    def batch_decode_paged_kv(
+        Q_handle: T.handle,
+        pages_handle: T.handle,
+        page_table_indptr_handle: T.handle,
+        page_table_values_handle: T.handle,
+        var_length_info: T.handle, # [b] when sliding window = False, or otherwise [3, b]
+        k_rope_pos_offset_handle: T.handle,
+        q_rope_position_handle: T.handle,
+        output_handle: T.handle,
+        lse_handle: T.handle,
+        rotary_mode: T.int32,
+        rope_scale: T.float32,
+        rope_theta: T.float32,
+        sm_scale: T.float32,
+    ):
+        T.func_attr({"tirx.is_scheduled": True, "global_symbol": global_symbol})
+        B = T.int32()
+        nnz_pages = T.int32()
+        max_num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        page_indptr_elem_offset = T.int32()
+        page_values_elem_offset = T.int32()
+        k_rope_pos_offset_elem_offset = T.int32()
+        q_rope_position_elem_offset = T.int32()
+        length_info_elem_offset = T.int32()
+
+        Q = T.match_buffer(Q_handle, (B, H_qo, D), qkv_dtype)
+        pages = T.match_buffer(pages_handle, (max_num_pages, 2, H_kv, page_size, D), qkv_dtype, elem_offset=pages_elem_offset)
+        page_table_indptr = T.match_buffer(page_table_indptr_handle, (B + 1,), "int32", elem_offset=page_indptr_elem_offset)
+        page_table_values = T.match_buffer(page_table_values_handle, (nnz_pages,), "int32", elem_offset=page_values_elem_offset)
+        k_rope_pos_offset = T.match_buffer(k_rope_pos_offset_handle, (B,), "int32", elem_offset=k_rope_pos_offset_elem_offset)
+        q_rope_position = T.match_buffer(q_rope_position_handle, (B,), "int32", elem_offset=q_rope_position_elem_offset)
+        output = T.match_buffer(output_handle, (B, H_qo, D), qkv_dtype)
+        lse = T.match_buffer(lse_handle, (B, H_qo), "float32")  # pylint: disable=unused-variable
+        length_info = _declare_length_info(var_length_info, B, sliding_window, length_info_elem_offset)
+
+        for bx in T.thread_binding(B, thread="blockIdx.x"):
+            for fused_by_bz in T.thread_binding(H_kv * gdz, thread="blockIdx.y"):
+                for ty in T.thread_binding(bdy, thread="threadIdx.y"):
+                    for tx in T.thread_binding(bdx, thread="threadIdx.x"):
+                        for tz in T.thread_binding(bdz, thread="threadIdx.z"):
+                            with T.sblock("attn"):
+                                Q_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
+                                kv_chunk_len = T.sblock_alloc_buffer((1,), "int32", scope="local")
+                                K_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
+                                V_smem = T.sblock_alloc_buffer((bdz * bdy * tile_size_per_bdx, D), qkv_dtype, scope="shared")
+                                O_allreduce = T.sblock_alloc_buffer((bdz, bdy, D), "float32", scope="shared")
+                                md_allreduce = T.sblock_alloc_buffer((bdz, bdy, 2), "float32", scope="shared")
+                                S_reduce_local = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                t0 = T.sblock_alloc_buffer((1,), "float32", scope="local")
+
+                                S_local = T.sblock_alloc_buffer((bdy * tile_size_per_bdx), "float32", scope="local")
+                                QK_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+                                V_local = T.sblock_alloc_buffer((VEC_SIZE,), qkv_dtype, scope="local")
+                                m_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                d_prev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                exp_mprev = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                exp_otherm = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                other_o = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+                                st_m = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                st_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
+                                O_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
+
+                                by: T.let[T.int32] = fused_by_bz % H_kv
+                                bz: T.let[T.int32] = fused_by_bz // H_kv
+                                batch_idx: T.let[T.int32] = bx
+                                cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[batch_idx]
+                                cur_page_indptr_end: T.let[T.int32] = page_table_indptr[batch_idx + 1]
+                                kv_chunk_len[0] = T.if_then_else(
+                                    cur_page_indptr_begin != cur_page_indptr_end,
+                                    _get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, batch_idx, length_info, sliding_window),
+                                    0
+                                )
+
+                                # init states
+                                st_m[0] = -5e4
+                                st_d[0] = 1.0
+                                for vec in T.vectorized(VEC_SIZE):
+                                    O_local[vec] = 0.0
+
+                                # load q
+                                for vec in T.vectorized(VEC_SIZE):
+                                    Q_local[vec] = T.if_then_else(
+                                        rotary_mode == 1,
+                                        _rope(Q, q_rope_position[batch_idx], head_dim, rope_theta, rope_scale, (bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
+                                        Q[bx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec]
+                                    )
+
+                                for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_size_per_bdx * bdy * bdz)):
+                                    tile_start_s: T.let[T.int32()] = (tz * bdy + ty) * tile_size_per_bdx  # type: ignore
+                                    tile_start_g: T.let[T.int32()] = ((iterator * bdz + tz) * bdy + ty) * tile_size_per_bdx  # type: ignore
+                                    # load KV from global memory to shared memory
+                                    for j in T.serial(tile_size_per_bdx):
+                                        with T.sblock("KV_load"):
+                                            T.reads()
+                                            T.writes()
+                                            row_g: T.let[T.int32()] = tile_start_g + j  # type: ignore
+                                            if row_g < kv_chunk_len[0]:
+                                                seq_offset: T.let[T.int32()] = _get_seq_offset(row_g, batch_idx, length_info, sliding_window)  # type: ignore
+                                                page_no: T.let[T.int32()] = page_table_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)]  # type: ignore
+                                                page_offset: T.let[T.int32()] = T.floormod(seq_offset, page_size)  # type: ignore
+                                                for vec in T.vectorized(VEC_SIZE):
+                                                    K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = T.if_then_else(
+                                                        rotary_mode == 1,
+                                                        _rope(pages, k_rope_pos_offset[batch_idx] + row_g, head_dim, rope_theta, rope_scale, (page_no, 0, by, page_offset, tx * VEC_SIZE + vec), qkv_dtype, rope_scaling),
+                                                        pages[page_no, 0, by, page_offset, tx * VEC_SIZE + vec]
+                                                    )
+                                                    V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = pages[page_no, 1, by, page_offset, tx * VEC_SIZE + vec]
+                                            else:
+                                                for vec in T.vectorized(VEC_SIZE):
+                                                    K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
+                                                    V_smem[tile_start_s + j, tx * VEC_SIZE + vec] = 0.0
+                                    T.tvm_storage_sync("shared")
+                                    # compute QK
+                                    m_prev[0] = st_m[0]
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        # compute S = Q * K * sm_scale
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            QK_local[vec] = T.cast(Q_local[vec], "float32") * T.cast(K_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec], "float32") * sm_scale * math.log2(math.exp(1))
+                                        S_reduce_local[0] = 0
+                                        for vec in T.unroll(VEC_SIZE):
+                                            S_reduce_local[0] += QK_local[vec]
+
+                                        with T.sblock("block_cross_thread"):
+                                            T.reads(S_reduce_local[0])
+                                            T.writes(t0[0])
+                                            T.attr(
+                                                T.comm_reducer(lambda x0, y0: x0 + y0, [T.float32(0)]),
+                                                "reduce_scope",
+                                                T.int32(0),
+                                            )
+                                            T.tvm_thread_allreduce(T.uint32(1), S_reduce_local[0], True, t0[0], tx, dtype="void")
+
+                                        S_local[j] = -5e4
+                                        if (iterator * bdz + tz) * bdy * tile_size_per_bdx + j < kv_chunk_len[0]:
+                                            S_local[j] = t0[0]
+                                        # update st_m
+                                        st_m[0] = T.max(st_m[0], S_local[j])
+
+                                    # update st_d, st_O
+                                    o_scale: T.let[T.float32] = T.exp2(m_prev[0] - st_m[0])
+                                    st_d[0] *= o_scale
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        S_local[j] = T.exp2(S_local[j] - st_m[0])
+                                        st_d[0] += S_local[j]
+                                    for j in T.vectorized(VEC_SIZE):
+                                        O_local[j] *= o_scale
+
+                                    # load V from shared memory to local memory
+                                    # compute O
+                                    for j in T.serial(bdy * tile_size_per_bdx):
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            V_local[vec] = V_smem[tz * bdy * tile_size_per_bdx + j, tx * VEC_SIZE + vec]
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            O_local[vec] += T.cast(V_local[vec], "float32") * S_local[j]
+
+                                if bdz > 1:
+                                    # allreduce over bdz
+                                    for vec in T.vectorized(VEC_SIZE):
+                                        O_allreduce[tz, ty, tx * VEC_SIZE + vec] = O_local[vec]
+                                    md_allreduce[tz, ty, 0] = st_m[0]
+                                    md_allreduce[tz, ty, 1] = st_d[0]
+                                    T.tvm_storage_sync("shared")
+
+                                    st_m[0] = -5e4
+                                    st_d[0] = 1.0
+                                    for vec in T.vectorized(VEC_SIZE):
+                                        O_local[vec] = 0.0
+
+                                    for j in T.serial(bdz):
+                                        m_prev[0] = st_m[0]
+                                        d_prev[0] = st_d[0]
+                                        other_m[0] = md_allreduce[j, ty, 0]
+                                        other_d[0] = md_allreduce[j, ty, 1]
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            other_o[vec] = O_allreduce[j, ty, tx * VEC_SIZE + vec]
+                                        st_m[0] = T.max(st_m[0], other_m[0])
+                                        st_d[0] = d_prev[0] * T.exp2(m_prev[0] - st_m[0]) + other_d[0] * T.exp2(other_m[0] - st_m[0])
+                                        exp_mprev[0] = T.exp2(m_prev[0] - st_m[0])
+                                        exp_otherm[0] = T.exp2(other_m[0] - st_m[0])
+                                        for vec in T.vectorized(VEC_SIZE):
+                                            O_local[vec] = O_local[vec] * exp_mprev[0] + other_o[vec] * exp_otherm[0]
+
+                                # normalize O
+                                for vec in T.vectorized(VEC_SIZE):
+                                    O_local[vec] /= st_d[0]
+
+                                # store O to global memory
+                                for vec in T.vectorized(VEC_SIZE):
+                                    output[batch_idx, by * GROUP_SIZE + bz * bdy + ty, tx * VEC_SIZE + vec] = O_local[vec]
+
+                                # store lse to global memory
+                                lse[batch_idx, by * GROUP_SIZE + bz * bdy + ty] = st_m[0] + T.log2(st_d[0])
+    # pylint: enable=too-many-branches
+    return batch_decode_paged_kv
+
+
+def _merge_state_inplace_cpu(v_dtype):
+    @T.prim_func(s_tir=True)
+    def merge_state_inplace_cpu(
+        v: T.handle,
+        s: T.handle,
+        v_other: T.handle,
+        s_other: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        N = T.int32()
+        H = T.int32()
+        D = T.int32()
+
+        V = T.match_buffer(v, (N, H, D), v_dtype)
+        S = T.match_buffer(s, (N, H), "float32")
+        V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
+        S_other = T.match_buffer(s_other, (N, H), "float32")
+
+        for n in T.serial(N):
+            for h in T.serial(H):
+                with T.sblock("merge"):
+                    s_val = _var_cpu("float32")
+                    s_other_val = _var_cpu("float32")
+                    s_max = _var_cpu("float32")
+                    scale = _var_cpu("float32")
+                    other_scale = _var_cpu("float32")
+
+                    s_val[0] = S[n, h]
+                    s_other_val[0] = S_other[n, h]
+                    s_max[0] = T.max(s_val[0], s_other_val[0])
+                    s_val[0] = T.exp2(s_val[0] - s_max[0])
+                    s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
+                    scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
+                    other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
+                    for d in T.serial(D):
+                        V[n, h, d] = V[n, h, d] * scale[0] + V_other[n, h, d] * other_scale[0]
+                    S[n, h] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
+
+    return merge_state_inplace_cpu
+
+
+def _merge_state_inplace(num_heads, head_dim, v_dtype, target: Target, global_symbol: str | None = None):
+    v_dtype_bytes = 2
+    VEC_SIZE = min(max(8 // v_dtype_bytes, head_dim // 32), 4)
+    bdx = head_dim // VEC_SIZE
+    bdy = num_heads
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    while bdx * bdy > max_num_threads_per_block and bdy > 1:
+        bdy //= 2
+    gdy = num_heads // bdy
+    check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def merge_state_inplace(
+        v: T.handle,
+        s: T.handle,
+        v_other: T.handle,
+        s_other: T.handle,
+    ):
+        T.func_attr({"tirx.is_scheduled": True})
+        N = T.int32()
+        H = T.int32()
+        D = T.int32()
+
+        V = T.match_buffer(v, (N, H, D), v_dtype)
+        S = T.match_buffer(s, (N, H), "float32")
+        V_other = T.match_buffer(v_other, (N, H, D), v_dtype)
+        S_other = T.match_buffer(s_other, (N, H), "float32")
+
+        for bx in T.thread_binding(N, thread="blockIdx.x"):
+            for by in T.thread_binding(gdy, thread="blockIdx.y"):
+                for ty in T.thread_binding(bdy, thread="threadIdx.y"):
+                    for tx in T.thread_binding(bdx, thread="threadIdx.x"):
+                        with T.sblock("merge"):
+                            s_val = _var("float32")
+                            s_other_val = _var("float32")
+                            s_max = _var("float32")
+                            scale = _var("float32")
+                            other_scale = _var("float32")
+
+                            v_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
+                            v_other_vec = T.sblock_alloc_buffer((VEC_SIZE,), v_dtype, scope="local")
+
+                            s_val[0] = S[bx, ty + by * bdy]
+                            s_other_val[0] = S_other[bx, ty + by * bdy]
+                            s_max[0] = T.max(s_val[0], s_other_val[0])
+                            s_val[0] = T.exp2(s_val[0] - s_max[0])
+                            s_other_val[0] = T.exp2(s_other_val[0] - s_max[0])
+                            scale[0] = s_val[0] / (s_val[0] + s_other_val[0])
+                            other_scale[0] = s_other_val[0] / (s_val[0] + s_other_val[0])
+
+                            # load v
+                            for vec in T.vectorized(VEC_SIZE):
+                                v_vec[vec] = V[bx, ty + by * bdy, tx * VEC_SIZE + vec]
+                            # load v_other
+                            for vec in T.vectorized(VEC_SIZE):
+                                v_other_vec[vec] = V_other[bx, ty + by * bdy, tx * VEC_SIZE + vec]
+
+                            # merge
+                            for vec in T.serial(VEC_SIZE):
+                                v_vec[vec] = v_vec[vec] * scale[0] + v_other_vec[vec] * other_scale[0]
+
+                            # store v
+                            for vec in T.vectorized(VEC_SIZE):
+                                V[bx, ty + by * bdy, tx * VEC_SIZE + vec] = v_vec[vec]
+
+                            # store s
+                            S[bx, ty + by * bdy] = T.log2(s_val[0] + s_other_val[0]) + s_max[0]
+
+    func = merge_state_inplace
+    if global_symbol:
+        func = func.with_attr("global_symbol", global_symbol)
+    return func

python/tvm/relax/frontend/nn/llm/_kernel_common.py

@@ -0,0 +1,569 @@
+# 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.
+# ruff: noqa: E501, E731, RUF005
+# fmt: off
+
+"""Shared TIR helpers used by KV-cache / attention kernels in this package.
+
+This module consolidates constructs reused by the prefill/decode/paged/tree
+attention kernels so each kernel file can focus on its own specialised logic.
+
+Contents:
+- Thread-limit checks (``get_max_num_threads_per_block``, ``check_thread_limits``)
+- KV-cache enums (``AttnKind``, ``RopeMode``)
+- Small TVMScript helpers (``_var``, ``_var_cpu``, ``_causal_mask``, ``_rope``)
+- Length-info accessors for sliding-window-aware indexing
+- Buffer allocators for the tiled online-softmax state used by every prefill kernel
+- ``_make_prefill_macros`` — the ``@T.macro`` bundle invoked by the prefill kernels
+- Tiling config (``_get_prefill_kernel_config``) and scheduling (``_schedule_prefill_kernel``)
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import enum
+import math
+from typing import Any
+
+import tvm
+from tvm import s_tir, tirx
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from .position_embedding import switch_rope_freq_func
+
+
+def _var(dtype):
+    return T.sblock_alloc_buffer((1,), dtype, scope="local")
+
+
+def _var_cpu(dtype):
+    return T.sblock_alloc_buffer((1,), dtype)
+
+
+def get_max_num_threads_per_block(target: Target) -> int:
+    """
+    max(max_num_threads, max_threads_per_block); if latter does not exist, return max_num_threads.
+    We add this method since some targets have both fields and `max_threads_per_block` is larger.
+    """
+    max_num_threads = int(target.attrs["max_num_threads"])
+    max_threads_per_block = target.attrs.get("max_threads_per_block", None)
+    if max_threads_per_block is None:
+        return max_num_threads
+    return max(max_num_threads, max_threads_per_block)
+
+
+def check_thread_limits(target: Target, bdx: int, bdy: int, bdz: int, gdz: int):
+    """
+    Check whether max num threads exceeded given a target.
+
+    Parameters
+    ----------
+    bdx: threadIdx.x
+    bdy: threadIdx.y
+    bdz: threadIdx.z
+    gdz: blockIdx.z
+    """
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+
+    assert bdx * bdy * bdz <= max_num_threads_per_block, (
+        f"{target.kind} max num threads exceeded: {bdx}*{bdy}*{bdz}>{max_num_threads_per_block}"
+    )
+
+    if target.kind.name == "webgpu":
+        # https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxcomputeworkgroupsizez
+        assert bdz <= 64, f"webgpu's threadIdx.z cannot exceed 64, but got bdz={bdz}"
+        assert gdz == 1, f"webgpu's blockIdx.z should be 1, but got gdz={gdz}"
+
+
+class AttnKind(enum.IntEnum):
+    """The attention kind class.
+    MHA denotes multi-head attention, multi-query attention or grouped query attention.
+    MLA denotes multi-head latent attention.
+    """
+
+    MHA = 0
+    MLA = 1
+    MHA_SLIDING = 3
+
+
+class RopeMode(enum.IntEnum):
+    """The RoPE mode of the Paged KV cache.
+    If it is none, the KV cache will not apply RoPE to q and k.
+    If it is normal, RoPE will be applied to k before adding k to cache.
+    Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
+    """
+
+    NONE = 0
+    NORMAL = 1
+    INLINE = 2
+
+
+def _rope(buffer: T.Buffer, offset: tirx.Var, rotary_dim: int, theta: tirx.Var, scale: tirx.Var, indices: tuple[tirx.Var, ...], qkv_dtype: str, rope_scaling: dict[str, Any]):
+    d = indices[-1]
+    cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(offset * scale, d, rotary_dim, theta, "float32")
+    cos = cos_freq * buffer[indices].astype("float32")
+    sin = sin_freq * tirx.if_then_else(
+        d < rotary_dim // 2,
+        -buffer[indices[:-1] + (d + rotary_dim // 2,)],
+        buffer[indices[:-1] + (d - rotary_dim // 2,)],
+    ).astype("float32")
+    expr = (cos + sin).astype(qkv_dtype)
+    for var, value in var_map.items():
+        expr = tirx.Let(var, value, expr)
+    return expr
+
+
+def _causal_mask(causal, row, col, kv_len, qo_len):
+    return T.if_then_else(
+        causal > 0,
+        col < kv_len - qo_len + row + 1,
+        col < kv_len,
+    )
+
+
+def _declare_length_info(var_length_info, batch_size, sliding_window, elem_offset):
+    return (
+        T.match_buffer(var_length_info, (3, batch_size), "int32", elem_offset=elem_offset)
+        if sliding_window
+        else T.match_buffer(var_length_info, (batch_size,), "int32", elem_offset=elem_offset)
+    )
+
+
+def _get_kv_chunk_len(num_pages, page_size, seq_id, length_info, sliding_window):
+    if not sliding_window:
+        return (num_pages - 1) * page_size + length_info[seq_id]
+    # ((num_pages - 1) * page_size + last_page_len) - sliding_window_offset + sink_size
+    return (num_pages - 1) * page_size + length_info[0, seq_id] - length_info[1, seq_id] + length_info[2, seq_id]
+
+
+def _get_seq_offset(pos, seq_id, length_info, sliding_window):
+    if not sliding_window:
+        return pos
+    # pos if pos < sink_size else pos - sink_size + sliding_window_offset
+    return T.if_then_else(
+        pos < length_info[2, seq_id],
+        pos,
+        pos - length_info[2, seq_id] + length_info[1, seq_id],
+    )
+
+
+def _alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps):
+    """Allocate the shared/local online-softmax working state used by every tiled prefill kernel.
+
+    Returns ``(S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new)``.
+    """
+    S_smem = T.sblock_alloc_buffer((tile_x, tile_z), "float32", scope="shared")
+    S_local = T.sblock_alloc_buffer((tile_x, tile_z), "float32", scope="local")
+    m_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    m_prev_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    d_smem = T.sblock_alloc_buffer((tile_x,), "float32", scope="shared")
+    md_shape = (math.ceil(tile_x / (bdx * num_warps)),)
+    m_new = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    m_prev = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    d_new = T.sblock_alloc_buffer(md_shape, "float32", scope="local")
+    return S_smem, S_local, m_smem, m_prev_smem, d_smem, m_new, m_prev, d_new
+
+
+def _alloc_mha_qkvo_buffers(tile_x, tile_z, d_qk, d_v, dtype):
+    """Allocate Q/K/V shared + O local buffers for standard MHA/GQA prefill kernels."""
+    Q_smem = T.sblock_alloc_buffer((tile_x, d_qk), dtype, scope="shared")
+    K_smem = T.sblock_alloc_buffer((tile_z, d_qk), dtype, scope="shared")
+    V_smem = T.sblock_alloc_buffer((tile_z, d_v), dtype, scope="shared")
+    O_local = T.sblock_alloc_buffer((tile_x, d_v), "float32", scope="local")
+    return Q_smem, K_smem, V_smem, O_local
+
+
+def _alloc_mla_qkvo_buffers(tile_x, tile_z, d_qk, d_latent, dtype):
+    """Allocate Q + combined KV shared + O local for MLA prefill (V reuses the KV buffer)."""
+    Q_smem = T.sblock_alloc_buffer((tile_x, d_qk), dtype, scope="shared")
+    KV_smem = T.sblock_alloc_buffer((tile_z, d_qk), dtype, scope="shared")
+    O_local = T.sblock_alloc_buffer((tile_x, d_latent), "float32", scope="local")
+    return Q_smem, KV_smem, O_local
+
+
+def _alloc_tile_walk_state():
+    """Return (tile_id, batch_idx, batch_tiles, batch_rows, iterator, kv_chunk_len) int32 scalars for the paged/ragged/MLA tile-walk state machine."""
+    return _var("int32"), _var("int32"), _var("int32"), _var("int32"), _var("int32"), _var("int32")
+
+
+def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_size):
+    """Build @T.macro helpers shared across tiled online-softmax prefill kernels.
+
+    Parameters
+    ----------
+    tile_x : int  # query/output row tile
+    tile_y : int  # QK reduction dim (head_dim for MHA, d_qk for MLA/ragged)
+    tile_z : int  # key/value column tile
+    tile_o : int  # output/V column dim (d for MHA/sequence, d_v for ragged, d_latent for MLA)
+    """
+    @T.macro
+    def init_states(
+        m_smem: T.Buffer, d_smem: T.Buffer, O_local: T.Buffer, ty: T.int32, tx: T.int32,
+    ):
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                m_smem[row] = -5e4
+                d_smem[row] = 1.0
+        for li, lj in T.grid(tile_x, tile_o):
+            with T.sblock("O_init"):
+                i, j = T.axis.remap("SS", [li, lj])
+                O_local[i, j] = 0.0
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def compute_s_gemm(
+        Q_smem: T.Buffer, K_smem: T.Buffer, S_local: T.Buffer, S_smem: T.Buffer, sm_scale: T.float32,
+    ):
+        with T.sblock():
+            for li, lj, lk in T.grid(tile_x, tile_z, tile_y):
+                with T.sblock("S_gemm"):
+                    i, j, k = T.axis.remap("SSR", [li, lj, lk])
+                    with T.init():
+                        S_local[i, j] = 0.0
+                    S_local[i, j] += T.cast(Q_smem[i, k], "float32") * T.cast(K_smem[j, k], "float32") * sm_scale * math.log2(math.exp(1))
+        T.tvm_storage_sync("shared")
+        for li, lj in T.grid(tile_x, tile_z):
+            with T.sblock("S_store"):
+                i, j = T.axis.remap("SS", [li, lj])
+                S_smem[i, j] = S_local[i, j]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def softmax_update_causal(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        causal: T.int32, kv_len: T.int32, qo_len: T.int32,
+    ):
+        # Phase 1: compute m_new = max(masked S over kv tile), d_new = d_prev * exp2(m_prev - m_new)
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    for j in T.serial(tile_z):
+                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        # Phase 2: exp-and-scale S_smem; masked-out entries use -inf
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    # predicate sits inside loop so sync stays outside conditional branches
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        # Phase 3: d_new += sum(S_smem[row, :]); write m/d/m_prev back to smem
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def compute_o_gemm(
+        S_smem: T.Buffer, V_smem: T.Buffer, O_local: T.Buffer,
+        m_prev_smem: T.Buffer, m_smem: T.Buffer,
+    ):
+        with T.sblock():
+            for li, lj, lk in T.grid(tile_x, tile_o, tile_z):
+                with T.sblock("O_gemm"):
+                    i, j, k = T.axis.remap("SSR", [li, lj, lk])
+                    with T.init():
+                        O_local[i, j] *= T.exp2(m_prev_smem[i] - m_smem[i])
+                    O_local[i, j] += S_smem[i, k] * T.cast(V_smem[k, j], "float32")
+
+    @T.macro
+    def paged_store_output_lse(
+        output: T.Buffer, lse: T.Buffer, O_local: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer,
+        q_indptr: T.Buffer, b_idx: T.int32, by: T.int32, LH_start: T.int32,
+    ):
+        """Paged-style (q_indptr-based) O_store + lse_store epilogue.
+
+        Used by paged prefill, ragged prefill and MLA prefill. MLA passes ``by=0`` so
+        the ``by * group_size`` term drops to zero at compile time.
+        """
+        for li, lj in T.grid(tile_x, tile_o):
+            with T.sblock("O_store"):
+                i, j = T.axis.remap("SS", [li, lj])
+                cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
+                cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
+                if cur_L < q_indptr[b_idx + 1]:
+                    output[cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
+        for li in T.grid(tile_x):
+            with T.sblock("lse_store"):
+                i = T.axis.remap("S", [li])
+                cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
+                cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
+                if cur_L < q_indptr[b_idx + 1]:
+                    lse[cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
+
+    @T.macro
+    def advance_tile_batch(
+        tile_id: T.Buffer, batch_idx: T.Buffer, batch_tiles: T.Buffer, batch_rows: T.Buffer,
+        q_indptr: T.Buffer, batch_size: T.int32,
+    ):
+        """Advance tile_id/batch_idx past exhausted batches.
+
+        After the loop, either batch_idx[0] >= batch_size (all tiles consumed) or
+        tile_id[0] < batch_tiles[0] (the current batch still has work to do).
+        """
+        while tile_id[0] >= batch_tiles[0] and batch_idx[0] < batch_size:
+            tile_id[0] -= batch_tiles[0]
+            batch_idx[0] += 1
+            if batch_idx[0] < batch_size:
+                b_idx: T.let[T.int32] = batch_idx[0]
+                batch_rows[0] = (q_indptr[b_idx + 1] - q_indptr[b_idx]) * group_size
+                batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
+
+    @T.macro
+    def softmax_update_valid_length(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        valid_len: T.int32, qo_len: T.int32, kv_len: T.int32,
+    ):
+        # Same three-phase online softmax as softmax_update_causal but with a
+        # per-batch right-padding mask in place of causal masking.
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    for j in T.serial(tile_z):
+                        if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    @T.macro
+    def softmax_update_causal_padded_left(
+        S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
+        m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
+        ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
+        valid_len: T.int32, qo_len: T.int32, kv_len: T.int32,
+    ):
+        # Three-phase online softmax with left-padding + causal mask. Real
+        # queries occupy [qo_len - valid_len, qo_len); real keys occupy
+        # [kv_len - valid_len, kv_len). Causal keeps
+        # col <= row + (kv_len - qo_len) within those valid suffixes.
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update1"):
+                    m_prev[i] = m_smem[row]
+                    m_new[i] = m_smem[row]
+                    row_: T.let[T.int32] = (LH_start + row) // group_size
+                    pad_q: T.let[T.int32] = qo_len - valid_len
+                    pad_kv: T.let[T.int32] = kv_len - valid_len
+                    for j in T.serial(tile_z):
+                        col_: T.let[T.int32] = L_kv_start + j
+                        if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
+                            m_new[i] = T.max(m_new[i], S_smem[row, j])
+                    d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            with T.sblock("update"):
+                for j in T.serial(tile_z):
+                    if row < tile_x:
+                        row_: T.let[T.int32] = (LH_start + row) // group_size
+                        pad_q: T.let[T.int32] = qo_len - valid_len
+                        pad_kv: T.let[T.int32] = kv_len - valid_len
+                        col_: T.let[T.int32] = L_kv_start + j
+                        if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
+                            S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
+                        else:
+                            S_smem[row, j] = T.exp2(-5e4 - m_new[i])
+        for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
+            row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
+            if row < tile_x:
+                with T.sblock("update"):
+                    for j in T.serial(tile_z):
+                        d_new[i] += S_smem[row, j]
+                    m_smem[row] = m_new[i]
+                    d_smem[row] = d_new[i]
+                    m_prev_smem[row] = m_prev[i]
+        T.tvm_storage_sync("shared")
+
+    return init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, softmax_update_valid_length, advance_tile_batch, paged_store_output_lse, softmax_update_causal_padded_left
+
+
+def _get_prefill_kernel_config(h_kv, h_q, d, dtype, target: Target):
+    NUM_BLKS = 16
+    LOAD_VEC = 8 // ((DataType(dtype).bits + 7) // 8)  # 8 bytes
+    group_size = h_q // h_kv
+
+    bdx = 32
+    num_warps = 4
+    tile_x, tile_y, tile_z = (
+        64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+        d,
+        64 // ((DataType(dtype).bits + 7) // 8) // max(d // 128, 1),
+    )
+    original_tile_y = tile_y
+    original_tile_z = tile_z
+    while (tile_x * tile_z) % (bdx * num_warps) != 0:
+        tile_z += original_tile_z
+    while (tile_x * tile_y) % (bdx * num_warps) != 0:
+        tile_y += original_tile_y
+
+    # Otherwise we would exceed maxComputeWorkgroupStorageSize
+    if (
+        target.kind.name == "webgpu"
+        and ((d + 127) // 128) * ((DataType(dtype).bits + 15) // 16) >= 4
+    ):
+        tile_z = 8
+        num_warps = 2
+    if target.kind.name == "opencl" and (
+        ("android" in str(target.host)) or ("adreno" in str(target.attrs))
+    ):
+        LOAD_VEC = 16 // ((DataType(dtype).bits + 7) // 8)  # 16 bytes
+        NUM_BLKS = group_size * 8
+
+    check_thread_limits(target, bdx=bdx, bdy=num_warps, bdz=1, gdz=1)
+
+    return NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z
+
+
+def _schedule_prefill_kernel(sch: s_tir.Schedule, load_vec, bdx, num_warps, tile_x, tile_y, tile_z, transform_k_load: bool, merged_qk_load: bool) -> tvm.s_tir.Schedule:
+    get_extent = lambda *lps: [int(sch.get(lp).extent) for lp in lps]
+
+    def get_vecsize(extent):
+        return min(load_vec, (extent & ~(extent - 1)))
+
+    def getxy_vecsize(x, y, t):
+        assert (x * y) % t == 0
+        return min(get_vecsize(y), get_vecsize(x * y // t))
+
+    def get_tile_size(x, y, t):
+        cnt = (x * y) // t
+        assert (x * y) % t == 0
+        tile_y = math.ceil(math.sqrt(cnt))
+        while (cnt % tile_y != 0 or y % tile_y != 0 or x % (cnt // tile_y) != 0) and tile_y <= cnt:
+            tile_y += 1
+        assert tile_y <= cnt
+        tile_x = cnt // tile_y
+        return tile_x, tile_y
+
+    def apply_to_qkv_load(sch: s_tir.Schedule, block):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        x_extent, y_extent = get_extent(loop_x, loop_y)
+        vec_size = getxy_vecsize(x_extent, y_extent, bdx * num_warps)
+        yo, yv = sch.split(loop_y, [None, vec_size])
+        yo_extent = y_extent // vec_size
+        tile_x, tile_y = get_tile_size(x_extent, yo_extent, (bdx * num_warps))
+        xo, xi = sch.split(loop_x, [tile_x, None])
+        yo, yi = sch.split(yo, [tile_y, None])
+        sch.reorder(xi, yi, xo, yo)
+        t = sch.fuse(xi, yi)
+        ty, tx = sch.split(t, [num_warps, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.vectorize(yv)
+
+    def apply_to_so_ewise(sch: s_tir.Schedule, block, tile):
+        loop_x, loop_y = sch.get_loops(block)[-2:]
+        xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+        yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+        sch.reorder(xo, yo, xi, yi)
+        yiv_extent = get_vecsize(tile[1])
+        yio, yiv = sch.split(yi, [None, yiv_extent])
+        sch.unroll(yio)
+        sch.vectorize(yiv)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    def apply_to_gemm(sch: s_tir.Schedule, block, tile, r_len=16, k_major=False):
+        loop_x, loop_y, loop_z = sch.get_loops(block)[-3:]
+        xo, xi = sch.split(loop_x, factors=[None, tile[0]])
+        yo, yi = sch.split(loop_y, factors=[None, tile[1]])
+        sch.reorder(xo, yo, xi, yi)
+        t = sch.fuse(xo, yo)
+        ty, tx = sch.split(t, factors=[None, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+        ko, ki = sch.split(loop_z, factors=[None, r_len])
+        if k_major:
+            sch.reorder(ko, xi, yi, ki)
+        else:
+            sch.reorder(ko, ki, xi, yi)
+        yiv_extent = get_vecsize(tile[1])
+        yio, yiv = sch.split(yi, [None, yiv_extent])
+        sch.unroll(yio)
+        sch.vectorize(yiv)
+        sch.unroll(xi)
+        sch.decompose_reduction(block, ty)
+
+    def apply_to_md(sch, block):
+        loop = sch.get_loops(block)[-1]
+        _, ty, tx = sch.split(loop, factors=[None, num_warps, bdx])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+
+    if transform_k_load and not merged_qk_load:
+        sch.transform_layout("K_load", ("write", 0), lambda i, j: (j, i))
+    tile_s = get_tile_size(tile_x, tile_z, bdx * num_warps)
+    tile_o = get_tile_size(tile_x, tile_y, bdx * num_warps)
+    apply_to_gemm(sch, sch.get_sblock("S_gemm"), tile_s, k_major=True)
+    apply_to_gemm(sch, sch.get_sblock("O_gemm"), tile_o, k_major=False)
+    apply_to_so_ewise(sch, sch.get_sblock("S_store"), tile_s)
+    apply_to_so_ewise(sch, sch.get_sblock("O_init"), tile_o)
+    apply_to_so_ewise(sch, sch.get_sblock("O_store"), tile_o)
+    apply_to_qkv_load(sch, sch.get_sblock("Q_load"))
+    if not merged_qk_load:
+        apply_to_qkv_load(sch, sch.get_sblock("K_load"))
+        apply_to_qkv_load(sch, sch.get_sblock("V_load"))
+    else:
+        apply_to_qkv_load(sch, sch.get_sblock("KV_load"))
+    apply_to_md(sch, sch.get_sblock("lse_store"))
+    return sch

python/tvm/relax/frontend/nn/llm/_page_kernels.py

@@ -0,0 +1,293 @@
+# 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.
+# ruff: noqa: E501
+# fmt: off
+
+"""TIR kernels that operate on paged KV-cache storage (without doing attention).
+
+This module contains:
+- Append helpers that transpose/write new K/V tokens into the paged layout
+  (``_kv_cache_transpose_append`` and its MLA variant).
+- Debug helpers that extract K/V from the paged layout for inspection
+  (``_kv_cache_debug_get_kv``, ``_kv_cache_debug_get_kv_mla``).
+- Copy helpers used by the cache runtime for forking/sharing pages
+  (``_copy_single_page``, ``_copy_single_page_mla``, ``_copy_single_page_cpu``).
+- Compact helpers that reorganise pages after removals
+  (``_compact_kv_copy``, ``_compact_kv_copy_cpu``).
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+from tvm.script import tirx as T
+from tvm.target import Target
+
+from ._kernel_common import get_max_num_threads_per_block
+
+
+def _kv_cache_transpose_append(num_key_value_heads, head_dim, dtype, page_size: int = 16):
+    """Return the TIR function that appends new k/v data to PagedKVCache."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_transpose_append(
+        var_pages: T.handle,
+        var_k_data: T.handle,
+        var_v_data: T.handle,
+        var_position_map: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        ntoken = T.Var("num_tokens_excluding_cache", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_key_value_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+        k_data = T.match_buffer(var_k_data, (ntoken, num_key_value_heads, head_dim), dtype)
+        v_data = T.match_buffer(var_v_data, (ntoken, num_key_value_heads, head_dim), dtype)
+        position_map = T.match_buffer(var_position_map, (ntoken,), "int32", elem_offset=position_map_elem_offset)
+        for global_pos, h, f in T.grid(ntoken, num_key_value_heads, head_dim):
+            if position_map[global_pos] != T.int32(-1):
+                with T.sblock("k_transpose_append"):
+                    vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
+                    T.reads(position_map[vgpos], k_data[vgpos, vh, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, 0, vh, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos]  # type: ignore
+                    pages[T.floordiv(position, page_size), 0, vh, T.floormod(position, page_size), vf] = k_data[vgpos, vh, vf]
+                with T.sblock("v_transpose_append"):
+                    vgpos, vh, vf = T.axis.remap("SSS", [global_pos, h, f])
+                    T.reads(position_map[vgpos], v_data[vgpos, vh, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, 1, vh, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos] # type: ignore[name-defined,no-redef]
+                    pages[T.floordiv(position, page_size), 1, vh, T.floormod(position, page_size), vf] = v_data[vgpos, vh, vf]
+
+    return tir_kv_cache_transpose_append
+
+
+def _kv_cache_transpose_append_mla(d_qk: int, dtype, page_size: int = 16):
+    """Return the TIR function that appends new compressed KV data to PagedKVCache for MLA."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_transpose_append_mla(
+        var_pages: T.handle,
+        var_kv_data: T.handle,
+        var_position_map: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        ntoken = T.Var("num_tokens_excluding_cache", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, d_qk), dtype, elem_offset=pages_elem_offset)
+        kv_data = T.match_buffer(var_kv_data, (ntoken, d_qk), dtype)
+        position_map = T.match_buffer(var_position_map, (ntoken,), "int32", elem_offset=position_map_elem_offset)
+        for global_pos, f in T.grid(ntoken, d_qk):
+            if position_map[global_pos] != T.int32(-1):
+                with T.sblock("k_transpose_append"):
+                    vgpos, vf = T.axis.remap("SS", [global_pos, f])
+                    T.reads(position_map[vgpos], kv_data[vgpos, vf])
+                    T.writes(pages[position_map[vgpos] // page_size, position_map[vgpos] % page_size, vf])
+                    position: T.int32 = position_map[vgpos]  # type: ignore
+                    pages[T.floordiv(position, page_size), T.floormod(position, page_size), vf] = kv_data[vgpos, vf]
+
+    return tir_kv_cache_transpose_append_mla
+
+
+def _kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, head_dim, dtype):
+    """Return the TIR function that fetches the k/v data on given positions and layer."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_debug_get_kv(
+        var_pages: T.handle,
+        var_position_map: T.handle,
+        var_k_data: T.handle,
+        var_v_data: T.handle,
+        layer_id: T.int64,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        seqlen = T.Var("num_tokens_including_cache", "int64")
+        page_size = T.Var("page_size", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_key_value_heads, page_size, head_dim), dtype,elem_offset=pages_elem_offset)
+        position_map = T.match_buffer(var_position_map, (seqlen,), "int32", elem_offset=position_map_elem_offset)
+        k_data = T.match_buffer(var_k_data, (num_hidden_layers, seqlen, num_key_value_heads, head_dim), dtype)
+        v_data = T.match_buffer(var_v_data, (num_hidden_layers, seqlen, num_key_value_heads, head_dim), dtype)
+        for p, h, d in T.grid(seqlen, num_key_value_heads, head_dim):
+            with T.sblock("copy0"):
+                vp, vh, vd = T.axis.remap("SSS", [p, h, d])
+                T.reads(position_map[vp], pages[position_map[vp] // page_size, 0:2, vh, position_map[vp] % page_size, vd])
+                T.writes(k_data[layer_id, vp, vh, vd], v_data[layer_id, vp, vh, vd])
+                position: T.int32 = position_map[vp] # type: ignore[name-defined]
+                k_data[layer_id, vp, vh, vd] = pages[T.floordiv(position, page_size), 0, vh, T.floormod(position, page_size), vd]
+                v_data[layer_id, vp, vh, vd] = pages[T.floordiv(position, page_size), 1, vh, T.floormod(position, page_size), vd]
+
+    return tir_kv_cache_debug_get_kv
+
+
+def _kv_cache_debug_get_kv_mla(num_hidden_layers, d_qk, dtype):
+    """Return the TIR function that fetches the k/v data on given positions and layer."""
+
+    @T.prim_func(s_tir=True)
+    def tir_kv_cache_debug_get_kv_mla(
+        var_pages: T.handle,
+        var_position_map: T.handle,
+        var_compressed_kv_with_k_pe_data: T.handle,
+        layer_id: T.int64,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        seqlen = T.Var("num_tokens_including_cache", "int64")
+        page_size = T.Var("page_size", "int64")
+        num_pages = T.int64()
+        pages_elem_offset = T.int64()
+        position_map_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, d_qk), dtype, elem_offset=pages_elem_offset)
+        position_map = T.match_buffer(var_position_map, (seqlen,), "int32", elem_offset=position_map_elem_offset)
+        compressed_kv_with_k_pe_data = T.match_buffer(var_compressed_kv_with_k_pe_data, (num_hidden_layers, seqlen, d_qk), dtype)
+        for p, d in T.grid(seqlen, d_qk):
+            with T.sblock("copy0"):
+                vp, vd = T.axis.remap("SS", [p, d])
+                T.reads(position_map[vp], pages[position_map[vp] // page_size, position_map[vp] % page_size, vd])
+                T.writes(compressed_kv_with_k_pe_data[layer_id, vp, vd])
+                position: T.int32 = position_map[vp] # type: ignore[name-defined]
+                compressed_kv_with_k_pe_data[layer_id, vp, vd] = pages[T.floordiv(position, page_size), T.floormod(position, page_size), vd]
+
+    return tir_kv_cache_debug_get_kv_mla
+
+
+def _copy_single_page(num_heads, page_size, head_dim, dtype, target: Target):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+
+        for b in T.thread_binding((copy_length * num_heads * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+            for t in T.thread_binding(tx, thread="threadIdx.x"):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * num_heads * head_dim)
+                    vh = T.axis.spatial(num_heads, T.Cast("int32", (b * tx + t) // (copy_length * head_dim)))
+                    vp = T.axis.spatial(copy_length, (b * tx + t) % (copy_length * head_dim) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, 0, vh, vp, vd] = pages[src_page_id, 0, vh, vp, vd]
+                    pages[tgt_page_id, 1, vh, vp, vd] = pages[src_page_id, 1, vh, vp, vd]
+
+    return copy_single_page
+
+
+def _copy_single_page_mla(page_size, head_dim, dtype, target: Target):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page_mla(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+
+        for b in T.thread_binding((copy_length * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+            for t in T.thread_binding(tx, thread="threadIdx.x"):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * head_dim)
+                    vp = T.axis.spatial(copy_length, (b * tx + t) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, vp, vd] = pages[src_page_id, vp, vd]
+
+    return copy_single_page_mla
+
+
+def _copy_single_page_cpu(num_heads, page_size, head_dim, dtype):
+    tx = 1
+
+    @T.prim_func(s_tir=True)
+    def copy_single_page_cpu(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype)
+
+        for b in T.serial((copy_length * num_heads * head_dim + tx - 1) // tx):
+            for t in T.serial(tx):
+                with T.sblock("copy"):
+                    T.where(b * tx + t < copy_length * num_heads * head_dim)
+                    vh = T.axis.spatial(num_heads, T.Cast("int32", (b * tx + t) // (copy_length * head_dim)))
+                    vp = T.axis.spatial(copy_length, (b * tx + t) % (copy_length * head_dim) // head_dim)
+                    vd = T.axis.spatial(head_dim, T.Cast("int32", (b * tx + t) % head_dim))
+                    pages[tgt_page_id, 0, vh, vp, vd] = pages[src_page_id, 0, vh, vp, vd]
+                    pages[tgt_page_id, 1, vh, vp, vd] = pages[src_page_id, 1, vh, vp, vd]
+
+    return copy_single_page_cpu
+
+
+def _compact_kv_copy(num_heads, head_dim, dtype, target: Target, page_size: int = 16):
+    tx = get_max_num_threads_per_block(target)
+
+    @T.prim_func(s_tir=True)
+    def compact_kv_copy(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        total_copy_length = T.int32()
+        copy_length_indptr_elem_offset = T.int32()
+        copy_src_dst_pos_elem_offset = T.int32()
+        pages_elem_offset = T.int64()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype, elem_offset=pages_elem_offset)
+        copy_length_indptr = T.match_buffer(var_copy_length_indptr, (batch_size + 1,), "int32", elem_offset=copy_length_indptr_elem_offset)
+        copy_src_dst_pos = T.match_buffer(var_copy_src_dst_pos, (2, total_copy_length), "int32", elem_offset=copy_src_dst_pos_elem_offset)
+
+        with T.sblock("root"):
+            for bhd_o in T.thread_binding((batch_size * num_heads * head_dim + tx - 1) // tx, thread="blockIdx.x"):
+                for bhd_i in T.thread_binding(tx, thread="threadIdx.x"):
+                    b: T.int32 = (bhd_o * tx + bhd_i) // (num_heads * head_dim)
+                    h: T.int32 = (bhd_o * tx + bhd_i) // head_dim % num_heads
+                    d: T.int32 = (bhd_o * tx + bhd_i) % head_dim
+                    if (bhd_o * tx + bhd_i) < batch_size * num_heads * head_dim:
+                        for i in T.serial(copy_length_indptr[b + 1] - copy_length_indptr[b]):
+                            src_pos: T.int32 = copy_src_dst_pos[0, copy_length_indptr[b] + i]
+                            dst_pos: T.int32 = copy_src_dst_pos[1, copy_length_indptr[b] + i]
+                            pages[dst_pos // page_size, 0, h, dst_pos % page_size, d] = pages[src_pos // page_size, 0, h, src_pos % page_size, d]
+                            pages[dst_pos // page_size, 1, h, dst_pos % page_size, d] = pages[src_pos // page_size, 1, h, src_pos % page_size, d]
+
+    return compact_kv_copy
+
+
+def _compact_kv_copy_cpu(num_heads, head_dim, dtype, page_size: int = 16):
+    tx = 8
+
+    @T.prim_func(s_tir=True)
+    def compact_kv_copy_cpu(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
+        T.func_attr({"tirx.is_scheduled": True})
+        num_pages = T.int32()
+        total_copy_length = T.int32()
+        copy_length_indptr_elem_offset = T.int32()
+        copy_src_dst_pos_elem_offset = T.int32()
+        pages = T.match_buffer(var_pages, (num_pages, 2, num_heads, page_size, head_dim), dtype)
+        copy_length_indptr = T.match_buffer(var_copy_length_indptr, (batch_size + 1,), "int32", elem_offset=copy_length_indptr_elem_offset)
+        copy_src_dst_pos = T.match_buffer(var_copy_src_dst_pos, (2, total_copy_length), "int32", elem_offset=copy_src_dst_pos_elem_offset)
+
+        with T.sblock("root"):
+            for bhd_o in T.serial((batch_size * num_heads * head_dim + tx - 1) // tx):
+                for bhd_i in T.serial(tx):
+                    b: T.int32 = (bhd_o * tx + bhd_i) // (num_heads * head_dim)
+                    h: T.int32 = (bhd_o * tx + bhd_i) // head_dim % num_heads
+                    d: T.int32 = (bhd_o * tx + bhd_i) % head_dim
+                    if (bhd_o * tx + bhd_i) < batch_size * num_heads * head_dim:
+                        for i in T.serial(copy_length_indptr[b + 1] - copy_length_indptr[b]):
+                            src_pos: T.int32 = copy_src_dst_pos[0, copy_length_indptr[b] + i]
+                            dst_pos: T.int32 = copy_src_dst_pos[1, copy_length_indptr[b] + i]
+                            pages[dst_pos // page_size, 0, h, dst_pos % page_size, d] = pages[src_pos // page_size, 0, h, src_pos % page_size, d]
+                            pages[dst_pos // page_size, 1, h, dst_pos % page_size, d] = pages[src_pos // page_size, 1, h, src_pos % page_size, d]
+
+    return compact_kv_copy_cpu

python/tvm/relax/frontend/nn/llm/kv_cache.py

@@ -0,0 +1,686 @@
+# 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.
+# ruff: noqa: E501, RUF012
+# fmt: off
+
+"""Attention KV cache modeling.
+
+This module exposes the public ``PagedKVCache`` classes (``FlashInferPagedKVCache``
+and ``TIRPagedKVCache``). The kernel factories that build the underlying TIR
+functions are split across sibling private modules:
+
+- ``_kernel_common``: shared helpers (enums, RoPE, mask, tile allocators,
+  ``@T.macro`` bundle, tiling config, scheduling).
+- ``_page_kernels``: page management (append, debug, copy, compact).
+- ``_prefill_kernels``: prefill attention kernels (paged/ragged/MLA/dense).
+- ``_decode_kernels``: decode attention kernels and state-merge helpers.
+
+The private-named kernel factories are re-exported from this module so the
+test suite can continue to import them via ``tvm.relax.frontend.nn.llm.kv_cache``.
+"""
+
+# pylint: disable=too-many-statements,too-many-arguments,invalid-name,line-too-long
+import math
+from typing import Any, Literal
+
+import tvm
+from tvm import relax as rx
+from tvm import tirx
+from tvm.relax.frontend.nn import Object, Tensor
+from tvm.target import Target
+
+# Re-export enums + kernel factories so existing ``from kv_cache import ...``
+# users (test suite, tree_attn.py, mlc-llm, etc.) continue to work after the
+# split. These names are referenced in ``__all__`` below to signal to linters
+# that the imports are intentional public API (not dead code).
+from ._decode_kernels import (
+    _attention_decode,
+    _attention_decode_cpu,
+    _merge_state_inplace,
+    _merge_state_inplace_cpu,
+)
+from ._kernel_common import AttnKind, RopeMode
+from ._page_kernels import (
+    _compact_kv_copy,
+    _compact_kv_copy_cpu,
+    _copy_single_page,
+    _copy_single_page_cpu,
+    _copy_single_page_mla,
+    _kv_cache_debug_get_kv,
+    _kv_cache_debug_get_kv_mla,
+    _kv_cache_transpose_append,
+    _kv_cache_transpose_append_mla,
+)
+from ._prefill_kernels import (
+    _attention_prefill,
+    _attention_prefill_cpu,
+    _attention_prefill_mla,
+    _attention_prefill_ragged,
+    _attention_prefill_ragged_cpu,
+    _attention_sequence_prefill,
+    _attention_sequence_prefill_with_mask,
+)
+from .position_embedding import llama_rope_with_position_map
+from .tree_attn import (
+    tree_attn,
+    tree_attn_cpu,
+    tree_attn_with_paged_kv_cache,
+    tree_attn_with_paged_kv_cache_cpu,
+)
+
+__all__ = [
+    "AttnKind",
+    "FlashInferPagedKVCache",
+    "PagedKVCache",
+    "RopeMode",
+    "TIRPagedKVCache",
+    "_attention_decode",
+    "_attention_decode_cpu",
+    "_attention_prefill",
+    "_attention_prefill_cpu",
+    "_attention_prefill_mla",
+    "_attention_prefill_ragged",
+    "_attention_prefill_ragged_cpu",
+    "_attention_sequence_prefill",
+    "_attention_sequence_prefill_with_mask",
+    "_compact_kv_copy",
+    "_compact_kv_copy_cpu",
+    "_copy_single_page",
+    "_copy_single_page_cpu",
+    "_copy_single_page_mla",
+    "_kv_cache_debug_get_kv",
+    "_kv_cache_debug_get_kv_mla",
+    "_kv_cache_transpose_append",
+    "_kv_cache_transpose_append_mla",
+    "_merge_state_inplace",
+    "_merge_state_inplace_cpu",
+    "llama_rope_with_position_map",
+    "tree_attn",
+    "tree_attn_cpu",
+    "tree_attn_with_paged_kv_cache",
+    "tree_attn_with_paged_kv_cache_cpu",
+]
+
+
+class PagedKVCache(Object):  # pylint: disable=too-few-public-methods
+    """The Paged KV Cache used in LLM batching for efficient attention computation."""
+
+    extern_mods: list[tvm.runtime.Module] = []
+
+    def attention_with_fused_qkv(
+        self,
+        layer_id: int,
+        qkv: Tensor,
+        num_qo_heads: int,
+        sm_scale: float,
+    ) -> Tensor:
+        """Compute attention with the given fused q/k/v data and in-cache k/v data
+        on the specified layer. Rotary position embeddings are applied to k/v
+        within this function.
+
+        - For prefill, the input qkv and output tensor have shape
+        (1, total_seq_len) for the first two dimensions.
+        - For decode, the input qkv and output tensor have shape
+        (batch_size, 1) for the first two dimensions.
+        - The input qkv have `2 * num_qo_heads + num_kv_heads` at the third dim.
+        - The output tensor have `num_qo_heads` at the third dim.
+        - The input qkv and output tensor have `head_dim` at the last dim.
+        """
+        # pylint: disable=protected-access
+        b, s, _, d = qkv._expr.ty.shape
+        qkv = qkv.reshape(b * s, qkv.shape[2], d)
+        return Tensor(
+            _expr=rx.BlockBuilder.current().emit(
+                rx.call_dps_packed(
+                    "vm.builtin.attention_kv_cache_attention_with_fused_qkv",
+                    [
+                        self._expr,
+                        rx.prim_value(layer_id),  # type: ignore[arg-type]
+                        rx.prim_value(sm_scale),
+                        qkv._expr,
+                    ],
+                    out_ty=rx.TensorType((b * s, num_qo_heads, d), qkv.dtype),
+                )
+            )
+        ).reshape(b, s, num_qo_heads, d)
+
+    def self_attention(  # pylint: disable=too-many-locals
+        self,
+        layer_id: int,
+        q: Tensor,
+        k: Tensor,
+        v: Tensor,
+        sm_scale: float,
+    ) -> tuple[Tensor, Tensor]:
+        """Fine-grained API that computes ragged self attention with Q/K/V data."""
+        # pylint: disable=protected-access
+        b, s, h_qo, d_qk = q._expr.ty.shape
+        _, _, h_kv, d_v = v._expr.ty.shape
+        q = q.reshape(b * s, h_qo, d_qk)
+        k = k.reshape(b * s, h_kv, d_qk)
+        v = v.reshape(b * s, h_kv, d_v)
+        bb = rx.BlockBuilder.current()
+        attn_results = bb.emit(
+            rx.call_dps_packed(
+                "vm.builtin.attention_kv_cache_self_attention",
+                [
+                    self._expr,
+                    rx.prim_value(layer_id),  # type: ignore[arg-type]
+                    rx.prim_value(sm_scale),
+                    q._expr,
+                    k._expr,
+                    v._expr,
+                ],
+                out_ty=[
+                    rx.TensorType((b * s, h_qo, d_v), q.dtype),
+                    rx.TensorType((b * s, h_qo), "float32"),
+                ],
+            )
+        )
+        assert isinstance(attn_results.ty, rx.TupleType)
+        assert len(attn_results.ty.fields) == 2
+        o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, d_v)
+        lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
+        return o, lse
+
+    def cross_attention(
+        self,
+        layer_id: int,
+        q: Tensor,
+        v_head_dim: int,
+        sm_scale: float,
+    ) -> tuple[Tensor, Tensor]:
+        """Fine-grained API that computes paged cross attention with Q and in-cache KV data."""
+        # pylint: disable=protected-access
+        b, s, h_qo, d_qk = q._expr.ty.shape
+        q = q.reshape(b * s, h_qo, d_qk)
+        bb = rx.BlockBuilder.current()
+        attn_results = bb.emit(
+            rx.call_dps_packed(
+                "vm.builtin.attention_kv_cache_cross_attention",
+                [
+                    self._expr,
+                    rx.prim_value(layer_id),  # type: ignore[arg-type]
+                    rx.prim_value(sm_scale),
+                    q._expr,
+                ],
+                out_ty=[
+                    rx.TensorType((b * s, h_qo, v_head_dim), q.dtype),
+                    rx.TensorType((b * s, h_qo), "float32"),
+                ],
+            )
+        )
+        assert isinstance(attn_results.ty, rx.TupleType)
+        assert len(attn_results.ty.fields) == 2
+        o = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 0))).reshape(b, s, h_qo, v_head_dim)
+        lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
+        return o, lse
+
+    def append_mla_kv(self, layer_id: int, kv: Tensor) -> "PagedKVCache":
+        """Fine-grained API that appends the MLA K/V data to KV cache."""
+        # pylint: disable=protected-access
+        b, s, _, d_qk = kv._expr.ty.shape
+        kv = kv.reshape(b * s, d_qk)
+        return PagedKVCache(
+            _expr=rx.call_pure_packed(
+                "vm.builtin.attention_kv_cache_append_mla_kv",
+                self._expr,
+                rx.prim_value(layer_id),  # type: ignore[arg-type]
+                kv._expr,
+                ty_args=rx.AnyType(),
+            ),
+            _name="paged_kv_cache",
+        )
+
+    def merge_attn_output_inplace(
+        self,
+        o_self_attn: Tensor,
+        lse_self_attn: Tensor,
+        o_cross_attn: Tensor,
+        lse_cross_attn: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        """Fine-grained API that merges the attention output from two sources.
+        The first two tensors will be inplace updated.
+        """
+        # pylint: disable=protected-access
+        b, s, h_qo, d_v = o_self_attn._expr.ty.shape
+        o_self_attn = o_self_attn.reshape(b * s, h_qo, d_v)
+        lse_self_attn = lse_self_attn.reshape(b * s, h_qo)
+        o_cross_attn = o_cross_attn.reshape(b * s, h_qo, d_v)
+        lse_cross_attn = lse_cross_attn.reshape(b * s, h_qo)
+        bb = rx.BlockBuilder.current()
+        merge_results = bb.emit(
+            rx.call_pure_packed(
+                "vm.builtin.attention_kv_cache_merge_attn_output_inplace",
+                self._expr,
+                o_self_attn._expr,
+                lse_self_attn._expr,
+                o_cross_attn._expr,
+                lse_cross_attn._expr,
+                ty_args=rx.TupleType(
+                    [o_self_attn._expr.ty, lse_self_attn._expr.ty]
+                ),
+            )
+        )
+        assert isinstance(merge_results.ty, rx.TupleType)
+        assert len(merge_results.ty.fields) == 2
+        o_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 0))).reshape(
+            b, s, h_qo, d_v
+        )
+        lse_self_attn = Tensor(_expr=bb.emit(rx.TupleGetItem(merge_results, 1))).reshape(b, s, h_qo)
+        return o_self_attn, lse_self_attn
+
+    def get_query_positions(self, total_length: tirx.Expr) -> Tensor:
+        """Get the in-sequence positions of each slot in the query,
+        which are needed for applying positional embeddings in some models.
+
+        Parameters
+        ----------
+        total_length : tirx.Expr
+            The summed-up total sequence length of queries in
+            the batch being forwarded.
+
+        Returns
+        -------
+        q_positions : Tensor
+            The in-sequence query positions, in shape `(total_length,)`
+        """
+        return Tensor(
+            _expr=rx.BlockBuilder.current().emit(
+                rx.call_pure_packed(
+                    "vm.builtin.attention_kv_cache_get_query_positions",
+                    self._expr,
+                    ty_args=rx.TensorType((total_length,), "int32"),
+                )
+            )
+        )
+
+    # pylint: enable=protected-access
+
+
+def _prepare_yarn_rope_scaling(rope_scaling: dict[str, Any] | None, rope_theta: float | None) -> dict[str, Any] | None:
+    """Ensure Yarn-specific scaling configs include the theta metadata."""
+    if rope_scaling is None:
+        return None
+    if rope_scaling.get("rope_type") != "yarn":
+        return rope_scaling
+
+    rope_scaling_updated = dict(rope_scaling)
+    if "inv_theta_log_scale" not in rope_scaling_updated and rope_theta is not None:
+        theta_value = float(rope_theta)
+        rope_scaling_updated["inv_theta_log_scale"] = 1.0 / (2 * math.log(theta_value))
+    return rope_scaling_updated
+
+
+class FlashInferPagedKVCache(PagedKVCache):  # pylint: disable=too-few-public-methods
+    """Paged KV cache using FlashInfer (CUDA) kernels."""
+
+    def __init__(  # pylint: disable=too-many-locals
+        self,
+        attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
+        max_batch_size: tirx.Var,
+        max_total_seq_len: tirx.Var,
+        prefill_chunk_size: tirx.Var,
+        page_size: tirx.Var,
+        support_sliding_window: tirx.Var,
+        layer_partition: rx.ShapeExpr,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        mla_original_qk_head_dim: int,
+        mla_original_v_head_dim: int,
+        rope_mode: RopeMode,
+        rope_scale: int,
+        rope_theta: int,
+        rope_scaling: dict[str, Any],
+        rope_ext_factors: rx.Expr,
+        rotary_dim: int,
+        enable_disaggregation: bool,
+        dtype: str,
+        target: Target,
+        name: str = "paged_kv_cache",
+    ) -> None:
+        """Create a paged KV cache object with FlashInfer kernels.
+
+        Parameters
+        ----------
+        max_batch_size : tirx.Var
+            The maximum allowed batch size of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        max_total_seq_len : tirx.Var
+            The maximum allowed total sequence length of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        prefill_chunk_size : tirx.Var
+            The maximum total sequence length in a prefill.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        page_size : tirx.Var
+            The size (a.k.a. number of tokens) of each page.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        support_sliding_window : tirx.Var
+            0 or 1, denoting whether the KV cache supports sliding window.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        layer_partition : rx.ShapeExpr
+            The KV cache layer partition for pipeline stages.
+            It is an indptr array, denoting the starting layer of each pipeline stage.
+        rope_mode : RopeMode
+            The RoPE mode of the Paged KV cache.
+            If it is normal, RoPE will be applied to k before adding k to cache.
+            Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
+        rope_scale : int
+            The scale of rotary position embedding.
+        rope_theta : int
+            The base of rotary position embedding.
+        rope_scaling: Dict[str, Any]
+            The RoPE scaling information dict.
+        rope_ext_factors: rx.Expr
+            The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
+        rotary_dim : int
+            The number of dimensions in the embedding that RoPE is applied to.
+        enable_disaggregation : bool
+            Whether to enable disaggregation in the KV cache.
+        """
+        assert rope_mode != RopeMode.INLINE, "FlashInfer RoPE does not support inline mode."
+        rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
+
+        attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
+        if attn_kind_single == "mha_sliding":
+            attn_kind_single = "mha"
+        flashinfer_prefill_mods = rx.backend.cuda.flashinfer.gen_flashinfer_prefill_module(
+            dtype_q=dtype,
+            dtype_kv=dtype,
+            dtype_o=dtype,
+            qk_head_dim=(qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim),
+            v_head_dim=(v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim),
+            enable_inline_rope=False,
+            return_static_libs=True,
+        )
+        flashinfer_decode_mods = (
+            rx.backend.cuda.flashinfer.gen_flashinfer_decode_module(
+                dtype_q=dtype,
+                dtype_kv=dtype,
+                dtype_o=dtype,
+                qk_head_dim=qk_head_dim,
+                v_head_dim=v_head_dim,
+                enable_inline_rope=False,
+                return_static_libs=True,
+            )
+            if attn_kind_single == "mha"
+            else []
+        )
+        flashinfer_mla_mods = (
+            rx.backend.cuda.flashinfer.gen_flashinfer_mla_module(
+                dtype_q=dtype,
+                dtype_kv=dtype,
+                dtype_o=dtype,
+                head_dim_ckv=v_head_dim,
+                head_dim_kpe=qk_head_dim - v_head_dim,
+                return_static_libs=True,
+            )
+            if attn_kind_single == "mla"
+            else []
+        )
+        self.extern_mods = flashinfer_prefill_mods + flashinfer_decode_mods + flashinfer_mla_mods
+
+        bb = rx.BlockBuilder.current()
+        mha_functions = (
+            [
+                rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
+                rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
+            ]
+            if attn_kind_single == "mha"
+            else [rx.Tuple([]) for _ in range(6)]
+        )
+        ragged_prefill_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan")]) if attn_kind_single == "mha" else rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_ragged_run"), rx.ExternFunc("batch_prefill_plan"), rx.prim_value(mla_original_qk_head_dim), rx.prim_value(mla_original_v_head_dim)])
+        mla_function = rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_mla_run"), rx.ExternFunc("batch_mla_plan")] if attn_kind_single == "mla" else [])
+        attn_merge_functions = [
+            bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
+        ]
+        if attn_kind_single == "mla":
+            attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
+
+        if isinstance(attn_kind, list):
+            attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
+        else:
+            attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
+
+        args = [
+            rx.ShapeExpr(
+                [
+                    max_batch_size,
+                    max_total_seq_len,
+                    prefill_chunk_size,
+                    page_size,
+                    support_sliding_window,
+                ]
+            ),
+            layer_partition,
+            rx.prim_value(num_attention_heads),
+            rx.prim_value(num_key_value_heads),
+            rx.prim_value(qk_head_dim),
+            rx.prim_value(v_head_dim),
+            rx.ShapeExpr(attn_kind),
+            rx.prim_value(enable_disaggregation),
+            rx.prim_value(rope_mode),
+            rx.prim_value(rope_scale),
+            rx.prim_value(rope_theta),
+            rope_ext_factors,
+            rx.op.zeros((), dtype),
+            bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
+            bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
+            ragged_prefill_function,
+            *mha_functions,
+            mla_function,
+            rx.Tuple(attn_merge_functions),
+            bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
+            bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
+            bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
+            bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
+        ]
+        super().__init__(
+            _expr=rx.call_pure_packed(
+                "vm.builtin.paged_attention_kv_cache_create",
+                *args,
+                ty_args=rx.AnyType(),
+            ),
+            _name=name,
+        )
+
+
+class TIRPagedKVCache(PagedKVCache):  # pylint: disable=too-few-public-methods
+    """Paged KV cache using TIR kernels."""
+
+    def __init__(  # pylint: disable=too-many-locals
+        self,
+        attn_kind: Literal["mha", "mla"] | list[Literal["mha", "mla", "mha_sliding"]],
+        max_batch_size: tirx.Var,
+        max_total_seq_len: tirx.Var,
+        prefill_chunk_size: tirx.Var,
+        page_size: tirx.Var,
+        support_sliding_window: tirx.Var,
+        layer_partition: rx.ShapeExpr,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        mla_original_qk_head_dim: int,
+        mla_original_v_head_dim: int,
+        rope_mode: RopeMode,
+        rope_scale: int,
+        rope_theta: int,
+        rope_scaling: dict[str, Any],
+        rope_ext_factors: rx.Expr,
+        rotary_dim: int,
+        enable_disaggregation: bool,
+        dtype: str,
+        target: Target,
+        name: str = "paged_kv_cache",
+    ) -> None:
+        """Create a paged KV cache object with TIR kernels.
+
+        Parameters
+        ----------
+        max_batch_size : tirx.Var
+            The maximum allowed batch size of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        max_total_seq_len : tirx.Var
+            The maximum allowed total sequence length of the KV cache.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        prefill_chunk_size : tirx.Var
+            The maximum total sequence length in a prefill.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        page_size : tirx.Var
+            The size (a.k.a. number of tokens) of each page.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        support_sliding_window : tirx.Var
+            0 or 1, denoting whether the KV cache supports sliding window.
+            It is a symbolic variable whose concrete value is specified
+            at runtime.
+        layer_partition : rx.ShapeExpr
+            The KV cache layer partition for pipeline stages.
+            It is an indptr array, denoting the starting layer of each pipeline stage.
+        rope_mode : RopeMode
+            The RoPE mode of the Paged KV cache.
+            If it is normal, RoPE will be applied to k before adding k to cache.
+            Otherwise, RoPE will be applied to q/k in attention kernel on-the-fly.
+        rope_scale : int
+            The scale of rotary position embedding.
+        rope_theta : int
+            The base of rotary position embedding.
+        rope_scaling: Dict[str, Any]
+            The RoPE scaling information dict.
+        rope_ext_factors: rx.Expr
+            The RoPE extension factors when "longrope" mode RoPE scaling is enabled.
+        rotary_dim : int
+            The number of dimensions in the embedding that RoPE is applied to.
+        enable_disaggregation : bool
+            Whether to enable disaggregation in the KV cache.
+        target : Target
+            The target to build the model to.
+        """
+        rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
+        attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
+        if attn_kind_single == "mha_sliding":
+            attn_kind_single = "mha"
+        if isinstance(attn_kind, list):
+            attn_kind = [int(getattr(AttnKind, layer_kind.upper())) for layer_kind in attn_kind]
+        else:
+            attn_kind = [int(getattr(AttnKind, attn_kind.upper())) for _ in range(num_hidden_layers)]
+        bb = rx.BlockBuilder.current()
+        args = [
+            rx.ShapeExpr(
+                [
+                    max_batch_size,
+                    max_total_seq_len,
+                    prefill_chunk_size,
+                    page_size,
+                    support_sliding_window,
+                ]
+            ),
+            layer_partition,
+            rx.prim_value(num_attention_heads),
+            rx.prim_value(num_key_value_heads),
+            rx.prim_value(qk_head_dim),
+            rx.prim_value(v_head_dim),
+            rx.ShapeExpr(attn_kind),
+            rx.prim_value(enable_disaggregation),
+            rx.prim_value(rope_mode),
+            rx.prim_value(rope_scale),
+            rx.prim_value(rope_theta),
+            rope_ext_factors,
+            rx.op.zeros((), dtype),
+            bb.add_func(_kv_cache_transpose_append(num_key_value_heads, qk_head_dim, dtype), "kv_cache_transpose_append"),
+            bb.add_func(_kv_cache_transpose_append_mla(qk_head_dim, dtype), "kv_cache_transpose_append_mla"),
+        ]
+
+        if target.kind.name == "llvm":
+            if attn_kind_single == "mla":
+                raise ValueError("MLA is not supported in TIR kernels for now.")
+            args.extend(
+                [
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_prefill_cpu_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
+                    rx.Tuple([]),  # f_mla_prefill
+                    rx.Tuple([bb.add_func(_merge_state_inplace_cpu(dtype), "tir_attention_merge_state_cpu")]),
+                    bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
+                    bb.add_func(_copy_single_page_cpu(num_key_value_heads, page_size, qk_head_dim, dtype), "kv_cache_copy_single_page_cpu"),
+                    bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
+                    bb.add_func(_compact_kv_copy_cpu(num_key_value_heads, qk_head_dim, dtype), "kv_cache_compact_kv_copy_cpu"),
+                ]
+            )
+        else:
+            ragged_qk_head_dim = qk_head_dim if attn_kind_single == "mha" else mla_original_qk_head_dim
+            ragged_v_head_dim = v_head_dim if attn_kind_single == "mha" else mla_original_v_head_dim
+            args.append(rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged(num_key_value_heads if attn_kind_single == "mha" else num_attention_heads, num_attention_heads, ragged_qk_head_dim, ragged_v_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_ragged")]))
+            mha_functions = (
+                [
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
+                ]
+                if attn_kind_single == "mha"
+                else [rx.Tuple([]) for _ in range(6)]
+            )
+            mla_function = rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_mla(num_attention_heads, v_head_dim, qk_head_dim - v_head_dim, dtype, False, target), "tir_attention_prefill_mla")] if attn_kind_single == "mla" else [])
+            attn_merge_functions = [
+                bb.add_func(_merge_state_inplace(num_attention_heads, v_head_dim, dtype, target, "tir_attention_merge_state"), "tir_attention_merge_state"),
+            ]
+            if attn_kind_single == "mla":
+                attn_merge_functions.append(bb.add_func(_merge_state_inplace(num_attention_heads, mla_original_v_head_dim, dtype, target, "tir_attention_merge_state_mla"), "tir_attention_merge_state_mla"))
+            args.extend(mha_functions)
+            args.append(mla_function)
+            args.extend(
+                [
+                    rx.Tuple(attn_merge_functions),
+                    bb.add_func(llama_rope_with_position_map(rope_theta, rope_scale, qk_head_dim, num_attention_heads, num_key_value_heads, dtype, rope_scaling, rotary_dim), "tir_split_rotary"),
+                    bb.add_func(_copy_single_page(num_key_value_heads, page_size, qk_head_dim, dtype, target) if attn_kind_single == "mha" else _copy_single_page_mla(page_size, qk_head_dim, dtype, target), "kv_cache_copy_single_page"),
+                    bb.add_func(_kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, qk_head_dim, dtype), "kv_cache_debug_get_kv"),
+                    bb.add_func(_compact_kv_copy(num_key_value_heads, qk_head_dim, dtype, target), "kv_cache_compact_kv_copy"),
+                ]
+            )
+
+        super().__init__(
+            _expr=rx.call_pure_packed(
+                "vm.builtin.paged_attention_kv_cache_create",
+                *args,
+                ty_args=rx.AnyType(),
+            ),
+            _name=name,
+        )

python/tvm/relax/frontend/nn/llm/position_embedding.py

@@ -0,0 +1,894 @@
+# 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.
+
+"""Operators for positional embeddings, e.g. RoPE."""
+
+import math
+from collections.abc import Callable
+from functools import partial
+from typing import Any
+
+from tvm import tirx
+from tvm.relax.frontend.nn import Tensor, op
+from tvm.script import tirx as T
+
+# pylint: disable=invalid-name
+
+
+def rope_freq_default(s: tirx.Var, d: tirx.Var, d_range: int, theta: float, dtype: str):
+    """Compute the inverse frequency of RoPE and then return the cosine and sine of it.
+
+    Parameters
+    ----------
+    s : tirx.Var
+        The position index.
+
+    d : tirx.Var
+        The dimension index.
+
+    d_range : int
+        The maximum dimension index.
+
+    theta : float
+        The theta value in RoPE, which controls the frequency.
+
+    dtype : str
+        The data type of the output.
+
+    Returns
+    -------
+    cos_freq : Tensor
+        The cosine of the inverse frequency.
+
+    sin_freq : Tensor
+        The sine of the inverse frequency.
+
+    var_map: Dict[tirx.Var, tirx.Expr]
+        The common expression map.
+    """
+    freq = s / tirx.power(theta, d * 2 % d_range / tirx.const(d_range, "float32"))
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_gptj(s: tirx.Var, d: tirx.Var, d_range: int, theta: float, dtype: str):
+    """Compute the inverse frequency of RoPE for gptj RoPE scaling."""
+    freq = s / tirx.power(theta, 2 * (d // 2) % d_range / tirx.const(d_range, "float32"))
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def rope_freq_llama4(  # pylint: disable=too-many-arguments,too-many-locals
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    original_max_position_embeddings: float,
+):
+    """Compute the inverse frequency of RoPE for llama4 RoPE scaling."""
+    orig_freq = tirx.const(1, "float32") / tirx.power(
+        theta, 2 * (d // 2) / tirx.const(d_range, "float32")
+    )
+    orig_freq_var = tirx.Var("orig_freq", "float32")
+
+    llama4_inv_scaling_factor = 1.0 / factor
+
+    if high_freq_factor == low_freq_factor:
+        wavelength = tirx.const(2 * math.pi, "float32") / orig_freq_var
+        threshold_wavelen = tirx.const(
+            original_max_position_embeddings / low_freq_factor, "float32"
+        )
+
+        scaled_freq = tirx.if_then_else(
+            wavelength > threshold_wavelen, orig_freq_var / factor, orig_freq_var
+        )
+        smoothed_freq = s * scaled_freq
+
+    else:
+        # Original smooth interpolation logic
+        inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+
+        llama4_alpha = original_max_position_embeddings / (2 * math.pi) * inv_diff_freq_factor
+        llama4_beta = low_freq_factor * inv_diff_freq_factor
+        smooth = tirx.max(0.0, tirx.min(1.0, llama4_alpha * orig_freq_var - llama4_beta))
+        smoothed_freq = s * (
+            (1.0 - smooth) * orig_freq_var * llama4_inv_scaling_factor + smooth * orig_freq_var
+        )
+
+    smoothed_freq_var = tirx.Var("smoothed_freq", "float32")
+    cos_freq = tirx.cos(smoothed_freq_var).astype(dtype)
+    sin_freq = tirx.sin(smoothed_freq_var).astype(dtype)
+    return (
+        cos_freq,
+        sin_freq,
+        {smoothed_freq_var: smoothed_freq, orig_freq_var: orig_freq},
+    )
+
+
+def rope_freq_llama3(  # pylint: disable=too-many-arguments,too-many-locals
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    factor: float,
+    low_freq_factor: float,
+    high_freq_factor: float,
+    original_max_position_embeddings: float,
+):
+    """Compute the inverse frequency of RoPE for llama3 RoPE scaling."""
+    orig_freq = tirx.const(1, "float32") / tirx.power(
+        theta, d * 2 % d_range / tirx.const(d_range, "float32")
+    )
+    orig_freq_var = tirx.Var("orig_freq", "float32")
+    inv_diff_freq_factor = 1.0 / (high_freq_factor - low_freq_factor)
+    llama3_inv_scaling_factor = 1.0 / factor
+    llama3_alpha = original_max_position_embeddings / (2 * math.pi) * inv_diff_freq_factor
+    llama3_beta = low_freq_factor * inv_diff_freq_factor
+    smooth = tirx.max(0.0, tirx.min(1.0, llama3_alpha * orig_freq_var - llama3_beta))
+    smoothed_freq = s * (
+        (1.0 - smooth) * orig_freq_var * llama3_inv_scaling_factor + smooth * orig_freq_var
+    )
+    smoothed_freq_var = tirx.Var("smoothed_freq", "float32")
+    cos_freq = tirx.cos(smoothed_freq_var).astype(dtype)
+    sin_freq = tirx.sin(smoothed_freq_var).astype(dtype)
+    return (
+        cos_freq,
+        sin_freq,
+        {smoothed_freq_var: smoothed_freq, orig_freq_var: orig_freq},
+    )
+
+
+def rope_freq_longrope(  # pylint: disable=too-many-arguments
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float,
+    dtype: str,
+    max_position_embeddings: int,
+    original_max_position_embeddings: int,
+    ext_factors: T.Buffer | None = None,
+):
+    """Compute the inverse frequency of RoPE for longrope scaling."""
+    scale = max_position_embeddings / original_max_position_embeddings
+    scaling_factor = (
+        math.sqrt(1 + math.log(scale) / math.log(original_max_position_embeddings))
+        if scale > 1.0
+        else 1.0
+    )
+    divisor = tirx.power(theta, d * 2 % d_range / tirx.const(d_range, "float32"))
+    if ext_factors is not None:
+        divisor = ext_factors[d % (d_range // 2)] * divisor
+    freq = s / divisor
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = (tirx.cos(freq_var) * scaling_factor).astype(dtype)
+    sin_freq = (tirx.sin(freq_var) * scaling_factor).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def yarn_find_correction_dim(
+    num_rotations: int,
+    d: tirx.Var,
+    max_position_embeddings: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):
+    """Inverse dim formula to find dim based on number of rotations"""
+    return (
+        d * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) * inv_theta_log_scale
+    )
+
+
+def yarn_find_correction_range(
+    low_rot: int,
+    high_rot: int,
+    d: tirx.Var,
+    max_position_embeddings: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):
+    """Find the correction range based on the number of rotations"""
+    low = yarn_find_correction_dim(
+        low_rot, d, max_position_embeddings, inv_theta_log_scale=inv_theta_log_scale
+    )
+    high = yarn_find_correction_dim(
+        high_rot, d, max_position_embeddings, inv_theta_log_scale=inv_theta_log_scale
+    )
+    return tirx.max(low, 0), tirx.min(high, d - 1)
+
+
+def rope_freq_yarn(
+    s: tirx.Var,
+    d: tirx.Var,
+    d_range: int,
+    theta: float | tirx.Expr,
+    dtype: str,
+    original_max_position_embeddings: int,
+    scaling_factor: float,
+    beta_fast: int,
+    beta_slow: int,
+    inv_theta_log_scale: float | tirx.Expr | None = None,
+):  # pylint: disable=too-many-arguments, too-many-locals
+    """Compute the inverse frequency of RoPE for yarn RoPE scaling."""
+
+    exponent = d * 2 % d_range / tirx.const(d_range, "float32")
+    freq_power = tirx.power(theta, exponent)
+    freq_extra = tirx.const(1, "float32") / freq_power
+    freq_inter = tirx.const(1, "float32") / (scaling_factor * freq_power)
+
+    low, high = yarn_find_correction_range(
+        beta_fast,
+        beta_slow,
+        d_range,
+        original_max_position_embeddings,
+        inv_theta_log_scale=inv_theta_log_scale,
+    )
+    high = tirx.if_then_else(low == high, high + 0.001, high)
+    inv_freq_mask = tirx.const(1, "float32") - tirx.max(
+        tirx.min((d - low) / (high - low), 1.0), 0.0
+    ).astype("float32")
+    inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask
+    freq = s * inv_freq
+    freq_var = tirx.Var("freq", "float32")
+    cos_freq = tirx.cos(freq_var).astype(dtype)
+    sin_freq = tirx.sin(freq_var).astype(dtype)
+    return cos_freq, sin_freq, {freq_var: freq}
+
+
+def switch_rope_freq_func(rope_scaling: dict[str, Any]) -> Callable:
+    """Return the RoPE inverse frequency computation function based
+    on the given RoPE scaling.
+    """
+    if "rope_type" not in rope_scaling:
+        return rope_freq_default
+    if rope_scaling["rope_type"] == "gptj":
+        return rope_freq_gptj
+    if rope_scaling["rope_type"] == "llama3":
+        return partial(
+            rope_freq_llama3,
+            factor=rope_scaling["factor"],
+            low_freq_factor=rope_scaling["low_freq_factor"],
+            high_freq_factor=rope_scaling["high_freq_factor"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "llama4":
+        return partial(
+            rope_freq_llama4,
+            factor=rope_scaling["factor"],
+            low_freq_factor=rope_scaling["low_freq_factor"],
+            high_freq_factor=rope_scaling["high_freq_factor"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "longrope":
+        return partial(
+            rope_freq_longrope,
+            max_position_embeddings=rope_scaling["max_position_embeddings"],
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+        )
+    if rope_scaling["rope_type"] == "yarn":
+        inv_theta_log_scale = rope_scaling.get("inv_theta_log_scale")
+        assert inv_theta_log_scale is not None, "inv_theta_log_scale must be precomputed for YaRN"
+        return partial(
+            rope_freq_yarn,
+            original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+            scaling_factor=rope_scaling["factor"],
+            beta_fast=rope_scaling["beta_fast"],
+            beta_slow=rope_scaling["beta_slow"],
+            inv_theta_log_scale=inv_theta_log_scale,
+        )
+    raise ValueError(f"Unsupported RoPE scaling type: {rope_scaling['rope_type']}")
+
+
+# mypy: disable-error-code="attr-defined"
+
+
+def llama_rope(  # pylint: disable=too-many-arguments
+    qkv: Tensor,
+    total_seq_len: tirx.Var,
+    theta: float,
+    scale: float,
+    num_q_heads: int,
+    num_kv_heads: int,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Llama-style RoPE. Given a fused QKV tensor, it returns three tensors, Q, K, and V, where Q
+    and K are rotated by RoPE while V remains unchanged.
+
+    Parameters
+    ----------
+    qkv : Tensor
+        The fused QKV tensor of shape: [batch_size, seq_len, #q_heads + #kv_heads * 2, head_dim]
+
+    total_seq_len : tirx.Var
+        The total sequence length after being concatenated with KVCache. It is used to compute the
+        offset of RoPE.
+
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : Optional[int]
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+
+    Returns
+    -------
+    q : Tensor
+        The query tensor of shape [batch_size, seq_len, #q_heads, head_dim] w/ RoPE applied
+
+    k : Tensor
+        The key tensor of shape [batch_size, seq_len, #kv_heads, head_dim] w/ RoPE applied
+
+    v : Tensor
+        The value tensor of shape [batch_size, seq_len, #kv_heads, head_dim] w/o RoPE applied
+    """
+    _, _, fused_heads, head_dim = qkv.shape
+    assert fused_heads == num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    dtype = qkv.dtype
+    scale = tirx.const(scale, dtype)
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        b: tirx.Var,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        offset: tirx.Var,
+    ):
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            (s + offset) * scale, d, rotary_dim, theta, dtype
+        )
+        cos = cos_freq * x[b, s, h, d]
+        if rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[b, s, h, d + 1],
+                x[b, s, h, d - 1],
+            )
+        else:
+            sin = sin_freq * tirx.if_then_else(
+                d < rotary_dim // 2,
+                -x[b, s, h, d + rotary_dim // 2],
+                x[b, s, h, d - rotary_dim // 2],
+            )
+        expr = cos + sin
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(private=True, s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        total_seq_len: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        batch_size = T.int64()
+        seq_len = T.int64()
+        qkv = T.match_buffer(var_qkv, (batch_size, seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (batch_size, seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (batch_size, seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (batch_size, seq_len, num_kv_heads, head_dim), dtype)
+        for iters in T.grid(batch_size, seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                b, s, h, d = T.axis.remap("SSSS", iters)
+                if h < num_q_heads:
+                    q[b, s, h, d] = T.if_then_else(
+                        d < rotary_dim,
+                        _rope(qkv, b, s, h, d, total_seq_len - seq_len),
+                        qkv[b, s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[b, s, h - num_q_heads, d] = T.if_then_else(
+                        d < rotary_dim,
+                        _rope(qkv, b, s, h, d, total_seq_len - seq_len),
+                        qkv[b, s, h, d],
+                    )
+                else:
+                    v[b, s, h - (num_q_heads + num_kv_heads), d] = qkv[b, s, h, d]
+
+    b, s, _, _ = qkv.shape
+    return op.tensor_ir_op(  # pylint: disable=no-member
+        fused_rope,
+        "llama_rope",
+        args=[qkv, total_seq_len],
+        out=(
+            Tensor.placeholder((b, s, num_q_heads, head_dim), dtype),
+            Tensor.placeholder((b, s, num_kv_heads, head_dim), dtype),
+            Tensor.placeholder((b, s, num_kv_heads, head_dim), dtype),
+        ),
+    )
+
+
+def llama_rope_with_position_map(  # pylint: disable=too-many-arguments
+    theta: float,
+    scale: float,
+    head_dim: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    dtype: str,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+):
+    """Return the TIR function that computes Llama-style RoPE with q position map.
+
+    Parameters
+    ----------
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    head_dim : int
+        The number of features on each head.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    dtype : str
+        The dtype of qkv data.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : int
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+    """
+    fused_heads = num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    scale = tirx.const(scale, "float32")
+    is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
+    if is_longrope_scaling and "original_max_position_embeddings" in rope_scaling:
+        original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+    else:
+        original_max_position_embeddings = 0
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        pos: tirx.Var,
+        ext_factors: T.Buffer | None = None,
+    ):
+        kwargs = {}
+        if ext_factors:
+            kwargs["ext_factors"] = ext_factors
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            pos * scale, d, rotary_dim, theta, "float32", **kwargs
+        )
+        cos = cos_freq * x[s, h, d].astype("float32")
+        if "rope_type" in rope_scaling and rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        else:
+            sin = sin_freq * tirx.if_then_else(
+                d < rotary_dim // 2,
+                -x[s, h, d + rotary_dim // 2],
+                x[s, h, d - rotary_dim // 2],
+            ).astype("float32")
+        expr = (cos + sin).astype(dtype)
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        apply_rope: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int32()
+        position_map_elem_offset = T.int32()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        for iters in T.grid(seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                s, h, d = T.axis.remap("SSS", iters)
+                if h < num_q_heads:
+                    q[s, h, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[s, h - num_q_heads, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                else:
+                    v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    @T.prim_func(s_tir=True)
+    def fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        ext_factors: T.Buffer((rotary_dim,), "float32"),  # type: ignore
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int64()
+        position_map_elem_offset = T.int64()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        # long factors is the first half, short factors is the second half
+        long_factors = T.decl_buffer((rotary_dim // 2,), "float32", data=ext_factors.data)
+        short_factors = T.decl_buffer(
+            (rotary_dim // 2,),
+            "float32",
+            data=ext_factors.data,
+            elem_offset=(rotary_dim // 2),
+        )
+
+        if seq_len > original_max_position_embeddings:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+        else:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    if is_longrope_scaling:
+        return fused_rope_longrope_scaling
+    return fused_rope
+
+
+def llama4_rope_with_position_map(  # pylint: disable=too-many-arguments
+    theta: float,
+    scale: float,
+    head_dim: int,
+    num_q_heads: int,
+    num_kv_heads: int,
+    dtype: str,
+    rope_scaling: dict[str, Any],
+    rotary_dim: int | None = None,
+):
+    """Return the TIR function that computes Llama-style RoPE with q position map.
+
+    Parameters
+    ----------
+    theta : float
+        The theta value, or "base" in RoPE, which controls the frequency.
+
+    scale : float
+        The RoPE scaling factor.
+
+    head_dim : int
+        The number of features on each head.
+
+    num_q_heads : int
+        The number of query heads.
+
+    num_kv_heads : int
+        The number of key/value heads. It differs from `num_q_heads` in group-query attention.
+
+    dtype : str
+        The dtype of qkv data.
+
+    rope_scaling : Dict
+        The configuration of RoPE scaling.
+
+    rotary_dim : int
+        The number of dimensions in the embedding that RoPE is applied to. By default, the
+        rotary_dim is the same as head_dim.
+    """
+    fused_heads = num_q_heads + num_kv_heads * 2
+    if rotary_dim is None:
+        rotary_dim = head_dim
+    scale = tirx.const(scale, "float32")
+    is_longrope_scaling = rope_scaling.get("rope_type") == "longrope"
+    if is_longrope_scaling and "original_max_position_embeddings" in rope_scaling:
+        original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+    else:
+        original_max_position_embeddings = 0
+
+    def _rope(  # pylint: disable=too-many-arguments
+        x: T.Buffer,
+        s: tirx.Var,
+        h: tirx.Var,
+        d: tirx.Var,
+        pos: tirx.Var,
+        ext_factors: T.Buffer | None = None,
+    ):
+        kwargs = {}
+        if ext_factors:
+            kwargs["ext_factors"] = ext_factors
+        cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)(
+            pos * scale, d, rotary_dim, theta, "float32", **kwargs
+        )
+        cos = cos_freq * x[s, h, d].astype("float32")
+        if "rope_type" in rope_scaling and rope_scaling["rope_type"] == "gptj":
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        else:
+            # Data layout is different for llama4 vs llama3
+            sin = sin_freq * tirx.if_then_else(
+                d % 2 == 0,
+                -x[s, h, d + 1],
+                x[s, h, d - 1],
+            ).astype("float32")
+        expr = (cos + sin).astype(dtype)
+        for var, value in var_map.items():
+            expr = tirx.Let(var, value, expr)
+        return expr
+
+    @T.prim_func(private=True, s_tir=True)
+    def fused_rope(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        apply_rope: T.int64,
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int32()
+        position_map_elem_offset = T.int32()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        for iters in T.grid(seq_len, fused_heads, head_dim):
+            with T.sblock("llama_fused_rope"):
+                s, h, d = T.axis.remap("SSS", iters)
+                if h < num_q_heads:
+                    q[s, h, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                elif h < num_q_heads + num_kv_heads:
+                    k[s, h - num_q_heads, d] = T.if_then_else(
+                        apply_rope > 0 and d < rotary_dim,
+                        _rope(qkv, s, h, d, position_map[s]),
+                        qkv[s, h, d],
+                    )
+                else:
+                    v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    @T.prim_func(s_tir=True)
+    def fused_rope_longrope_scaling(  # pylint: disable=too-many-locals
+        var_qkv: T.handle,
+        var_position_map: T.handle,
+        var_q: T.handle,
+        var_k: T.handle,
+        var_v: T.handle,
+        ext_factors: T.Buffer((rotary_dim,), "float32"),  # type: ignore
+    ):
+        T.func_attr(
+            {
+                "op_pattern": 8,  # 2 means injective, 8 means opaque
+                "tirx.noalias": True,
+            }
+        )
+        seq_len = T.int64()
+        position_map_elem_offset = T.int64()
+        qkv = T.match_buffer(var_qkv, (seq_len, fused_heads, head_dim), dtype)
+        q = T.match_buffer(var_q, (seq_len, num_q_heads, head_dim), dtype)
+        k = T.match_buffer(var_k, (seq_len, num_kv_heads, head_dim), dtype)
+        v = T.match_buffer(var_v, (seq_len, num_kv_heads, head_dim), dtype)
+        position_map = T.match_buffer(
+            var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset
+        )
+        # long factors is the first half, short factors is the second half
+        long_factors = T.decl_buffer((rotary_dim // 2,), "float32", data=ext_factors.data)
+        short_factors = T.decl_buffer(
+            (rotary_dim // 2,),
+            "float32",
+            data=ext_factors.data,
+            elem_offset=(rotary_dim // 2),
+        )
+
+        if seq_len > original_max_position_embeddings:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                long_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+        else:
+            for iters in T.grid(seq_len, fused_heads, head_dim):
+                with T.sblock("llama_fused_rope"):
+                    s, h, d = T.axis.remap("SSS", iters)
+                    if h < num_q_heads:
+                        q[s, h, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    elif h < num_q_heads + num_kv_heads:
+                        k[s, h - num_q_heads, d] = T.if_then_else(
+                            d < rotary_dim,
+                            _rope(
+                                qkv,
+                                s,
+                                h,
+                                d,
+                                position_map[s],
+                                short_factors if is_longrope_scaling else None,
+                            ),
+                            qkv[s, h, d],
+                        )
+                    else:
+                        v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
+
+    if is_longrope_scaling:
+        return fused_rope_longrope_scaling
+    return fused_rope

python/tvm/relax/frontend/nn/spec.py

@@ -0,0 +1,257 @@
+# 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.
+"""Compilation specifications, for example, dynamic shape inputs."""
+
+import inspect
+import typing
+
+if typing.TYPE_CHECKING:
+    from .core import Module as nn_module_class
+
+ArgSpecType = typing.Union["Int", "Tensor"]
+MethodSpecType = typing.Union["MethodSpec", dict[str, ArgSpecType]]
+ModuleSpecType = typing.Union["ModuleSpec", dict[str, MethodSpecType]]
+SpecAny = typing.Union["Object", "Int", "Tensor", "Tuple"]
+
+
+class Int:  # pylint: disable=too-few-public-methods
+    """An integer input"""
+
+    def __init__(self) -> None:
+        pass
+
+    def __repr__(self) -> str:
+        return "int"
+
+
+class Tensor:  # pylint: disable=too-few-public-methods
+    """A tensor input with static ndim and dtype, but can have symbolic shapes."""
+
+    shape: list[int | str]
+    dtype: str
+
+    def __init__(self, shape: typing.Sequence[int | str], dtype: str) -> None:
+        self.shape = list(shape)
+        self.dtype = dtype
+
+    def __repr__(self) -> str:
+        shape = ", ".join(str(i) for i in self.shape)
+        return f"Tensor([{shape}], '{self.dtype}')"
+
+
+class Object:  # pylint: disable=too-few-public-methods
+    """An non-tensor opaque frontend object."""
+
+    object_type: type
+
+    def __init__(self, object_type: type) -> None:
+        self.object_type = object_type
+
+    def __repr__(self) -> str:
+        return "object"
+
+
+class Tuple:  # pylint: disable=too-few-public-methods
+    """A tuple input or a list input"""
+
+    name: str
+    elements: list[SpecAny] | tuple[SpecAny, ...]
+
+    def __init__(
+        self,
+        name: str,
+        elements: list[SpecAny] | tuple[SpecAny, ...],
+    ) -> None:
+        assert isinstance(elements, tuple | list), f"Unsupported container type: {type(elements)}"
+        self.name = name
+        self.elements = elements
+
+    def __repr__(self) -> str:
+        return self.elements.__repr__()
+
+
+class MethodSpec:
+    """A spec for a compiled method"""
+
+    method: typing.Callable
+    arg_names: list[str]
+    arg_specs: list[ArgSpecType]
+    param_mode: str  # "plain", "packed", "none"
+    effect_mode: str  # "plain", "packed", "none"
+
+    def __init__(  # pylint: disable=too-many-arguments
+        self,
+        method: typing.Callable,
+        arg_names: list[str],
+        arg_specs: list[ArgSpecType],
+        param_mode: str,
+        effect_mode: str,
+    ):
+        if param_mode not in ["plain", "packed", "none"]:
+            raise ValueError(f"Invalid param_mode: {param_mode}")
+        if effect_mode not in ["plain", "packed", "none"]:
+            raise ValueError(f"Invalid effect_mode: {effect_mode}")
+        self.method = method
+        self.arg_names = arg_names
+        self.arg_specs = arg_specs
+        self.param_mode = param_mode
+        self.effect_mode = effect_mode
+
+    def _repr(self, name: str) -> str:
+        args = ", ".join(
+            f"{name}: {spec}"
+            for name, spec in zip(
+                self.arg_names,
+                self.arg_specs,
+            )
+        )
+        return f"{name}({args})"
+
+    def __repr__(self) -> str:
+        return self._repr(name="MethodSpec")
+
+    @staticmethod
+    def from_raw(spec: MethodSpecType, method: typing.Callable) -> "MethodSpec":
+        """Create MethodSpec from raw python dictionaries.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            MethodSpec.from_raw(
+                spec={
+                    "inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
+                    "total_seq_len": "int",
+                },
+                method=module.prefill,
+            )
+        """
+        if isinstance(spec, MethodSpec):
+            return spec
+        config: dict[str, typing.Any] = spec.pop("$", {})  # type: ignore[assignment]
+        param_mode = config.get("param_mode", "plain")
+        effect_mode = config.get("effect_mode", "plain")
+        method_signature = inspect.signature(method)
+        arg_names = list(method_signature.parameters.keys())
+        arg_specs = []
+
+        def _convert_arg_spec(arg_spec, arg_name):
+            if arg_spec is Int or arg_spec is int:
+                return Int()
+            if isinstance(arg_spec, str) and arg_spec == "int":
+                return Int()
+            if isinstance(arg_spec, Int | Tensor | Object):
+                return arg_spec
+            if isinstance(arg_spec, tuple | list | Tuple):
+                return Tuple(
+                    arg_name,
+                    elements=type(arg_spec)(
+                        [
+                            _convert_arg_spec(arg_spec_i, f"{arg_name}_{i}")
+                            for i, arg_spec_i in enumerate(arg_spec)
+                        ]
+                    ),
+                )
+            raise TypeError(f"Invalid spec for argument {arg_name}: {arg_spec}")
+
+        for arg_name in arg_names:
+            if arg_name in spec:
+                arg_spec = spec[arg_name]
+                arg_spec = _convert_arg_spec(arg_spec, arg_name)
+                arg_specs.append(arg_spec)
+        return MethodSpec(
+            method,
+            arg_names,
+            arg_specs,
+            param_mode=param_mode,
+            effect_mode=effect_mode,
+        )
+
+    @staticmethod
+    def from_torch(args: list[typing.Any], method: typing.Callable) -> "MethodSpec":
+        """Converts a list of torch tensors to MethodSpec."""
+        from .torch import (  # pylint: disable=import-outside-toplevel
+            _method_spec_from_torch,
+        )
+
+        return _method_spec_from_torch(args, method)
+
+
+class ModuleSpec:
+    """A spec for a compiled nn.Module"""
+
+    module: "nn_module_class"
+    method_names: list[str]
+    method_specs: list[MethodSpec]
+
+    def __init__(
+        self,
+        module: "nn_module_class",
+        method_names: list[str],
+        method_specs: list[MethodSpec],
+    ) -> None:
+        self.module = module
+        self.method_names = method_names
+        self.method_specs = method_specs
+
+    @staticmethod
+    def from_raw(spec: ModuleSpecType, module: "nn_module_class") -> "ModuleSpec":
+        """Create ModuleSpec from raw python dictionaries.
+
+        Examples
+        --------
+        .. code-block:: python
+
+            ModuleSpec.from_raw(
+                spec={
+                    "prefill": {
+                        "inputs": spec.Tensor([batch_size, "seq_len"], "int32"),
+                        "total_seq_len": int,
+                    },
+                    "decode": {
+                        "inputs": spec.Tensor([batch_size, 1], "int32"),
+                        "total_seq_len": int,
+                    },
+                    "softmax_with_temperature": {
+                        "logits": spec.Tensor([1, 1, config.vocab_size], "float32"),
+                        "temperature": spec.Tensor([], "float32"),
+                    },
+                },
+                module=module,
+            )
+        """
+        if isinstance(spec, ModuleSpec):
+            return spec
+        method_names = list(spec.keys())
+        method_specs: list[MethodSpec] = []
+        for method_name in method_names:
+            method_spec = spec[method_name]
+            if isinstance(method_spec, MethodSpec):
+                pass
+            else:
+                method_spec = MethodSpec.from_raw(method_spec, getattr(module, method_name))
+            method_specs.append(method_spec)
+        return ModuleSpec(module, method_names, method_specs)
+
+    def __repr__(self) -> str:
+        return "ModuleSpec:\n" + "\n".join(
+            "  " + spec._repr(name)  # pylint: disable=protected-access
+            for name, spec in zip(
+                self.method_names,
+                self.method_specs,
+            )
+        )

python/tvm/relax/frontend/nn/subroutine.py

@@ -0,0 +1,188 @@
+# 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=too-many-lines,invalid-name,protected-access
+"""nn.Module mixin for subroutine dispatch"""
+
+import collections
+import contextlib
+import functools
+import inspect
+import re
+import typing
+
+import tvm_ffi
+
+from tvm import ir, relax
+from tvm.relax.frontend import nn
+
+
+def _camel_to_snake(name):
+    """Convert from CamelCase to snake_case"""
+
+    # Adapted from https://stackoverflow.com/a/1176023
+    name = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
+    name = re.sub("([a-z0-9])([A-Z])", r"\1_\2", name)
+    name = name.lower()
+    return name
+
+
+def _normalize_expr(block_builder, arg, as_relax_expr=False):
+    """Ensure that an argument is a relax.Expr with type"""
+    if isinstance(arg, tuple):
+        arg = relax.Tuple([_normalize_expr(block_builder, element) for element in arg])
+
+    if isinstance(arg, relax.Expr) and arg.ty.is_missing():
+        arg = block_builder.emit(arg)
+
+    if isinstance(arg, nn.Tensor) and as_relax_expr:
+        arg = arg._expr
+
+    return arg
+
+
+def _get_ty(arg):
+    if isinstance(arg, relax.Expr):
+        return arg.ty
+    elif isinstance(arg, nn.Tensor):
+        return arg._expr.ty
+    elif isinstance(arg, tuple | list | tvm_ffi.Array):
+        return relax.TupleType([_get_ty(field) for field in arg])
+    else:
+        raise TypeError(f"Cannot find type for {arg} of type {type(arg)}")
+
+
+class SubroutineMixin:
+    """A mixin that generates a
+
+    Contains common logic for `tvm.relax.frontend.nn.Module` and
+    `tvm.relax.testing.nn.Module`.
+    """
+
+    define_subroutine: bool = False
+
+    def __init_subclass__(cls):
+        """Update the cls.forward of subclasses"""
+        if hasattr(cls, "forward"):
+            is_wrapped = getattr(cls.forward, "_is_subroutine_mixin", False)
+            if not is_wrapped:
+                cls.forward = cls._subroutine_dispatch(cls.forward)
+
+    @classmethod
+    def _subroutine_dispatch(cls, old_forward):
+        @functools.wraps(old_forward)
+        def new_forward(self, *args, **kwargs):
+            if not self.define_subroutine:
+                return old_forward(self, *args, **kwargs)
+
+            block_builder = relax.BlockBuilder.current()
+            assert block_builder is not None, (
+                f"Class {type(self)} has cls.define_subroutines = True, "
+                "but is called outsdie of a block_builder environment.  "
+                "relax.BlockBuilder.current() is required "
+                "to determine where to generate the subroutine."
+            )
+
+            func_args = self._normalize_subroutine_args(block_builder, *args, **kwargs)
+            subroutine, is_nn_tensor_output = self._get_subroutine(
+                block_builder, old_forward, func_args
+            )
+            subroutine_args = [
+                arg._expr if isinstance(arg, nn.Tensor) else arg
+                for arg in [*func_args.values(), *self.parameters()]
+            ]
+
+            out = subroutine(*subroutine_args)
+
+            if is_nn_tensor_output:
+                if out.ty.is_missing():
+                    out = block_builder.emit(out, name_hint=f"{subroutine.name_hint}_output")
+                out = nn.Tensor(_expr=out)
+            return out
+
+        new_forward._is_subroutine_mixin = True
+
+        return new_forward
+
+    def _normalize_subroutine_args(
+        self, block_builder, *args, **kwargs
+    ) -> typing.OrderedDict[str, relax.Expr]:
+        signature = inspect.signature(self.forward)
+        bindings = signature.bind(*args, **kwargs)
+        func_args = collections.OrderedDict(
+            (name, _normalize_expr(block_builder, arg)) for name, arg in bindings.arguments.items()
+        )
+        return func_args
+
+    def _get_subroutine(
+        self,
+        block_builder,
+        old_forward: typing.Callable,
+        func_args: typing.OrderedDict[str, relax.Expr],
+    ) -> (ir.GlobalVar, bool):
+        cls = type(self)
+        if not hasattr(cls, "_gvar"):
+            cls._gvar = {}
+
+        model_params = [
+            param._expr if isinstance(param, nn.Tensor) else param for param in self.parameters()
+        ]
+
+        arg_ty = _get_ty([*func_args.values(), *model_params])
+        is_dataflow = block_builder.current_block_is_dataflow()
+        lookup_key = (
+            old_forward,
+            tvm_ffi.structural_hash(arg_ty, map_free_vars=True),
+            is_dataflow,
+        )
+
+        for cached_ty, cached_result in cls._gvar.get(lookup_key, []):
+            if tvm_ffi.structural_equal(cached_ty, arg_ty, map_free_vars=True):
+                return cached_result
+
+        func_name = _camel_to_snake(cls.__name__)
+        func_params = [relax.Var(name, ty) for name, ty in zip(func_args, arg_ty.fields)]
+        old_forward_args = [
+            nn.Tensor(_expr=param) if isinstance(old_arg, nn.Tensor) else param
+            for param, old_arg in zip(func_params, func_args.values())
+        ]
+
+        with block_builder.function(func_name, [*func_params, *model_params], private=True):
+            with contextlib.ExitStack() as stack:
+                if is_dataflow:
+                    stack.enter_context(block_builder.dataflow())
+
+                out = old_forward(self, *old_forward_args)
+                is_nn_tensor_output = isinstance(out, nn.Tensor)
+                if is_nn_tensor_output:
+                    out = out._expr
+
+                if is_dataflow:
+                    out = block_builder.emit_output(out)
+            gvar = block_builder.emit_func_output(out)
+
+        # The relax.Var instances in model_params, along with any
+        # tirx.Var instances in the type, appear in both the
+        # calling scope and as parameters for the subroutine.  To
+        # maintain SSA, replace all relax and TIR variables in the
+        # subroutine.
+        mod = block_builder.get()
+        mod.update_func(gvar, relax.utils.copy_with_new_vars(mod[gvar]))
+
+        result = (gvar, is_nn_tensor_output)
+        bucket = cls._gvar.setdefault(lookup_key, [])
+        bucket.append((arg_ty, result))
+        return result

python/tvm/relax/frontend/nn/torch.py

@@ -0,0 +1,136 @@
+# 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.
+"""PyTorch integration with nn.Module"""
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from tvm_ffi import Array, Shape
+
+from tvm.runtime import Tensor, _tensor
+from tvm.runtime.vm import VirtualMachine
+
+from . import core
+from . import spec as _spec
+
+
+class TorchModule:  # pylint: disable=too-few-public-methods
+    """A wrapper on top of TVM VirtualMachine that takes torch tensors as inputs and returns torch
+    tensors as outputs"""
+
+    spec: _spec.ModuleSpec
+    vm: VirtualMachine  # pylint: disable=invalid-name
+    params: list[Tensor]
+    effects: list[Any]
+
+    def __init__(  # pylint: disable=invalid-name
+        self,
+        spec: _spec.ModuleSpec,
+        vm: VirtualMachine,
+        params: list[Tensor],
+    ):
+        try:
+            self.effects = vm["_initialize_effect"]()
+        except AttributeError:
+            self.effects = None
+
+        self.spec = spec
+        self.vm = vm
+        self.params = params
+
+    def __getitem__(self, method_name: str) -> Callable:
+        def _find_method(method_name):
+            for key, value in zip(self.spec.method_names, self.spec.method_specs):
+                if method_name == key:
+                    return value
+            raise ValueError(f"Method `{method_name}` is not found in the module spec. {self.spec}")
+
+        method_spec = _find_method(method_name)
+        method = self.vm[method_name]
+
+        def _closure(*args):
+            if len(args) != len(method_spec.arg_names):
+                raise TypeError(
+                    f"Argument length mismatch. Expected {len(method_spec.arg_names)} arguments, "
+                    f"but got {len(args)} arguments. The spec is: {method_spec}"
+                )
+            args = [
+                _torch_to_tvm(arg_name, arg_spec, arg)
+                for arg_name, arg_spec, arg in zip(
+                    method_spec.arg_names, method_spec.arg_specs, args
+                )
+            ]
+            if self.effects is not None:
+                outputs, self.effects = method(*args, *self.effects, *self.params)
+            else:
+                outputs = method(*args, *self.params)
+            return _tvm_to_torch(outputs)
+
+        _closure.__name__ = method_name
+        return _closure
+
+
+def _tvm_to_torch(arg):
+    if isinstance(arg, list | tuple | Array):
+        return [_tvm_to_torch(i) for i in arg]
+    if isinstance(arg, _tensor.Tensor):
+        return torch.utils.dlpack.from_dlpack(arg)
+    if isinstance(arg, Shape):
+        return list(arg)
+    raise TypeError(f"Unsupported argument type: {type(arg)}")
+
+
+def _torch_to_tvm(arg_name, arg_spec, arg_torch):
+    if isinstance(arg_spec, _spec.Tensor):
+        if not isinstance(arg_torch, torch.Tensor):
+            raise TypeError(
+                f"Expected argument `{arg_name}` to be `torch.Tensor`, but got {type(arg_torch)}"
+            )
+        return core._from_dlpack(arg_torch)  # pylint: disable=protected-access
+    if isinstance(arg_spec, _spec.Int):
+        if not isinstance(arg_torch, int):
+            raise TypeError(
+                f"Expected argument `{arg_name}` to be `int`, but got {type(arg_torch)}"
+            )
+        return Shape([arg_torch])
+    if isinstance(arg_spec, _spec.Tuple):
+        return [
+            _torch_to_tvm(f"{arg_name}[{i}]", x, arg_torch[i])
+            for i, x in enumerate(arg_spec.elements)
+        ]
+    raise TypeError(f"Unsupported spec item type: {type(arg_spec)}")
+
+
+def _method_spec_from_torch(
+    args_torch: list[Any],
+    method: Callable,
+):
+    def _as_spec(arg_torch):
+        if isinstance(arg_torch, torch.Tensor):
+            _, dtype = str(arg_torch.dtype).rsplit(".", maxsplit=1)
+            return _spec.Tensor(shape=list(arg_torch.shape), dtype=dtype)
+        if isinstance(arg_torch, int):
+            return _spec.Int()
+        raise TypeError(f"Unsupported argument type: {type(arg_torch)}")
+
+    arg_names = list(inspect.signature(method).parameters.keys())
+    if len(arg_names) != len(args_torch):
+        raise TypeError(f"Expected {len(arg_names)} arguments, but got {len(args_torch)} arguments")
+    arg_specs = [_as_spec(i) for i in args_torch]
+    return _spec.MethodSpec(method, arg_names, arg_specs, param_mode="plain", effect_mode="plain")

python/tvm/relax/frontend/nn/visitor.py

@@ -0,0 +1,234 @@
+# 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 visitor and mutator infra for nn.Module."""
+
+from typing import Any
+
+from . import core as nn
+
+
+class Mutator:
+    """The mutator for nn.Module transform. Users can override the `visit_*` methods
+    to apply transform in different structures, or even override the `visit` method
+    to change the logic of traversal."""
+
+    def visit_module(self, name: str, node: nn.Module) -> Any:
+        """The base visiting method for mutation of nn.Module nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.Module
+            The current node of nn.Module to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_effect(self, name: str, node: nn.Parameter) -> Any:
+        """The base visiting method for mutation of nn.Parameter nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.Parameter
+            The current node of nn.Parameter to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_param(self, name: str, node: nn.Effect) -> Any:
+        """The base visiting method for mutation of nn.Effect nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.Effect
+            The current node of nn.Effect to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_moduledict(self, name: str, node: nn.ModuleDict) -> Any:
+        """The base visiting method for mutation of nn.ModuleDict nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ModuleDict
+            The current node of nn.ModuleDict to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_modulelist(self, name: str, node: nn.ModuleList) -> Any:
+        """The base visiting method for mutation of nn.ModuleList nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ModuleList
+            The current node of nn.ModuleList to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_parameterdict(self, name: str, node: nn.ParameterDict) -> Any:
+        """The base visiting method for mutation of nn.ParameterDict nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ParameterDict
+            The current node of nn.ParameterDict to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit_parameterlist(self, name: str, node: nn.ParameterList) -> Any:
+        """The base visiting method for mutation of nn.ParameterList nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : nn.ParameterList
+            The current node of nn.ParameterList to mutate.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+        return self.visit(name, node)
+
+    def visit(self, name: str, node: Any) -> Any:
+        """The base dispatching method for visiting of all nodes.
+
+        Parameters
+        ----------
+        name : str
+            The name of the current node in parent's attribute.
+
+        node : Any
+            The current node to visit.
+
+        Returns
+        ------
+        ret_node: Any
+            The new node to replace current node.
+        """
+
+        def _get_child_name(parent: str, child: str) -> str:
+            """Get the name of the child node/key given the parent's name."""
+            if parent == "":
+                # in the top level of the module
+                return child
+            else:
+                return f"{parent}.{child}"
+
+        if isinstance(node, nn.ParameterList):
+            for i in range(len(node)):
+                node[i] = self.visit_param(_get_child_name(name, str(i)), node[i])
+        elif isinstance(node, nn.ParameterDict):
+            for k, v in node.items():
+                node[k] = self.visit_param(_get_child_name(name, k), v)
+        elif isinstance(node, nn.ModuleList):
+            for i in range(len(node)):
+                if isinstance(node[i], nn.ParameterDict):
+                    node[i] = self.visit_parameterdict(_get_child_name(name, str(i)), node[i])
+                elif isinstance(node[i], nn.ParameterList):
+                    node[i] = self.visit_parameterlist(_get_child_name(name, str(i)), node[i])
+                elif isinstance(node[i], nn.ModuleDict):
+                    node[i] = self.visit_moduledict(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.ModuleList):
+                    node[i] = self.visit_modulelist(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.Module):
+                    node[i] = self.visit_module(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.Effect):
+                    node[i] = self.visit_effect(f"{name}.{i}", node[i])
+                elif isinstance(node[i], nn.Parameter):
+                    node[i] = self.visit_param(f"{name}.{i}", node[i])
+        elif isinstance(node, nn.ModuleDict):
+            for k, v in node.items():
+                if isinstance(v, nn.ParameterDict):
+                    node[k] = self.visit_parameterdict(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ParameterList):
+                    node[k] = self.visit_parameterlist(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ModuleDict):
+                    node[k] = self.visit_moduledict(_get_child_name(name, k), v)
+                elif isinstance(v, nn.ModuleList):
+                    node[k] = self.visit_modulelist(_get_child_name(name, k), v)
+                elif isinstance(v, nn.Module):
+                    node[k] = self.visit_module(_get_child_name(name, k), v)
+                elif isinstance(v, nn.Effect):
+                    node[k] = self.visit_effect(_get_child_name(name, k), v)
+                elif isinstance(v, nn.Parameter):
+                    node[k] = self.visit_param(_get_child_name(name, k), v)
+        else:
+            for key, value in node.__dict__.items():
+                if isinstance(value, nn.ParameterDict):
+                    setattr(node, key, self.visit_parameterdict(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ParameterList):
+                    setattr(node, key, self.visit_parameterlist(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ModuleDict):
+                    setattr(node, key, self.visit_moduledict(_get_child_name(name, key), value))
+                elif isinstance(value, nn.ModuleList):
+                    setattr(node, key, self.visit_modulelist(_get_child_name(name, key), value))
+                elif isinstance(value, nn.Module):
+                    setattr(node, key, self.visit_module(_get_child_name(name, key), value))
+                elif isinstance(value, nn.Effect):
+                    setattr(node, key, self.visit_effect(_get_child_name(name, key), value))
+                elif isinstance(value, nn.Parameter):
+                    setattr(node, key, self.visit_param(_get_child_name(name, key), value))
+        return node

python/tvm/relax/frontend/onnx/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+"""
+Tools for converting ONNX graphs into Relax graphs.
+"""
+
+from .onnx_frontend import from_onnx

python/tvm/relax/frontend/stablehlo/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+"""
+StableHLO Frontends for constructing Relax programs, with the model importers
+"""
+
+from .stablehlo_translator import from_stablehlo

python/tvm/relax/frontend/stablehlo/stablehlo_translator.py

@@ -0,0 +1,445 @@
+# 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=import-outside-toplevel, unused-argument
+
+"""StableHLO frontend of Relax."""
+
+from collections.abc import Callable
+from typing import Any
+
+import tvm
+from tvm import relax, tirx
+
+
+class StableHLOImporter:
+    """An importer from StableHLO to Relax."""
+
+    from jaxlib import mlir
+    from jaxlib.mlir.dialects import stablehlo
+
+    def __init__(self) -> None:
+        from jaxlib import mlir
+
+        self._nodes: dict[str | mlir.ir.Operation, relax.Expr] = {}
+        self.block_builder: relax.BlockBuilder = None
+        self.create_convert_map()
+
+    @staticmethod
+    def _convert_data_type(input_type):
+        """converts the data type from mlir to tvm."""
+        from jaxlib import mlir
+
+        if mlir.ir.ShapedType.isinstance(input_type):
+            input_type = mlir.ir.ShapedType(input_type).element_type
+
+        input_type = str(input_type)
+        if input_type == "f16":
+            return "float16"
+        elif input_type in ["f32", "F32Type"]:
+            return "float32"
+        elif input_type in ["f64", "F64Type"]:
+            return "float64"
+        elif input_type == "i1":
+            return "bool"
+        elif input_type == "i8":
+            return "int8"
+        elif input_type == "i16":
+            return "int16"
+        elif input_type == "i32":
+            return "int32"
+        elif input_type == "i64":
+            return "int64"
+        elif input_type == "ui8":
+            return "uint8"
+        elif input_type == "ui16":
+            return "uint16"
+        elif input_type == "ui32":
+            return "uint32"
+        elif input_type == "ui64":
+            return "uint64"
+        else:
+            raise NotImplementedError(f"input_type {input_type} is not handled yet")
+
+    def _attr2value(self, node) -> Any | list[Any]:
+        import numpy as np
+        from jaxlib import mlir
+
+        if mlir.ir.IntegerAttr.isinstance(node):
+            int_attr = mlir.ir.IntegerAttr(node)
+            return int_attr.value
+        if mlir.ir.FloatAttr.isinstance(node):
+            float_attr = mlir.ir.FloatAttr(node)
+            return float_attr.value
+        if mlir.ir.DenseIntElementsAttr.isinstance(node):
+            dense_attr = mlir.ir.DenseIntElementsAttr(node)
+        elif mlir.ir.DenseFPElementsAttr.isinstance(node):
+            dense_attr = mlir.ir.DenseFPElementsAttr(node)
+        else:
+            raise ValueError("Unsupported Attribute type: " + str(type(node)))
+        ret = []
+        for val in dense_attr:
+            ret.append(val)
+        shape = self.get_shape(node.type)
+        dtype = self._convert_data_type(node.type)
+        return np.asarray(ret, dtype).reshape(shape).tolist()
+
+    def retrieve_operands(self, node):
+        return self._retrieve_operands(node.operands)
+
+    def _retrieve_operands(self, node):
+        from jaxlib import mlir
+
+        # the operand is one of the inputs of FuncOp
+        if isinstance(node, mlir.ir.Operation):
+            return self._nodes[node]
+        if isinstance(node, tuple):
+            return tuple(self._retrieve_operands(x) for x in node)
+        if isinstance(node, list | mlir.ir.OpOperandList):
+            return [self._retrieve_operands(x) for x in node]
+        if isinstance(node, dict):
+            return {self._retrieve_operands(k): self._retrieve_operands(v) for k, v in node.items()}
+        if isinstance(node, mlir.ir.Value):
+            if isinstance(node.owner, mlir.ir.Block):
+                block_arg = mlir.ir.BlockArgument(node)
+                return self._nodes["arg" + str(block_arg.arg_number)]
+            return self._retrieve_operands(node.owner)
+        return node
+
+    def get_shape(self, inpt_type) -> list[Any]:
+        """Get the shape from Type like tensor<?x?xf32>"""
+        from jaxlib import mlir
+
+        shape_type = inpt_type
+        if isinstance(shape_type, mlir.ir.Type):
+            shape_type = mlir.ir.ShapedType(shape_type)
+        ret = []
+        for i in range(shape_type.rank):
+            # get_dim_size
+            if shape_type.is_dynamic_dim(i):
+                n = tirx.Var("n", "int64")
+                ret.append(n)
+            else:
+                ret.append(shape_type.get_dim_size(i))
+
+        return ret
+
+    @staticmethod
+    def _promote_binary_op_args(lhs, rhs):
+        if not isinstance(lhs, relax.Expr) and not isinstance(rhs, relax.Expr):
+            msg = "Both the lhs and the rhs are not expressions."
+            raise AssertionError(msg)
+        if isinstance(lhs, relax.Expr) and isinstance(rhs, relax.Expr):
+            return lhs, rhs
+        if isinstance(lhs, relax.Expr):
+            assert isinstance(lhs.ty, relax.TensorType)
+            return lhs, relax.const(rhs, lhs.ty.dtype)
+        assert isinstance(rhs.ty, relax.TensorType)
+        return relax.const(lhs, rhs.ty.dtype), rhs
+
+    def _call_binary_op(self, op, lhs, rhs):
+        lhs, rhs = StableHLOImporter._promote_binary_op_args(lhs, rhs)
+        return self.block_builder.emit(op(lhs, rhs))
+
+    def _add(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.add, lhs, rhs)
+        return lhs + rhs
+
+    def _maximum(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.maximum(lhs, rhs))
+
+    def _minimum(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.minimum(lhs, rhs))
+
+    def _divide(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.divide, lhs, rhs)
+        return lhs / rhs
+
+    def _multiply(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.multiply, lhs, rhs)
+        return lhs * rhs
+
+    def _subtract(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        if isinstance(lhs, relax.Var) or isinstance(rhs, relax.Var):
+            return self._call_binary_op(relax.op.subtract, lhs, rhs)
+        return lhs - rhs
+
+    def _broadcast_in_dim(self, node: mlir.ir.Operation) -> relax.Expr:
+        operands = self.retrieve_operands(node)
+        data = operands[0]
+        # broadcast_dims = self._attr2value(node.attributes["broadcast_dimensions"])
+        shape = self.get_shape(node.result.type)
+        # scalar
+        if len(shape) == 0:
+            return data
+        return self.block_builder.emit(relax.op.broadcast_to(data, shape))
+
+    def _const(self, node: mlir.ir.Operation) -> relax.Expr:
+        const_value = self._attr2value(node.attributes["value"])
+        dtype = self._convert_data_type(node.result.type)
+        return relax.const(const_value, dtype)
+
+    def _dot_general(self, node: mlir.ir.Operation) -> relax.Expr:
+        lhs, rhs = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.matmul(lhs, rhs))
+
+    def _convolution(self, node) -> relax.Expr:
+        from jaxlib import mlir
+
+        x, weight = self.retrieve_operands(node)
+        shaped_type = mlir.ir.ShapedType(node.result.type)
+        out_dtype = self._convert_data_type(shaped_type.element_type)
+        strides = self._attr2value(node.attributes["window_strides"])
+        padding = self._attr2value(node.attributes["padding"])
+        lhs_dilation = self._attr2value(node.attributes["lhs_dilation"])
+        rhs_dilation = self._attr2value(node.attributes["rhs_dilation"])
+        if len(lhs_dilation) > 0:
+            lhs_dilation = lhs_dilation[0]
+        if len(rhs_dilation) > 0:
+            rhs_dilation = rhs_dilation[0]
+        dilation = (lhs_dilation, rhs_dilation)
+        groups = self._attr2value(node.attributes["batch_group_count"])
+        conv2d = relax.op.nn.conv2d(
+            x,
+            weight,
+            strides=strides,
+            padding=padding[0],
+            dilation=dilation,
+            groups=groups,
+            data_layout="NHWC",
+            kernel_layout="HWIO",
+            out_dtype=out_dtype,
+        )
+
+        return self.block_builder.emit(conv2d)
+
+    def _reshape(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        if isinstance(data, list):
+            assert len(data) == 1
+            data = data[0]
+        new_shape = self.get_shape(node.result.type)
+        return self.block_builder.emit(relax.op.reshape(data, new_shape))
+
+    def _reduce(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        dimensions = self._attr2value(node.attributes["dimensions"])
+        if node.body is not None:
+            reducer_op = node.body.blocks[0].operations[0].OPERATION_NAME
+            assert reducer_op == "stablehlo.add", f"reducer {reducer_op} in reduce is not supported"
+        return self.block_builder.emit(relax.op.sum(data[0], axis=dimensions))
+
+    def _reduce_window(self, node: mlir.ir.Operation) -> relax.Expr:
+        operands = self.retrieve_operands(node)
+        window_dimensions = self._attr2value(node.attributes["window_dimensions"])
+        window_dilations = self._attr2value(node.attributes["window_dilations"])
+
+        if node.body is not None:
+            reducer_op = node.body.blocks[0].operations[0].OPERATION_NAME
+            assert reducer_op == "stablehlo.maximum", (
+                f"the reducer {reducer_op} in reduce_window is not supported"
+            )
+
+        pool_size = []
+        for i, window_dim in enumerate(window_dimensions):
+            if window_dim == 0:
+                pool_size.append(0)
+            else:
+                dilated_window_size = (window_dim - 1) * window_dilations[i] + 1
+                pool_size.append(dilated_window_size)
+        strides = self._attr2value(node.attributes["window_strides"])
+        # padding = self._attr2value(node.attributes["padding"])
+
+        # TODO (yongwww): Infer the layout automatically
+        layout = "NHWC"
+
+        ret = self.block_builder.emit(
+            relax.op.nn.max_pool2d(
+                operands[0],
+                pool_size=pool_size[1:3],  # HW
+                strides=strides[1:3],
+                padding=[1, 1],
+                dilation=window_dilations[1:3],
+                layout=layout,
+            )
+        )
+        return ret
+
+    def _rsqrt(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.rsqrt(data[0]))
+
+    def _sin(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sin(data[0]))
+
+    def _sinh(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sinh(data[0]))
+
+    def _cos(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.cos(data[0]))
+
+    def _cosh(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.cosh(data[0]))
+
+    def _sqrt(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.sqrt(data[0]))
+
+    def _round(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.round(data[0]))
+
+    def _exp(self, node: mlir.ir.Operation) -> relax.Expr:
+        data = self.retrieve_operands(node)
+        return self.block_builder.emit(relax.op.exp(data[0]))
+
+    def _return(self, node: mlir.ir.Operation) -> relax.Expr:
+        outputs = self.retrieve_operands(node)
+        return self.block_builder.emit_output(self.nodes[outputs])
+
+    def create_convert_map(self):
+        from jaxlib import mlir
+
+        self.convert_map: dict[str, Callable[[mlir.ir.Operation], relax.Var]] = {
+            "stablehlo.add": self._add,
+            "stablehlo.broadcast_in_dim": self._broadcast_in_dim,
+            "stablehlo.constant": self._const,
+            "stablehlo.convolution": self._convolution,
+            "stablehlo.cosine": self._cos,
+            "stablehlo.cosh": self._cosh,
+            "stablehlo.divide": self._divide,
+            "stablehlo.dot_general": self._dot_general,
+            "stablehlo.exponential": self._exp,
+            "stablehlo.maximum": self._maximum,
+            "stablehlo.minimum": self._minimum,
+            "stablehlo.multiply": self._multiply,
+            "stablehlo.reshape": self._reshape,
+            "stablehlo.reduce": self._reduce,
+            "stablehlo.reduce_window": self._reduce_window,
+            "stablehlo.round_nearest_afz": self._round,
+            "stablehlo.rsqrt": self._rsqrt,
+            "stablehlo.sine": self._sin,
+            "chlo.sinh": self._sinh,
+            "stablehlo.sqrt": self._sqrt,
+            "stablehlo.subtract": self._subtract,
+            "func.return": self._return,
+            "stablehlo.return": self._return,
+        }
+
+    def from_stablehlo(self, model, input_info: list[tuple[tuple[int], str]]) -> tvm.IRModule:
+        """Convert a StableHLO Module to a Relax program.
+
+        Parameters
+        ----------
+        model : mlir.ir.Module
+            The StableHLO Module to convert.
+
+        input_info : List[Tuple[Tuple[int], str]]
+            A list of shapes and data types of input tensors.
+
+        Returns
+        -------
+        output : tvm.IRModule
+            The result IRModule with entry function "main"
+        """
+        from jaxlib import mlir
+        from jaxlib.mlir.dialects import stablehlo
+
+        assert isinstance(model, mlir.ir.Module)
+        block: mlir.ir.Block = model.body.operations[0].regions[0].blocks[0]
+
+        # inputs of the function
+        inputs = []
+        for idx, arg in enumerate(block.arguments.types):
+            arg_shape = mlir.ir.ShapedType(arg)
+            ipt_shape = self.get_shape(arg_shape)
+            ipt_dtype = self._convert_data_type(arg_shape.element_type)
+            ipt_name = "arg" + str(idx)
+            ipt_var = relax.Var(f"arg{idx}", relax.TensorType(ipt_shape, ipt_dtype))
+            self._nodes[ipt_name] = ipt_var
+            inputs.append(ipt_var)
+
+        # TODO (yongwww): Handle mlir.ir.Module with multiple functions
+        # Initialize the block builder with a function and a dataflow block.
+        # Raise error if the input stablehlo op is impure
+        func_name = "main"
+        self.block_builder = relax.BlockBuilder()
+
+        with self.block_builder.function(name=func_name, params=inputs.copy()):
+            output = None
+            with self.block_builder.dataflow():
+                block = model.body.operations[0].regions[0].blocks[0]
+                for operation in block.operations:
+                    if isinstance(operation, mlir.dialects.func.ReturnOp | stablehlo.ReturnOp):
+                        operation = operation.operands[0].owner
+                        # TODO (yongwww): handle multiple outputs
+                        output = self.block_builder.emit_output(self._nodes[operation])
+                        break
+
+                    if isinstance(operation, mlir.ir.OpView):
+                        op_name = operation.operation.name
+                        assert op_name in self.convert_map, f"Unsupported operation {op_name}"
+                        self._nodes[operation] = self.convert_map[op_name](operation)
+                    else:
+                        raise ValueError(f"Unsupported op {operation}")
+            assert output is not None
+            self.block_builder.emit_func_output(output)
+
+        mod = self.block_builder.get()
+        return mod
+
+
+def from_stablehlo(
+    stablehlo_module,
+    input_info: list[tuple[tuple[int], str]] | None = None,
+) -> tvm.IRModule:
+    """Convert a StableHLO Module to a Relax program
+
+    Parameters
+    ----------
+    stablehlo_module : Union[str, mlir.ir.Module]
+        The StableHLO Module to convert.
+
+    input_info : List[Tuple[Tuple[int], str]]
+        A list of shapes and data types of input tensors.
+
+    Returns
+    -------
+    output : tvm.IRModule
+        The result IRModule with entry function "main"
+    """
+    from jax._src.interpreters import mlir as jax_mlir
+
+    if isinstance(stablehlo_module, str):
+        # TODO (yongwww): support the serialized bytecode format of StableHLO
+        # model using stablehlo.deserialize_portable_artifact(ir) if the python
+        # binding is ready
+        context = jax_mlir.make_ir_context()
+        stablehlo_module = jax_mlir.ir.Module.parse(stablehlo_module, context)
+    return StableHLOImporter().from_stablehlo(stablehlo_module, input_info)

python/tvm/relax/frontend/tflite/__init__.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.
+"""
+Tools for converting TFLite graphs into Relax graphs.
+"""
+
+from .tflite_frontend import from_tflite

python/tvm/relax/frontend/tflite/tflite_flexbuffer.py

@@ -0,0 +1,161 @@
+# 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, too-many-lines, import-outside-toplevel
+# pylint: disable=broad-exception-raised, use-list-literal
+"""Tensorflow lite frontend helper to parse custom options in Flexbuffer format."""
+
+import struct
+from enum import IntEnum
+
+
+class BitWidth(IntEnum):
+    """Flexbuffer bit width schema from flexbuffers.h"""
+
+    BIT_WIDTH_8 = 0
+    BIT_WIDTH_16 = 1
+    BIT_WIDTH_32 = 2
+    BIT_WIDTH_64 = 3
+
+
+class FlexBufferType(IntEnum):
+    """Flexbuffer type schema from flexbuffers.h"""
+
+    FBT_NULL = 0
+    FBT_INT = 1
+    FBT_UINT = 2
+    FBT_FLOAT = 3
+    # Types above stored inline, types below store an offset.
+    FBT_KEY = 4
+    FBT_STRING = 5
+    FBT_INDIRECT_INT = 6
+    FBT_INDIRECT_UINT = 7
+    FBT_INDIRECT_FLOAT = 8
+    FBT_MAP = 9
+    FBT_VECTOR = 10  # Untyped.
+    FBT_VECTOR_INT = 11  # Typed any size (stores no type table).
+    FBT_VECTOR_UINT = 12
+    FBT_VECTOR_FLOAT = 13
+    FBT_VECTOR_KEY = 14
+    FBT_VECTOR_STRING = 15
+    FBT_VECTOR_INT2 = 16  # Typed tuple (no type table, no size field).
+    FBT_VECTOR_UINT2 = 17
+    FBT_VECTOR_FLOAT2 = 18
+    FBT_VECTOR_INT3 = 19  # Typed triple (no type table, no size field).
+    FBT_VECTOR_UINT3 = 20
+    FBT_VECTOR_FLOAT3 = 21
+    FBT_VECTOR_INT4 = 22  # Typed quad (no type table, no size field).
+    FBT_VECTOR_UINT4 = 23
+    FBT_VECTOR_FLOAT4 = 24
+    FBT_BLOB = 25
+    FBT_BOOL = 26
+    FBT_VECTOR_BOOL = 36  # To Allow the same type of conversion of type to vector type
+
+
+class FlexBufferDecoder:
+    """
+    This implements partial flexbuffer deserialization to be able
+    to read custom options. It is not intended to be a general
+    purpose flexbuffer deserializer and as such only supports a
+    limited number of types and assumes the data is a flat map.
+    """
+
+    def __init__(self, buffer):
+        self.buffer = buffer
+
+    def indirect_jump(self, offset, byte_width):
+        """Helper function to read the offset value and jump"""
+        unpack_str = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[byte_width]
+        back_jump = struct.unpack(unpack_str, self.buffer[offset : offset + byte_width])[0]
+        return offset - back_jump
+
+    def decode_keys(self, end, size, byte_width):
+        """Decodes the flexbuffer type vector. Map keys are stored in this form"""
+        # Keys are strings here. The format is all strings separated by null, followed by back
+        # offsets for each of the string. For example, (str1)\0(str1)\0(offset1)(offset2) The end
+        # pointer is pointing at the end of all strings
+        keys = list()
+        for i in range(0, size):
+            offset_pos = end + i * byte_width
+            start_index = self.indirect_jump(offset_pos, byte_width)
+            str_size = self.buffer[start_index:].find(b"\0")
+            assert str_size != -1
+            s = self.buffer[start_index : start_index + str_size].decode("utf-8")
+            keys.append(s)
+        return keys
+
+    def decode_vector(self, end, size, byte_width):
+        """Decodes the flexbuffer vector"""
+        # Each entry in the vector can have different datatype. Each entry is of fixed length. The
+        # format is a sequence of all values followed by a sequence of datatype of all values. For
+        # example - (4)(3.56)(int)(float) The end here points to the start of the values.
+        # Each type byte contains: (type << 2) | bit_width, where bit_width determines actual size.
+        values = list()
+        for i in range(0, size):
+            value_type_pos = end + size * byte_width + i
+            value_type_packed = self.buffer[value_type_pos]
+            value_type = FlexBufferType(value_type_packed >> 2)
+            value_bit_width = BitWidth(value_type_packed & 3)
+            value_byte_width = 1 << value_bit_width
+            value_bytes = self.buffer[
+                end + i * byte_width : end + i * byte_width + value_byte_width
+            ]
+            if value_type == FlexBufferType.FBT_BOOL:
+                value = bool(value_bytes[0])
+            elif value_type == FlexBufferType.FBT_INT:
+                fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[value_byte_width]
+                value = struct.unpack(fmt, value_bytes)[0]
+            elif value_type == FlexBufferType.FBT_UINT:
+                fmt = {1: "<B", 2: "<H", 4: "<I", 8: "<Q"}[value_byte_width]
+                value = struct.unpack(fmt, value_bytes)[0]
+            elif value_type == FlexBufferType.FBT_FLOAT:
+                fmt = {4: "<f", 8: "<d"}[value_byte_width]
+                value = struct.unpack(fmt, value_bytes)[0]
+            else:
+                raise Exception
+            values.append(value)
+        return values
+
+    def decode_map(self, end, byte_width, parent_byte_width):
+        """Decodes the flexbuffer map and returns a dict"""
+        mid_loc = self.indirect_jump(end, parent_byte_width)
+        size_fmt = {1: "<b", 2: "<h", 4: "<i", 8: "<q"}[byte_width]
+        map_size = struct.unpack(size_fmt, self.buffer[mid_loc - byte_width : mid_loc])[0]
+
+        # Find keys
+        keys_offset = mid_loc - byte_width * 3
+        keys_end = self.indirect_jump(keys_offset, byte_width)
+        keys = self.decode_keys(keys_end, map_size, 1)
+
+        # Find values
+        values_end = self.indirect_jump(end, parent_byte_width)
+        values = self.decode_vector(values_end, map_size, byte_width)
+        return dict(zip(keys, values))
+
+    def decode(self):
+        """Decode the buffer. Decoding is partially implemented"""
+        root_end = len(self.buffer) - 1
+        root_byte_width = self.buffer[root_end]
+        root_end -= 1
+        root_packed_type = self.buffer[root_end]
+        root_end -= root_byte_width
+
+        root_type = FlexBufferType(root_packed_type >> 2)
+        byte_width = 1 << BitWidth(root_packed_type & 3)
+
+        if root_type == FlexBufferType.FBT_MAP:
+            return self.decode_map(root_end, byte_width, root_byte_width)
+        raise NotImplementedError("Flexbuffer Decoding is partially imlpemented.")

python/tvm/relax/frontend/torch/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+"""
+PyTorch Frontends for constructing Relax programs, with the model importers
+"""
+
+from .exported_program_translator import from_exported_program
+from .fx_translator import from_fx
+from .dynamo import relax_dynamo, dynamo_capture_subgraphs

python/tvm/relax/frontend/torch/dynamo.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.
+
+# pylint: disable=invalid-name, missing-function-docstring, not-callable
+# pylint: disable=import-outside-toplevel, unused-argument, use-list-literal
+# mypy: ignore-errors
+"""PyTorch Dynamo backend of Relax."""
+
+import functools
+
+import tvm_ffi
+
+import tvm
+from tvm.relax import build as relax_build
+
+from .fx_translator import from_fx
+
+
+def device_from_inputs(example_inputs):
+    for x in example_inputs:
+        if hasattr(x, "device"):
+            return x.device
+    return None
+
+
+def relax_dynamo(pipeline: tvm.transform.Pass | None = None):
+    """A helper function to create a relax backend.
+
+    Parameters
+    ----------
+    pipeline : Optional[tvm.transform.Pass]
+        The pipeline to be applied to the relax module before sent to build.
+
+    Returns
+    -------
+    backend : Callable[[torch.fx.GraphModule, List[torch.Tensor]], Callable]
+        The relax dynamo backend.
+    """
+
+    def _relax_backend(graph_module, example_inputs):
+        import torch  # type: ignore[import]
+
+        assert isinstance(graph_module, torch.fx.GraphModule)
+
+        def to_torch_tensor(nd_tensor):
+            """A helper function to transfer a Tensor to torch.tensor."""
+            if isinstance(nd_tensor, torch.Tensor):
+                # tvm-ffi #517 (Recursive DLPack container conversion) auto-converts
+                # ffi::Tensor items returned in containers back to torch.Tensor when
+                # the call site passed torch.Tensor inputs.
+                return nd_tensor
+            if isinstance(nd_tensor, tvm.runtime.Tensor):
+                return torch.from_numpy(nd_tensor.numpy())
+            elif isinstance(nd_tensor, tvm_ffi.Array):
+                return tuple(to_torch_tensor(x) for x in nd_tensor)
+            else:
+                raise ValueError(f"Unsupported type {type(nd_tensor)}")
+
+        graph_module.graph.eliminate_dead_code()
+
+        device = device_from_inputs(example_inputs)
+
+        assert len(example_inputs)
+
+        fake_inputs = []
+        if isinstance(example_inputs[0], torch._subclasses.fake_tensor.FakeTensor):
+            # Fake tensors
+            fake_inputs = example_inputs
+        else:
+            # Real tensors
+            for node in graph_module.graph.nodes:
+                if node.op != "placeholder":
+                    continue
+                if "grapharg" not in node.meta:
+                    continue
+                fake_tensor = node.meta["grapharg"].fake_tensor
+                if fake_tensor is None:
+                    continue
+                fake_inputs.append(fake_tensor)
+
+        input_info = []
+        shape_vars = {}
+        for tensor in fake_inputs:
+            shape = []
+            for s in tensor.shape:
+                if isinstance(s, torch.SymInt):
+                    if str(s) not in shape_vars:
+                        shape_vars[str(s)] = tvm.tirx.Var(str(s), "int64")
+                    shape.append(shape_vars[str(s)])
+                else:
+                    shape.append(s)
+            input_info.append((shape, tensor.dtype))
+
+        mod = from_fx(graph_module, input_info)
+
+        if device.type == "cuda":
+            dev = tvm.cuda(device.index)
+            target = tvm.target.Target("cuda")
+        else:
+            dev = tvm.cpu(0)
+            target = tvm.target.Target(llvm_target())
+
+        # invoke optimization pipeline.
+        if pipeline is None:
+            # get default pipeline
+            seq = tvm.relax.get_pipeline()
+        elif isinstance(pipeline, str):
+            # lookup by name
+            seq = tvm.relax.get_pipeline(pipeline)
+        else:
+            seq = pipeline
+
+        mod = mod.with_attr("target", target)
+        mod = seq(mod)
+
+        ex = relax_build(mod, target=target)
+
+        vm = tvm.relax.VirtualMachine(ex.mod, device=dev)
+
+        def exec_tvm(*i_args):
+            args = [a.contiguous() for a in i_args if isinstance(a, torch.Tensor)]
+            vm_args = list()
+            for arg in args:
+                if arg.requires_grad:
+                    arg = arg.detach()
+                if isinstance(arg, torch._subclasses.fake_tensor.FakeTensor):
+                    # Materialize a real (eager) Tensor
+                    arg = torch.randn(arg.shape, dtype=arg.dtype, device=device)
+                vm_args.append(arg)
+            outputs = vm["main"](*vm_args)
+            return to_torch_tensor(outputs)
+
+        return exec_tvm
+
+    return _relax_backend
+
+
+def dynamo_capture_subgraphs(model, *params, **kwargs) -> tvm.IRModule:
+    """Capture subgraphs of the PyTorch model using torch.compile into an IRModule.
+
+    Parameters
+    ----------
+    model : torch.nn.Module
+        The PyTorch model to be captured.
+
+    params : List[torch.Tensor]
+        The parameters of the PyTorch model.
+
+    keep_params_as_input : bool
+        Whether to keep model parameters as input variables of the captured Relax functions.
+
+    Returns
+    -------
+    output : ImporterOutput
+        The output of translation, including the translated IRModule.
+        If `keep_params_as_input` is true, the functions in the IRModule have an
+        attribute "params" that contains the weights of the input model. The
+        weights can be detached by `relax.frontend.detach_params`.
+    """
+    import torch  # type: ignore[import]
+    from torch import _dynamo as dynamo  # type: ignore[import]
+    from torch import fx  # type: ignore[import]
+
+    keep_params_as_input = "keep_params_as_input" in kwargs and kwargs["keep_params_as_input"]
+    kwargs.pop("keep_params_as_input", None)
+    mod = tvm.IRModule()
+
+    def _capture(graph_module: fx.GraphModule, example_inputs):
+        assert isinstance(graph_module, torch.fx.GraphModule)
+        input_info = [(tuple(tensor.shape), str(tensor.dtype)) for tensor in example_inputs]
+        mod_ = from_fx(
+            graph_module,
+            input_info,
+            keep_params_as_input=keep_params_as_input,
+            unwrap_unit_return_tuple=True,
+        )
+        new_name = f"subgraph_{len(mod.get_global_vars())}"
+        mod[new_name] = mod_["main"].with_attr("global_symbol", new_name)
+        return graph_module.forward
+
+    dynamo.reset()
+    compiled_model = torch.compile(model, backend=_capture)
+
+    with torch.no_grad():
+        compiled_model(*params, **kwargs)
+
+    return mod
+
+
+@functools.lru_cache(None)
+def llvm_target():
+    import platform
+    import subprocess
+
+    AVX512_TARGET = {"kind": "llvm", "mcpu": "skylake-avx512"}
+    AVX2_TARGET = {"kind": "llvm", "mcpu": "core-avx2"}
+    DEFAULT_TARGET = "llvm"
+
+    system = platform.system()
+
+    if system == "Linux":
+        try:
+            with open("/proc/cpuinfo") as f:
+                cpuinfo = f.read()
+            if "avx512" in cpuinfo:
+                return AVX512_TARGET
+            return AVX2_TARGET
+        except FileNotFoundError:
+            pass
+    elif system == "Darwin":
+        try:
+            result = subprocess.run(
+                ["sysctl", "-n", "machdep.cpu.features"],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+            if result.returncode == 0:
+                cpu_features = result.stdout.lower()
+                if "avx512" in cpu_features:
+                    return AVX512_TARGET
+                if "avx2" in cpu_features:
+                    return AVX2_TARGET
+        except (FileNotFoundError, subprocess.SubprocessError):
+            pass
+
+        if platform.machine() == "arm64":
+            return DEFAULT_TARGET
+
+    # Default fallback
+    return DEFAULT_TARGET

python/tvm/relax/ir/instrument.py

@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Common relax pass instrumentation across IR variants."""
+
+import tvm
+from tvm import relax
+
+
+@tvm.instrument.pass_instrument
+class WellFormedInstrument:
+    """An instrument that checks the input/output IRModule of the Pass
+    is well formed. It will skip specific passes, like Normalize.
+
+    Parameters
+    ----------
+    check_ty: bool
+
+        If True, validate the type in the module.  If False,
+        skip these checks.
+
+    validate_before_transform: bool
+
+        If True (default), perform a well-formed check before running
+        a transform.  If False, only perform the well-formed check
+        after running a transform.
+    """
+
+    def __init__(self, check_ty: bool = True, validate_before_transform: bool = True):
+        self.skip_pass_name = ["Normalize", "NormalizeGlobalVar", "ResolveGlobals"]
+        self.check_ty = check_ty
+        self.validate_before_transform = validate_before_transform
+
+    def run_before_pass(self, mod, pass_info):
+        if self.validate_before_transform:
+            self._check(mod, pass_info.name, "Before")
+
+    def run_after_pass(self, mod, pass_info):
+        self._check(mod, pass_info.name, "After")
+
+    def _check(self, mod, pass_name, name_prefix):
+        if pass_name not in self.skip_pass_name:
+            is_well_formed = relax.analysis.check_well_formed(mod, self.check_ty)
+            if not is_well_formed:
+                mod.show(name=f"{name_prefix}{pass_name}")
+            assert is_well_formed

python/tvm/relax/op/__init__.py

@@ -0,0 +1,240 @@
+# 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.
+# pylint: disable= redefined-builtin
+"""Relax core operators."""
+
+# Register operator gradient functions
+from . import _op_gradient, builtin, ccl, distributed, grad, image, memory, nn, op_attrs
+
+# Operators
+from .base import (
+    assert_op,
+    call_builtin_with_ctx,
+    call_dps_packed,
+    call_inplace_packed,
+    call_pure_packed,
+    call_py_func,
+    call_tir,
+    call_tir_inplace,
+    call_tir_with_grad,
+    hint_on_device,
+    invoke_closure,
+    invoke_pure_closure,
+    make_closure,
+    null_value,
+    print,
+    register_gradient,
+    shape_of,
+    shape_to_tensor,
+    size,
+    tensor_to_shape,
+    to_vdevice,
+)
+from .binary import (
+    add,
+    atan2,
+    bitwise_and,
+    bitwise_or,
+    bitwise_xor,
+    divide,
+    equal,
+    floor_divide,
+    log_add_exp,
+    floor_mod,
+    greater,
+    greater_equal,
+    left_shift,
+    less,
+    less_equal,
+    logical_and,
+    logical_or,
+    logical_xor,
+    maximum,
+    minimum,
+    mod,
+    multiply,
+    not_equal,
+    power,
+    right_shift,
+    subtract,
+)
+from .create import (
+    arange,
+    full,
+    full_like,
+    hamming_window,
+    ones,
+    ones_like,
+    eye,
+    eye_like,
+    tril,
+    triu,
+    zeros,
+    zeros_like,
+)
+from .datatype import astype, wrap_param
+from .index import dynamic_strided_slice, strided_slice, take
+from .linear_algebra import einsum, linear, matmul, outer
+from .manipulate import (
+    broadcast_to,
+    collapse_sum_like,
+    collapse_sum_to,
+    concat,
+    expand_dims,
+    flatten,
+    flip,
+    gather_elements,
+    gather_nd,
+    index_put,
+    index_tensor,
+    meshgrid,
+    layout_transform,
+    one_hot,
+    permute_dims,
+    repeat,
+    reshape,
+    reverse_sequence,
+    scatter_elements,
+    scatter_nd,
+    slice_scatter,
+    split,
+    squeeze,
+    stack,
+    tile,
+)
+from .mask import masked_fill
+from .qdq import dequantize, quantize
+from .sampling import multinomial_from_uniform
+from .search import argmax, argmin, where, bucketize
+from .set import nonzero, unique
+from .sorting import argsort, sort, topk
+from .statistical import cumprod, cumsum, max, mean, min, prod, std, sum, variance, median
+from .ternary import ewise_fma
+from .unary import (
+    abs,
+    acos,
+    acosh,
+    asin,
+    asinh,
+    atan,
+    atanh,
+    bitwise_not,
+    ceil,
+    clip,
+    cos,
+    cosh,
+    erf,
+    exp,
+    floor,
+    isfinite,
+    isinf,
+    isnan,
+    log,
+    logical_not,
+    negative,
+    round,
+    rsqrt,
+    sigmoid,
+    sign,
+    sin,
+    sinh,
+    sqrt,
+    square,
+    tan,
+    tanh,
+    trunc,
+)
+from .vision import (
+    all_class_non_max_suppression,
+    get_valid_counts,
+    multibox_transform_loc,
+    non_max_suppression,
+    roi_align,
+    roi_pool,
+)
+
+
+def _register_op_make():
+    # pylint: disable=import-outside-toplevel
+    from .. import expr
+    from tvm.ir import _tensor_expr_overload
+    from . import _ffi_api
+
+    expr._op_ffi_api = _ffi_api  # type: ignore
+
+    def _add(lhs, rhs):
+        if isinstance(lhs.ty, expr.tvm.relax.TupleType) and isinstance(rhs, tuple):
+            return tuple([*lhs, *rhs])
+        return expr._binary_op_helper(lhs, rhs, _ffi_api.add)
+
+    def _rhs(_lhs, rhs):
+        return expr._binary_rhs_helper(rhs)
+
+    def _getitem(value, index):
+        try:
+            return expr.TupleGetItem(value, index)
+        except RuntimeError as err:
+            if "Index out of bounds" in err.args[0]:
+                raise IndexError from err
+            raise
+
+    _tensor_expr_overload.astype = lambda lhs, dtype, _span=None: _ffi_api.astype(lhs, dtype)
+    _tensor_expr_overload.__call__ = lambda func, *args, attrs=None: expr.tvm.ir.Call(
+        func, args, attrs=attrs
+    )
+    _tensor_expr_overload.__getitem__ = _getitem
+    _tensor_expr_overload.__neg__ = lambda lhs: _ffi_api.negative(lhs)
+    _tensor_expr_overload.__lt__ = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.less)
+    _tensor_expr_overload.__le__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.less_equal
+    )
+    _tensor_expr_overload.__gt__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.greater
+    )
+    _tensor_expr_overload.__ge__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.greater_equal
+    )
+    _tensor_expr_overload.__add__ = _add
+    _tensor_expr_overload.__radd__ = _add
+    _tensor_expr_overload.__sub__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.subtract
+    )
+    _tensor_expr_overload.__rsub__ = _rhs
+    _tensor_expr_overload.__mul__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.multiply
+    )
+    _tensor_expr_overload.__rmul__ = _tensor_expr_overload.__mul__
+    _tensor_expr_overload.__div__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.divide
+    )
+    _tensor_expr_overload.__rdiv__ = _rhs
+    _tensor_expr_overload.__truediv__ = _tensor_expr_overload.__div__
+    _tensor_expr_overload.__rtruediv__ = _rhs
+    _tensor_expr_overload.__floordiv__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.floor_divide
+    )
+    _tensor_expr_overload.__rfloordiv__ = _rhs
+    _tensor_expr_overload.__mod__ = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.mod)
+    _tensor_expr_overload.__rmod__ = _rhs
+    _tensor_expr_overload.__pow__ = lambda lhs, rhs: expr._binary_op_helper(
+        lhs, rhs, _ffi_api.power
+    )
+    _tensor_expr_overload.__rpow__ = _rhs
+
+
+_register_op_make()

python/tvm/relax/op/_ffi_api.py

@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+"""FFI APIs for tvm.relax.op"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("relax.op", __name__)

python/tvm/relax/op/base.py

@@ -0,0 +1,880 @@
+# 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
+# pylint: disable=redefined-builtin
+# ruff: noqa: F821
+"""The base Relax operators."""
+
+from collections.abc import Callable
+
+import tvm_ffi
+
+import tvm
+import tvm.runtime
+from tvm.ir import Call
+from tvm.runtime import Object, ObjectConvertible
+
+from ..expr import Expr, ExternFunc, GlobalVar, ShapeExpr, StringImm, Var
+from ..type import TensorType, Type
+from ..utils import convert_to_expr
+from . import _ffi_api
+
+py_print = print  # pylint: disable=invalid-name
+
+
+def register_gradient(
+    op_name: str,
+    fgradient: Callable[[Var, Call, Var, "BlockBuilder"], list[Expr]] | None = None,
+    level: int = 10,
+):
+    """Register operator gradient function for a relax operator.
+
+    Parameters
+    ----------
+    op_name: str
+        The name of the op.
+
+    fgradient: function (orig_var: Var, orig_call: Call, output_grad: Var, ctx: BlockBuilder)
+         -> partials: List[Expr]
+        The gradient function being used.
+
+    level: int
+        The priority level
+    """
+    return tvm.ir.register_op_attr(op_name, "FPrimalGradient", fgradient, level)
+
+
+def null_value() -> Call:
+    """Create a call node that represents a null value object.
+
+    Returns
+    -------
+    ret: Call
+        The created call node.
+    """
+    return _ffi_api.null_value()  # type: ignore
+
+
+def _wrap_inline_arg_tuple(args) -> Expr:
+    """Helper function to wrap argument tuple
+
+    Normalize the arguments provided the functions that accept a tuple
+    of arguments, and require the tuple of arguments to be written
+    in-line.  If the arguments provided are a single relax expression,
+    and are not a reference to a relax tuple, then wrap them into an
+    in-line relax Tuple.
+
+    """
+    if isinstance(args, tuple | list):
+        return tvm.relax.Tuple([convert_to_expr(a) for a in args])
+    elif (
+        isinstance(args, Expr)
+        and not isinstance(args, tvm.relax.Tuple)
+        and (args.ty is None or not isinstance(args.ty, tvm.relax.TupleType))
+    ):
+        return tvm.relax.Tuple([args])
+    else:
+        return args
+
+
+def call_tir(
+    gvar: GlobalVar,
+    args: Expr,
+    out_ty: TensorType | list[TensorType],
+    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
+) -> Call:
+    """
+    Call a tirx.prim_func and return the output.
+
+    Parameters
+    ----------
+    gvar : GlobalVar
+        The GlobalVar referring to a tirx PrimFunc.
+
+    args : Expr
+        The input arguments.
+
+    out_ty : Union[TensorType, List[TensorType]]
+        The type information of the call_tir output.
+        It should be a single or a list of TensorType. Each one denotes the
+        type information of a returned tensor.
+
+    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
+        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
+
+    Returns
+    -------
+    ret: Call
+        A call node for the call_tir operator.
+    """
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(out_ty, list):
+        out_ty = [out_ty]
+
+    if isinstance(tir_vars, list | tuple):
+        tir_vars = ShapeExpr(tir_vars)
+
+    return _ffi_api.call_tir(gvar, args, out_ty, tir_vars)  # type: ignore
+
+
+def call_tir_with_grad(
+    gvar: GlobalVar,
+    args: Expr,
+    out_ty: TensorType | list[TensorType],
+    te_grad_name: str,
+    te_grad_kwargs: dict[str, Object] | None = None,
+    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
+) -> Call:
+    """
+    Call a tirx.prim_func and return the output. This intrinsic will bind a te gradient function
+    (refered by te_grad_name) to the call_tir_with_grad node. The te gradient function will be
+    called by the Gradient pass.
+
+    Parameters
+    ----------
+    gvar : GlobalVar
+        The GlobalVar referring to a tirx PrimFunc.
+
+    args : Expr
+        The input arguments.
+
+    out_ty : Union[TensorType, List[TensorType]]
+        The type information of the call_tir_with_grad output.
+        It should be a single or a list of TensorType. Each one denotes the
+        type information of a returned tensor.
+
+    te_grad_name : str
+        The registered name of the te gradient function associated with the call_tir_with_grad
+        node. Must be provided as a keyword argument.
+
+    te_grad_kwargs : Dict[str, Object], optional
+        The keyword arguments passed to the te gradient function.
+        Optionally provided as a keyword argument. Default: {}.
+
+    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
+        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
+
+    Returns
+    -------
+    ret: Call
+        A call node for the call_tir_with_grad operator.
+    """
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(out_ty, list):
+        out_ty = [out_ty]
+
+    if isinstance(tir_vars, list | tuple):
+        tir_vars = ShapeExpr(tir_vars)
+
+    if te_grad_kwargs is None:
+        te_grad_kwargs = {}
+
+    return _ffi_api.call_tir_with_grad(  # type: ignore
+        gvar, args, out_ty, te_grad_name, te_grad_kwargs, tir_vars
+    )
+
+
+def call_tir_inplace(
+    gvar: GlobalVar,
+    args: Expr,
+    inplace_indices: int | list[int],
+    out_ty: TensorType | list[TensorType],
+    tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
+) -> Call:
+    """
+    Call a TIR PrimFunc and return the result, doing the specified computations in-place
+    (based on the `inplace_indices` argument; outputs will alias the inputs
+    selected by in-place indices).
+
+    Warning: This operator is considered pure by the type system but actually mutates
+    the arguments specified by `inplace_indices`. This operator should not be used directly,
+    but rather should be inserted by passes that have checked whether it is safe to perform
+    operations in-place (i.e., none of the arguments specified as an output is aliased or is
+    live after calling call_tir_inplace).
+
+    Direct calls to this operator should be done for testing purposes only.
+
+    Parameters
+    ----------
+    gvar : GlobalVar
+        The GlobalVar referring to a TIR PrimFunc.
+
+    args : Expr
+        The input arguments.
+
+    inplace_indices : Union[int, List[int]]
+        Specify which arguments should be used for in-place computations.
+        If `inplace_indices` is a single integer, it will be made into a singleton list.
+        Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
+        will be an alias of `args[j]`.
+        If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
+        At least one member of `inplace_indices` must not be -1.
+
+    out_ty : Union[TensorType, List[TensorType]]
+        The type information of the call_tir_inplace output.
+        It should be a single `TensorType` or a list of `TensorType`.
+        Each one denotes the type information of a returned tensor.
+        If a list of `TensorType` is given, the result will be a tuple of `TensorType`.
+
+    tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
+        ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
+
+    Returns
+    -------
+    ret: Call
+        A call node for the call_tir operator.
+    """
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(inplace_indices, list):
+        inplace_indices = [inplace_indices]
+
+    if not isinstance(out_ty, list):
+        out_ty = [out_ty]
+
+    if isinstance(tir_vars, list | tuple):
+        tir_vars = ShapeExpr(tir_vars)
+
+    return _ffi_api.call_tir_inplace(  # type: ignore
+        gvar,
+        args,
+        inplace_indices,
+        out_ty,
+        tir_vars,
+    )
+
+
+def call_dps_packed(
+    func: str | Expr,
+    args: Expr,
+    out_ty: TensorType | list[TensorType],
+) -> Call:
+    """
+    Call a destination-passing-style packed function and return the output.
+
+    Note: The called function is assumed to be _pure_ (other than modifying the designated
+    output arguments). If the function _does_ result in other side effects, then the compiler
+    may end up removing, reordering, or repeating those effects--no guarantees can be made.
+
+    Parameters
+    ----------
+    func : Union[str, Expr]
+        The destination-passing-style function, can be ExternFunc.
+
+    args : Expr
+        The input arguments.
+
+    out_ty : Union[TensorType, List[TensorType]]
+        The type information of the call_dps_packed output.
+        It should be a single or a list of TensorType. Each one denotes the
+        type information of a returned tensor.
+
+    Returns
+    -------
+    ret: Call
+        A call node for the call_dps_packed operator.
+    """
+    if isinstance(func, str):
+        func = ExternFunc(func)
+
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(out_ty, list):
+        out_ty = [out_ty]
+
+    return _ffi_api.call_dps_packed(func, args, out_ty)  # type: ignore
+
+
+def call_py_func(
+    func_name: str,
+    args: Expr,
+    out_ty: TensorType | list[TensorType],
+) -> Call:
+    """
+    Call a Python function and return the output.
+
+    Parameters
+    ----------
+    func_name : str
+        The name of the Python function to call. This should correspond to a function
+        in the IRModule's pyfuncs attribute.
+
+    args : Expr
+        The input arguments.
+
+    out_ty : Union[TensorType, List[TensorType]]
+        The type information of the call_py_func output.
+        It should be a single or a list of TensorType. Each one denotes the
+        type information of a returned tensor.
+
+    Returns
+    -------
+    ret: Call
+        A call node for the call_py_func operator.
+    """
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(out_ty, list):
+        out_ty = [out_ty]
+
+    return _ffi_api.call_py_func(func_name, args, out_ty)  # type: ignore
+
+
+def call_builtin_with_ctx(
+    func: str | Expr,
+    args: Expr,
+    *,
+    ty_args: Type | list[Type] | None = None,
+) -> Call:
+    """Call a builtin function func.
+
+    Parameters
+    ----------
+    func : Expr
+        The builtin function to be called.
+
+    args : Expr
+        The input arguments.
+
+    ty_args: Optional[Union[Type, List[Type]]]
+        The type arguments to the call node.
+
+    Returns
+    -------
+    ret: Call
+        The created call node.
+    """
+    if isinstance(func, str):
+        func = ExternFunc(func)
+
+    args = _wrap_inline_arg_tuple(args)
+
+    if ty_args is not None and not isinstance(ty_args, list | tuple):
+        ty_args = [ty_args]
+
+    return _ffi_api.call_builtin_with_ctx(  # type: ignore
+        func,
+        args,
+        ty_args,  # type: ignore
+    )
+
+
+def make_closure(
+    func: Expr,
+    args: Expr,
+) -> Object:
+    """
+    Create a closure with free variables and return the closure.
+
+    Parameters
+    ----------
+    func : Expr
+        The closure, can be ExternFunc or PrimFunc.
+
+    args : Expr
+        The input arguments.
+
+
+    Returns
+    -------
+    ret: Object
+        The VMClosure.
+    """
+
+    args = _wrap_inline_arg_tuple(args)
+
+    return _ffi_api.make_closure(func, args)  # type: ignore
+
+
+def invoke_closure(
+    closure: Expr,
+    args: Expr,
+    ty_args: list[Type] | Type,
+) -> Call:
+    """
+    Invoke a closure.
+
+    Parameters
+    ----------
+    closure : Expr
+        The VMClosure object.
+
+    args : Expr
+        The input arguments.
+
+    type_args: Union[List[Type], Type]
+        The type information arguments of the CallNode
+
+    Returns
+    -------
+    ret: Call
+        A call to `invoke_closure`.
+    """
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(ty_args, list | tuple):
+        ty_args = [ty_args]
+
+    return _ffi_api.invoke_closure(closure, args, ty_args)  # type: ignore
+
+
+def render_object(val: tvm.Object) -> str:
+    """
+    Given a TVM Object, renders it in string form. Used for Relax printing and assertions.
+
+    Parameters
+    ----------
+    val: tvm.Object
+        An object to render
+
+    Returns
+    -------
+    ret: str
+        A string representing the value, ideally human-readable
+    """
+    if isinstance(val, tvm.runtime.Tensor):
+        return str(val)
+    if isinstance(val, tvm_ffi.Array):
+        fields = ", ".join([render_object(val[i]) for i in range(len(val))])
+        return f"({fields})"
+    return str(val)
+
+
+@tvm.register_global_func("relax.run.shape_to_tensor")
+def relax_shape_to_tensor(shape_tuple: tvm_ffi.Shape) -> tvm.runtime.Tensor:
+    """
+    Takes a Shape and convert it to Tensor.
+
+    Parameters
+    ----------
+    shape_tuple: tvm_ffi.Shape
+        Shape tuple that we want to convert to Tensor at runtime
+    """
+    return tvm.runtime.tensor([int(v) for v in shape_tuple])
+
+
+@tvm.register_global_func("relax.run.print")
+def relax_print(format_str: str, *format_args: tvm.Object) -> None:
+    """
+    Takes a list of values to print, formats with the given format string.
+    If the format string is empty, simply prints.
+
+    Call from TVM script like this:
+    `relax.print(value1, value2, ..., valueN, format=format_str)`
+    or
+    `relax.print(value1, value2, ..., valueN) # format_str defaults to ""`
+
+    Parameters
+    ----------
+    format_str: str
+        The last argument is a Python-style format string for printing the value
+
+    format_args: List[Object]
+        The values to print.
+    """
+    val_strs = map(render_object, format_args)
+    if format_str == "":
+        py_print(*val_strs)
+    else:
+        py_print(format_str.format(*val_strs))
+
+
+def print(*values: list[Expr], format: str | Expr = "") -> Expr:
+    """Print op to print the values
+
+    Parameters
+    ----------
+    values : List[Expr]
+        The values to print.
+
+    format: Union[str, Expr]
+        The format string or StringImm.
+
+    Returns
+    -------
+    result : Expr
+        A relax Call, which will print the value during runtime.
+    """
+    if isinstance(format, str):
+        format = StringImm(format)
+
+    return _ffi_api.print(values, format)  # type: ignore # pylint: disable=no-member
+
+
+@tvm.register_global_func("relax.run.assert_op")
+def relax_assert_op(condition: tvm.Object, format_str: str, *format_args: tvm.Object) -> None:
+    """
+    A variadic function. The first value serves as the assertion condition:
+    If the condition is true, then the operator does nothing.
+    If the condition is false, then the operator raises an assertion error.
+
+    Arguments after the first value serve as format arguments for the error message;
+    the last argument must be a format string for the error message (empty by default).
+    If the format string is the empty string, then the error message will simply include
+    a comma-separated list of the format arguments.
+    The condition argument is not included in the format string.
+
+    Parameters
+    ----------
+    condition: tvm.Object
+        The assertion condition. Must be a boolean scalar.
+
+    format_str: str
+        The last argument is a Python-style format string for printing the value
+
+    format_args: List[tvm.Object]
+        Values used for formatting the string.
+    """
+    if not isinstance(format_str, str):
+        raise ValueError(
+            f"The format string argument to assert must be a string, given {type(format_str)})"
+        )
+
+    if isinstance(condition, bool | int):
+        val = condition
+    elif isinstance(condition, tvm.runtime.Tensor):
+        # may happen if the original program had unknown shape or dtype for the tensor's type
+        dtype = condition.dtype
+        if dtype != "bool":
+            raise ValueError(f"The condition must be a bool scalar, but given a {dtype} tensor")
+        shape = condition.shape
+        if len(shape) != 0:
+            raise ValueError(f"The condition must be a scalar, but it has a shape of {shape}")
+
+        val = condition.numpy()
+
+    else:
+        # should be guaranteed by the type system
+        raise ValueError(
+            f"The condition for relax assert must be a bool, int, or Tensor, "
+            f"but received a {type(condition)}."
+        )
+
+    if not val:
+        error_message = "Assertion Failed"
+        if format_args or format_str != "":
+            rendered = map(render_object, format_args)
+            if format_str != "":
+                error_message = format_str.format(*rendered)
+            else:
+                error_message = ", ".join(rendered)
+        raise AssertionError(error_message)
+
+
+def assert_op(
+    condition: Expr,
+    format_args: Expr | list[Expr] | None = None,
+    format: str | Expr = "",
+) -> Expr:
+    """
+    Create a call to Relax's assert_op operation (`assert` is reserved in Python,
+    so the name must be distinct).
+
+    Parameters
+    ----------
+    condition: Expr
+        The assertion condition.
+
+    format_args: Optional[Union[Expr, List[Expr]]]
+        Format arguments for the error message if the condition fails.
+
+    format: Union[str, Expr]
+        The format string or StringImm for the error message.
+
+    Returns
+    -------
+    result : Expr
+        A Call to the Relax assert operation.
+    """
+    if not isinstance(condition, Expr):
+        condition = tvm.relax.prim_value(condition)
+
+    if format_args is None:
+        format_args = []
+    elif isinstance(format_args, Expr):
+        format_args = [format_args]
+
+    if isinstance(format, str):
+        format = StringImm(format)
+
+    return _ffi_api.assert_op(condition, format_args, format)  # type: ignore
+
+
+def shape_of(expr: Expr) -> Expr:
+    """Get shape of a tensor.
+
+    Parameters
+    ----------
+    expr : Expr
+        The input Expr.
+
+    Returns
+    -------
+    result : Expr
+        A relax Call, which gets the shape of the input
+    """
+    return _ffi_api.shape_of(expr)  # type: ignore # pylint: disable=no-member
+
+
+def size(expr: Expr) -> Expr:
+    """Get the total number of elements in a tensor.
+
+    Parameters
+    ----------
+    expr : Expr
+        The input tensor.
+
+    Returns
+    -------
+    result : Expr
+        A scalar tensor of dtype int64 containing the total number of elements.
+    """
+    return _ffi_api.size(expr)  # type: ignore # pylint: disable=no-member
+
+
+def tensor_to_shape(expr: Expr) -> Expr:
+    """Convert tensor to shape expr.
+    Parameters
+    ----------
+    expr : Expr
+        The input Expr
+    Returns
+    -------
+    result : Expr
+        A relax Call, which transforms the tensor values to the shape
+    """
+    return _ffi_api.tensor_to_shape(expr)  # type: ignore # pylint: disable=no-member
+
+
+def shape_to_tensor(expr: Expr) -> Expr:
+    """Convert shape to tensor expr.
+    Parameters
+    ----------
+    expr : Expr
+        The input Expr
+    Returns
+    -------
+    result : Expr
+        A relax Call, which transforms the shape values to the tensor
+    """
+    return _ffi_api.shape_to_tensor(expr)  # type: ignore # pylint: disable=no-member
+
+
+def call_inplace_packed(
+    func: str | ExternFunc | GlobalVar,
+    *args: Expr,
+    inplace_indices: int | list[int],
+    ty_args: Type | list[Type],
+) -> Expr:
+    """
+    Construct a call to a packed function that consumes some of its arguments "in-place"
+    and returns the mutated arguments (aliased), but should be considered to be otherwise pure.
+    The `inplace_indices` argument indicates which of the outputs are mutated arguments.
+
+    The resulting call will have the same semantics as calling the packed function directly.
+
+    Note: This should be used for cases when the user knows that calling the packed function
+    with these arguments will **in reality** not cause any other side effects.
+    If it is used for a call that **does** result in other side effects, then the compiler
+    may end up removing, reordering, or repeating that call, with no guarantees
+    made about any side effects from the callee.
+
+    Warning: This operator as treated as pure by the type system even though it *is* performing
+    side effects (mutating some arguments). It is therefore incumbent upon the user to ensure
+    that it is being used safely (viz., that mutated arguments are not live after the mutation,
+    that they do not alias values live after the mutation).
+
+    Parameters
+    ----------
+    func : Union[str, ExternFunc]
+      The name (global symbol) for a PackedFunc or an ExternFunc node.
+
+    args: Expr
+      The arguments for the PackedFunc.
+
+    inplace_indices : Union[int, List[int]]
+      Specify which arguments should be used for in-place computations.
+      If `inplace_indices` is a single integer, it will be made into a singleton list.
+      Suppose `inplace_indices[i] = j`, where `j >= 0`. Then the `i`th output
+      will be an alias of `args[j]`.
+      If `inplace_indices[i] = -1`, then the `i`th output will be a freshly allocated tensor.
+      At least one member of `inplace_indices` must not be -1.
+
+    ty_args: Union[Type, List[Type]]
+        The list of type information arguments (giving the type information for the returned value).
+
+    Returns
+    -------
+    result : Expr
+      A Relax call, corresponding to
+      `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
+    """
+    if isinstance(func, ExternFunc):
+        func = func.global_symbol
+
+    op = ExternFunc(func)
+    args = tuple(convert_to_expr(a) for a in args)
+    if ty_args is None:
+        raise ValueError("R.call_pure_packed is required to have type_args")
+    if isinstance(ty_args, tuple):  # type: ignore
+        ty_args = list(ty_args)
+    elif not isinstance(ty_args, list):
+        ty_args = [ty_args]
+    if not isinstance(inplace_indices, list):
+        inplace_indices = [inplace_indices]
+
+    return _ffi_api.call_inplace_packed(op, args, inplace_indices, ty_args)  # type: ignore # pylint: disable=no-member
+
+
+def call_pure_packed(
+    func: str | ExternFunc | GlobalVar,
+    *args: Expr,
+    ty_args: Type | list[Type],
+) -> Expr:
+    """
+    Construct a call to a packed function that should be treated as pure,
+    even though packed calls are normally not treated as pure.
+
+    The resulting call will have the same semantics as calling the packed function directly.
+
+    Note: This should be used for cases when the user knows that calling the packed function
+    with these arguments will **in reality** not cause any side effects.
+    If it is used for a call that **does** result in side effects, then the compiler
+    may end up removing, reordering, or repeating that call, with no guarantees
+    made about any side effects from the callee.
+
+    Parameters
+    ----------
+    func : Union[str, ExternFunc]
+      The name (global symbol) for a PackedFunc or an ExternFunc node.
+
+    args: Expr
+      The arguments for the PackedFunc.
+
+    ty_args: Union[Type, List[Type]]
+        The list of type information arguments (giving the type information for the returned value).
+
+    Returns
+    -------
+    result : Expr
+      A Relax call, corresponding to
+      `call_pure_packed(ExternFunc(func), args, DictAttrs(kwargs), ty_args)`
+    """
+    if isinstance(func, ExternFunc):
+        func = func.global_symbol
+
+    op = ExternFunc(func)
+    args = tuple(convert_to_expr(a) for a in args)
+
+    if ty_args is None:
+        raise ValueError("R.call_pure_packed is required to have type_args")
+
+    if isinstance(ty_args, tuple):  # type: ignore
+        ty_args = list(ty_args)
+    elif not isinstance(ty_args, list):
+        ty_args = [ty_args]
+
+    ty_args = [
+        (ty() if callable(ty) else ty.asobject() if isinstance(ty, ObjectConvertible) else ty)
+        for ty in ty_args
+    ]
+
+    # note: if we need attributes, we can also take them here
+
+    return _ffi_api.call_pure_packed(op, args, None, ty_args)  # type: ignore # pylint: disable=no-member
+
+
+def invoke_pure_closure(
+    closure: Expr,
+    args: Expr,
+    ty_args: list[Type] | Type,
+) -> Call:
+    """
+    Invoke a closure and indicate to the compiler that it is pure.
+
+    Note: This should be used for cases when the user knows that calling the closure
+    with these arguments will **in reality** not cause any side effects.
+    If it is used for a call that _does_ result in side effects, then the compiler
+    may end up removing, reordering, or repeating that call, with no guarantees
+    made about any side effects from the callee.
+
+    Parameters
+    ----------
+    closure : Expr
+        The VMClosure object.
+
+    args : Expr
+        The input arguments.
+
+    type_args: Union[List[Type], Type]
+        The type information arguments of the CallNode
+
+    Returns
+    -------
+    ret: Call
+        A call to `invoke_pure_closure`.
+    """
+    args = _wrap_inline_arg_tuple(args)
+
+    if not isinstance(ty_args, list | tuple):
+        ty_args = [ty_args]
+
+    return _ffi_api.invoke_pure_closure(closure, args, ty_args)  # type: ignore
+
+
+def to_vdevice(data, dst_vdevice) -> Expr:
+    """Copy data to the destination device. This
+    operator helps data transferring between difference devices for
+    heterogeneous execution.
+
+    Parameters
+    ----------
+    data : Expr
+        The tensor to be copied.
+
+    dst_device : VDevice
+        The destination device where the data is copied to.
+
+    Returns
+    -------
+    result : Expr
+        The copied result.
+    """
+    return _ffi_api.to_vdevice(data, dst_vdevice)  # type: ignore
+
+
+def hint_on_device(data, dst_vdevice, memory_scope="global") -> Expr:
+    """It provides a hint specifying the device on which the input data should be executed.
+    This hint is utilized by RealizeVDevice to propagate the virtual device."
+
+    Parameters
+    ----------
+    data : Expr
+        The tensor to be copied.
+
+    dst_device : Device
+        The destination device where the data is supposed to be executed.
+
+    memory_scope: String
+       Memory scope of buffer on target device.
+
+    Returns
+    -------
+    result : Expr
+        The result.
+    """
+    return _ffi_api.hint_on_device(data, dst_vdevice, memory_scope)  # type: ignore

python/tvm/relax/op/binary.py

@@ -0,0 +1,484 @@
+# 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=redefined-builtin, invalid-name
+"""Relax binary arithmetic and comparison operators."""
+
+from ..expr import Expr
+from . import _ffi_api
+
+###################### Arithmetic operators ######################
+
+
+def add(x1: Expr, x2: Expr) -> Expr:
+    """Addition with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : Expr
+        The first input tensor.
+    x2 : Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : Expr
+        The computed result.
+
+    Examples
+    --------
+    .. code:: python
+
+      bb = relax.BlockBuilder()
+      a = relax.Var("a", relax.TensorType(shape=(2, 3), dtype="float32"))
+      b = relax.Var("b", relax.TensorType(shape=(2, 1), dtype="float32"))
+      c = bb.normalize(relax.op.add(a, b))  # c has TensorType(shape=(2, 3), dtype="float32")
+    """
+    return _ffi_api.add(x1, x2)  # type: ignore
+
+
+def divide(x1: Expr, x2: Expr) -> Expr:
+    """Division with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.divide(x1, x2)  # type: ignore
+
+
+def floor_divide(x1: Expr, x2: Expr) -> Expr:
+    """Floor division with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.floor_divide(x1, x2)  # type: ignore
+
+
+def log_add_exp(x1: Expr, x2: Expr) -> Expr:
+    """
+    Compute the log of the sum of exponentials of the inputs, element-wise.
+
+    Parameters
+    ----------
+    x1 : Expr
+        The first input tensor.
+    x2 : Expr
+        The second input tensor.
+
+    Returns
+    -------
+    Expr
+        The element-wise log-sum-exp of `x1` and `x2`.
+    """
+    return _ffi_api.log_add_exp(x1, x2)
+
+
+def multiply(x1: Expr, x2: Expr) -> Expr:
+    """Multiplication with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : Expr
+        The first input tensor.
+    x2 : Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : Expr
+        The computed result.
+    """
+    return _ffi_api.multiply(x1, x2)  # type: ignore
+
+
+def power(x1: Expr, x2: Expr):
+    """Power with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.power(x1, x2)  # type: ignore
+
+
+def atan2(x1: Expr, x2: Expr) -> Expr:
+    """Atan2 with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor (y-coordinates).
+    x2 : relax.Expr
+        The second input tensor (x-coordinates).
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.atan2(x1, x2)  # type: ignore
+
+
+def subtract(x1: Expr, x2: Expr) -> Expr:
+    """Subtraction with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.subtract(x1, x2)  # type: ignore
+
+
+def mod(x1: Expr, x2: Expr) -> Expr:
+    """Modulo with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : Expr
+        The first input tensor.
+    x2 : Expr
+        The second input tensor.
+    """
+    return _ffi_api.mod(x1, x2)  # type: ignore
+
+
+def floor_mod(x1: Expr, x2: Expr) -> Expr:
+    """Floor modulo with numpy-style broadcasting.
+
+    Parameters
+    ----------
+    x1 : Expr
+        The first input tensor.
+    x2 : Expr
+        The second input tensor.
+    """
+    return _ffi_api.floor_mod(x1, x2)  # type: ignore
+
+
+###################### Comparison operators ######################
+
+
+def equal(x1: Expr, x2: Expr) -> Expr:
+    """Broadcasted element-wise test for (lhs == rhs).
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.equal(x1, x2)  # type: ignore
+
+
+def greater(x1: Expr, x2: Expr) -> Expr:
+    """Broadcasted element-wise test for (lhs > rhs).
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.greater(x1, x2)  # type: ignore
+
+
+def greater_equal(x1: Expr, x2: Expr) -> Expr:
+    """Broadcasted element-wise test for (lhs >= rhs).
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.greater_equal(x1, x2)  # type: ignore
+
+
+def less(x1: Expr, x2: Expr) -> Expr:
+    """Broadcasted element-wise test for (lhs < rhs).
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.less(x1, x2)  # type: ignore
+
+
+def less_equal(x1: Expr, x2: Expr) -> Expr:
+    """Broadcasted element-wise test for (lhs <= rhs).
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.less_equal(x1, x2)  # type: ignore
+
+
+def not_equal(x1: Expr, x2: Expr) -> Expr:
+    """Broadcasted element-wise test for (lhs != rhs).
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.not_equal(x1, x2)  # type: ignore
+
+
+def maximum(x1: Expr, x2: Expr) -> Expr:
+    """Element-wise maximum
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.maximum(x1, x2)
+
+
+def minimum(x1: Expr, x2: Expr) -> Expr:
+    """Element-wise minimum
+
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.minimum(x1, x2)
+
+
+###################### Logical operators ######################
+
+
+def logical_and(x1: Expr, x2: Expr) -> Expr:
+    """Logical AND
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.logical_and(x1, x2)
+
+
+def logical_or(x1: Expr, x2: Expr) -> Expr:
+    """Logical OR
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.logical_or(x1, x2)
+
+
+def logical_xor(x1: Expr, x2: Expr) -> Expr:
+    """Logical XOR
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.logical_xor(x1, x2)
+
+
+###################### Bitwise operators ######################
+
+
+def bitwise_and(x1: Expr, x2: Expr) -> Expr:
+    """Bitwise AND
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.bitwise_and(x1, x2)
+
+
+def bitwise_or(x1: Expr, x2: Expr) -> Expr:
+    """Bitwise OR
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.bitwise_or(x1, x2)
+
+
+def bitwise_xor(x1: Expr, x2: Expr) -> Expr:
+    """Bitwise XOR
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The first input tensor.
+    x2 : relax.Expr
+        The second input tensor.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.bitwise_xor(x1, x2)
+
+
+def left_shift(x1: Expr, x2: Expr) -> Expr:
+    """Bitwise Shift Left
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The input tensor to be shifted.
+    x2 : relax.Expr
+        The number of positions to shift.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.left_shift(x1, x2)
+
+
+def right_shift(x1: Expr, x2: Expr) -> Expr:
+    """Bitwise Shift Right
+    Parameters
+    ----------
+    x1 : relax.Expr
+        The input tensor to be shifted.
+    x2 : relax.Expr
+        The number of positions to shift.
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.right_shift(x1, x2)

python/tvm/relax/op/builtin/__init__.py

@@ -0,0 +1,19 @@
+# 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 builtin operators."""
+
+from .builtin import alloc_tensor, stop_lift_params