chore: import upstream snapshot with attribution

python/.gitignore

@@ -0,0 +1,4 @@
+build
+dist
+*.cpp
+requirements/*.txt

python/tvm/__init__.py

@@ -0,0 +1,126 @@
+# 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, wildcard-import
+"""TVM: Open Deep Learning Compiler Stack."""
+
+import multiprocessing
+import sys
+import os
+
+# ffi module must load first
+from tvm_ffi import register_object, register_global_func, get_global_func
+
+# top-level alias
+from .libinfo import __version__
+from .base import _RUNTIME_ONLY
+
+# tvm.runtime
+from .runtime import Object
+from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal
+from .runtime._tensor import vpi, rocm, ext_dev, hexagon
+from .runtime import DataType, DataTypeCode
+
+# tvm.error
+from . import error
+
+# tvm.ir
+from .ir import IRModule
+from .ir import transform
+from .ir import instrument
+from . import ir
+
+# tvm.script — must be imported before any dialect package so that
+# tvm.script.register_dialect is reachable when dialect __init__.py files run.
+from . import script
+
+# tvm.tirx — registers itself via tvm.script.register_dialect in its __init__
+from . import tirx
+
+# tvm.backend — owns backend Python load hooks
+from . import backend
+
+# tvm.target
+from . import target
+
+# tvm.te
+from . import te
+
+# tvm.driver
+from .driver import build, compile
+
+# others
+from . import arith
+
+# support infra
+from . import support
+
+# Side-effect imports: register CUDA/ROCm FFI callbacks at TVM startup
+from .support import rocm as _rocm, nvcc as _nvcc
+
+# Relax contain modules that are only available in compiler package
+# Do not import them if TVM is built with runtime only
+if not _RUNTIME_ONLY:
+    # tvm.relax — registers itself via tvm.script.register_dialect in its __init__
+    from . import relax
+
+# NOTE: This file should be python2 compatible so we can
+# raise proper error message when user run the package using
+# an older version of the python
+
+
+def _should_print_backtrace():
+    in_pytest = "PYTEST_CURRENT_TEST" in os.environ
+    tvm_backtrace = os.environ.get("TVM_BACKTRACE", "0")
+
+    try:
+        tvm_backtrace = bool(int(tvm_backtrace))
+    except ValueError:
+        raise ValueError(f"invalid value for TVM_BACKTRACE {tvm_backtrace}, please set to 0 or 1.")
+
+    return in_pytest or tvm_backtrace
+
+
+def tvm_wrap_excepthook(exception_hook):
+    """Wrap given excepthook with TVM additional work."""
+
+    def wrapper(exctype, value, trbk):
+        """Clean subprocesses when TVM is interrupted."""
+        if exctype is error.DiagnosticError and not _should_print_backtrace():
+            # TODO(@jroesch): consider moving to C++?
+            print("note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.")
+        else:
+            exception_hook(exctype, value, trbk)
+
+        if hasattr(multiprocessing, "active_children"):
+            # pylint: disable=not-callable
+            for p in multiprocessing.active_children():
+                p.terminate()
+
+    return wrapper
+
+
+sys.excepthook = tvm_wrap_excepthook(sys.excepthook)
+
+# Autoload loads built-in and out-of-tree backends. Out-of-tree extensions opt
+# into being loaded automatically at ``import tvm`` time by declaring an entry
+# point in the ``tvm.backends`` group:
+# [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
+# Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
+from .backend._autoload_backends import _autoload_backends
+
+_autoload_backends()

python/tvm/arith/__init__.py

@@ -0,0 +1,40 @@
+# 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.
+"""Integer bound analysis, simplification and pattern detection."""
+
+from .int_set import (
+    IntSet,
+    IntervalSet,
+    PresburgerSet,
+    estimate_region_lower_bound,
+    estimate_region_strict_bound,
+    estimate_region_upper_bound,
+)
+from .analyzer import ModularSet, ConstIntBound, Analyzer, ProofStrength, Extension, CompareResult
+from .bound import deduce_bound
+from .pattern import detect_linear_equation, detect_clip_bound
+from .int_solver import solve_linear_equations, solve_linear_inequalities
+from .iter_affine_map import IterMapExpr, IterMark, IterSplitExpr, IterSumExpr
+from .iter_affine_map import (
+    detect_iter_map,
+    iter_map_simplify,
+    normalize_iter_map_to_expr,
+    normalize_to_iter_sum,
+    subspace_divide,
+    inverse_affine_iter_map,
+)

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

python/tvm/arith/analyzer.py

@@ -0,0 +1,529 @@
+# 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
+"""Arithmetic data structure and utility"""
+
+import enum
+
+import tvm_ffi
+
+from tvm import ir, tirx
+from tvm.arith import IntSet
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+class ProofStrength(enum.IntEnum):
+    """Proof strength of the analysis"""
+
+    DEFAULT = 0
+    SYMBOLIC_BOUND = 1
+
+
+class CompareResult(enum.IntEnum):
+    """Result of a transitive comparison.
+
+    Values must match the C++ ``arith::CompareResult`` enum.
+    """
+
+    INCONSISTENT = 0
+    EQ = 1
+    LT = 2
+    LE = 3
+    GT = 4
+    GE = 5
+    NE = 6
+    UNKNOWN = 7
+
+
+class Extension(enum.Flag):
+    """Extensions enabled for RewriteSimplifier
+
+    Values should match `RewriteSimplifier::Extensions`
+    """
+
+    NoExtensions = 0
+    TransitivelyProveInequalities = 1 << 0
+    ConvertBooleanToAndOfOrs = 1 << 1
+    ApplyConstraintsToBooleanBranches = 1 << 2
+    ComparisonOfProductAndSum = 1 << 3
+
+
+@tvm_ffi.register_object("arith.ModularSet")
+class ModularSet(Object):
+    """Represent range of (coeff * x + base) for x in Z"""
+
+    def __init__(self, coeff, base):
+        self.__init_handle_by_constructor__(_ffi_api.ModularSet, coeff, base)
+
+
+@tvm_ffi.register_object("arith.ConstIntBound")
+class ConstIntBound(Object):
+    """Represent constant integer bound
+
+    Parameters
+    ----------
+    min_value : int
+        The minimum value of the bound.
+
+    max_value : int
+        The maximum value of the bound.
+    """
+
+    POS_INF = (1 << 63) - 1
+    NEG_INF = -POS_INF
+
+    def __init__(self, min_value, max_value):
+        self.__init_handle_by_constructor__(_ffi_api.ConstIntBound, min_value, max_value)
+
+
+class ConstraintScope:
+    """Constraint scope.
+
+    Parameters
+    ----------
+    fenter : function
+        A function that will be called to create an enter context.
+
+    Note
+    ----
+    Do not create object directly, use Analyzer.constraint_scope
+    """
+
+    def __init__(self, fenter):
+        self._fenter = fenter
+        self._fexit = None
+
+    def __enter__(self):
+        self._fexit = self._fenter()
+
+    def __exit__(self, ptype, value, trace):
+        self._fexit()
+
+
+@tvm_ffi.register_object("arith.Analyzer")
+class Analyzer(Object):
+    """Integer arithmetic analyzer
+
+    This is a stateful analyzer class that can be used to perform
+    various symbolic integer analysis. The same analyzer instance can
+    be passed to FFI APIs to share accumulated facts across calls.
+    """
+
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_api.Analyzer)
+
+    @property
+    def is_z3_enabled(self) -> bool:
+        """Whether this build includes the Z3 backend (``USE_Z3=ON``).
+
+        The Z3-specific methods (:py:meth:`get_smtlib2`, :py:meth:`get_z3_stats`,
+        :py:meth:`set_z3_timeout_ms`, :py:meth:`set_z3_rlimit`) only work when
+        this is ``True``.
+        """
+        return bool(_ffi_api.AnalyzerIsZ3Enabled(self))
+
+    def _check_z3_enabled(self) -> None:
+        if not self.is_z3_enabled:
+            raise RuntimeError(
+                "The Z3 backend is not available in this build. "
+                "Rebuild TVM with USE_Z3=ON to use Z3-specific Analyzer APIs."
+            )
+
+    def get_smtlib2(self, expr: tirx.Expr | None = None) -> str:
+        """Get the current Z3 problem in SMT-LIB2 format.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``), since there is no
+            solver state to export. Use :py:attr:`is_z3_enabled` to check first.
+
+        Parameters
+        ----------
+        expr : Optional[Expr]
+            The expression to prove. If provided, its negation is added to the problem.
+        """
+        self._check_z3_enabled()
+        return _ffi_api.AnalyzerGetSMTLIB2(self, expr)
+
+    def set_z3_timeout_ms(self, timeout_ms: int) -> None:
+        """Set Z3 timeout in milliseconds.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
+        Parameters
+        ----------
+        timeout_ms : int
+            The timeout in milliseconds.
+        """
+        self._check_z3_enabled()
+        _ffi_api.AnalyzerSetZ3TimeoutMs(self, timeout_ms)
+
+    def set_z3_rlimit(self, rlimit: int) -> None:
+        """Set Z3 resource limit.
+
+        The resource limit gives deterministic solver budgeting (unlike a wall
+        clock timeout). A value of ``0`` disables the limit.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
+        Parameters
+        ----------
+        rlimit : int
+            The resource limit.
+        """
+        self._check_z3_enabled()
+        _ffi_api.AnalyzerSetZ3RLimit(self, rlimit)
+
+    def get_z3_stats(self) -> str:
+        """Get Z3 solver statistics.
+
+        Raises
+        ------
+        RuntimeError
+            If TVM was built without Z3 (``USE_Z3=OFF``).
+
+        Returns
+        -------
+        stats : str
+            The Z3 statistics.
+        """
+        self._check_z3_enabled()
+        return _ffi_api.AnalyzerGetZ3Stats(self)
+
+    def const_int_bound(self, expr: tirx.Expr) -> ConstIntBound:
+        """Find constant integer bound for expr.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        bound : ConstIntBound
+            The result bound
+        """
+        return _ffi_api.AnalyzerConstIntBound(self, expr)
+
+    def const_int_bound_is_bound(self, var: tirx.Var) -> bool:
+        """Check if a variable is bound to a range.
+
+        Parameters
+        ----------
+        var : tvm.tirx.Var
+            The variable.
+
+        Returns
+        -------
+        result : bool
+            Whether the variable is bound to a range.
+        """
+        return _ffi_api.AnalyzerConstIntBoundIsBound(self, var)
+
+    def modular_set(self, expr: tirx.Expr) -> ModularSet:
+        """Find a modular set that expr belongs to.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        result : ModularSet
+            The result.
+        """
+        return _ffi_api.AnalyzerModularSet(self, expr)
+
+    def simplify(self, expr: tirx.Expr, steps: int = 2) -> tirx.Expr:
+        """Simplify expression via both rewrite and canonicalization.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+        steps : The simplification runs in the order of
+                rewrite_simplify (step 1) -> canonical_simplify (step 2) ->
+                rewrite_simplify (step 3) -> canonical_simplify (step 4) -> ...
+                param steps controls how many steps to run.
+                Default is 2, i.e., rewrite_simplify + canonical_simplify.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerSimplify(self, expr, steps)
+
+    def clone(self) -> "Analyzer":
+        """Return a deep copy of this analyzer with independent state.
+
+        The returned analyzer carries the same accumulated facts (variable
+        bounds, modular sets, bindings, integer-set domains, literal
+        constraints and transitive comparisons) as this one, but owns its own
+        state: binding or simplifying on either analyzer afterwards does not
+        affect the other. Unlike copying the handle, this is a true deep copy.
+
+        Do not call this while a constraint scope is active on this analyzer.
+
+        Returns
+        -------
+        result : Analyzer
+            A new analyzer holding an independent copy of the facts.
+        """
+        return _ffi_api.AnalyzerClone(self)
+
+    def rewrite_simplify(self, expr: tirx.Expr) -> tirx.Expr:
+        """Simplify expression via rewriting rules.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerRewriteSimplify(self, expr)
+
+    @property
+    def rewrite_simplify_stats(self):
+        return _ffi_api.AnalyzerGetRewriteSimplifyStats(self)
+
+    def reset_rewrite_simplify_stats(self):
+        _ffi_api.AnalyzerResetRewriteSimplifyStats(self)
+
+    def canonical_simplify(self, expr: tirx.Expr) -> tirx.Expr:
+        """Simplify expression via canonicalization.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerCanonicalSimplify(self, expr)
+
+    def int_set(self, expr: tirx.Expr, dom_map: dict[tirx.Var, IntSet] | None = None) -> IntSet:
+        """Compute a symbolic IntSet that covers expr for all values in dom_map.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        dom_map : Optional[Dict[tvm.tirx.Var, tvm.arith.IntSet]]
+            The domain for variables to be relaxed.  When omitted, the analyzer
+            uses the domains of the variables already bound to it.
+
+        Returns
+        -------
+        result : IntSet
+            The result.
+        """
+        return _ffi_api.AnalyzerIntSet(self, expr, dom_map)
+
+    def can_prove(self, expr: tirx.Expr, strength: ProofStrength = ProofStrength.DEFAULT) -> bool:
+        """Check whether we can prove expr to be true.
+
+        Parameters
+        ----------
+        expr : Expr
+            The expression.
+
+        strength: ProofStrength
+            The proof strength. When TVM is built with Z3 (``USE_Z3=ON``), the
+            optional Z3 fallback is only consulted at ``SYMBOLIC_BOUND`` or
+            higher, after the native analyzers fail to prove the predicate.
+
+        Returns
+        -------
+        result : Expr
+            The result.
+        """
+        return _ffi_api.AnalyzerCanProve(self, expr, strength)
+
+    def set_maximum_rewrite_steps(self, maximum: int) -> None:
+        """Set the maximum allowed number of rewrite-simplify steps.
+
+        When a positive limit is set, the simplifier raises an exception once
+        it exceeds that number of rewrite steps.  This is useful for guarding
+        against performance regressions in tests.
+
+        Parameters
+        ----------
+        maximum : int
+            The maximum number of rewrite steps, or a non-positive value to
+            allow an unlimited number of steps.
+        """
+        _ffi_api.AnalyzerSetMaximumRewriteSteps(self, maximum)
+
+    def bind(
+        self,
+        var: tirx.Var,
+        expr: tirx.Expr | ir.Range,
+        allow_override: bool = False,
+    ) -> None:
+        """Bind a variable to the expression.
+
+        Parameters
+        ----------
+        var : tvm.tirx.Var
+            The variable.
+
+        expr : Union[tirx.Expr, ir.Range]
+            The expression or the range to bind to.
+
+        allow_override : bool
+            Whether to allow overriding an existing binding for the variable.
+        """
+        return _ffi_api.AnalyzerBind(self, var, expr, allow_override)
+
+    def constraint_scope(self, constraint: tirx.Expr) -> ConstraintScope:
+        """Create a constraint scope.
+
+        Parameters
+        ----------
+        constraint : Expr
+            The constraint expression.
+
+        returns
+        -------
+        scope : ConstraintScope
+            The constraint scope
+
+        Examples
+        --------
+        .. code-block:: python
+
+          x = te.var("x")
+          analyzer = tvm.arith.Analyzer()
+          with analyzer.constraint_scope(x % 3 == 0):
+              # constraint in effect
+              assert analyzer.modular_set(x).coeff == 3
+          # constraint no longer in effect
+          assert analyzer.modular_set(x).coeff != 3
+        """
+
+        def _fenter():
+            return _ffi_api.AnalyzerEnterConstraintContext(self, constraint)
+
+        return ConstraintScope(_fenter)
+
+    def update(
+        self, var: tirx.Var, info: ConstIntBound | ModularSet | IntSet, override: bool = False
+    ) -> None:
+        """Update information about var.
+
+        Parameters
+        ----------
+        var : tvm.tirx.Var
+            The variable.
+
+        info : Union[ConstIntBound, ModularSet, IntSet]
+            Related information.  A ``ConstIntBound`` updates the constant
+            integer bound, a ``ModularSet`` updates the modular set, and an
+            ``IntSet`` updates the integer-set domain of ``var``.
+
+        override : bool
+            Whether allow override.
+        """
+        if isinstance(info, ConstIntBound):
+            _ffi_api.AnalyzerConstIntBoundUpdate(self, var, info, override)
+        elif isinstance(info, ModularSet):
+            _ffi_api.AnalyzerModularSetUpdate(self, var, info, override)
+        elif isinstance(info, IntSet):
+            _ffi_api.AnalyzerIntSetUpdate(self, var, info, override)
+        else:
+            raise TypeError(f"Do not know how to handle type {type(info)}")
+
+    def can_prove_equal(self, lhs: tirx.Expr, rhs: tirx.Expr) -> bool:
+        """Whether we can prove that lhs == rhs
+
+        Parameters
+        ----------
+        lhs: Expr
+            The left-hand side of the comparison
+
+        rhs: Expr
+            The right-hand side of the comparison
+
+        Returns
+        -------
+        result: bool
+            Whether we can prove that lhs == rhs
+        """
+        return _ffi_api.AnalyzerCanProveEqual(self, lhs, rhs)
+
+    def try_compare(
+        self, lhs: tirx.Expr, rhs: tirx.Expr, propagate_inequalities: bool = True
+    ) -> CompareResult:
+        """Compare lhs and rhs using previously provided known comparisons.
+
+        Parameters
+        ----------
+        lhs : Expr
+            The left-hand side of the comparison.
+
+        rhs : Expr
+            The right-hand side of the comparison.
+
+        propagate_inequalities : bool
+            If true, attempt to find a sequence of transitive inequalities that
+            allow lhs and rhs to be compared.
+
+        Returns
+        -------
+        result : CompareResult
+            The most specific result that can be proven about the comparison.
+            Returns ``CompareResult.UNKNOWN`` when nothing can be proven.
+        """
+        return CompareResult(_ffi_api.AnalyzerTryCompare(self, lhs, rhs, propagate_inequalities))
+
+    @property
+    def enabled_extensions(self) -> Extension:
+        """Return the currently enabled extensions"""
+        value = _ffi_api.AnalyzerGetEnabledExtensions(self)
+        return Extension(value)
+
+    @enabled_extensions.setter
+    def enabled_extensions(self, flags: int | Extension):
+        """Enable extensions for the analyzer
+
+        Parameters
+        ----------
+        flags: Union[int,Extension]
+
+            The extensions to enable.
+        """
+        flags = Extension(flags).value
+        _ffi_api.AnalyzerSetEnabledExtensions(self, flags)

python/tvm/arith/bound.py

@@ -0,0 +1,40 @@
+# 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.
+"""Bound deduction."""
+
+from . import _ffi_api
+
+
+def deduce_bound(var, cond, hint_map, relax_map):
+    """Deduce the bound of the target variable in the cond.
+
+    Parameters
+    ----------
+    var : tvm.tirx.Var
+        The target variable to be deduced.
+
+    cond : Expr
+        The condition
+
+    hint_map : Map[tvm.tirx.Var, IntSet]
+        Domain of variables used to help deduction.
+
+    relax_map : Map[tvm.tirx.Var, IntSet]
+        The fomain of the variables to be relaxed
+        using the provided domain.
+    """
+    return _ffi_api.DeduceBound(var, cond, hint_map, relax_map)

python/tvm/arith/int_set.py

@@ -0,0 +1,213 @@
+# 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.
+"""Integer set."""
+
+import tvm_ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("ir.IntSet")
+class IntSet(Object):
+    """Represent a set of integer in one dimension."""
+
+    def is_nothing(self):
+        """Whether the set represent nothing"""
+        return _ffi_api.IntSetIsNothing(self)
+
+    def is_everything(self):
+        """Whether the set represent everything"""
+        return _ffi_api.IntSetIsEverything(self)
+
+    @staticmethod
+    def vector(vec):
+        """Construct an integer set that covers the vector expr
+
+        Parameters
+        ----------
+        vec : Expr
+            The vector expression.
+
+        Returns
+        -------
+        rset : IntSet
+            The result set.
+        """
+        return _ffi_api.intset_vector(vec)
+
+    @staticmethod
+    def single_point(point):
+        """Construct a point set.
+
+        Parameters
+        ----------
+        point : Expr
+            The vector expression.
+
+        Returns
+        -------
+        rset : IntSet
+            The result set.
+        """
+        return _ffi_api.intset_single_point(point)
+
+
+@tvm_ffi.register_object("arith.IntervalSet")
+class IntervalSet(IntSet):
+    """Represent set of continuous interval [min_value, max_value]
+
+    Parameters
+    ----------
+    min_value : Expr
+        The minimum value in the interval.
+
+    max_value : Expr
+        The maximum value in the interval.
+    """
+
+    def __init__(self, min_value, max_value):
+        self.__init_handle_by_constructor__(_ffi_api.IntervalSet, min_value, max_value)
+
+
+@tvm_ffi.register_object("arith.PresburgerSet")
+class PresburgerSet(IntSet):
+    """Represent of Presburger Set"""
+
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_api.PresburgerSet)
+
+
+def estimate_region_lower_bound(region, var_dom, predicate, analyzer=None):
+    """Analyze the region with affine map, given the domain of variables and their predicate
+    Some subregion may be discarded during the lower-bound analysis.
+
+    Parameters
+    ----------
+    region : List[Range]
+        The region to be analyzed.
+
+    var_dom : Dict[tvm.tirx.Var, Range]
+        The ranges of the variables
+
+    predicate : Expr
+        The predicate for the affine map
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    ----------
+    region_int_set : Optional[List[IntSet]]
+        None if the detection fails, or an array of IntSets as the result of analysis
+    """
+    return _ffi_api.EstimateRegionLowerBound(region, var_dom, predicate, analyzer)
+
+
+def estimate_region_strict_bound(region, var_dom, predicate, analyzer=None):
+    """Analyze the region with affine map, given the domain of variables and their predicate
+    The result should be strict, i.e. no region is discarded or relaxed.
+
+    Parameters
+    ----------
+    region : List[Range]
+        The region to be analyzed.
+
+    var_dom : Dict[tvm.tirx.Var, Range]
+        The ranges of the variables
+
+    predicate : Expr
+        The predicate for the affine map
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    ----------
+    region_int_set : Optional[List[IntSet]]
+        None if the detection fails, or an array of IntSets as the result of analysis
+    """
+    return _ffi_api.EstimateRegionStrictBound(region, var_dom, predicate, analyzer)
+
+
+def estimate_region_upper_bound(region, var_dom, predicate, analyzer=None):
+    """Analyze the region with affine map, given the domain of variables and their predicate
+    Relaxation of the region may be used in upper-bound analysis,
+    i.e. some extra region may be added to the result.
+
+    Parameters
+    ----------
+    region : List[Range]
+        The region to be analyzed.
+
+    var_dom : Dict[tvm.tirx.Var, Range]
+        The ranges of the variables
+
+    predicate : Expr
+        The predicate for the affine map
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    ----------
+    region_int_set : List[IntSet]
+        an array of IntSets as the result of analysis
+    """
+    return _ffi_api.EstimateRegionUpperBound(region, var_dom, predicate, analyzer)
+
+
+def pos_inf():
+    """Returns the symbolic positive infinity
+
+    Returns
+    ----------
+    pos_inf : tvm.tirx.Var
+        A symbolic var that indicates positive infinity
+    """
+    return _ffi_api.PosInf()
+
+
+def neg_inf():
+    """Returns the symbolic positive infinity
+
+    Returns
+    ----------
+    neg_inf : tvm.tirx.Var
+        A symbolic var that indicates positive infinity
+    """
+    return _ffi_api.NegInf()
+
+
+def union_lower_bound(sets):
+    """Create a lower-bound of union set, where some of the segments may be dropped
+
+    Parameters
+    ----------
+    sets : List[IntSet]
+        The sets to be combined
+
+    Returns
+    ----------
+    union_lower_bound : List[IntSet]
+        An N-dimensional integer set, the lower bound of the union
+    """
+    return _ffi_api.UnionLowerBound(sets)

python/tvm/arith/int_solver.py

@@ -0,0 +1,183 @@
+# 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.
+"""integer constraints data structures and solvers"""
+
+import tvm_ffi
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("arith.IntGroupBounds")
+class IntGroupBounds(Object):
+    """Represent integer grouped bounds which are classified into
+       lower bounds (include), upper bounds (include) and equalities.
+
+    Parameters
+    ----------
+    coef : tvm.ir.Expr
+        The coefficient. Must be integer type.
+        coef * var >= lower
+        coef * var == equal
+        coef * var >= upper
+    lower : List[tvm.ir.Expr]
+        the lower bounds (include)
+    equal : List[tvm.ir.Expr]
+        equalities
+    upper : List[tvm.ir.Expr]
+        the upper bounds (include)
+    """
+
+    def __init__(self, coef, lower, equal, upper):
+        self.__init_handle_by_constructor__(_ffi_api.IntGroupBounds, coef, lower, equal, upper)
+
+    @staticmethod
+    def from_range(rng):
+        """Construct a IntGroupedBounds by Range.
+
+        Parameters
+        ----------
+        rng : tvm.ir.Range
+
+
+        Returns
+        -------
+        ret : Range
+            The constructed range.
+        """
+        return _ffi_api.IntGroupBounds_from_range(rng)
+
+    def find_best_range(self):
+        """Return the best range from the grouped bounds.
+        None if (-inf, +inf).
+        """
+        return _ffi_api.IntGroupBounds_FindBestRange(self)
+
+
+@tvm_ffi.register_object("arith.IntConstraints")
+class IntConstraints(Object):
+    """Represent a set of integer constraints including variables, their ranges and
+       the relations between them (either equations or inequalities)
+
+    Parameters
+    ----------
+    variables : List[tvm.tirx.Var]
+        The variables in the constraints. Must be integers
+    ranges    : Map[tvm.tirx.Var, tvm.ir.Range]
+        The ranges of the variables.
+    relations : List[tvm.ir.Expr]
+        The relations between the variables (either equations or inequalities)
+    """
+
+    def __init__(self, variables, ranges, relations):
+        self.__init_handle_by_constructor__(_ffi_api.IntConstraints, variables, ranges, relations)
+
+
+@tvm_ffi.register_object("arith.IntConstraintsTransform")
+class IntConstraintsTransform(Object):
+    """We can have different set of variables to represent the same integer constraints.
+       For example, the following two constrains are equivalent,
+       {a + b = 0 | a >= 0, b >= 0} and
+       {m - n = 0 | m >= 0, n <= 0}
+       This data structure represents the transformation
+       between two equivalent integer constraints.
+       In the above example,
+       src        : {a + b = 0 | a >= 0, b >= 0}
+       dst        : {m - n = 0 | m >= 0, n <= 0}
+       src_to_dst : {a -> m, b -> -n}
+       dst_to_src : {m -> a, n -> -b}
+
+    Parameters
+    ----------
+    src : arith.IntConstraints
+        source integer constraints, e.g., {a + b = 0 | a >= 0, b >= 0}
+    dst : arith.IntConstraints
+        integer constraints equivalent to the source, e.g., {m - n = 0 | m >= 0, n <= 0}
+    src_to_dst : Map[tvm.tirx.Var, tvm.ir.Expr]
+        mapping from variables in the src to the variables in the dst,
+                e.g., {a -> m, b -> -n}
+    dst_to_src : Map[tvm.tirx.Var, tvm.ir.Expr]
+        mapping from variables in the dst to the variables in the src,
+        e.g., {m -> a, n -> -b}
+    """
+
+    def __init__(self, src, dst, src_to_dst, dst_to_src):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IntConstraintsTransform, src, dst, src_to_dst, dst_to_src
+        )
+
+
+def solve_linear_equations(equations, variables=None, ranges=None):
+    """Solve linear equations.
+
+    Parameters
+    ----------
+    equations: List[tvm.ir.Expr] or IntConstraints
+        The equations of the variables
+    variables : Optional[List[tvm.tirx.Var]]
+        The variables in the system.
+    ranges    : Optional[Map[tvm.tirx.Var, tvm.ir.Range]]
+        The ranges of the variables.
+
+    Returns
+    -------
+    int_constraints_transform : IntConstraintsTransform
+        New integer constraints, with less variables (if the problem is NOT of full rank),
+        or no variable (if the problem is of full rank),
+        or an empty integer constraints (if the problem is unsolvable).
+        It also provides the ranges of the variables in the new system,
+        as well as inequalities inferred from the problem.
+        You can get the mapping from the original variables to the solution via
+        int_constraints_transform.src_to_dst.
+    """
+    if isinstance(equations, IntConstraints):
+        return _ffi_api.SolveLinearEquations(equations)
+    return _ffi_api.SolveLinearEquations(variables, ranges, equations)
+
+
+def solve_linear_inequalities(equations, variables=None, ranges=None, deskew_range=False):
+    """Solve linear inequalities.
+
+    Parameters
+    ----------
+    equations   : List[tvm.ir.Expr] or IntConstraints
+        The inequalities of the variables
+    variables   : Optional[List[tvm.tirx.Var]]
+        The variables in the system.
+    ranges      : Optional[Map[tvm.tirx.Var, tvm.ir.Range]]
+        The ranges of the variables.
+    deskew_range: Optional[bool]
+        Whether deskew the result ranges to be started from zero.
+        Default false.
+
+    Returns
+    -------
+    ret_ranges: IntConstraints or IntConstraintsTransform
+        The result ranges for each variables.
+        Constrains that cannot be transformed to Range will be stored in IntConstraints.relations.
+        If deskew_range is set (=True), the result ranges will be deskewed to be started from zero.
+        New variables are created accordingly therefore IntConstraintsTransform is returned.
+    """
+    solver = (
+        _ffi_api.SolveInequalitiesDeskewRange if deskew_range else _ffi_api.SolveInequalitiesToRange
+    )
+    if isinstance(equations, IntConstraints):
+        assert variables is None
+        assert ranges is None
+        return solver(equations)
+    return solver(variables, ranges, equations)

python/tvm/arith/iter_affine_map.py

@@ -0,0 +1,375 @@
+# 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.
+"""Iterator (quasi)affine mapping patterns."""
+
+from enum import IntEnum
+
+import tvm_ffi
+
+from tvm.ir import Expr
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("arith.IterMapExpr")
+class IterMapExpr(Expr):
+    """Base class of all IterMap expressions."""
+
+
+@tvm_ffi.register_object("arith.IterMark")
+class IterMark(Object):
+    """Mark the source as an iterator in [0, extent).
+
+    Parameters
+    ----------
+    source : Expr.
+        The source expression.
+
+    extent : Expr
+        The extent of the iterator.
+    """
+
+    def __init__(self, source, extent):
+        self.__init_handle_by_constructor__(_ffi_api.IterMark, source, extent)
+
+
+@tvm_ffi.register_object("arith.IterSplitExpr")
+class IterSplitExpr(IterMapExpr):
+    """Split of an iterator.
+
+    result = floormod(floordiv(source, lower_factor), extent) * scale
+
+    Parameters
+    ----------
+    source : IterMark
+        The source marked iterator.
+
+    lower_factor : Expr
+        The lower factor to split the domain.
+
+    extent : Expr
+        The extent of the split.
+
+    scale : Expr
+        Additional scale to the split.
+    """
+
+    def __init__(self, source, lower_factor, extent, scale):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IterSplitExpr, source, lower_factor, extent, scale
+        )
+
+
+@tvm_ffi.register_object("arith.IterSumExpr")
+class IterSumExpr(IterMapExpr):
+    """Fuse multiple iterators by summing them with scaling.
+
+    result = sum(args) + base
+
+    Parameters
+    ----------
+    args : List[IterSplitExpr]
+        The input to the sum expression.
+
+    base : Expr
+        The base offset.
+    """
+
+    def __init__(self, args, base):
+        self.__init_handle_by_constructor__(_ffi_api.IterSumExpr, args, base)
+
+
+@tvm_ffi.register_object("arith.IterMapResult")
+class IterMapResult(Object):
+    """Result of iter map detection."""
+
+
+class IterMapLevel(IntEnum):
+    """Possible kinds of iter mapping check level."""
+
+    Bijective = 0
+    Surjective = 1
+    NoCheck = 3
+
+    @staticmethod
+    def from_str(name: str):
+        """Helper to create level enum from string"""
+        if name is None:
+            return IterMapLevel.NoCheck
+        name = name.lower()
+        if name == "bijective":
+            check_level = IterMapLevel.Bijective
+        elif name == "surjective":
+            check_level = IterMapLevel.Surjective
+        elif name == "nocheck":
+            check_level = IterMapLevel.NoCheck
+        else:
+            raise ValueError(f"Unknown check level {name}")
+        return check_level
+
+
+def detect_iter_map(
+    indices,
+    input_iters,
+    predicate=True,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+    analyzer=None,
+):
+    """Detect if indices can be written as mapped iters from input iters
+
+    Parameters
+    ----------
+    indices : List[Expr]
+        The input indices
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of each input iterators.
+
+    predicate : Expr
+        The predicate constraints on the input iterators
+
+    check_level : Union[str, IterMapLevel]
+        Checking level of iteration mapping
+
+    simplify_trivial_iterators: bool
+        If true, iterators with extent of 1 will be replaced with a
+        constant value.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    -------
+    results : IterMapResult
+        The iter map matching result.
+        The result's .indices is empty array if no match can be found.
+
+    """
+    if isinstance(check_level, str):
+        check_level = IterMapLevel.from_str(check_level)
+    elif check_level is None:
+        check_level = IterMapLevel.NoCheck
+    return _ffi_api.DetectIterMap(
+        indices, input_iters, predicate, check_level, simplify_trivial_iterators, analyzer
+    )
+
+
+def normalize_to_iter_sum(index, input_iters, analyzer=None):
+    """Normalize expr to iter sum.
+
+    The normalized result ensures that
+    each scale is in the form of (symbol_prod) * cscale
+    It will also sort in desc order by cscale then len(symbol_prod).
+
+    Parameters
+    ----------
+    index : Expr
+        The input index
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of each input iterators.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    -------
+    iter_sum: IterSumExpr
+        The result iter sum
+
+    Note
+    ----
+    This function does best effort detection, so some undetected
+    part can go into iter_sum.base
+
+    This function is useful to decide the stride multiplier and
+    division factor in buffer access patterns.
+    """
+    return _ffi_api.NormalizeToIterSum(index, input_iters, analyzer)
+
+
+def iter_map_simplify(
+    indices,
+    input_iters,
+    predicate=True,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+    analyzer=None,
+):
+    """Simplify the indices using iter map detection.
+
+    Parameters
+    ----------
+    indices : List[Expr]
+        The input indices
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of each input iterators.
+
+    predicate : Expr
+        The predicate constraints on the input iterators
+
+    check_level : Union[str, IterMapLevel]
+        Checking level of iteration mapping
+
+    simplify_trivial_iterators: bool
+        If true, iterators with extent of 1 will be replaced with a
+        constant value.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    -------
+    results : IterMapResult
+        The iter map matching result.
+        The result's .indices is empty array if no match can be found.
+
+    """
+    if isinstance(check_level, str):
+        check_level = IterMapLevel.from_str(check_level)
+    elif check_level is None:
+        check_level = IterMapLevel.NoCheck
+    return _ffi_api.IterMapSimplify(
+        indices, input_iters, predicate, check_level, simplify_trivial_iterators, analyzer
+    )
+
+
+def normalize_iter_map_to_expr(expr):
+    """Given an IterMapExpr, transform it to normal Expr
+
+    Parameters
+    ----------
+    expr : IterMapExpr
+        the input IterMapExpr
+
+    Returns
+    -------
+    result : Expr
+        the corresponding normal Expr
+    """
+    return _ffi_api.NormalizeIterMapToExpr(expr)
+
+
+def subspace_divide(
+    bindings,
+    input_iters,
+    sub_iters,
+    predicate=True,
+    check_level=IterMapLevel.Surjective,
+    simplify_trivial_iterators=True,
+    analyzer=None,
+):
+    """Detect if bindings can be written as
+    ``[a_0*e_0 + b_0 + c_0, a_1*e_1 + b_1, ..., a_n*e_n + b_n]``
+
+    where::
+
+        a = some-quasi-affine-iter-map(input_iters set_minus sub_iters)
+        b = some-quasi-affine-iter-map(sub_iters)
+        c is constant symbols
+        e is the extent of b
+
+    For example::
+
+        z*12 + y*3 + x + c = (z*4+y)*3 + x
+        bindings = [z*12 + y*3 + x + c]
+        input_iters = [z, y, x]
+        sub_iter = [x]
+        Then the result will be [a, b] where
+        a = [z*4 + y]
+        b = [x]
+
+    Parameters
+    ----------
+    bindings : List[Expr]
+        The input bindings
+
+    input_iters : Map[tvm.tirx.Var, Range]
+        The domain of input iterator, which is the basis of the whole space
+
+    sub_iters : Array[tvm.tirx.Var]
+        The subset of input_iters, which is the basis of the subspace
+
+    predicate : Expr
+        The predicate constraints on the input iterators
+
+    check_level : Union[str, IterMapLevel]
+        Checking level of iteration mapping
+
+    simplify_trivial_iterators: bool
+        If true, iterators with extent of 1 will be replaced with a
+        constant value.
+
+    analyzer : Optional[tvm.arith.Analyzer]
+        The analyzer to use.  When provided, its accumulated bindings and
+        constraints are reused; otherwise a fresh analyzer is created.
+
+    Returns
+    -------
+    results : List[List[Expr]]
+        The result list has length ``len(bindings) + 1``.
+
+        - ``[0, len(bindings))``: The iter map matching result.
+          The inner list is of length 2. The first expr is the basis
+          of the quotient space. The second expr is the basis of the subspace.
+        - ``len(bindings)``: the predicate of outer space and inner space.
+        - Empty array if no match can be found.
+    """
+    if isinstance(check_level, str):
+        check_level = IterMapLevel.from_str(check_level)
+    return _ffi_api.SubspaceDivide(
+        bindings,
+        input_iters,
+        sub_iters,
+        predicate,
+        check_level,
+        simplify_trivial_iterators,
+        analyzer,
+    )
+
+
+def inverse_affine_iter_map(iter_map, outputs):
+    """Apply the inverse of the affine transformation to the outputs.
+    Similar to the back-propagation, starting from the outputs, it visits the DAG of the expressions
+    in reverse topology order and applies the inverse of the affine transformation until it reaches
+    the input. The affine iter map is required to be bijective.
+
+    For example, iter_map = [l0 // 16, l0 % 16], outputs = [output_0, output_1],
+    the affine transformation specified by `iter_map` will be applied to `outputs` and the result
+    will be {l0: ((output_0*16) + output_1)}.
+
+    See also :any:`detect_iter_map`.
+
+    Parameters
+    ----------
+    iter_map : List[IterSumExpr]
+        The bijective affine iter map.
+    outputs : List[Expr]
+        The outputs of the affine transformation.
+
+    Returns
+    -------
+    results : Map[tvm.tirx.Var, Expr]
+        The map from the input to the transformed result.
+    """
+    return _ffi_api.InverseAffineIterMap(iter_map, outputs)

python/tvm/arith/pattern.py

@@ -0,0 +1,61 @@
+# 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.
+"""Detect common patterns."""
+
+from . import _ffi_api
+
+
+def detect_linear_equation(expr, var_list):
+    """Match `expr = sum_{i=0}^{n-1} var[i] * coeff[i] + coeff[n]`
+
+    Where coeff[i] and base are invariant of var[j] for all i and j.
+
+    Parameters
+    ----------
+    expr : Expr
+        The expression to be matched.
+
+    var_list : List[tvm.tirx.Var]
+        A list of variables.
+
+    Returns
+    -------
+    coeff : List[Expr]
+        A list of co-efficients if the match is successful.
+        An empty list if the match failed.
+    """
+    return _ffi_api.DetectLinearEquation(expr, var_list)
+
+
+def detect_clip_bound(expr, var_list):
+    """Detect if expression corresponds to clip bound of the vars
+
+    Parameters
+    ----------
+    expr : Expr
+        The expression to be matched.
+
+    var_list : List[tvm.tirx.Var]
+        A list of variables.
+
+    Returns
+    -------
+    coeff : List[Expr]
+        `concat([min_value[i], max_value[i]] for i, v in enumerate(var_list))`
+        An empty list if the match failed.
+    """
+    return _ffi_api.DetectClipBound(expr, var_list)

python/tvm/backend/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+"""Backend-owned Python modules and load hooks."""
+
+from __future__ import annotations
+
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)  # type: ignore[name-defined]
+
+from .loader import is_loaded, load
+
+__all__ = ["is_loaded", "load"]

python/tvm/backend/_autoload_backends.py

@@ -0,0 +1,88 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Autoload built-in and out-of-tree backend libraries and registration hooks."""
+
+from __future__ import annotations
+
+import os
+import warnings
+from importlib.metadata import entry_points
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.backend.loader import load
+from tvm.base import _LOADED_LIBS
+
+_BUILTIN_BACKENDS = (
+    "cuda",
+    "metal",
+    "rocm",
+    "trn",
+    "opencl",
+    "vulkan",
+    "webgpu",
+    "hexagon",
+    "adreno",
+)
+
+_AUTO_LOAD_DONE = False
+
+
+def _load_builtin_backends() -> None:
+    """Load all in-tree backend Python hooks."""
+    for name in _BUILTIN_BACKENDS:
+        load(name)
+
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Load libtvm_runtime_extra if available for registration side effects.
+        _LOADED_LIBS["tvm_runtime_extra"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_extra",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    return None
+
+
+def _autoload_backends() -> None:
+    """Load built-in backends and invoke out-of-tree backend entry points."""
+    global _AUTO_LOAD_DONE
+    if _AUTO_LOAD_DONE:
+        return
+    _AUTO_LOAD_DONE = True
+
+    if os.environ.get("TVM_DEVICE_BACKEND_AUTOLOAD", "1") == "0":
+        return
+
+    from tvm import _RUNTIME_ONLY  # pylint: disable=import-outside-toplevel
+
+    if not _RUNTIME_ONLY:
+        _load_builtin_backends()
+
+    # Out-of-tree extensions opt into being loaded automatically at ``import tvm`` time
+    # by declaring an entry point in the ``tvm.backends`` group:
+    # [project.entry-points."tvm.backends"] tvm_foo = "tvm_foo:_autoload".
+    # Autoload can be disabled via ``TVM_DEVICE_BACKEND_AUTOLOAD=0``.
+    for entry_pt in entry_points(group="tvm.backends"):
+        try:
+            entry_pt.load()()
+        except Exception as e:  # pylint: disable=broad-except
+            warnings.warn(f"Failed to autoload tvm backend '{entry_pt.name}': {e}")

python/tvm/backend/adreno/__init__.py

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Adreno-owned backend hooks."""
+
+from importlib import import_module
+
+_LAZY_SUBMODULES = {"target_tags"}
+
+
+def register_backend():
+    """Register Adreno-owned Python semantics."""
+    import tvm.backend as backend  # pylint: disable=import-outside-toplevel
+
+    backend.load("opencl")
+    backend.load("vulkan")
+    import_module(f"{__name__}.target_tags")
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["register_backend", "target_tags"]

python/tvm/backend/adreno/target_tags.py

@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Qualcomm Adreno GPU target tags."""
+
+from tvm.target import register_tag
+
+register_tag(
+    "qcom/adreno-opencl",
+    {
+        "kind": "opencl",
+        "device": "adreno",
+        "keys": ["adreno", "opencl", "gpu"],
+    },
+)
+
+register_tag(
+    "qcom/adreno-opencl-clml",
+    {
+        "kind": "opencl",
+        "device": "adreno",
+        "keys": ["adreno", "opencl", "gpu", "clml"],
+    },
+)
+
+register_tag(
+    "qcom/adreno-opencl-texture",
+    {
+        "kind": "opencl",
+        "device": "adreno",
+        "keys": ["adreno", "opencl", "gpu", "texture"],
+    },
+)
+
+register_tag(
+    "qcom/adreno-vulkan",
+    {
+        "kind": "vulkan",
+        "device": "adreno",
+        "keys": ["adreno", "vulkan", "gpu"],
+    },
+)
+
+register_tag(
+    "qcom/adreno-vulkan-texture",
+    {
+        "kind": "vulkan",
+        "device": "adreno",
+        "keys": ["adreno", "vulkan", "gpu", "texture"],
+    },
+)

python/tvm/backend/cuda/__init__.py

@@ -0,0 +1,103 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CUDA-owned TIRx modules."""
+
+from importlib import import_module
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.base import _LOADED_LIBS
+
+_LAZY_SUBMODULES = {"lang", "op", "operator", "script", "target_tags"}
+
+
+def _detect_target_from_device(dev):
+    from tvm.target import Target  # pylint: disable=import-outside-toplevel
+
+    return Target(
+        {
+            "kind": "cuda",
+            "max_shared_memory_per_block": dev.max_shared_memory_per_block,
+            "max_threads_per_block": dev.max_threads_per_block,
+            "thread_warp_size": dev.warp_size,
+            "arch": "sm_" + dev.compute_version.replace(".", ""),
+        }
+    )
+
+
+def register_backend():
+    """Register CUDA-owned Python semantics."""
+    from tvm.target.detect_target import register_device_target_detector
+    from tvm.tirx.script.builder import ir as builder_ir  # pylint: disable=import-outside-toplevel
+
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Runtime sidecars only need registration side effects; libtvm_runtime is global.
+        _LOADED_LIBS["tvm_runtime_cuda"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_cuda",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    register_device_target_detector("cuda", _detect_target_from_device)
+    for name, namespace in script_namespaces().items():
+        builder_ir.register_script_namespace(name, namespace)
+
+    import_module(f"{__name__}.operator.intrinsics")
+    import_module(f"{__name__}.operator.tile_primitive")
+    import_module(f"{__name__}.target_tags")
+
+
+def script_namespaces(**_):
+    """Return CUDA-owned TVMScript namespaces."""
+    from .script import (  # pylint: disable=import-outside-toplevel
+        CUDANamespace,
+        NVSHMEMNamespace,
+        PTXNamespace,
+    )
+
+    return {
+        "cuda": CUDANamespace(),
+        "nvshmem": NVSHMEMNamespace(),
+        "ptx": PTXNamespace(),
+    }
+
+
+def script_namespace(**kwargs):
+    """Return the CUDA TVMScript namespace object."""
+    return script_namespaces(**kwargs)["cuda"]
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "lang",
+    "op",
+    "operator",
+    "register_backend",
+    "script",
+    "script_namespace",
+    "script_namespaces",
+    "target_tags",
+]

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

@@ -0,0 +1,70 @@
+# 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.
+"""CUDA-specific TIRx language helpers."""
+
+from importlib import import_module
+
+__all__ = [
+    "BaseTileScheduler",
+    "ClusterPersistentScheduler2D",
+    "FlashAttentionLPTScheduler",
+    "FlashAttentionLinearScheduler",
+    "GroupMajor3D",
+    "IndexedTripleTileScheduler",
+    "MBarrier",
+    "Pipeline",
+    "PipelineState",
+    "RankAwareGroupMajorTileScheduler",
+    "SMEMPool",
+    "SmemDescriptor",
+    "TCGen05Bar",
+    "TMABar",
+    "TMEMPool",
+    "TMEMStages",
+    "WarpRole",
+    "WarpgroupRole",
+]
+
+_HELPER_MODULES = {
+    "MBarrier": ".pipeline",
+    "Pipeline": ".pipeline",
+    "PipelineState": ".pipeline",
+    "BaseTileScheduler": ".tile_scheduler",
+    "ClusterPersistentScheduler2D": ".tile_scheduler",
+    "FlashAttentionLPTScheduler": ".tile_scheduler",
+    "FlashAttentionLinearScheduler": ".tile_scheduler",
+    "GroupMajor3D": ".tile_scheduler",
+    "IndexedTripleTileScheduler": ".tile_scheduler",
+    "RankAwareGroupMajorTileScheduler": ".tile_scheduler",
+    "SMEMPool": ".alloc_pool",
+    "TCGen05Bar": ".pipeline",
+    "TMABar": ".pipeline",
+    "TMEMPool": ".alloc_pool",
+    "TMEMStages": ".alloc_pool",
+    "WarpRole": ".warp_role",
+    "WarpgroupRole": ".warp_role",
+    "SmemDescriptor": ".smem_desc",
+}
+
+
+def __getattr__(name: str):
+    module_name = _HELPER_MODULES.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(module_name, __name__), name)
+    globals()[name] = value
+    return value

python/tvm/backend/cuda/lang/alloc_pool.py

@@ -0,0 +1,529 @@
+# 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.
+"""SMEM and TMEM bump-allocator pools for TIRX kernels."""
+
+from __future__ import annotations
+
+import functools
+import operator
+
+from tvm import DataType
+from tvm.tirx.layout import S, TCol, TileLayout, TLane
+
+# ---------------------------------------------------------------------------
+# ir_builder helpers — imported lazily to avoid circular deps at module level
+# ---------------------------------------------------------------------------
+
+_ir = None
+
+
+def _get_ir():
+    global _ir
+    if _ir is None:
+        from tvm.tirx.script.builder import ir as _mod
+
+        _ir = _mod
+    return _ir
+
+
+def _get_frame():
+    from tvm.tirx.script.builder import frame
+
+    return frame
+
+
+# ---------------------------------------------------------------------------
+# Shared utilities
+# ---------------------------------------------------------------------------
+
+_POOL_UNSET = object()
+
+
+def _default_tmem_layout(rows, cols):
+    return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)])
+
+
+def _emit_stmt(expr):
+    ir = _get_ir()
+    ir.add_to_parent(ir.evaluate(expr))
+
+
+def _shape_product(shape):
+    return functools.reduce(operator.mul, shape, 1)
+
+
+def _auto_swizzle_mode(dtype):
+    """Select the default MMA swizzle mode for a shared-memory allocation."""
+    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+    del dtype
+    return SwizzleMode.SWIZZLE_128B_ATOM
+
+
+def _swizzle_atom_bytes(swizzle_mode):
+    """Return the row width (in bytes) of one swizzle atom for *swizzle_mode*."""
+    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+    return {
+        SwizzleMode.SWIZZLE_NONE: 0,
+        SwizzleMode.SWIZZLE_32B_ATOM: 32,
+        SwizzleMode.SWIZZLE_64B_ATOM: 64,
+        SwizzleMode.SWIZZLE_128B_ATOM: 128,
+    }[swizzle_mode]
+
+
+def _suggest_swizzle_for_row_bytes(row_bytes):
+    """Pick the largest valid swizzle mode whose atom row fits within *row_bytes*."""
+
+    for atom_bytes, mode in (
+        (128, "SWIZZLE_128B_ATOM"),
+        (64, "SWIZZLE_64B_ATOM"),
+        (32, "SWIZZLE_32B_ATOM"),
+    ):
+        if row_bytes >= atom_bytes and row_bytes % atom_bytes == 0:
+            return mode
+    return "SWIZZLE_NONE"
+
+
+def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
+    """Validate that *shape* / *dtype* / *swizzle_mode* are mutually compatible.
+
+    ``mma_shared_layout`` tiles a swizzle atom of shape ``[8, swizzle_bytes / dtype_bytes]``
+    over the last two logical dimensions of *shape*. If the row width or row count of
+    the request is smaller than (or not a multiple of) the atom, the underlying
+    ``Layout.tile_to`` lowers to a ``floordiv``/``floormod`` by zero and raises an
+    opaque internal "Divide by zero" diagnostic from ``tile_tile_ops.cc``. Catch the
+    misconfiguration here so callers see *what* is wrong and *how* to fix it.
+
+    Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
+    """
+    from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+    if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
+        return
+
+    if len(shape) < 2:
+        raise ValueError(
+            f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; "
+            f"swizzled MMA layouts tile over the last two dims (rows, cols). "
+            f"Use swizzle_mode='none' for 1-D allocations."
+        )
+
+    # Only validate concrete int dims; symbolic dims fall through (the analyzer
+    # in C++ will still ICHECK on them, but at least we don't false-positive).
+    rows = shape[-2]
+    cols = shape[-1]
+    if not (isinstance(rows, int) and isinstance(cols, int)):
+        return
+
+    dtype_bytes = DataType(dtype).bits // 8
+    if dtype_bytes == 0:
+        # Sub-byte dtype (e.g. float4); ``cols`` is already in element units, so
+        # use a fractional check expressed via bits.
+        col_bits = cols * DataType(dtype).bits
+        atom_bits = _swizzle_atom_bytes(swizzle_mode) * 8
+        if col_bits < atom_bits or col_bits % atom_bits != 0:
+            row_bytes = col_bits // 8 if col_bits % 8 == 0 else col_bits / 8
+            atom_bytes = _swizzle_atom_bytes(swizzle_mode)
+            suggestion = _suggest_swizzle_for_row_bytes(col_bits // 8 if col_bits >= 8 else 0)
+            raise ValueError(
+                f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
+                f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
+                f"swizzle atom selected by {swizzle_mode.name}. "
+                f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
+                f"to a multiple of "
+                f"{(atom_bits + DataType(dtype).bits - 1) // DataType(dtype).bits} elements."
+            )
+    else:
+        row_bytes = cols * dtype_bytes
+        atom_bytes = _swizzle_atom_bytes(swizzle_mode)
+        if row_bytes < atom_bytes or row_bytes % atom_bytes != 0:
+            suggestion = _suggest_swizzle_for_row_bytes(row_bytes)
+            min_cols = atom_bytes // dtype_bytes
+            raise ValueError(
+                f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces "
+                f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B "
+                f"swizzle atom selected by {swizzle_mode.name}. "
+                f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] "
+                f"to a multiple of {min_cols} elements (>= {atom_bytes}B at {dtype})."
+            )
+
+    # Atom rows is always 8 (see ``mma_atom_shape`` in tma_utils.py).
+    atom_rows = 8
+    if rows < atom_rows or rows % atom_rows != 0:
+        raise ValueError(
+            f"alloc_mma shape={tuple(shape)} has shape[-2]={rows}, but the "
+            f"{swizzle_mode.name} atom requires shape[-2] to be a positive "
+            f"multiple of {atom_rows}. Use swizzle_mode='none', or widen shape[-2] "
+            f"to a multiple of {atom_rows}."
+        )
+
+
+# ---------------------------------------------------------------------------
+# TMEMStages
+# ---------------------------------------------------------------------------
+
+
+def _meta_class(cls):
+    """Apply @meta_class decorator from ir_builder."""
+    return _get_ir().meta_class(cls)
+
+
+@_meta_class
+class TMEMStages:
+    """Parse-time staged view over a TMEM buffer.
+
+    Parameters
+    ----------
+    buf : Buffer
+        The underlying TMEM buffer (e.g. f32 or f16 view).
+    col_start : int
+        First column of stage 0 in *buf*'s column space.
+    width : int
+        Number of columns per stage.
+    stages : int
+        Number of pipeline stages (default 1).
+    stride : int or None
+        Column distance between consecutive stages.  When *None* (default),
+        equals *width* (stages are packed back-to-back).
+    """
+
+    def __init__(self, buf, col_start, width, stages=1, stride=None):
+        self.buf = buf
+        self.col_start = col_start
+        self.width = width
+        self.stages = stages
+        self.stride = width if stride is None else stride
+
+    def _stage_base(self, stage):
+        return self.col_start + stage * self.stride
+
+    def __getitem__(self, item):
+        if isinstance(item, tuple):
+            assert len(item) == 2, "TMEMStages expects region[stage] or region[stage, start:stop]"
+            stage, col_slice = item
+            assert isinstance(col_slice, slice), "TMEMStages tuple indexing requires a slice"
+            base = self._stage_base(stage)
+            start = 0 if col_slice.start is None else col_slice.start
+            stop = self.width if col_slice.stop is None else col_slice.stop
+            return self.buf[:, base + start : base + stop : col_slice.step]
+        base = self._stage_base(item)
+        return self.buf[:, base : base + self.width]
+
+
+# ---------------------------------------------------------------------------
+# TMEMPool
+# ---------------------------------------------------------------------------
+
+
+@_meta_class
+class TMEMPool:
+    """Bump allocator over TMEM columns."""
+
+    def __init__(
+        self,
+        pool,
+        total_cols=512,
+        *,
+        cta_group=1,
+        alloc_warp=0,
+        dealloc_warp=None,
+        tmem_addr=None,
+        sync_after_alloc=True,
+    ):
+        # tcgen05 alloc/dealloc are warp-uniform PTX instructions: every lane
+        # in the chosen warp must participate, and exactly one warp in the
+        # CTA must execute them. The pool emits its own
+        # ``if warp_id() == target_warp: tcgen05.alloc(...)``
+        # guard, using the cta->warp scope id ``T.warp_id()``.
+        # NOTE: synccheck currently false-deadlocks on kernels that declare a
+        # second warp-scope id (cpusim binds only one warp var); the generated
+        # CUDA is equivalent to ``thread_rank() // 32 == target_warp``.
+        self.pool = pool
+        self.total_cols = total_cols
+        self.cta_group = cta_group
+        self.alloc_warp = alloc_warp
+        self.dealloc_warp = alloc_warp if dealloc_warp is None else dealloc_warp
+        self.sync_after_alloc = sync_after_alloc
+        self.offset = 0
+        self.max_offset = 0
+        self._committed = False
+        self._deallocated = False
+        self._addr_buf = pool.alloc([1], "uint32", align=4) if tmem_addr is None else tmem_addr
+
+    def _addr_slot(self):
+        try:
+            return self._addr_buf[0]
+        except TypeError:
+            return self._addr_buf
+
+    @property
+    def addr(self):
+        return self._addr_slot()
+
+    def _emit_warp_guard(self, target_warp, emit):
+        from tvm.script import tirx as T
+
+        warp_id = T.warp_id()
+        with T.If(warp_id == target_warp):
+            with T.Then():
+                emit()
+
+    def _resolve_cols(self, shape, dtype, cols, layout=None):
+        if cols is not None:
+            return cols
+        bits = DataType(dtype).bits
+        if layout is not None:
+            # span("TCol") is in *element* (buffer dtype) units; one TMEM cell
+            # holds 32 bits regardless of the element type.
+            tcol_elems = int(layout.span("TCol"))
+            tcol_bits = tcol_elems * bits
+            assert tcol_bits % 32 == 0, (
+                f"layout TCol span={tcol_elems} elems x {bits}b is not 32-bit aligned"
+            )
+            return tcol_bits // 32
+        assert len(shape) == 2, "TMEMPool.alloc() requires cols= for non-2D TMEM buffers"
+        total_bits = _shape_product(shape) * bits
+        rows = shape[0]
+        assert total_bits % (32 * rows) == 0, (
+            f"Cannot infer TMEM columns from shape={shape}, dtype={dtype!r}; "
+            "please pass cols= explicitly"
+        )
+        return total_bits // (32 * rows)
+
+    def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None):
+        """Allocate a TMEM buffer.
+
+        Parameters
+        ----------
+        shape, dtype, cols
+            Standard buffer shape / dtype / column count.
+        layout
+            Explicit ``TileLayout``. Mutually exclusive with ``datapath``.
+        datapath : str | None
+            Optional tcgen05 datapath letter (``"D"`` for M=128 full datapath,
+            ``"F"`` for M=64 non-``.ws`` scattered). When provided, the buffer's
+            layout is derived from ``tmem_datapath_layout(datapath, *shape)``
+            so the row index reflects the *physical* TMEM lane occupation
+            (PTX ISA §9.7.16.10.5). The downstream ``.16x*b`` / ``.32x32b``
+            dispatches structurally check this layout to catch mismatched
+            atoms (e.g. a ``.16x*b`` M=128 read against a Layout F buffer).
+            Defaults to ``None``, which means Layout D's identity row→lane
+            mapping — keep this for shape ``(128, X)`` buffers that hold
+            an M=128 MMA accumulator.
+        """
+        from tvm.tirx.layout import tmem_datapath_layout
+
+        if layout is not None and datapath is not None:
+            raise ValueError("TMEMPool.alloc: pass at most one of layout= and datapath=")
+        if datapath is not None:
+            assert len(shape) == 2, "TMEMPool.alloc: datapath= requires a 2-D shape"
+            layout = tmem_datapath_layout(datapath, shape[0], shape[1])
+
+        ir = _get_ir()
+        cols = self._resolve_cols(shape, dtype, cols, layout)
+        col_start = self.offset
+        col_end = col_start + cols
+        assert col_end <= self.total_cols, f"TMEM overflow: {col_end} > {self.total_cols}"
+        if layout is None:
+            assert len(shape) == 2, "TMEMPool.alloc() requires layout= for non-2D TMEM buffers"
+            layout = _default_tmem_layout(shape[0], shape[1])
+        res = ir.decl_buffer(shape, dtype, scope="tmem", allocated_addr=col_start, layout=layout)
+        self.offset = col_end
+        self.max_offset = max(self.max_offset, self.offset)
+        return res
+
+    def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1):
+        """Allocate a tcgen05 block-scaled SF TMEM buffer with an inferred layout.
+
+        ``shape`` last two dims are ``(rows, SF_K * sf_reuse)`` (the last dim is
+        what gemm dispatch iterates over). When ``shape`` has 3 dims, the first
+        is treated as a pipe-depth outer.
+        """
+        from tvm.backend.cuda.operator.tile_primitive.gemm_async.tcgen05 import sf_tmem_layout
+
+        if len(shape) == 2:
+            pipe_depth, rows, last = None, shape[0], shape[1]
+        elif len(shape) == 3:
+            pipe_depth, rows, last = shape[0], shape[1], shape[2]
+        else:
+            raise ValueError(
+                f"alloc_sf expects 2D (rows, SF_K*sf_reuse) or 3D "
+                f"(pipe_depth, rows, SF_K*sf_reuse); got shape={shape}"
+            )
+        assert last % sf_reuse == 0, (
+            f"alloc_sf: shape last dim {last} must be divisible by sf_reuse={sf_reuse}"
+        )
+        SF_K = last // sf_reuse
+        layout = sf_tmem_layout(
+            rows=rows, SF_K=SF_K, sf_per_mma=sf_per_mma, sf_reuse=sf_reuse, pipe_depth=pipe_depth
+        )
+        return self.alloc(shape, dtype, layout=layout)
+
+    def move_base_to(self, col):
+        self.offset = col
+        self.max_offset = max(self.max_offset, self.offset)
+
+    def commit(self):
+        assert not self._committed, "TMEMPool.commit() can only be called once"
+        from tvm.script import tirx as T
+
+        def emit_alloc():
+            _emit_stmt(
+                T.ptx.tcgen05.alloc(
+                    T.address_of(self.addr), n_cols=self.total_cols, cta_group=self.cta_group
+                )
+            )
+            if self.sync_after_alloc:
+                _emit_stmt(T.cuda.warp_sync())
+
+        self._emit_warp_guard(self.alloc_warp, emit_alloc)
+        self._committed = True
+
+    def dealloc(self):
+        assert self._committed, "TMEMPool.dealloc() called before commit()"
+        assert not self._deallocated, "TMEMPool.dealloc() can only be called once"
+        self._deallocated = True
+        from tvm.script import tirx as T
+
+        def emit_dealloc():
+            _emit_stmt(T.ptx.tcgen05.relinquish_alloc_permit(cta_group=self.cta_group))
+            _emit_stmt(
+                T.ptx.tcgen05.dealloc(self.addr, n_cols=self.total_cols, cta_group=self.cta_group)
+            )
+
+        self._emit_warp_guard(self.dealloc_warp, emit_dealloc)
+
+
+# ---------------------------------------------------------------------------
+# SMEMPool
+# ---------------------------------------------------------------------------
+
+
+@_meta_class
+class SMEMPool:
+    """Bump allocator over a contiguous shared memory region.
+
+    Parameters
+    ----------
+    ptr : Var or None, optional
+        If omitted, an ``alloc_buffer([0], "uint8", scope="shared.dyn")`` is
+        created automatically and ``commit()`` must be called after all
+        allocations to emit the size annotation.
+        If a ``Var`` is provided, the caller manages the backing buffer and
+        ``commit()`` is a no-op.
+    """
+
+    def __init__(self, ptr=_POOL_UNSET):
+        ir = _get_ir()
+        if ptr is _POOL_UNSET:
+            self.buf = ir.alloc_buffer([0], "uint8", scope="shared.dyn")
+            self.ptr = self.buf.data
+            self._owns_buffer = True
+        else:
+            self.buf = None
+            self.ptr = ptr
+            self._owns_buffer = False
+        self.offset = 0
+        self.max_offset = 0
+
+    def alloc(
+        self,
+        shape,
+        dtype="float32",
+        strides=None,
+        scope="shared.dyn",
+        align=0,
+        buffer_type="",
+        axis_separators=None,
+        layout="default",
+    ):
+        ir = _get_ir()
+        if align > 0:
+            self.offset = (self.offset + align - 1) // align * align
+        res = ir.decl_buffer(
+            shape,
+            dtype,
+            data=self.ptr,
+            strides=strides,
+            byte_offset=self.offset,
+            scope=scope,
+            align=align,
+            buffer_type=buffer_type,
+            axis_separators=axis_separators,
+            layout=layout,
+        )
+        # Advance in bits then round up to bytes so sub-byte dtypes (e.g.
+        # float4_e2m1fn = 4 bits) still bump the cursor instead of leaving it
+        # at 0 (bits // 8) and silently overlapping the next allocation.
+        self.offset += (_shape_product(shape) * DataType(dtype).bits + 7) // 8
+        if self._owns_buffer:
+            self.max_offset = max(self.max_offset, self.offset)
+        return res
+
+    def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024):
+        """Allocate MMA-compatible shared memory with an inferred swizzle layout."""
+        from tvm.backend.cuda.operator.tile_primitive.tma_utils import (
+            SwizzleMode,
+            mma_shared_layout,
+        )
+
+        if isinstance(swizzle_mode, str):
+            if swizzle_mode == "auto":
+                swizzle_mode = _auto_swizzle_mode(dtype)
+            elif swizzle_mode == "none":
+                swizzle_mode = SwizzleMode.SWIZZLE_NONE
+            else:
+                raise ValueError(
+                    f"Unsupported swizzle_mode={swizzle_mode!r}; expected 'auto', 'none', "
+                    "or SwizzleMode"
+                )
+        _validate_mma_alloc_shape(shape, dtype, swizzle_mode)
+        layout = mma_shared_layout(dtype, swizzle_mode, shape)
+        return self.alloc(shape, dtype, align=align, layout=layout)
+
+    def move_base_to(self, offset):
+        self.offset = offset
+        if self._owns_buffer:
+            self.max_offset = max(self.max_offset, self.offset)
+
+    def commit(self, size=None):
+        """Emit pool size annotation into the IR.
+
+        Must be called after all ``alloc()`` / ``move_base_to()`` calls.
+
+        Parameters
+        ----------
+        size : int, optional
+            Explicit shared memory size in bytes.  When *None* (the default),
+            the high-water mark ``max_offset`` tracked by the allocator is used.
+        """
+        if not self._owns_buffer:
+            return
+        ir = _get_ir()
+        frame_mod = _get_frame()
+        resolved = size if size is not None else self.max_offset
+        assert resolved >= self.max_offset, (
+            f"Specified smem size ({resolved}) is smaller than "
+            f"the pool high-water mark ({self.max_offset})"
+        )
+        attr_frame = ir.attr(self.ptr, "tirx.pool_max_bytes", resolved)
+        if isinstance(attr_frame, frame_mod.AttrFrame):
+            from functools import partial
+
+            attr_frame.add_callback(partial(attr_frame.__exit__, None, None, None))
+            attr_frame.__enter__()

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

@@ -0,0 +1,251 @@
+# 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.
+"""Reusable pipeline state and mbarrier helpers for SM100 kernels.
+
+These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
+instances are automatically treated as meta values inside @T.prim_func.
+"""
+
+from tvm.script import tirx as T
+
+
+@T.meta_class
+class PipelineState:
+    """Tracks stage and phase for a software-pipelined ring buffer.
+
+    This class does not know anything about full/empty barriers. Use it when
+    the kernel manually waits/signals barriers, or when the stage/phase drives
+    a ring not wrapped in a ``Pipeline``.
+
+    Parameters
+    ----------
+    depth : int
+        Number of stages in the ring.
+    phase : int, optional
+        Initial phase. Omit when initialization should happen later.
+    """
+
+    def __init__(self, depth: int, phase=None):
+        self.stage = T.local_scalar("int32")
+        self.phase = T.local_scalar("int32")
+        self.depth = depth
+        if phase is not None:
+            self.init(phase)
+
+    @T.inline
+    def init(self, phase):
+        self.stage = 0
+        self.phase = phase
+
+    @T.inline
+    def advance(self):
+        if self.depth > 1:
+            self.stage = self.stage + 1
+            if self.stage == self.depth:
+                self.stage = 0
+                self.phase = self.phase ^ 1
+        else:
+            self.phase = self.phase ^ 1
+
+
+@T.meta_class
+class MBarrier:
+    """Mbarrier wrapper with regular ``mbarrier.arrive``.
+
+    Parameters
+    ----------
+    pool : SMEMPool
+        Shared memory pool allocator.
+    depth : int
+        Number of barrier slots (one per pipeline stage).
+    phase_offset : int
+        XORed into the phase bit on every ``wait`` / ``arrive``.
+    leader : Expr, optional
+        Boolean predicate selecting the single thread that runs
+        ``mbarrier.init``. Defaults to ``T.cuda.thread_rank() == 0`` --
+        thread 0 of the enclosing CTA, which always picks exactly one
+        thread regardless of which scope_id vars the caller declared.
+        Override only when you want a different CTA-local thread to do
+        the init.
+
+        Note: the default deliberately avoids ``T.warp_id()`` /
+        ``T.lane_id()``. Those introduce deferred ``cta->warp`` /
+        ``warp->thread`` ScopeIdDefs that the verifier cannot pin down
+        unless the kernel header declares the full warp/lane chain (e.g. a
+        single-CTA DSMEM kernel that only declares ``thread_id``). It also
+        avoids the synccheck false-deadlock on kernels that declare a
+        second warp-scope id. The generated CUDA is equivalent.
+    """
+
+    def __init__(self, pool, depth, phase_offset=0, leader=None):
+        self.buf = pool.alloc((depth,), "uint64", align=8)
+        self.depth = depth
+        self.phase_offset = phase_offset
+        self.leader = leader if leader is not None else (T.cuda.thread_rank() == 0)
+
+    @T.inline
+    def init(self, count):
+        if self.leader:
+            for i in T.unroll(self.depth):
+                T.ptx.mbarrier.init(self.buf.ptr_to([i]), count)
+
+    @T.inline
+    def wait(self, stage, phase):
+        # Blocks: ``mbarrier.try_wait`` loops internally until the phase flips,
+        # so this returns only once the barrier has completed.
+        T.ptx.mbarrier.try_wait(self.buf.ptr_to([stage]), phase ^ self.phase_offset)
+
+    @T.inline
+    def arrive(self, stage, cta_id=None, pred=None, count=None):
+        # Default: local-CTA arrive — emits the simple
+        # ``mbarrier.arrive.shared.b64`` form. To arrive on a remote
+        # CTA's mbarrier in a cluster kernel, callers must pass
+        # ``cta_id=`` explicitly (e.g. ``bar.arrive(stage, cta_id=0)``)
+        # or use ``MBarrier.remote_view(rank).arrive(stage)``. Defaulting
+        # the cross-CTA path was both surprising (``bar.arrive(stage)``
+        # silently ``mapa`` ed across the cluster) and a per-call cost
+        # of ~3 PTX ops on every single-CTA kernel.
+        #
+        # ``count`` (cross-CTA path only) emits the explicit arrival-count
+        # operand, i.e. ``mbarrier.arrive.shared::cluster.b64 _, [addr], count``.
+        # When ``None`` the implicit count-of-1 form is emitted. Passing
+        # ``count=1`` is semantically identical but spells the count explicitly.
+        if cta_id is None:
+            T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
+        else:
+            actual_pred = True if pred is None else pred
+            T.ptx.mbarrier.arrive(
+                self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred, count=count
+            )
+
+    def ptr_to(self, idx):
+        return self.buf.ptr_to(idx)
+
+    def remote_view(self, rank):
+        """Create a view of this barrier mapped to another CTA's shared memory.
+
+        Arrive-only: the returned view is built with ``object.__new__`` and
+        never copies ``self.leader``, so calling ``.init()`` on it would fail.
+        Use it solely to ``arrive`` on a remote CTA's mbarrier.
+        """
+        from tvm.ir import PointerType, PrimType
+        from tvm.tirx import Var as TIRVar
+
+        expr = T.reinterpret("handle", T.ptx.map_shared_rank(self.buf.ptr_to([0]), rank))
+        ptr = TIRVar("remote_mbar_ptr", PointerType(PrimType("uint64")))
+        T.Bind(expr, var=ptr)
+        buf = T.decl_buffer([self.depth], "uint64", data=ptr, scope="shared")
+        remote = object.__new__(type(self))
+        remote.buf = buf
+        remote.depth = self.depth
+        remote.phase_offset = self.phase_offset
+        return remote
+
+
+class TMABar(MBarrier):
+    """Barrier signaled by TMA (mbarrier.arrive.expect_tx).
+
+    When ``tx_count`` is None, falls back to a remote mbarrier.arrive
+    (matching MBarrier.arrive defaults).
+    """
+
+    @T.inline
+    def arrive(self, stage, tx_count=None, cta_id=None, pred=None):
+        # NOTE: this arrive() kwarg set intentionally differs from
+        # MBarrier.arrive (hardware necessity, LSP-incompatible by design).
+        # ``tx_count``: TMA byte count for ``mbarrier.arrive.expect_tx``.
+        # ``cta_id`` / ``pred``: forwarded to the underlying
+        # ``mbarrier.arrive`` (cluster path) when set; otherwise the
+        # arrive is local-CTA only. See ``MBarrier.arrive`` for the
+        # full default-local rationale.
+        if tx_count is not None:
+            T.ptx.mbarrier.arrive.expect_tx(self.buf.ptr_to([stage]), tx_count)
+        elif cta_id is None:
+            T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
+        else:
+            actual_pred = True if pred is None else pred
+            T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred)
+
+
+class TCGen05Bar(MBarrier):
+    """Barrier signaled by ``tcgen05`` commit.
+
+    The caller is responsible for ensuring only one thread issues the
+    commit, e.g. by wrapping the call in ``if T.ptx.elect_sync():``.
+    """
+
+    @T.inline
+    def arrive(self, stage, cta_group=1, cta_mask=None):
+        # NOTE: this arrive() kwarg set intentionally differs from
+        # MBarrier.arrive (hardware necessity, LSP-incompatible by design).
+        if cta_mask is None and cta_group == 1:
+            T.ptx.tcgen05.commit(self.buf.ptr_to([stage]))
+        else:
+            T.ptx.tcgen05.commit(self.buf.ptr_to([stage]), cta_group=cta_group, cta_mask=cta_mask)
+
+
+# Barrier-type tags accepted by Pipeline's ``full=`` / ``empty=`` arguments.
+_BAR_KINDS = {"tma": TMABar, "tcgen05": TCGen05Bar, "mbar": MBarrier}
+
+
+@T.meta_class
+class Pipeline:
+    """A full/empty mbarrier pair for a software-pipelined data flow.
+
+    Pass barrier-type tags and ``Pipeline`` constructs and ``init``\\ s the
+    barriers itself. Tags: ``"tma"`` (TMABar), ``"tcgen05"`` (TCGen05Bar),
+    ``"mbar"`` (MBarrier). The barrier type and arrival count of each event
+    stay explicit at the call site -- e.g. ``Pipeline(pool, n, full="tma",
+    empty="tcgen05", init_empty=NUM_CONSUMER)``.
+
+    Both signals are required: a ``Pipeline`` is a *pair*. For a one-way event
+    (a pure "X happened" signal with no slot to recycle) use a bare barrier
+    (``TMABar``/``TCGen05Bar``/``MBarrier``) directly -- it has no empty side.
+
+    Parameters
+    ----------
+    pool : SMEMPool
+        Shared memory pool allocator.
+    stages : int
+        Number of pipeline stages (barrier slots).
+    full, empty : str
+        Barrier-type tag for the full / empty signal (see above).
+    init_full, init_empty : int
+        Expected arrival count for the full / empty barrier.
+    empty_phase_offset : int
+        XORed into the empty barrier's phase bit on every wait / arrive.
+    leader : Expr, optional
+        Propagated to both barriers; defaults to thread 0 of the CTA.
+    """
+
+    def __init__(
+        self,
+        pool,
+        stages,
+        *,
+        full,
+        empty,
+        init_full=1,
+        init_empty=1,
+        empty_phase_offset=0,
+        leader=None,
+    ):
+        self.stages = stages
+        self.full = _BAR_KINDS[full](pool, stages, leader=leader)
+        self.full.init(init_full)
+        self.empty = _BAR_KINDS[empty](pool, stages, phase_offset=empty_phase_offset, leader=leader)
+        self.empty.init(init_empty)

python/tvm/backend/cuda/lang/smem_desc.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""SMEM matrix descriptor helper for tcgen05 / wgmma."""
+
+from tvm.backend.cuda.operator.tile_primitive.common import smem_desc_add_16B_offset
+from tvm.script import tirx as T
+
+
+@T.meta_class
+class SmemDescriptor:
+    """Encoded once via :meth:`init`, reused via :meth:`add_16B_offset`."""
+
+    def __init__(self):
+        self._buf = T.alloc_local([1], "uint64")
+
+    @property
+    def desc(self):
+        return self._buf[0]
+
+    @T.inline
+    def init(self, smem_ptr, ldo, sdo, swizzle):
+        T.ptx.tcgen05.encode_matrix_descriptor(
+            T.address_of(self._buf[0]), smem_ptr, ldo, sdo, swizzle
+        )
+
+    def add_16B_offset(self, offset):
+        return smem_desc_add_16B_offset(self._buf[0], offset)
+
+    def make_lo_uniform(self):
+        """Broadcast the lower 32 bits to all warp lanes via ``__shfl_sync``."""
+        func_name = "smem_desc_make_lo_uniform"
+        source_code = f"""
+__forceinline__ __device__ void {func_name}(uint64_t* desc) {{
+    SmemDescriptor* d = reinterpret_cast<SmemDescriptor*>(desc);
+    d->lo = __shfl_sync(0xffffffff, d->lo, 0);
+}}
+"""
+        return T.cuda.func_call(
+            func_name, T.address_of(self._buf[0]), source_code=source_code, return_type="void"
+        )

python/tvm/backend/cuda/lang/tile_scheduler.py

@@ -0,0 +1,945 @@
+# 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.
+"""Reusable tile scheduler helpers for TIR tests/kernels.
+
+These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
+instances are automatically treated as meta values inside @T.prim_func.
+"""
+
+from tvm.backend.cuda.lang.pipeline import Pipeline, PipelineState
+from tvm.script import tirx as T
+
+
+@T.meta_class
+class BaseTileScheduler:
+    """Base class for tile schedulers with common state and macros."""
+
+    def __init__(self, prefix: str):
+        self.m_idx = T.local_scalar("int32")
+        self.n_idx = T.local_scalar("int32")
+        self.linear_idx = T.local_scalar("int32")
+
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        # To be implemented by subclasses
+        pass
+
+    @T.inline
+    def init(self, linear_init):
+        self.linear_idx = linear_init
+        self.update_current_m_n_idx(linear_init)
+
+    @T.inline
+    def next_tile(self, step):
+        self.linear_idx = self.linear_idx + step
+        self.update_current_m_n_idx(self.linear_idx)
+
+    def valid(self, total_tiles):
+        return self.linear_idx < total_tiles
+
+
+class ClusterPersistentScheduler2D(BaseTileScheduler):
+    """
+    Tile scheduler for cluster-based persistent kernels.
+
+    Distributes a 2D tile grid across persistent clusters using group-major ordering
+    for L2 cache locality. Each cluster starts at its cluster_id and strides by
+    num_clusters to process tiles.
+
+    Tile Ordering (group-major for L2 locality):
+    - Tiles are grouped into "L2 groups" of `l2_group_size` rows
+    - Within a group, tiles are visited in column-major order within the group
+    - Groups are processed in row-major order
+
+    Example with 4x4 tiles, l2_group_size=2:
+        Group 0 (rows 0-1):  0  2  4  6
+                             1  3  5  7
+        Group 1 (rows 2-3):  8 10 12 14
+                             9 11 13 15
+
+    Serpentine Mode (serpentine=True):
+    - Uses CUTLASS-style 2D block swizzle with serpentine traversal
+    - Grid is divided into swizzle_size x swizzle_size blocks
+    - Within each block, tiles are visited in row-major order
+    - Blocks are traversed in serpentine order (even block-rows forward, odd backward)
+    - This provides better L2 locality by reusing both A and B tiles
+
+    Example with 4x4 tiles, swizzle_size=2, serpentine=True:
+        Block layout:
+          Block(0,0)  Block(0,1)
+          Block(1,0)  Block(1,1)
+
+        Tile numbering with serpentine:
+               n=0  n=1  n=2  n=3
+          m=0   0    1   14   15
+          m=1   2    3   12   13
+          m=2   4    5   10   11
+          m=3   6    7    8    9
+
+        Traversal: Block(0,0) -> Block(1,0) -> Block(1,1) -> Block(0,1)
+                   (serpentine: down in col 0, then up in col 1)
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for TIR variable names
+    num_m_tiles : int | T.ExprLike
+        Total number of tiles in M dimension (can be runtime expression)
+    num_n_tiles : int
+        Total number of tiles in N dimension
+    num_clusters : int
+        Number of persistent clusters (determines stride)
+    l2_group_size : int
+        Number of M-tile rows per L2 locality group (default: 8)
+        When serpentine=True, this is used as swizzle_size for 2D blocks
+    cluster_m : int
+        Cluster dimension in M for hierarchical scheduling (default: 1)
+    cluster_n : int
+        Cluster dimension in N for hierarchical scheduling (default: 1)
+    serpentine : bool
+        If True, use CUTLASS-style 2D block swizzle with serpentine traversal (default: False)
+
+    Attributes
+    ----------
+    m_idx : T.local_scalar
+        Current M tile index (output)
+    n_idx : T.local_scalar
+        Current N tile index (output)
+    work_idx : T.local_scalar
+        Global work item index for this cluster
+    tile_count : T.local_scalar
+        Number of tiles processed by this cluster so far
+
+    Usage
+    -----
+    ```python
+    scheduler = ClusterPersistentScheduler2D(
+        "sched", num_m_tiles=M_TILES, num_n_tiles=N_TILES,
+        num_clusters=NUM_CLUSTERS, l2_group_size=8
+    )
+    scheduler.init(cluster_id)  # cluster_id = cta_idx // CLUSTER_SIZE
+
+    while scheduler.valid():
+        m = T.meta_var(scheduler.m_idx)  # current M tile
+        n = T.meta_var(scheduler.n_idx)  # current N tile
+        # ... process tile (m, n) ...
+        scheduler.next_tile()
+    ```
+
+    Examples
+    --------
+    Example 1: Basic persistent kernel
+    ```
+    num_m_tiles=4, num_n_tiles=4, num_clusters=3, l2_group_size=2
+    cluster_m=1, cluster_n=1 (default, no tile subdivision)
+
+    Group-major tile numbering (l2_group_size=2):
+           n=0  n=1  n=2  n=3
+      m=0   0    2    4    6   ┐ L2 group 0
+      m=1   1    3    5    7   ┘
+      m=2   8   10   12   14   ┐ L2 group 1
+      m=3   9   11   13   15   ┘
+
+    Work distribution (cluster starts at cluster_id, strides by num_clusters=3):
+      cluster 0: work_idx 0,3,6,9,12,15  -> tiles 0,3,6,9,12,15
+      cluster 1: work_idx 1,4,7,10,13    -> tiles 1,4,7,10,13
+      cluster 2: work_idx 2,5,8,11,14    -> tiles 2,5,8,11,14
+
+    Tile grid (which cluster handles each tile):
+           n=0  n=1  n=2  n=3
+      m=0   C0   C2   C1   C0   ┐ L2 group 0
+      m=1   C1   C0   C2   C1   ┘
+      m=2   C2   C1   C0   C2   ┐ L2 group 1
+      m=3   C0   C2   C1   C0   ┘
+
+    Tile sequence per cluster (in execution order):
+      cluster 0: (0,0)->(1,1)->(0,3)->(2,0)->(2,3)->(3,3)
+      cluster 1: (1,0)->(0,2)->(1,3)->(2,1)->(3,2)
+      cluster 2: (0,1)->(1,2)->(2,0)->(3,1)->(2,3)
+    ```
+
+    Example 2: 2SM GEMM (typical B200 config)
+    ```
+    M=1024, N=512, CTA_M=128, MMA_N=128, CLUSTER_M=2, CLUSTER_N=1
+    => M_TILES=8, N_TILES=4
+    => CLUSTER_M_TILES=4, CLUSTER_N_TILES=4 (scheduler at cluster granularity)
+
+    Scheduler params:
+      num_m_tiles=4, num_n_tiles=4, num_clusters=74, l2_group_size=8
+      cluster_m=1, cluster_n=1
+
+    Key: Scheduler outputs CLUSTER-level tiles.
+         All CTAs in same cluster get SAME (m_idx, n_idx) from scheduler.
+         CTAs differentiate via cluster_rank (computed OUTSIDE scheduler):
+           cluster_rank = cta_idx % CLUSTER_SIZE
+           cb_m = cluster_rank % CLUSTER_M   # 0 or 1 for 2SM
+           cb_n = cluster_rank // CLUSTER_M  # 0 for 2SM
+
+    Final CTA tile:
+      cta_m = m_idx * CLUSTER_M + cb_m
+      cta_n = n_idx * CLUSTER_N + cb_n
+
+    Example: cluster 5 gets scheduler tile (1,2)
+      CTA rank=0 (cb_m=0): actual tile (2,2)
+      CTA rank=1 (cb_m=1): actual tile (3,2)
+    ```
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        num_m_tiles,
+        num_n_tiles: int,
+        num_clusters: int,
+        l2_group_size: int = 8,
+        cluster_m: int = 1,
+        cluster_n: int = 1,
+        serpentine: bool = False,
+    ):
+        super().__init__(prefix)
+        self._num_m_tiles = num_m_tiles
+        self._num_n_tiles = num_n_tiles
+        self._num_clusters = num_clusters
+        self._l2_group_size = l2_group_size
+        self._cluster_m = cluster_m
+        self._cluster_n = cluster_n
+        self._serpentine = serpentine
+
+        # Rename internal state for clarity
+        self.work_idx = self.linear_idx  # alias: global work item index
+        self.tile_count = T.local_scalar("int32")
+        self.tile_idx = self.tile_count  # alias for backward compatibility
+
+        is_static_m = isinstance(num_m_tiles, int)
+
+        # Number of tile columns after accounting for cluster_n
+        n_tile_cols = (num_n_tiles + cluster_n - 1) // cluster_n
+        self._N_TILE_COLS = n_tile_cols
+
+        if is_static_m:
+            self._M_TILE_ROWS = (num_m_tiles + cluster_m - 1) // cluster_m
+            self._FULL_GROUPS = self._M_TILE_ROWS // l2_group_size
+        else:
+            # Dynamic expressions for runtime M
+            self._M_TILE_ROWS = T.truncdiv(self._num_m_tiles + self._cluster_m - 1, self._cluster_m)
+            self._FULL_GROUPS = T.truncdiv(self._M_TILE_ROWS, self._l2_group_size)
+
+        self._TAIL_ROWS = self._M_TILE_ROWS - self._FULL_GROUPS * l2_group_size
+        self._TOTAL_TILES = self._M_TILE_ROWS * n_tile_cols * cluster_m * cluster_n
+
+        # For serpentine mode: precompute block counts
+        if serpentine:
+            self._N_BLOCKS = n_tile_cols // l2_group_size  # full blocks in N
+            self._M_BLOCKS = (
+                self._M_TILE_ROWS // l2_group_size
+                if is_static_m
+                else T.truncdiv(self._M_TILE_ROWS, l2_group_size)
+            )
+            self._BLOCK_SIZE = l2_group_size * l2_group_size  # tiles per block
+            self._FULL_BLOCK_TILES = self._M_BLOCKS * self._N_BLOCKS * self._BLOCK_SIZE
+            # Residual tiles (not covered by full blocks)
+            self._RESIDUAL_N = n_tile_cols - self._N_BLOCKS * l2_group_size
+            self._RESIDUAL_M = self._M_TILE_ROWS - self._M_BLOCKS * l2_group_size
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, work_idx):
+        """Convert global work index to (m_idx, n_idx) tile coordinates."""
+        CLUSTER_M = T.meta_var(self._cluster_m)
+        CLUSTER_N = T.meta_var(self._cluster_n)
+
+        # Extract hierarchical cluster-local offsets
+        cluster_m_offset = T.meta_var(work_idx % CLUSTER_M)
+        t = T.meta_var(work_idx // CLUSTER_M)
+        cluster_n_offset = T.meta_var(t % CLUSTER_N)
+        tile_linear = T.meta_var(t // CLUSTER_N)
+
+        @T.inline
+        def set_tile_coords(tile_row, tile_col):
+            self.m_idx = tile_row * CLUSTER_M + cluster_m_offset
+            self.n_idx = tile_col * CLUSTER_N + cluster_n_offset
+
+        if self._serpentine:
+            self._update_serpentine(tile_linear, set_tile_coords)
+        else:
+            self._update_group_major(tile_linear, set_tile_coords)
+
+    def _update_group_major(self, tile_linear, set_tile_coords):
+        """Group-major ordering with parse-time pruning of statically-dead branches.
+
+        The TIR script parser does not constant-fold ``if False: ...``, so a
+        Python-literal ``FULL_GROUPS == 0`` would otherwise produce
+        ``T.bitwise_and(T.bool(False), tile_linear < 0)`` IR plus the dead
+        then-leg.  Branch in plain Python here and only invoke the inline
+        emitter that can actually fire.
+        """
+        full_zero = isinstance(self._FULL_GROUPS, int) and self._FULL_GROUPS == 0
+        tail_zero = isinstance(self._TAIL_ROWS, int) and self._TAIL_ROWS == 0
+        if full_zero and tail_zero:
+            self._gm_emit_zero(set_tile_coords)
+        elif full_zero:
+            self._gm_emit_tail_only(tile_linear, set_tile_coords)
+        elif tail_zero:
+            self._gm_emit_full_only(tile_linear, set_tile_coords)
+        else:
+            self._gm_emit_full_and_tail(tile_linear, set_tile_coords)
+
+    @T.inline
+    def _gm_emit_zero(self, set_tile_coords):
+        set_tile_coords(0, 0)
+
+    @T.inline
+    def _gm_emit_full_only(self, tile_linear, set_tile_coords):
+        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
+        GROUP_SIZE = T.meta_var(self._l2_group_size)
+        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
+        if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
+            group_id: T.let = tile_linear // GROUP_SPAN
+            within_group: T.let = tile_linear % GROUP_SPAN
+            tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
+            tile_col: T.let = within_group // GROUP_SIZE
+            set_tile_coords(tile_row, tile_col)
+        else:
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def _gm_emit_tail_only(self, tile_linear, set_tile_coords):
+        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
+        TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
+        GROUP_SIZE = T.meta_var(self._l2_group_size)
+        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
+        if TAIL_ROWS > 0:
+            rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
+            tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
+            tile_col: T.let = rem // TAIL_ROWS
+            set_tile_coords(tile_row, tile_col)
+        else:
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def _gm_emit_full_and_tail(self, tile_linear, set_tile_coords):
+        FULL_GROUPS = T.meta_var(self._FULL_GROUPS)
+        TAIL_ROWS = T.meta_var(self._TAIL_ROWS)
+        GROUP_SIZE = T.meta_var(self._l2_group_size)
+        GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS)
+        if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN):
+            group_id: T.let = tile_linear // GROUP_SPAN
+            within_group: T.let = tile_linear % GROUP_SPAN
+            tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE)
+            tile_col: T.let = within_group // GROUP_SIZE
+            set_tile_coords(tile_row, tile_col)
+        elif TAIL_ROWS > 0:
+            rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN
+            tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS)
+            tile_col: T.let = rem // TAIL_ROWS
+            set_tile_coords(tile_row, tile_col)
+        else:
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def _update_serpentine(self, tile_linear, set_tile_coords):
+        """CUTLASS-style 2D block swizzle with serpentine traversal.
+
+        Algorithm:
+        1. Divide grid into swizzle_size x swizzle_size blocks
+        2. Within each block, visit tiles in row-major order
+        3. Blocks are traversed column by column (along N)
+        4. Within each column of blocks, use serpentine:
+           - Even columns: top to bottom
+           - Odd columns: bottom to top
+
+        This maximizes L2 reuse for both A and B matrices.
+        """
+        S = T.meta_var(self._l2_group_size)  # swizzle_size
+        M_BLOCKS = T.meta_var(self._M_BLOCKS)
+        N_BLOCKS = T.meta_var(self._N_BLOCKS)
+        BLOCK_SIZE = T.meta_var(self._BLOCK_SIZE)  # S * S
+        FULL_BLOCK_TILES = T.meta_var(self._FULL_BLOCK_TILES)
+        M_TILE_ROWS = T.meta_var(self._M_TILE_ROWS)
+        T.meta_var(self._N_TILE_COLS)
+        RESIDUAL_N = T.meta_var(self._RESIDUAL_N)
+        RESIDUAL_M = T.meta_var(self._RESIDUAL_M)
+
+        # Check if we're in the full block region
+        if (M_BLOCKS > 0) & (N_BLOCKS > 0) & (tile_linear < FULL_BLOCK_TILES):
+            # Which block (in linear order along columns of blocks)
+            block_linear: T.let = tile_linear // BLOCK_SIZE
+            within_block: T.let = tile_linear % BLOCK_SIZE
+
+            # Block column and row
+            block_col: T.let = block_linear // M_BLOCKS
+            block_row_raw: T.let = block_linear % M_BLOCKS
+
+            # Serpentine: odd columns go bottom-to-top
+            block_row: T.let = T.Select(
+                block_col % 2 == 0,
+                block_row_raw,
+                M_BLOCKS - 1 - block_row_raw
+            )
+
+            # Position within block (row-major within block)
+            local_row: T.let = within_block // S
+            local_col: T.let = within_block % S
+
+            tile_row: T.let = block_row * S + local_row
+            tile_col: T.let = block_col * S + local_col
+            set_tile_coords(tile_row, tile_col)
+
+        elif RESIDUAL_N > 0:
+            # Residual tiles in the rightmost partial column of blocks
+            # These are tiles where n >= N_BLOCKS * S
+            rem: T.let = tile_linear - FULL_BLOCK_TILES
+
+            # First handle the right residual strip (full M height, partial N width)
+            right_strip_tiles: T.let = M_TILE_ROWS * RESIDUAL_N
+            if rem < right_strip_tiles:
+                # Row-major within the right strip
+                tile_row: T.let = rem // RESIDUAL_N
+                tile_col: T.let = N_BLOCKS * S + (rem % RESIDUAL_N)
+                set_tile_coords(tile_row, tile_col)
+            elif RESIDUAL_M > 0:
+                # Bottom residual strip (already covered in right strip overlap)
+                # This handles corner case - shouldn't normally reach here
+                # as right strip already covers full M height
+                set_tile_coords(0, 0)
+            else:
+                set_tile_coords(0, 0)
+
+        elif RESIDUAL_M > 0:
+            # Bottom residual strip only (no right residual)
+            rem: T.let = tile_linear - FULL_BLOCK_TILES
+            bottom_strip_tiles: T.let = RESIDUAL_M * (N_BLOCKS * S)
+            if rem < bottom_strip_tiles:
+                tile_row: T.let = M_BLOCKS * S + (rem % RESIDUAL_M)
+                tile_col: T.let = rem // RESIDUAL_M
+                set_tile_coords(tile_row, tile_col)
+            else:
+                set_tile_coords(0, 0)
+        else:
+            # Fallback
+            set_tile_coords(0, 0)
+
+    @T.inline
+    def init(self, cluster_id):
+        """Initialize scheduler for a given cluster.
+
+        Parameters
+        ----------
+        cluster_id : int
+            The cluster's index (typically cta_idx // CLUSTER_SIZE)
+        """
+        self.linear_idx = cluster_id
+        self.tile_count = 0
+        self.update_current_m_n_idx(cluster_id)
+
+    @T.inline
+    def next_tile(self):
+        """Advance to the next tile for this cluster."""
+        self.linear_idx = self.linear_idx + self._num_clusters
+        self.tile_count = self.tile_count + 1
+        self.update_current_m_n_idx(self.linear_idx)
+
+    @T.inline
+    def next_tile_stride(self, stride: int):
+        """Advance by a custom stride (for non-standard scheduling)."""
+        self.linear_idx = self.linear_idx + stride
+        self.tile_count = self.tile_count + 1
+        self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        """Check if this cluster has more tiles to process."""
+        return self.linear_idx < self._TOTAL_TILES
+
+
+class GroupMajor3D(BaseTileScheduler):
+    """
+    3D grouped-row scheduler (M,N,K) with tail handling on M.
+
+    Args
+    ----
+    prefix: str
+    m_tiles: int | T Expr   # tiles along M (static or runtime)
+    n_tiles: int                # tiles along N (static)
+    k_tiles: int                # tiles along K (static)
+    group_rows: int             # rows per group along M
+    step: int = 1               # default stride for next_tile()
+    """
+
+    def __init__(
+        self, prefix: str, m_tiles, n_tiles: int, k_tiles: int, group_rows: int, step: int = 1
+    ):
+        super().__init__(prefix)
+        self._step = step
+        self.tile_idx = T.local_scalar("int32")
+        self.k_idx = T.local_scalar("int32")
+
+        # ---- constants / primexprs baked once ----
+        self._G = group_rows
+        self._N = n_tiles
+        self._K = k_tiles
+
+        if isinstance(m_tiles, int):
+            self._GROUPS = m_tiles // group_rows
+            self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
+            self._SAFE_FINAL_ROWS = max(self._FINAL_ROWS, 1)
+            self._GROUP_SIZE = group_rows * n_tiles * k_tiles
+            self._TOTAL = m_tiles * n_tiles * k_tiles
+        else:
+            self._GROUPS = T.truncdiv(m_tiles, group_rows)
+            self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
+            self._SAFE_FINAL_ROWS = T.max(self._FINAL_ROWS, 1)
+            self._GROUP_SIZE = self._G * self._N * self._K
+            self._TOTAL = m_tiles * n_tiles * k_tiles
+
+        # handy composites used in macro
+        self._FULL_BOUND = self._GROUPS * self._GROUP_SIZE
+        self._HAS_FULL = self._GROUPS > 0
+        self._HAS_TAIL = self._FINAL_ROWS > 0
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        # full-group formulas
+        full_m: T.let = T.floordiv(linear_idx, self._GROUP_SIZE) * self._G + T.floormod(
+            linear_idx, self._G
+        )
+        full_n: T.let = T.floormod(T.floordiv(linear_idx, self._G), self._N)
+        full_k: T.let = T.floordiv(T.floormod(linear_idx, self._GROUP_SIZE), self._G * self._N)
+
+        # tail formulas (relative to FULL_BOUND)
+        # Use _SAFE_FINAL_ROWS (max(FINAL_ROWS, 1)) to avoid divide-by-zero when there is no tail
+        rem: T.let = linear_idx - self._FULL_BOUND
+        tail_m: T.let = self._GROUPS * self._G + T.floormod(rem, self._SAFE_FINAL_ROWS)
+        tail_n: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS) % self._N
+        tail_k: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS * self._N)
+
+        # choose phase
+        if self._HAS_FULL & (linear_idx < self._FULL_BOUND):
+            self.m_idx = full_m
+            self.n_idx = full_n
+            self.k_idx = full_k
+        elif self._HAS_TAIL:
+            self.m_idx = tail_m
+            self.n_idx = tail_n
+            self.k_idx = tail_k
+        else:
+            self.m_idx = 0
+            self.n_idx = 0
+            self.k_idx = 0
+
+    @T.inline
+    def init(self, linear_init):
+        self.linear_idx = linear_init
+        self.tile_idx = 0
+        self.update_current_m_n_idx(linear_init)
+
+    @T.inline
+    def next_tile(self):
+        self.linear_idx = self.linear_idx + self._step
+        self.tile_idx = self.tile_idx + 1
+        self.update_current_m_n_idx(self.linear_idx)
+
+    @T.inline
+    def next_tile_stride(self, stride: int):
+        self.linear_idx = self.linear_idx + stride
+        self.tile_idx = self.tile_idx + 1
+        self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        return self.linear_idx < self._TOTAL
+
+
+class RankAwareGroupMajorTileScheduler(BaseTileScheduler):
+    """
+    Group-major scheduler that applies a rank-aware remapping (remote rows first).
+    Kept as a thin adapter because it depends on NVSHMEM rank at device-side.
+    """
+
+    def __init__(
+        self, prefix: str, m_clusters: int, n_clusters: int, group_size: int, world_size: int
+    ):
+        super().__init__(prefix)
+        self._m_clusters = m_clusters
+        self._n_clusters = n_clusters
+        self._group_size = group_size
+        self._world_size = world_size
+
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        my_rank: T.let = T.nvshmem.my_pe()
+        remote_m_clusters: T.let = self._m_clusters - self._m_clusters // self._world_size
+        group_rows: T.let = (remote_m_clusters // self._group_size) * self._group_size
+        final_rows: T.let = remote_m_clusters - group_rows
+        group_repeat: T.let = self._group_size * self._n_clusters
+        if linear_idx < group_rows * self._n_clusters and group_rows > 0:
+            self.m_idx = (
+                (linear_idx // group_repeat) * self._group_size
+                + (linear_idx % self._group_size)
+                + (my_rank + 1) * self._m_clusters // self._world_size
+            ) % self._m_clusters
+            self.n_idx = (linear_idx % group_repeat) // self._group_size
+        elif linear_idx < remote_m_clusters * self._n_clusters:
+            remainder_idx: T.let = linear_idx - group_rows * self._n_clusters
+            self.m_idx = (
+                group_rows
+                + remainder_idx % final_rows
+                + (my_rank + 1) * self._m_clusters // self._world_size
+            ) % self._m_clusters
+            self.n_idx = remainder_idx // final_rows
+        else:
+            remainder_idx: T.let = linear_idx - remote_m_clusters * self._n_clusters
+            self.m_idx = (
+                remote_m_clusters
+                + remainder_idx % (self._m_clusters // self._world_size)
+                + (my_rank + 1) * self._m_clusters // self._world_size
+            ) % self._m_clusters
+            self.n_idx = remainder_idx // (self._m_clusters // self._world_size)
+
+    @T.inline
+    def next_tile(self, stride: int):
+        self.linear_idx = self.linear_idx + stride
+        self.update_current_m_n_idx(self.linear_idx)
+
+    def valid(self):
+        return self.linear_idx < self._m_clusters * self._n_clusters
+
+
+class IndexedTripleTileScheduler(BaseTileScheduler):
+    """Scheduler that maps linear_idx to (b_idx, h_idx, q_idx) via index lists."""
+
+    def __init__(self, prefix: str, b_indices, h_indices, q_indices, tiles_indptr):
+        super().__init__(prefix)
+        self.b_indices = b_indices
+        self.h_indices = h_indices
+        self.q_indices = q_indices
+        self.tiles_indptr = tiles_indptr
+        self.q_idx = T.local_scalar("int32")
+        self.h_idx = T.local_scalar("int32")
+        self.b_idx = T.local_scalar("int32")
+        self.linear_lim = T.local_scalar("int32")
+
+    @T.inline
+    def _load(self):
+        self.q_idx = self.q_indices[self.linear_idx]
+        self.h_idx = self.h_indices[self.linear_idx]
+        self.b_idx = self.b_indices[self.linear_idx]
+
+    @T.inline
+    def init(self, sm):
+        self.linear_idx = self.tiles_indptr[sm]
+        self.linear_lim = self.tiles_indptr[sm + 1]
+        self._load()
+
+    @T.inline
+    def next_tile(self):
+        self.linear_idx = self.linear_idx + 1
+        self._load()
+
+    def valid(self):
+        return self.linear_idx < self.linear_lim
+
+
+class FlashAttentionLinearScheduler(BaseTileScheduler):
+    """Linear 3D scheduler for flash attention (batch, head, m_block).
+
+    Used for non-causal attention with simple linear decomposition.
+    Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:
+        batch = linear_idx // (num_heads * num_m_blocks)
+        head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks
+        m_block = linear_idx % num_m_blocks
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for TIR variable names
+    num_batches : int
+        Number of batches
+    num_heads : int
+        Number of KV heads
+    num_m_blocks : int
+        Number of Q blocks (M dimension tiles)
+    num_ctas : int
+        Number of CTAs for persistent kernel stride
+    """
+
+    def __init__(
+        self, prefix: str, num_batches: int, num_heads: int, num_m_blocks: int, num_ctas: int
+    ):
+        super().__init__(prefix)
+        self._num_batches = num_batches
+        self._num_heads = num_heads
+        self._num_m_blocks = num_m_blocks
+        self._num_ctas = num_ctas
+        self._total_tasks = num_batches * num_heads * num_m_blocks
+
+        # Output indices
+        self.batch_idx = T.local_scalar("int32")
+        self.head_idx = T.local_scalar("int32")
+        self.m_block_idx = T.local_scalar("int32")
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        """Convert linear index to (batch, head, m_block) coordinates."""
+        NUM_HEADS = T.meta_var(self._num_heads)
+        NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
+        HEAD_M_PRODUCT = T.meta_var(NUM_HEADS * NUM_M_BLOCKS)
+
+        self.batch_idx = linear_idx // HEAD_M_PRODUCT
+        self.head_idx = (linear_idx % HEAD_M_PRODUCT) // NUM_M_BLOCKS
+        self.m_block_idx = linear_idx % NUM_M_BLOCKS
+
+    @T.inline
+    def init(self, cta_id):
+        """Initialize scheduler with CTA ID."""
+        self.linear_idx = cta_id
+        self.update_current_m_n_idx(cta_id)
+
+    @T.inline
+    def next_tile(self):
+        """Advance to next tile by striding by num_ctas."""
+        self.linear_idx = self.linear_idx + self._num_ctas
+        self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        """Check if there are more tiles to process."""
+        return self.linear_idx < self._total_tasks
+
+
+class FlashAttentionLPTScheduler(BaseTileScheduler):
+    """LPT scheduler with L2 swizzle for causal flash attention.
+
+    Processes high-work Q blocks (with more KV blocks to attend to) first using
+    Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle
+    for better cache locality across batch*head dimensions.
+
+    The LPT aspect comes from reversing m_block order: lower Q blocks have more
+    KV blocks to process due to causal masking, so processing them first balances load.
+
+    The scheduler is only applied to non-persistent kernels.
+
+    L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.
+
+    Parameters
+    ----------
+    prefix : str
+        Prefix for TIR variable names
+    num_batches : int
+        Number of batches
+    num_heads : int
+        Number of KV heads
+    num_m_blocks : int
+        Number of Q blocks (M dimension tiles)
+    num_ctas : int
+        Number of CTAs (should equal total_tasks for causal)
+    l2_swizzle : int
+        L2 swizzle factor for cache locality
+    """
+
+    def __init__(
+        self,
+        prefix: str,
+        num_batches: int,
+        num_heads: int,
+        num_m_blocks: int,
+        l2_swizzle: int,
+        num_ctas: int | None = None,
+    ):
+        super().__init__(prefix)
+        self._num_batches = num_batches
+        self._num_heads = num_heads
+        self._num_m_blocks = num_m_blocks
+        self._l2_swizzle = l2_swizzle
+        self._num_ctas = num_ctas
+        self._total_tasks = num_batches * num_heads * num_m_blocks
+
+        # Derived constants for L2 swizzle
+        self._num_hb = num_batches * num_heads
+        self._l2_major = l2_swizzle * num_m_blocks
+        self._num_hb_quotient = self._num_hb // l2_swizzle
+
+        # Output indices
+        self.batch_idx = T.local_scalar("int32")
+        self.head_idx = T.local_scalar("int32")
+        self.m_block_idx = T.local_scalar("int32")
+
+    # fmt: off
+    @T.inline
+    def update_current_m_n_idx(self, linear_idx):
+        """Convert linear index to (batch, head, m_block) with LPT + L2 swizzle."""
+        L2_SWIZZLE = T.meta_var(self._l2_swizzle)
+        L2_MAJOR = T.meta_var(self._l2_major)
+        NUM_HB_QUOTIENT = T.meta_var(self._num_hb_quotient)
+        NUM_HB = T.meta_var(self._num_hb)
+        NUM_HEADS = T.meta_var(self._num_heads)
+        NUM_M_BLOCKS = T.meta_var(self._num_m_blocks)
+
+        # L2 swizzle decomposition
+        bidhb: T.let = linear_idx // L2_MAJOR
+        l2_mod: T.let = linear_idx % L2_MAJOR
+
+        # Handle residual section (last partial swizzle group)
+        num_hb_remainder: T.let = T.max(NUM_HB % L2_SWIZZLE, 1)
+        m_block_raw: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod // L2_SWIZZLE, l2_mod // num_hb_remainder)  # noqa: E501
+        bidhb_residual: T.let = T.Select(bidhb < NUM_HB_QUOTIENT, l2_mod % L2_SWIZZLE, l2_mod % num_hb_remainder)  # noqa: E501
+        bidhb_actual: T.let = bidhb * L2_SWIZZLE + bidhb_residual
+
+        self.batch_idx = bidhb_actual // NUM_HEADS
+        self.head_idx = bidhb_actual % NUM_HEADS
+
+        # LPT: Reverse block order so high-work blocks are processed first
+        self.m_block_idx = (NUM_M_BLOCKS - 1) - m_block_raw
+
+    @T.inline
+    def init(self, cta_id):
+        """Initialize scheduler with CTA ID."""
+        self.linear_idx = cta_id
+        self.update_current_m_n_idx(cta_id)
+
+    @T.inline
+    def next_tile(self):
+        """Advance to the next tile.
+
+        Single-tile mode (``num_ctas=None``, the default): each CTA owns one
+        task; terminate. Persistent mode (``num_ctas=N``): stride by N, like
+        :class:`FlashAttentionLinearScheduler`, while keeping the LPT + L2
+        swizzle index mapping.
+        """
+        if self._num_ctas is None:
+            self.linear_idx = self._total_tasks
+        else:
+            self.linear_idx = self.linear_idx + self._num_ctas
+            self.update_current_m_n_idx(self.linear_idx)
+    # fmt: on
+
+    def valid(self):
+        """Check if there are more tiles to process."""
+        return self.linear_idx < self._total_tasks
+
+
+class _CLCWorker(ClusterPersistentScheduler2D):
+    """Per-role CLC handle: IS-A ClusterPersistentScheduler2D (so m_idx / n_idx work as
+    usual) plus the role-local barrier phase and handshake. A coord-free role (e.g. an
+    MMA warp consuming whatever a loader staged) arms the loop with reset() not init().
+    """
+
+    def __init__(self, clc, prefix):
+        super().__init__(
+            prefix,
+            num_m_tiles=clc._num_m_tiles,
+            num_n_tiles=clc._num_n_tiles,
+            num_clusters=clc._num_m_tiles * clc._num_n_tiles,
+            l2_group_size=clc._l2_group_size,
+        )
+        self._clc = clc
+        self._sa = PipelineState(1, 0)
+        self._done = T.local_scalar("int32")
+        self._nxt = T.local_scalar("uint32")
+
+    @T.inline
+    def reset(self):
+        self._done = 0
+
+    @T.inline
+    def init(self, cluster_id):
+        # Explicit base call: TVMScript's parser has no zero-arg super().
+        ClusterPersistentScheduler2D.init(self, cluster_id)
+        self._done = 0
+
+    def valid(self):
+        return self._done == 0
+
+    @T.inline
+    def consume(self):
+        # Single-elected-thread scope: wait for the handle, decode, release the slot.
+        self._clc.sched_arr.full.wait(0, self._sa.phase)
+        self._sa.advance()
+        self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
+        self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)
+
+    @T.inline
+    def consume_wg(self, wg_id, warp_id, lane_id):
+        # Warpgroup scope: all threads decode; one elected lane releases the slot.
+        self._clc.sched_arr.full.wait(0, self._sa.phase)
+        self._sa.advance()
+        self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0]))
+        T.cuda.warpgroup_sync(wg_id + 1)
+        if (warp_id == 0) & (lane_id == 0):
+            self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True)
+
+    @T.inline
+    def advance_coords(self):
+        if self._nxt != 0xFFFFFFFF:
+            self.update_current_m_n_idx(self._nxt // self._clc._cta_group)
+
+    @T.inline
+    def mark_done_if_drained(self):
+        if self._nxt == 0xFFFFFFFF:
+            self._done = 1
+
+
+@T.meta_class
+class ClusterLaunchControlScheduler:
+    """Blackwell Cluster Launch Control (CLC) tile scheduler.
+
+    A scheduler warp runs ``run_scheduler`` (issues ``try_cancel`` to steal the next
+    cluster); worker roles each take a ``worker()`` handle and pull the stolen tile
+    through the shared smem handshake. Owns the CLC smem: the 16B response handle, the
+    arrival barrier (handle ready), and the finished barrier (slot consumed;
+    ``finish_arrivals`` arrivals per round). Tile-coord mapping is delegated to
+    ``ClusterPersistentScheduler2D`` (group-major L2 ordering).
+    """
+
+    def __init__(self, pool, num_m_tiles, num_n_tiles, l2_group_size, cta_group, finish_arrivals):
+        self._num_m_tiles = num_m_tiles
+        self._num_n_tiles = num_n_tiles
+        self._l2_group_size = l2_group_size
+        self._cta_group = cta_group
+        self.sched_arr = Pipeline(pool, 1, full="tma", empty="mbar", init_empty=1)
+        self.sched_fin = Pipeline(pool, 1, full="mbar", empty="mbar", init_empty=finish_arrivals)
+        self.clc_handle = pool.alloc((4,), "uint32", align=16)
+        self._s_done = T.local_scalar("int32")
+        self._s_nxt = T.local_scalar("uint32")
+
+    def worker(self, prefix):
+        return _CLCWorker(self, prefix)
+
+    @T.inline
+    def run_scheduler(self, cbx):
+        # cta0 drives try_cancel; both CTAs expect_bytes + consume the handle so the
+        # finished-barrier count is met and the slot can be reissued.
+        if T.ptx.elect_sync():
+            sa = PipelineState(1, 0)
+            sf = PipelineState(1, 1)
+            self._s_done = 0
+            while self._s_done == 0:
+                if cbx == 0:
+                    self.sched_fin.empty.wait(0, sf.phase)
+                    sf.advance()
+                    T.ptx.clc_try_cancel(
+                        T.address_of(self.clc_handle[0]), T.address_of(self.sched_arr.full.buf[0])
+                    )
+                self.sched_arr.full.arrive(0, 16)  # expect_bytes for the 16B handle
+                self.sched_arr.full.wait(0, sa.phase)
+                sa.advance()
+                self._s_nxt = T.ptx.clc_query_cancel(T.address_of(self.clc_handle[0]))
+                self.sched_fin.empty.arrive(0, cta_id=0, pred=True)
+                if self._s_nxt == 0xFFFFFFFF:
+                    self._s_done = 1

python/tvm/backend/cuda/lang/warp_role.py

@@ -0,0 +1,144 @@
+# 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.
+"""Warp role helpers for SM100 kernels.
+
+Simplifies the common pattern of dispatching warps to named roles
+with register budgets.
+
+Example::
+
+    # Declare roles
+    tma_warp = WarpRole(warp_id, 1, regs=48)
+    store_warp = WarpRole(warp_id, 2, regs=48)
+    mma_warp = WarpRole(warp_id, 0, regs=232, increase=True)
+
+    # Use with context manager
+    with tma_warp:
+        # TMA load code
+    with store_warp:
+        # TMA store code
+    with mma_warp:
+        # MMA compute code
+"""
+
+from tvm.script import tirx as T
+
+
+class WarpRole:
+    """A warp-level role that guards a block of code by warp_id comparison
+    with optional register budget.
+
+    Generates::
+
+        if <warp_id_var> == <warp_id_val>:
+            T.ptx.setmaxnreg(<increase>, <regs>)  # if regs specified
+            <user code>
+
+    The ``if`` guard narrows the active set to the single warp; individual
+    tile-primitive calls inside ``<user code>`` carry their own exec scope via
+    a scope-namespace prefix (e.g. ``Tx.warp.copy(...)``).
+
+    Parameters
+    ----------
+    warp_id_var : Var
+        The warp_id variable (from ``T.warp_id(...)``).
+    warp_id_val : int
+        Which warp index this role corresponds to.
+    regs : int, optional
+        Register budget (passed to ``T.ptx.setmaxnreg``).
+        If None, no setmaxnreg is emitted.
+    increase : bool
+        Direction for ``setmaxnreg`` (default False = decrease).
+    """
+
+    def __init__(self, warp_id_var, warp_id_val, regs=None, increase=False):
+        self.warp_id_var = warp_id_var
+        self.warp_id_val = warp_id_val
+        self.regs = regs
+        self.increase = increase
+
+    def __enter__(self):
+        self._if_frame = T.If(self.warp_id_var == self.warp_id_val)
+        self._if_frame.__enter__()
+        self._then_frame = T.Then()
+        self._then_frame.__enter__()
+        if self.regs is not None:
+            T.evaluate(T.ptx.setmaxnreg(self.increase, self.regs))
+        return self
+
+    def __exit__(self, *exc):
+        self._then_frame.__exit__(*exc)
+        self._if_frame.__exit__(*exc)
+        return False
+
+
+class WarpgroupRole:
+    """A warpgroup-level role that guards by wg_id comparison,
+    with optional register budget.
+
+    Generates (single wg_id)::
+
+        if <wg_id_var> == <wg_id_val>:
+            T.ptx.setmaxnreg(<increase>, <regs>)  # if regs specified
+            <user code>
+
+    Generates (range of wg_ids, e.g. ``wg_id_val=(0, 2)``)::
+
+        if 0 <= <wg_id_var> and <wg_id_var> < 2:
+            T.ptx.setmaxnreg(<increase>, <regs>)
+            <user code>
+
+    The ``if`` guard narrows the active set to the target warpgroup(s);
+    individual tile-primitive calls inside ``<user code>`` carry their own exec
+    scope via a scope-namespace prefix (e.g. ``Tx.wg.copy(...)``).
+
+    Parameters
+    ----------
+    wg_id_var : Var
+        The warpgroup_id variable (from ``T.warpgroup_id(...)``).
+    wg_id_val : int or tuple[int, int]
+        Which warpgroup index (int) or range ``(start, stop)`` this role
+        corresponds to.
+    regs : int, optional
+        Register budget.
+    increase : bool
+        Direction for ``setmaxnreg`` (default False = decrease).
+    """
+
+    def __init__(self, wg_id_var, wg_id_val, regs=None, increase=False):
+        self.wg_id_var = wg_id_var
+        self.wg_id_val = wg_id_val
+        self.regs = regs
+        self.increase = increase
+
+    def __enter__(self):
+        if isinstance(self.wg_id_val, tuple):
+            start, stop = self.wg_id_val
+            self._if_frame = T.If(start <= self.wg_id_var and self.wg_id_var < stop)
+        else:
+            self._if_frame = T.If(self.wg_id_var == self.wg_id_val)
+        self._if_frame.__enter__()
+        self._then_frame = T.Then()
+        self._then_frame.__enter__()
+        if self.regs is not None:
+            T.evaluate(T.ptx.setmaxnreg(self.increase, self.regs))
+        return self
+
+    def __exit__(self, *exc):
+        self._then_frame.__exit__(*exc)
+        self._if_frame.__exit__(*exc)
+        return False

python/tvm/backend/cuda/operator/__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.
+"""CUDA backend operator registrations and helpers."""
+
+__all__ = ["intrinsics", "tile_primitive"]

python/tvm/backend/cuda/operator/intrinsics/__init__.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=unused-import
+"""CUDA HW intrinsic codegens, grouped by feature domain.
+
+- ``mma`` / ``wgmma`` / ``tcgen05`` — matrix-multiply hardware (Volta+/Hopper/Blackwell).
+- ``cp_async`` — cp.async + cp.async.bulk + cp.async.bulk.tensor (TMA), incl. TMA address helpers.
+- ``sync`` — barriers, fences, mbarrier, cluster.barrier, warp vote, elect, sync helpers.
+- ``math`` — packed-f32x2 arithmetic, exp2/rcp/reduce3, warp/CTA reductions.
+- ``memory`` — typed copies, ldg, ld.global.acquire, atomics, type conversions, address casts.
+- ``nvshmem`` — NVSHMEM RMA / signal / collective.
+- ``misc`` — register-allocation control, profiler timer, debug helpers (printf / trap).
+
+Plus the support modules:
+
+- ``header`` — CUDA header generator and helper-tag table.
+- ``registry`` — codegen registry.
+- ``types`` — PTX dtype enum.
+- ``utils`` — small parsing / validation helpers.
+"""
+
+# Import op modules to register their codegen functions.
+from . import cp_async, math, memory, misc, mma, nvshmem, sync, tcgen05, wgmma
+from .header import TAGS, header_generator
+from .registry import CODEGEN_REGISTRY, get_codegen, register_codegen
+from .types import PTXDataType
+
+__all__ = [
+    "CODEGEN_REGISTRY",
+    "TAGS",
+    "PTXDataType",
+    "get_codegen",
+    "header_generator",
+    "register_codegen",
+]

python/tvm/backend/cuda/operator/intrinsics/_schema.py

@@ -0,0 +1,144 @@
+# 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
+"""Thin device-helper registration for TIRx intrinsic codegens.
+
+Exposes one entry point: :func:`device_intrinsic`. Given an op name plus
+``(helper_name, c_signature, body)`` (each a string or a ``(*args) -> str``
+callable), it:
+
+* wraps the body in
+  ``__forceinline__ __device__ <ret> <name><sig> { <body> }``,
+* registers a codegen function under the op name so
+  ``call_intrin("", "tirx.<op_name>", *args)`` resolves to a call to that
+  helper. TVM Op registration is static C++ only.
+
+Args passed to the codegen are split into ``(forward_args, attr_args)``:
+the trailing ``n_attrs`` are attrs (consumed by the ``helper_name`` /
+``c_signature`` / ``body`` callables but **not** forwarded to the helper),
+and the rest are operand args (forwarded). The default ``n_attrs=0`` means
+every arg is forwarded — appropriate for fixed-arity ops with literal
+``c_signature`` and ``helper_name``.
+
+Coerce / validate attrs explicitly inside the callables — there is no
+``Choice`` / ``Bool`` / ``IntAttr`` machinery; just call ``parse_str`` /
+``int`` / ``bool`` on the raw arg as needed.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from tvm.backend.cuda.op import cuda_func_call
+from tvm.backend.cuda.operator.intrinsics.registry import register_codegen
+
+# C primitive type → TVM dtype string. Used when the caller specifies a
+# non-void ``return_type`` but no explicit ``tvm_return_type`` — the helper
+# knows the TVM-side dtype from the C return type.
+_C_TO_TVM_DTYPE = {
+    "float": "float32",
+    "double": "float64",
+    "uint32_t": "uint32",
+    "int32_t": "int32",
+    "uint64_t": "uint64",
+    "int64_t": "int64",
+    "uint16_t": "uint16",
+    "int16_t": "int16",
+    "unsigned long long": "uint64",
+    "long long": "int64",
+    "unsigned short": "uint16",
+    "bool": "bool",
+    "unsigned int": "uint32",
+    "int": "int32",
+}
+
+
+def device_intrinsic(
+    op_name: str,
+    *,
+    helper_name: str | Callable | None = None,
+    c_signature: str | Callable = "()",
+    body: str | Callable,
+    n_attrs: int = 0,
+    return_type: str | Callable = "void",
+    tvm_return_type: str | Callable | None = None,
+    templated: bool = False,
+    extra_deps: tuple = (),
+) -> None:
+    """Register a CUDA device-helper intrinsic.
+
+    Parameters
+    ----------
+    op_name :
+        Registry key — ``call_intrin("", "tirx.<op_name>", ...)`` resolves
+        here. Also used as the default helper name (``tvm_builtin_<op_name>``)
+        when ``helper_name`` is not provided.
+    helper_name :
+        Literal C function name, OR ``(*args) -> str`` to compute it from
+        attr values. Defaults to ``f"tvm_builtin_{op_name}"``.
+    c_signature :
+        Literal C parameter list including outer parens (``"(int x, int y)"``),
+        OR ``(*args) -> str`` to compute it from attr values. Defaults to
+        ``"()"``.
+    body :
+        Literal C body string (already indented), OR ``(*args) -> str``.
+    n_attrs :
+        Number of trailing args that are attrs (consumed by ``helper_name``
+        / ``c_signature`` / ``body`` callables, NOT forwarded to the helper
+        as call arguments). The first ``len(args) - n_attrs`` args are the
+        operand args forwarded to the helper.
+    return_type :
+        C return type. Default ``"void"``. Either a literal string or
+        ``(*args) -> str`` when the helper return type depends on attrs.
+    tvm_return_type :
+        TVM dtype for the call result, when the helper has a non-void
+        return. Either a literal string (``"int32"``) or ``(*args) -> str``.
+        If omitted and ``return_type`` is non-void, it is auto-derived from
+        the ``_C_TO_TVM_DTYPE`` table.
+    templated :
+        Prefix the helper with ``template <typename T>``.
+    extra_deps :
+        Helper-tag list (e.g. ``("get_tmem_addr",)``) forwarded as the second
+        element of the codegen result so the header generator emits the
+        prerequisite snippets.
+    """
+    if helper_name is None:
+        helper_name = f"tvm_builtin_{op_name}"
+    extra_deps = tuple(extra_deps)
+
+    def codegen(*args):
+        forward = args if n_attrs == 0 else args[:-n_attrs]
+        name = helper_name(*args) if callable(helper_name) else helper_name
+        sig = c_signature(*args) if callable(c_signature) else c_signature
+        body_str = body(*args) if callable(body) else body
+        ret_type = return_type(*args) if callable(return_type) else return_type
+        prefix = "template <typename T>\n" if templated else ""
+        source_code = (
+            f"\n{prefix}__forceinline__ __device__ {ret_type} {name}{sig} {{\n{body_str}\n}}\n"
+        )
+        kwargs = {"source_code": source_code}
+        if tvm_return_type is not None:
+            kwargs["return_type"] = (
+                tvm_return_type(*args) if callable(tvm_return_type) else tvm_return_type
+            )
+        elif ret_type != "void":
+            kwargs["return_type"] = _C_TO_TVM_DTYPE.get(ret_type, ret_type)
+        result = cuda_func_call(name, *forward, **kwargs)
+        return (result, list(extra_deps)) if extra_deps else result
+
+    codegen.__name__ = f"codegen_{op_name}"
+    register_codegen(op_name)(codegen)

python/tvm/backend/cuda/operator/intrinsics/cp_async.py

@@ -0,0 +1,924 @@
+# 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, too-many-arguments, too-many-locals, too-many-positional-arguments
+"""PTX cp.async / cp.async.bulk / cp.async.bulk.tensor intrinsics.
+
+Each PTX form table entry is registered as one ``device_intrinsic``.
+User-facing wrappers in ``tvm.tirx.op`` keep their v1 signatures;
+``register_codegen`` dispatchers below decode the (cp_size, fill_mode,
+predicate) / (dim, cta_mask, tile_mode) arguments to pick the right form.
+Bodies are hand-written ``asm volatile(...)`` strings.  The file is grouped
+as cp.async, cp.async.bulk.tensor, cp.async.bulk non-TMA, and CUDA
+compatibility helpers.
+"""
+
+import tvm
+from tvm.backend.cuda.op import cuda_func_call
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+from .utils import parse_str
+
+_PREFETCH_CHOICES = ("", "64", "128", "256")
+_DIM_CHOICES = (1, 2, 3, 4, 5)
+_TILE_MODE_CHOICES = ("tile", "tile_gather4")
+
+
+def _safe(s):
+    return s.replace("::", "_").replace(".", "_")
+
+
+# =============================================================================
+# cp.async forms from the PTX Syntax block.
+#
+# Includes commit/wait plus the non-bulk shared/global copy forms.
+# =============================================================================
+device_intrinsic(
+    "ptx_cp_async_commit_group",
+    helper_name="tvm_builtin_ptx_cp_async_commit_group",
+    body='    asm volatile("cp.async.commit_group;");',
+)
+device_intrinsic(
+    "ptx_cp_async_wait_group",
+    n_attrs=1,
+    helper_name=lambda n: f"tvm_builtin_ptx_cp_async_wait_group_{int(n)}",
+    body=lambda n: f'    asm volatile("cp.async.wait_group {int(n)};");',
+)
+
+
+# cp.async non-bulk copy forms:
+#   Form 1: cp.async.ca.shared.global ... [dst], [src], cp-size{, src-size}{, cache-policy}
+#   Form 2: cp.async.cg.shared.global ... [dst], [src], 16{, src-size}{, cache-policy}
+#   Form 3: cp.async.ca.shared.global ... [dst], [src], cp-size{, ignore-src}{, cache-policy}
+#   Form 4: cp.async.cg.shared.global ... [dst], [src], 16{, ignore-src}{, cache-policy}
+
+
+def _cp_async_modifier_str(has_cache_hint, prefetch_size):
+    s = ""
+    if has_cache_hint:
+        s += ".L2::cache_hint"
+    if prefetch_size:
+        s += f".L2::{prefetch_size}B"
+    return s
+
+
+def _make_form_parts(ca_or_cg, fixed_cp_size, extra):
+    """Build a parts callable for one of the cp.async PTX forms.
+
+    Args layout: (dst, src [, extra_int], cache_policy, has_cache, prefetch_size [, cp_size_attr])
+    Forwarded operands: dst, src [, extra_int], cache_policy.
+    Trailing attrs: has_cache, prefetch_size [, cp_size if .ca].
+    """
+    n_op = 3 if extra is not None else 2
+    n_attrs = 2 if fixed_cp_size is not None else 3
+    extra_in_name = f"_with_{extra}" if extra is not None else ""
+
+    def _parts(*args):
+        # Operand args (forwarded) come first, then attr args.
+        attr_args = args[-n_attrs:]
+        has_cache = _bool_attr(attr_args[0])
+        prefetch_size = parse_str(attr_args[1])
+        cp_size = fixed_cp_size if fixed_cp_size is not None else int(attr_args[2])
+        modifier = _cp_async_modifier_str(has_cache, prefetch_size)
+        cache_operand = ', "l"(cache_policy)' if has_cache else ""
+        # name parts
+        name_cache = "_cache_hint" if has_cache else ""
+        name_prefetch = f"_prefetch_{prefetch_size}" if prefetch_size else ""
+        name = (
+            f"tvm_builtin_ptx_cp_async_{ca_or_cg}_{cp_size}"
+            f"{name_cache}{name_prefetch}{extra_in_name}"
+        )
+        sig = (
+            "(void* dst, void* src"
+            + (f", int {extra}" if extra else "")
+            + ", unsigned long long cache_policy)"
+        )
+        instr_base = f"cp.async.{ca_or_cg}.shared.global{modifier}"
+        if extra is None:
+            cache_arg = ", %2" if has_cache else ""
+            body = (
+                "    unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
+                f'    asm volatile("{instr_base} [%0], [%1], {cp_size}{cache_arg};\\n"\n'
+                f'                 :: "r"(dst_addr), "l"(src){cache_operand} : "memory");'
+            )
+        else:
+            cache_arg = ", %3" if has_cache else ""
+            body = (
+                "    unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
+                f'    asm volatile("{instr_base} [%0], [%1], {cp_size}, %2{cache_arg};\\n"\n'
+                f'                 :: "r"(dst_addr), "l"(src), "r"({extra})'
+                f'{cache_operand} : "memory");'
+            )
+        return name, sig, body
+
+    return _parts, n_op + n_attrs - n_op  # n_attrs
+
+
+def _register_nb_form(op_name, ca_or_cg, fixed_cp_size, extra):
+    parts_fn, n_attrs = _make_form_parts(ca_or_cg, fixed_cp_size, extra)
+    n_op = 3 if extra is not None else 2
+    sig_static = (
+        "(void* dst, void* src"
+        + (f", int {extra}" if extra else "")
+        + ", unsigned long long cache_policy)"
+    )
+    device_intrinsic(
+        f"ptx_cp_async_{op_name}",
+        n_attrs=n_attrs,
+        c_signature=sig_static,  # static — depends on `extra` not on attrs
+        helper_name=lambda *a, fn=parts_fn: fn(*a)[0],
+        body=lambda *a, fn=parts_fn: fn(*a)[2],
+    )
+    return n_op
+
+
+# Form 1: .ca + src-size (cp-size ∈ {4, 8}). src-size is required when present.
+_register_nb_form("ca_src_size", "ca", fixed_cp_size=None, extra="src_size")
+# Form 2: .cg + src-size (cp-size = 16).
+_register_nb_form("cg_src_size", "cg", fixed_cp_size=16, extra="src_size")
+# Form 3: .ca + ignore-src.
+_register_nb_form("ca_ignore_src", "ca", fixed_cp_size=None, extra="ignore_src")
+# Form 4: .cg + ignore-src.
+_register_nb_form("cg_ignore_src", "cg", fixed_cp_size=16, extra="ignore_src")
+# Plain degenerate of forms 1+2 with optional src-size omitted.
+_register_nb_form("ca", "ca", fixed_cp_size=None, extra=None)
+_register_nb_form("cg", "cg", fixed_cp_size=16, extra=None)
+
+
+def _make_setp_at_p_helper(ca_or_cg, cp_size, has_cache, prefetch):
+    """Wrapper convenience: ``setp+@p`` around a form 1/2 cp.async (predicate-
+    gated skip with dst untouched on false). Not a PTX form — emitted directly
+    here as a one-off helper rather than a separate device_intrinsic."""
+    modifier = _cp_async_modifier_str(has_cache, prefetch)
+    cache_arg = ", %4" if has_cache else ""
+    cache_operand = ', "l"(cache_policy)' if has_cache else ""
+    func_name = (
+        f"tvm_builtin_ptx_cp_async_{cp_size}"
+        + ("_cache_hint" if has_cache else "")
+        + (f"_prefetch_{prefetch}" if prefetch else "")
+        + "_predicate"
+    )
+    body = (
+        "  unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
+        "  __asm__ __volatile__(\n"
+        '    "{\\n"\n'
+        '    " .reg .pred p;\\n"\n'
+        '    " setp.eq.u32 p, %3, 1;\\n"\n'
+        f'    " @p cp.async.{ca_or_cg}.shared.global{modifier}'
+        f' [%0], [%1], %2{cache_arg};\\n"\n'
+        '    "}\\n"\n'
+        f'    :: "r"(dst_addr), "l"(src), "n"({cp_size}), "r"(predicate){cache_operand}\n'
+        "  );"
+    )
+    source_code = (
+        f"\n__forceinline__ __device__ void {func_name}"
+        "(void* dst, void* src, int predicate, unsigned long long cache_policy) {\n"
+        f"{body}\n"
+        "}\n"
+    )
+    return func_name, source_code
+
+
+@register_codegen("ptx_cp_async")
+def codegen_ptx_cp_async(*args):
+    """Map the wrapper API to the 4 PTX form table entries.
+
+    Accepts three call shapes (sorted by arity):
+
+    * 5 args ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size)`` —
+      the legacy form emitted by
+      ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``.
+      Offsets are folded into the pointers via ``tvm_access_ptr`` (in
+      bytes; offsets are pre-scaled by the pass) and the call is
+      forwarded with default cache / predicate / fill_mode.
+    * 6 args ``(dst_ptr, dst_offset, src_ptr, src_offset, cp_size,
+      predicate)`` — same as 5-arg form with an explicit predicate,
+      zero-filling the destination when the predicate is false.
+    * 8 args ``(dst_ptr, src_ptr, cp_size, cache_policy, has_cache_hint,
+      prefetch_size, predicate, fill_mode)`` — the fork-native wrapper
+      API.
+
+    The three resulting form_kinds:
+
+    * ``fill_mode == "zero"`` -> form 1/2 (src-size = predicate ? cp_size : 0)
+    * ``predicate != -1`` and no fill_mode -> form 1/2 wrapped in setp+@p
+      (wrapper convenience; not a PTX form)
+    * else -> form 1/2 with src-size omitted (the "plain" degenerate)
+    """
+    from tvm.tirx.op import if_then_else
+
+    if len(args) in (5, 6):
+        # Legacy InjectPTXAsyncCopy emission: (dst_ptr, dst_off, src_ptr,
+        # src_off, cp_size [, predicate]). Offsets are element indices into
+        # the typed buffers (the pass uses index_factor=1 except for the
+        # shared.dyn-merged byte-buffer path). Emit a C helper that scales
+        # the offset by the buffer element size, then runs cp.async.
+        #
+        # PTX plain form for both .ca and .cg is just
+        # ``cp.async.<v>.shared.global [dst], [src], cp_size;`` — three
+        # operands, no trailing src-size / cache-policy.
+        from tvm import DataType
+
+        dst_ptr_in, dst_offset, src_ptr_in, src_offset, cp_size = args[:5]
+        predicate = args[5] if len(args) == 6 else -1
+        cp_size_v = int(cp_size)
+        ca_or_cg = "cg" if cp_size_v == 16 else "ca"
+
+        # Recover the per-side element dtype from each pointer's type
+        # type (Var has ty = PointerType(PrimType(dtype))).
+        # InjectPTXAsyncCopy emits offsets in element-units of each side's
+        # buffer dtype (dst gets dst_offset * src_elem_size only when dst is a
+        # merged shared.dyn byte buffer, in which case dst_elem_dtype is uint8
+        # and the resulting scale-by-1 is a no-op).
+        def _elem_bytes(ptr):
+            ta = getattr(ptr, "ty", None)
+            if ta is None or getattr(ta, "element_type", None) is None:
+                return 1
+            et = ta.element_type
+            if not hasattr(et, "dtype"):
+                return 1
+            bits = DataType(str(et.dtype)).bits
+            assert bits % 8 == 0, f"non-byte element dtype: {et.dtype}"
+            return bits // 8
+
+        dst_elem_bytes = _elem_bytes(dst_ptr_in)
+        src_elem_bytes = _elem_bytes(src_ptr_in)
+        has_predicate = not (
+            (isinstance(predicate, int) and predicate == -1)
+            or (hasattr(predicate, "value") and int(predicate.value) == -1)
+        )
+
+        def _scale(n):
+            return "" if n == 1 else f" * {n}"
+
+        dst_scale = _scale(dst_elem_bytes)
+        src_scale = _scale(src_elem_bytes)
+        if has_predicate:
+            func_name = (
+                f"ptx_cp_async_legacy_pred_{ca_or_cg}_{cp_size_v}_{dst_elem_bytes}_{src_elem_bytes}"
+            )
+            if cp_size_v == 4:
+                zero_fill = '    " @!p st.shared.u32 [%0], {%4};\\n"\n'
+            elif cp_size_v == 8:
+                zero_fill = '    " @!p st.shared.v2.u32 [%0], {%4, %4};\\n"\n'
+            elif cp_size_v == 16:
+                zero_fill = '    " @!p st.shared.v4.u32 [%0], {%4, %4, %4, %4};\\n"\n'
+            else:
+                raise ValueError(f"unsupported legacy predicated cp.async size: {cp_size_v}")
+            body = (
+                f"  uint8_t* dst_p = (uint8_t*)dst + dst_off{dst_scale};\n"
+                f"  uint8_t* src_p = (uint8_t*)src + src_off{src_scale};\n"
+                "  unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n"
+                "  __asm__ __volatile__(\n"
+                '    "{\\n"\n'
+                '    " .reg .pred p;\\n"\n'
+                '    " setp.eq.u32 p, %3, 1;\\n"\n'
+                f'    " @p cp.async.{ca_or_cg}.shared.global'
+                ' [%0], [%1], %2;\\n"\n'
+                f"{zero_fill}"
+                '    "}\\n"\n'
+                f'    :: "r"(dst_addr), "l"(src_p), "n"({cp_size_v}), "r"(predicate), "r"(0)\n'
+                "  );"
+            )
+            source_code = (
+                f"\n__forceinline__ __device__ void {func_name}"
+                "(void* dst, int dst_off, void* src, int src_off, int predicate) {\n"
+                f"{body}\n"
+                "}\n"
+            )
+            return cuda_func_call(
+                func_name,
+                dst_ptr_in,
+                dst_offset,
+                src_ptr_in,
+                src_offset,
+                predicate,
+                source_code=source_code,
+            )
+        # No predicate — plain cp.async.
+        func_name = f"ptx_cp_async_legacy_{ca_or_cg}_{cp_size_v}_{dst_elem_bytes}_{src_elem_bytes}"
+        body = (
+            f"  uint8_t* dst_p = (uint8_t*)dst + dst_off{dst_scale};\n"
+            f"  uint8_t* src_p = (uint8_t*)src + src_off{src_scale};\n"
+            "  unsigned int dst_addr = __cvta_generic_to_shared(dst_p);\n"
+            f'  asm volatile("cp.async.{ca_or_cg}.shared.global'
+            ' [%0], [%1], %2;"\n'
+            f'    :: "r"(dst_addr), "l"(src_p), "n"({cp_size_v}));'
+        )
+        source_code = (
+            f"\n__forceinline__ __device__ void {func_name}"
+            "(void* dst, int dst_off, void* src, int src_off) {\n"
+            f"{body}\n"
+            "}\n"
+        )
+        return cuda_func_call(
+            func_name,
+            dst_ptr_in,
+            dst_offset,
+            src_ptr_in,
+            src_offset,
+            source_code=source_code,
+        )
+    elif len(args) == 8:
+        (
+            dst_ptr,
+            src_ptr,
+            cp_size,
+            cache_policy,
+            has_cache_hint,
+            prefetch_size,
+            predicate,
+            fill_mode,
+        ) = args
+    else:
+        raise ValueError(f"ptx_cp_async codegen expects 5/6/8 args, got {len(args)}")
+
+    cp_size_v = int(cp_size)
+    ca_or_cg = "cg" if cp_size_v == 16 else "ca"
+    pref = "" if int(prefetch_size) == -1 else str(int(prefetch_size))
+    fill = parse_str(fill_mode)
+    has_cache = _bool_attr(has_cache_hint)
+    has_predicate = not (
+        (isinstance(predicate, int) and predicate == -1)
+        or (hasattr(predicate, "value") and int(predicate.value) == -1)
+    )
+
+    if fill == "zero":
+        src_size = if_then_else(predicate != 0, cp_size_v, 0)
+        op = f"tirx.ptx_cp_async_{ca_or_cg}_src_size"
+        if cp_size_v == 16:
+            args = [dst_ptr, src_ptr, src_size, cache_policy, has_cache, pref]
+        else:
+            args = [dst_ptr, src_ptr, src_size, cache_policy, has_cache, pref, cp_size_v]
+        result = CODEGEN_REGISTRY[op](args)
+        return result[0] if isinstance(result, tuple) else result
+
+    if has_predicate:
+        func_name, source_code = _make_setp_at_p_helper(ca_or_cg, cp_size_v, has_cache, pref)
+        return cuda_func_call(
+            func_name, dst_ptr, src_ptr, predicate, cache_policy, source_code=source_code
+        )
+
+    # Plain — form 1/2 with src-size omitted.
+    op = f"tirx.ptx_cp_async_{ca_or_cg}"
+    if cp_size_v == 16:
+        args = [dst_ptr, src_ptr, cache_policy, has_cache, pref]
+    else:
+        args = [dst_ptr, src_ptr, cache_policy, has_cache, pref, cp_size_v]
+    result = CODEGEN_REGISTRY[op](args)
+    return result[0] if isinstance(result, tuple) else result
+
+
+CODEGEN_REGISTRY["tirx.ptx.cp_async_raw"] = CODEGEN_REGISTRY["tirx.ptx.cp_async"]
+
+
+# =============================================================================
+# cp.async.bulk.tensor (TMA) — one device_intrinsic per arity variant of each
+# PTX form. Per-dim coord operands materialise via the ``c_signature`` callable.
+# =============================================================================
+
+
+def _is_sm100_or_higher():
+    target = tvm.target.Target.current()
+    if target is None:
+        return False
+    arch = target.arch[3:]
+    if not arch[-1].isdigit():
+        arch = arch[:-1]
+    return int(arch) >= 100
+
+
+def _resolve_cta_group_str(cta_group):
+    if cta_group == 2 or (cta_group != -1 and _is_sm100_or_higher()):
+        return f".cta_group::{cta_group}"
+    return ""
+
+
+def _coord_template(coord_count, start_slot):
+    inner = ", ".join(f"%{start_slot + i}" for i in range(coord_count))
+    return f"{{{inner}}}"
+
+
+def _coord_constraints(coord_count):
+    return ", ".join(f'"r"(coord{i})' for i in range(coord_count))
+
+
+def _coord_sig(n):
+    return ", ".join(f"int coord{i}" for i in range(n))
+
+
+# PTX cp.async.bulk.tensor global -> shared::cluster form:
+#   cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism
+#       {.multicast}{.cta_group}{.level::cache_hint}
+#       [dstMem], [tensorMap, tensorCoords], [mbar]{, im2colInfo}
+#       {, ctaMask} {, cache-policy}
+#   .dst = {.shared::cluster}; .src = {.global}
+#   .completion_mechanism = {.mbarrier::complete_tx::bytes}
+#   .multicast = {.multicast::cluster}
+#   .cta_group = {.cta_group::1, .cta_group::2}
+#   .load_mode = {.tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128}
+#   .level::cache_hint = {.L2::cache_hint}
+# This registration supports tile/tile::gather4 modes; ctaMask is only used
+# when the optional ``.multicast::cluster`` modifier is enabled.
+def _g2cluster_parts(*args):
+    attrs = args[-6:]
+    dim = int(attrs[0])
+    cta_group = int(attrs[1])
+    has_cache = _bool_attr(attrs[2])
+    tile_mode = parse_str(attrs[3])
+    bar_is_addr = _bool_attr(attrs[4])
+    multicast = _bool_attr(attrs[5])
+    coord_count = 5 if tile_mode == "tile_gather4" else dim
+    bar_type = "unsigned int bar_addr" if bar_is_addr else "void* bar"
+    sig = (
+        f"(void* dst, {bar_type}, unsigned long long tensormap_addr, "
+        "uint16_t cta_mask, unsigned long long cache_policy"
+        + (", " + _coord_sig(coord_count) if coord_count else "")
+        + ")"
+    )
+    name = (
+        f"ptx_cp_async_bulk_tensor_g2cluster_{tile_mode}_{dim}d"
+        f"{'_multicast' if multicast else ''}"
+        f"{'_cache_hint' if has_cache else ''}{'_bar_addr' if bar_is_addr else ''}"
+    )
+    tile_modifier = ".tile::gather4" if tile_mode == "tile_gather4" else ""
+    cta_group_str = _resolve_cta_group_str(cta_group)
+    multicast_inst = ".multicast::cluster" if multicast else ""
+    cache_inst = ".L2::cache_hint" if has_cache else ""
+    mask_arg = ',\n          "h"(cta_mask)' if multicast else ""
+    cache_arg = ',\n          "l"(cache_policy)' if has_cache else ""
+    mask_slot = ", %3" if multicast else ""
+    cache_slot = ", %4" if multicast and has_cache else ", %3" if has_cache else ""
+    coord_start = 5 if multicast and has_cache else 4 if multicast or has_cache else 3
+    coord_tpl = _coord_template(coord_count, coord_start)
+    instr = (
+        f"cp.async.bulk.tensor.{dim}d.shared::cluster.global{tile_modifier}"
+        f".mbarrier::complete_tx::bytes{multicast_inst}"
+        f"{cta_group_str}{cache_inst}"
+    )
+    bar_addr_decl = (
+        "" if bar_is_addr else "    unsigned int bar_addr = __cvta_generic_to_shared(bar);\n"
+    )
+    body = (
+        "    unsigned int dst_addr = __cvta_generic_to_shared(dst);\n"
+        f"{bar_addr_decl}"
+        "    asm volatile(\n"
+        f'        "{instr} [%0], [%1, {coord_tpl}], [%2]{mask_slot}{cache_slot};"\n'
+        "        :\n"
+        f'        : "r"(dst_addr), "l"(tensormap_addr), "r"(bar_addr){mask_arg}{cache_arg},\n'
+        f"          {_coord_constraints(coord_count)}\n"
+        '        : "memory"\n'
+        "    );"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_tensor_g2cluster",
+    n_attrs=6,
+    helper_name=lambda *a: _g2cluster_parts(*a)[0],
+    c_signature=lambda *a: _g2cluster_parts(*a)[1],
+    body=lambda *a: _g2cluster_parts(*a)[2],
+)
+
+
+# PTX cp.async.bulk.tensor shared::cta -> global form:
+#   cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism
+#       {.level::cache_hint}
+#       [tensorMap, tensorCoords], [srcMem] {, cache-policy}
+#   .dst = {.global}; .src = {.shared::cta}
+#   .completion_mechanism = {.bulk_group}
+#   .load_mode = {.tile, .tile::scatter4, .im2col_no_offs}
+#   .level::cache_hint = {.L2::cache_hint}
+# This registration supports tile mode; cache-policy is a real operand.
+def _s2g_parts(*args):
+    attrs = args[-2:]
+    dim = int(attrs[0])
+    has_cache = _bool_attr(attrs[1])
+    sig = (
+        "(void* src, unsigned long long tensormap_addr, unsigned long long cache_policy"
+        + (", " + _coord_sig(dim) if dim else "")
+        + ")"
+    )
+    name = f"ptx_cp_async_bulk_tensor_shared_to_global_{dim}d{'_cache_hint' if has_cache else ''}"
+    cache_inst = ".L2::cache_hint" if has_cache else ""
+    cache_arg = ', "l"(cache_policy)' if has_cache else ""
+    cache_slot = ", %2" if has_cache else ""
+    coord_start = 3 if has_cache else 2
+    coord_tpl = _coord_template(dim, coord_start)
+    instr = f"cp.async.bulk.tensor.{dim}d.global.shared::cta.tile.bulk_group{cache_inst}"
+    body = (
+        "    unsigned int src_addr = __cvta_generic_to_shared(src);\n"
+        "    asm volatile(\n"
+        f'        "{instr} [%0, {coord_tpl}], [%1]{cache_slot};"\n'
+        "        :\n"
+        f'        : "l"(tensormap_addr), "r"(src_addr){cache_arg},\n'
+        f"          {_coord_constraints(dim)}\n"
+        '        : "memory"\n'
+        "    );"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_tensor_s2g",
+    n_attrs=2,
+    helper_name=lambda *a: _s2g_parts(*a)[0],
+    c_signature=lambda *a: _s2g_parts(*a)[1],
+    body=lambda *a: _s2g_parts(*a)[2],
+)
+
+
+# PTX cp.async.bulk.prefetch.tensor form:
+#   cp.async.bulk.prefetch.tensor.dim.L2.src{.load_mode}{.level::cache_hint}
+#       [tensorMap, tensorCoords] {, im2colInfo} {, cache-policy}
+#   .src = {.global}
+#   .load_mode = {.tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128}
+#   .level::cache_hint = {.L2::cache_hint}
+# This registration supports tile mode; cache-policy is a real operand.
+def _prefetch_parts(*args):
+    attrs = args[-2:]
+    dim = int(attrs[0])
+    has_cache = _bool_attr(attrs[1])
+    sig = (
+        "(unsigned long long tensormap_addr, unsigned long long cache_policy"
+        + (", " + _coord_sig(dim) if dim else "")
+        + ")"
+    )
+    name = (
+        f"ptx_cp_async_bulk_tensor_global_to_cluster_prefetch_{dim}d"
+        f"{'_cache_hint' if has_cache else ''}"
+    )
+    cache_inst = ".L2::cache_hint" if has_cache else ""
+    cache_arg = ', "l"(cache_policy)' if has_cache else ""
+    cache_slot = ", %1" if has_cache else ""
+    coord_start = 2 if has_cache else 1
+    coord_tpl = _coord_template(dim, coord_start)
+    instr = f"cp.async.bulk.prefetch.tensor.{dim}d.L2.global.tile{cache_inst}"
+    body = (
+        "    asm volatile(\n"
+        f'        "{instr} [%0, {coord_tpl}]{cache_slot};"\n'
+        "        :\n"
+        f'        : "l"(tensormap_addr){cache_arg},\n'
+        f"          {_coord_constraints(dim)}\n"
+        '        : "memory"\n'
+        "    );"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_tensor_prefetch",
+    n_attrs=2,
+    helper_name=lambda *a: _prefetch_parts(*a)[0],
+    c_signature=lambda *a: _prefetch_parts(*a)[1],
+    body=lambda *a: _prefetch_parts(*a)[2],
+)
+
+
+# PTX cp.reduce.async.bulk.tensor shared::cta -> global form:
+#   cp.reduce.async.bulk.tensor.dim.dst.src.redOp{.load_mode}.completion_mechanism
+#       {.level::cache_hint}
+#       [tensorMap, tensorCoords], [srcMem] {, cache-policy}
+#   .dst = {.global}; .src = {.shared::cta}
+#   .completion_mechanism = {.bulk_group}
+#   .redOp = {.add, .min, .max, .inc, .dec, .and, .or, .xor}
+#   .level::cache_hint = {.L2::cache_hint}
+# This registration supports tile mode; redOp is syntax, cache-policy is an operand.
+def _reduce_parts(*args):
+    attrs = args[-3:]
+    dim = int(attrs[0])
+    has_cache = _bool_attr(attrs[1])
+    red_op = parse_str(attrs[2])
+    sig = (
+        "(void* src, unsigned long long tensormap_addr, unsigned long long cache_policy"
+        + (", " + _coord_sig(dim) if dim else "")
+        + ")"
+    )
+    name = (
+        f"ptx_cp_async_bulk_tensor_shared_to_global_reduce_{dim}d"
+        f"{'_cache_hint' if has_cache else ''}"
+    )
+    cache_inst = ".L2::cache_hint" if has_cache else ""
+    cache_arg = ', "l"(cache_policy)' if has_cache else ""
+    cache_slot = ", %2" if has_cache else ""
+    coord_start = 3 if has_cache else 2
+    coord_tpl = _coord_template(dim, coord_start)
+    instr = (
+        f"cp.reduce.async.bulk.tensor.{dim}d.global.shared::cta"
+        f".{red_op}.tile.bulk_group{cache_inst}"
+    )
+    body = (
+        "    unsigned int src_addr = __cvta_generic_to_shared(src);\n"
+        "    asm volatile(\n"
+        f'        "{instr} [%0, {coord_tpl}], [%1]{cache_slot};"\n'
+        "        :\n"
+        f'        : "l"(tensormap_addr), "r"(src_addr){cache_arg},\n'
+        f"          {_coord_constraints(dim)}\n"
+        '        : "memory"\n'
+        "    );"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_tensor_reduce",
+    n_attrs=3,
+    helper_name=lambda *a: _reduce_parts(*a)[0],
+    c_signature=lambda *a: _reduce_parts(*a)[1],
+    body=lambda *a: _reduce_parts(*a)[2],
+)
+
+
+# User-facing dispatchers for tensor global -> shared::cluster.  The same
+# backend root handles the optional ``.multicast::cluster`` modifier.
+
+
+def _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode):
+    cta_mask, cta_group, cache_policy, has_cache, *rest = args
+    coord_count = 5 if tile_mode == "tile_gather4" else int(dim)
+    if len(rest) == coord_count + 1:
+        bar_is_addr = _bool_attr(rest[0])
+        coords = rest[1:]
+    else:
+        bar_is_addr = False
+        coords = rest
+    is_unicast = isinstance(cta_mask, tvm.tirx.IntImm) and bin(int(cta_mask)).count("1") <= 1
+    cg = int(cta_group)
+    op = "tirx.ptx_cp_async_bulk_tensor_g2cluster"
+    call_args = [
+        dst_ptr,
+        bar,
+        tensormap,
+        cta_mask,
+        cache_policy,
+        *coords,
+        int(dim),
+        cg,
+        has_cache,
+        tile_mode,
+        bar_is_addr,
+        int(not is_unicast),
+    ]
+    result = CODEGEN_REGISTRY[op](call_args)
+    return result[0] if isinstance(result, tuple) else result
+
+
+@register_codegen("ptx_cp_async_bulk_tensor_global_to_cluster")
+def codegen_g2c(dim, dst_ptr, bar, tensormap, *args):
+    return _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode="tile")
+
+
+@register_codegen("ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster")
+def codegen_g2c_gather4(dim, dst_ptr, bar, tensormap, *args):
+    return _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode="tile_gather4")
+
+
+@register_codegen("ptx_cp_async_bulk_tensor_shared_to_global")
+def codegen_s2g(dim, src_ptr, tensormap, *args):
+    cache_policy, has_cache, *coords = args
+    result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_s2g"](
+        [src_ptr, tensormap, cache_policy, *coords, int(dim), has_cache]
+    )
+    return result[0] if isinstance(result, tuple) else result
+
+
+@register_codegen("ptx_cp_async_bulk_tensor_global_to_cluster_prefetch")
+def codegen_prefetch(dim, tensormap, *args):
+    cache_policy, has_cache, *coords = args
+    result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_prefetch"](
+        [tensormap, cache_policy, *coords, int(dim), has_cache]
+    )
+    return result[0] if isinstance(result, tuple) else result
+
+
+@register_codegen("ptx_cp_async_bulk_tensor_shared_to_global_reduce")
+def codegen_reduce(dim, src_ptr, tensormap, *args):
+    cache_policy, has_cache, red_op, *coords = args
+    result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_reduce"](
+        [src_ptr, tensormap, cache_policy, *coords, int(dim), has_cache, red_op]
+    )
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# cp.async.bulk non-TMA forms from the PTX Syntax block. Each form is one
+# device_intrinsic; optional PTX modifiers are attrs, not separate fixed ops.
+# =============================================================================
+device_intrinsic(
+    "ptx_cp_async_bulk_commit_group",
+    helper_name="ptx_cp_async_bulk_tensor_commit_group",
+    body='    asm volatile("cp.async.bulk.commit_group;");',
+)
+
+
+def _ptx_cp_async_bulk_wait_group_parts(n, read):
+    n = int(n)
+    read_b = bool(int(read)) if hasattr(read, "value") else bool(read)
+    return (
+        f"ptx_cp_async_bulk_wait_group{'_read' if read_b else ''}_{n}",
+        f'    asm volatile("cp.async.bulk.wait_group{".read" if read_b else ""} {n};");',
+    )
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_wait_group",
+    n_attrs=2,
+    helper_name=lambda n, read: _ptx_cp_async_bulk_wait_group_parts(n, read)[0],
+    body=lambda n, read: _ptx_cp_async_bulk_wait_group_parts(n, read)[1],
+)
+
+
+def _bool_attr(value):
+    return bool(int(value)) if hasattr(value, "value") else bool(value)
+
+
+def _bulk_cache_operand_constraint(has_cache):
+    return ', "l"(cache_policy)' if has_cache else ""
+
+
+def _bulk_cache_operand_suffix(has_cache):
+    return ".L2::cache_hint" if has_cache else ""
+
+
+# PTX cp.async.bulk global -> shared::cta form:
+#   cp.async.bulk.dst.src.completion_mechanism{.level::cache_hint}{.ignore_oob}
+#       [dstMem], [srcMem], size{, ignoreBytesLeft, ignoreBytesRight}, [mbar] {, cache-policy}
+#   .dst = {.shared::cta}; .src = {.global}
+#   .completion_mechanism = {.mbarrier::complete_tx::bytes}
+#   .level::cache_hint = {.L2::cache_hint}
+def _bulk_g2s_cta_parts(*args):
+    has_cache = _bool_attr(args[-2])
+    ignore_oob = _bool_attr(args[-1])
+    instr = (
+        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes"
+        f"{_bulk_cache_operand_suffix(has_cache)}{'.ignore_oob' if ignore_oob else ''}"
+    )
+    if ignore_oob:
+        asm_args = (
+            '"r"(dst), "l"(src_ptr), "r"(num_bytes), "r"(ignore_bytes_left), '
+            '"r"(ignore_bytes_right), "r"(mbarrier)'
+        )
+        operands = "%2, %3, %4, [%5]"
+        cache_slot = ", %6" if has_cache else ""
+    else:
+        asm_args = '"r"(dst), "l"(src_ptr), "r"(num_bytes), "r"(mbarrier)'
+        operands = "%2, [%3]"
+        cache_slot = ", %4" if has_cache else ""
+    body = (
+        "    unsigned int dst = (unsigned int)__cvta_generic_to_shared(dst_ptr);\n"
+        "    unsigned int mbarrier = (unsigned int)__cvta_generic_to_shared(mbarrier_ptr);\n"
+        f'    asm volatile("{instr} [%0], [%1], {operands}{cache_slot};"\n'
+        "                 :\n"
+        f"                 : {asm_args}{_bulk_cache_operand_constraint(has_cache)}\n"
+        '                 : "memory");'
+    )
+    name = (
+        "tvm_builtin_ptx_cp_async_bulk_g2s_cta"
+        f"{'_cache_hint' if has_cache else ''}{'_ignore_oob' if ignore_oob else ''}"
+    )
+    return name, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_g2s_cta",
+    n_attrs=2,
+    helper_name=lambda *a: _bulk_g2s_cta_parts(*a)[0],
+    c_signature=(
+        "(void* dst_ptr, void* src_ptr, unsigned int num_bytes, "
+        "unsigned int ignore_bytes_left, unsigned int ignore_bytes_right, "
+        "void* mbarrier_ptr, unsigned long long cache_policy)"
+    ),
+    body=lambda *a: _bulk_g2s_cta_parts(*a)[1],
+)
+
+
+# PTX cp.async.bulk global -> shared::cluster form:
+#   cp.async.bulk.dst.src.completion_mechanism{.multicast}{.level::cache_hint}
+#       [dstMem], [srcMem], size, [mbar] {, ctaMask} {, cache-policy}
+#   .dst = {.shared::cluster}; .src = {.global}
+#   .completion_mechanism = {.mbarrier::complete_tx::bytes}
+#   .level::cache_hint = {.L2::cache_hint}
+#   .multicast = {.multicast::cluster}
+def _bulk_g2s_cluster_parts(*args):
+    has_cache = _bool_attr(args[-2])
+    multicast = _bool_attr(args[-1])
+    instr = (
+        "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
+        f"{'.multicast::cluster' if multicast else ''}{_bulk_cache_operand_suffix(has_cache)}"
+    )
+    cta_constraint = ', "h"(cta_mask)' if multicast else ""
+    mask_slot = ", %4" if multicast else ""
+    cache_slot = ", %5" if multicast and has_cache else ", %4" if has_cache else ""
+    body = (
+        "    unsigned int dst = (unsigned int)__cvta_generic_to_shared(dst_ptr);\n"
+        "    unsigned int mbarrier = (unsigned int)__cvta_generic_to_shared(mbarrier_ptr);\n"
+        f'    asm volatile("{instr} [%0], [%1], %2, [%3]'
+        f'{mask_slot}{cache_slot};"\n'
+        "                 :\n"
+        '                 : "r"(dst), "l"(src_ptr), "r"(num_bytes), "r"(mbarrier)'
+        f"{cta_constraint}{_bulk_cache_operand_constraint(has_cache)}\n"
+        '                 : "memory");'
+    )
+    name = (
+        "tvm_builtin_ptx_cp_async_bulk_g2s_cluster"
+        f"{'_multicast' if multicast else ''}{'_cache_hint' if has_cache else ''}"
+    )
+    return name, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_g2s_cluster",
+    n_attrs=2,
+    helper_name=lambda *a: _bulk_g2s_cluster_parts(*a)[0],
+    c_signature=(
+        "(void* dst_ptr, void* src_ptr, unsigned int num_bytes, "
+        "void* mbarrier_ptr, unsigned short cta_mask, unsigned long long cache_policy)"
+    ),
+    body=lambda *a: _bulk_g2s_cluster_parts(*a)[1],
+)
+
+
+# PTX cp.async.bulk shared::cta -> shared::cluster form:
+#   cp.async.bulk.dst.src.completion_mechanism [dstMem], [srcMem], size, [mbar]
+#   .dst = {.shared::cluster}; .src = {.shared::cta}
+#   .completion_mechanism = {.mbarrier::complete_tx::bytes}
+device_intrinsic(
+    "ptx_cp_async_bulk_s2s_cluster",
+    helper_name="tvm_builtin_ptx_cp_async_bulk_s2s_cluster",
+    c_signature="(uint64_t dst, void* src, int size, uint64_t mbar)",
+    body=r"""    unsigned int dst_addr = static_cast<unsigned int>(dst);
+    unsigned int src_addr = __cvta_generic_to_shared(src);
+    unsigned int mbar_addr = static_cast<unsigned int>(mbar);
+    asm volatile(
+        "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"
+        " [%0], [%1], %2, [%3];"
+        :
+        : "r"(dst_addr), "r"(src_addr), "r"(size), "r"(mbar_addr)
+        : "memory");""",
+)
+
+
+@register_codegen("ptx_cp_async_bulk_shared_to_cluster")
+def codegen_ptx_cp_async_bulk_shared_to_cluster(dst_ptr, src_ptr, size, mbar):
+    result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_s2s_cluster"]([dst_ptr, src_ptr, size, mbar])
+    return result[0] if isinstance(result, tuple) else result
+
+
+# PTX cp.async.bulk shared::cta -> global form:
+#   cp.async.bulk.dst.src.completion_mechanism{.level::cache_hint}{.cp_mask}
+#       [dstMem], [srcMem], size {, cache-policy} {, byteMask}
+#   .dst = {.global}; .src = {.shared::cta}
+#   .completion_mechanism = {.bulk_group}
+#   .level::cache_hint = {.L2::cache_hint}
+def _bulk_s2g_parts(*args):
+    has_cache = _bool_attr(args[-2])
+    cp_mask = _bool_attr(args[-1])
+    if cp_mask and not has_cache:
+        raise ValueError("cp.async.bulk shared::cta -> global .cp_mask requires .L2::cache_hint")
+    instr = f"cp.async.bulk.global.shared::cta.bulk_group{_bulk_cache_operand_suffix(has_cache)}"
+    if cp_mask:
+        instr += ".cp_mask"
+    cache_slot = ", %3" if has_cache else ""
+    mask_slot = ", %4" if cp_mask else ""
+    mask_constraint = ', "r"(byte_mask)' if cp_mask else ""
+    body = (
+        "    unsigned int src = (unsigned int)__cvta_generic_to_shared(src_ptr);\n"
+        f'    asm volatile("{instr} [%0], [%1], %2'
+        f'{cache_slot}{mask_slot};"\n'
+        "                 :\n"
+        '                 : "l"(dst_ptr), "r"(src), "r"(num_bytes)'
+        f"{_bulk_cache_operand_constraint(has_cache)}{mask_constraint}\n"
+        '                 : "memory");'
+    )
+    name = (
+        "tvm_builtin_ptx_cp_async_bulk_s2g"
+        f"{'_cache_hint' if has_cache else ''}{'_cp_mask' if cp_mask else ''}"
+    )
+    return name, body
+
+
+device_intrinsic(
+    "ptx_cp_async_bulk_s2g",
+    n_attrs=2,
+    helper_name=lambda *a: _bulk_s2g_parts(*a)[0],
+    c_signature=(
+        "(void* dst_ptr, void* src_ptr, unsigned int num_bytes, "
+        "unsigned int byte_mask, unsigned long long cache_policy)"
+    ),
+    body=lambda *a: _bulk_s2g_parts(*a)[1],
+)

python/tvm/backend/cuda/operator/intrinsics/header.py

@@ -0,0 +1,809 @@
+# 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=line-too-long
+"""CUDA header generator for codegen.
+
+The header generator is used to generate the header for the CUDA code.
+It's controlled by the predefined tags.
+The tags are used to identify the utility functions/classes necessary for the codegen.
+"""
+
+import tvm_ffi
+
+TAGS = {
+    "cuda",
+    "cuda/barrier",
+    "cooperative_groups",
+    "fp16",
+    "bf16",
+    "fp8",
+    "fp6",
+    "fp4",
+    "int8",
+    "math_constants",
+    "mma",
+    "warp_shuffle",
+    "cast_smem_ptr_to_int",
+    "get_tmem_addr",
+    "gmma_descriptor",
+    "smem_descriptor",
+    "instr_descriptor",
+    "instr_descriptor_block_scaled",
+    "get_time_stamp",
+    "nvshmem",
+    "elect_one_sync",
+}
+
+
+@tvm_ffi.register_global_func("tirx.intrinsics.cuda.header_generator")
+def header_generator(tags):
+    """Generate the header for the CUDA code."""
+    for tag in tags:
+        if tag not in TAGS:
+            raise ValueError(f"Invalid tag: {tag}")
+
+    header = ""
+    if "nvshmem" in tags:
+        header += R"""
+#include <nvshmem.h>
+#include <nvshmemx.h>
+"""
+
+    if "cuda/barrier" in tags or "cooperative_groups" in tags:
+        header += (
+            R"""
+#include <cuda/barrier>
+#include <cooperative_groups.h>
+"""
+            + "\n"
+        )
+
+    # NVRTC has no host C++ stdlib and no <cuda.h>. Branch on __CUDACC_RTC__ so
+    # the same emitted source compiles under both nvcc (offline) and NVRTC
+    # (runtime) without any post-processing in tvm.support.nvcc.
+    header += """
+#ifdef __CUDACC_RTC__
+  #include <cuda/std/cstdint>
+  using cuda::std::uint8_t;
+  using cuda::std::uint16_t;
+  using cuda::std::uint32_t;
+  using cuda::std::uint64_t;
+  using cuda::std::int8_t;
+  using cuda::std::int16_t;
+  using cuda::std::int32_t;
+  using cuda::std::int64_t;
+
+  #include <cuda/std/type_traits>
+  namespace std {
+    using cuda::std::is_same;
+    using cuda::std::is_same_v;
+    using cuda::std::is_integral;
+    using cuda::std::is_signed;
+    using cuda::std::is_unsigned;
+    using cuda::std::is_floating_point;
+    using cuda::std::enable_if;
+    using cuda::std::conditional;
+  }
+
+  // NVRTC uses asm/volatile instead of __asm__/__volatile__ (gcc extension).
+  #ifndef __asm__
+  #define __asm__ asm
+  #endif
+  #ifndef __volatile__
+  #define __volatile__ volatile
+  #endif
+#else
+  #include <cstdint>
+  #include <type_traits>
+  #include <cuda.h>
+#endif
+"""
+
+    if "fp16" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 530)
+#include <cuda_fp16.h>
+__device__ half max(half a, half b)
+{
+  return __hgt(__half(a), __half(b)) ? a : b;
+}
+__device__ half min(half a, half b)
+{
+  return __hlt(__half(a), __half(b)) ? a : b;
+}
+#endif // __CUDA_ARCH__ >= 530
+
+// Pack two half values.
+static inline __device__ __host__ unsigned
+__pack_half2(const half x, const half y) {
+  unsigned v0 = *((unsigned short *)&x);
+  unsigned v1 = *((unsigned short *)&y);
+  return (v1 << 16) | v0;
+}
+
+#define CUDA_UNSUPPORTED_HALF_MATH_BINARY(HALF_MATH_NAME, FP32_MATH_NAME) \
+static inline __device__ __host__ half HALF_MATH_NAME(half x, half y) {   \
+  float tmp_x = __half2float(x);                                          \
+  float tmp_y = __half2float(y);                                          \
+  float result = FP32_MATH_NAME(tmp_x, tmp_y);                            \
+  return __float2half(result);                                            \
+}
+
+#define CUDA_UNSUPPORTED_HALF_MATH_UNARY(HALF_MATH_NAME, FP32_MATH_NAME) \
+static inline __device__ __host__ half HALF_MATH_NAME(half x) {          \
+  float tmp_x = __half2float(x);                                         \
+  float result = FP32_MATH_NAME(tmp_x);                                  \
+  return __float2half(result);                                           \
+}
+
+// Some fp16 math functions are not supported in cuda_fp16.h,
+// so we define them here to make sure the generated CUDA code
+// is valid.
+#if defined(__CUDA_ARCH__)
+#if (__CUDA_ARCH__ >= 530)
+CUDA_UNSUPPORTED_HALF_MATH_BINARY(hpow, powf)
+#if ((__CUDACC_VER_MAJOR__ < 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ < 8)))
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(htanh, tanhf)
+#endif
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(htan, tanf)
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(hatan, atanf)
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(herf, erf)
+#else
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(hexp, exp)
+#endif
+#endif
+
+#undef CUDA_UNSUPPORTED_HALF_MATH_BINARY
+#undef CUDA_UNSUPPORTED_HALF_MATH_UNARY
+"""
+
+    if "bf16" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+#include <cuda_bf16.h>
+__device__ nv_bfloat16 max(nv_bfloat16 a, nv_bfloat16 b)
+{
+  return __hgt(a, b) ? a : b;
+}
+__device__ nv_bfloat16 min(nv_bfloat16 a, nv_bfloat16 b)
+{
+  return __hlt(a, b) ? a : b;
+}
+#endif // __CUDA_ARCH__ >= 800
+// Pack two bfloat16 values.
+static inline __device__ __host__ unsigned
+__pack_nv_bfloat162(const nv_bfloat16 x, const nv_bfloat16 y) {
+  unsigned v0 = *((unsigned short *)&x);
+  unsigned v1 = *((unsigned short *)&y);
+  return (v1 << 16) | v0;
+}
+
+// Some bfp16 math functions are not supported in cuda_bfp16.h,
+// so we define them here to make sure the generated CUDA code
+// is valid.
+#define CUDA_UNSUPPORTED_HALF_MATH_BINARY(HALF_MATH_NAME, FP32_MATH_NAME) \
+static inline __device__ __host__ nv_bfloat16 HALF_MATH_NAME(nv_bfloat16 x, nv_bfloat16 y) {   \
+  float tmp_x = __bfloat162float(x);                                      \
+  float tmp_y = __bfloat162float(y);                                      \
+  float result = FP32_MATH_NAME(tmp_x, tmp_y);                            \
+  return __float2bfloat16(result);                                        \
+}
+
+#define CUDA_UNSUPPORTED_HALF_MATH_UNARY(HALF_MATH_NAME, FP32_MATH_NAME) \
+static inline __device__ __host__ nv_bfloat16 HALF_MATH_NAME(nv_bfloat16 x) {          \
+  float tmp_x = __bfloat162float(x);                                     \
+  float result = FP32_MATH_NAME(tmp_x);                                  \
+  return __float2bfloat16(result);                                       \
+}
+
+CUDA_UNSUPPORTED_HALF_MATH_BINARY(hpow, powf)
+#if ((__CUDACC_VER_MAJOR__ < 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ < 8)))
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(htanh, tanhf)
+#endif
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(htan, tanf)
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(hatan, atanf)
+CUDA_UNSUPPORTED_HALF_MATH_UNARY(herf, erf)
+
+#undef CUDA_UNSUPPORTED_HALF_MATH_BINARY
+#undef CUDA_UNSUPPORTED_HALF_MATH_UNARY
+"""
+
+    if "fp8" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
+#include <cuda_fp8.h>
+using fp8_e4_t = __nv_fp8_e4m3;
+using fp8_e4x2_t = __nv_fp8x2_e4m3;
+using fp8_e4x4_t = __nv_fp8x4_e4m3;
+struct fp8_e4x8_t {
+ fp8_e4_t data[8];
+};
+struct fp8_e4x16_t {
+ fp8_e4_t data[16];
+};
+using fp8_e5_t = __nv_fp8_e5m2;
+using fp8_e5x2_t = __nv_fp8x2_e5m2;
+using fp8_e5x4_t = __nv_fp8x4_e5m2;
+struct fp8_e5x8_t {
+ fp8_e5_t data[8];
+};
+struct fp8_e5x16_t {
+ fp8_e5_t data[16];
+};
+using fp8_e8_t = __nv_fp8_e8m0;
+using fp8_e8x2_t = __nv_fp8x2_e8m0;
+using fp8_e8x4_t = __nv_fp8x4_e8m0;
+struct fp8_e8x8_t {
+ fp8_e8_t data[8];
+};
+struct fp8_e8x16_t {
+ fp8_e8_t data[16];
+};
+#endif // __CUDA_ARCH__ >= 890
+"""
+
+    if "fp6" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+#include <cuda_fp6.h>
+using fp6_e2_t = __nv_fp6_e2m3;
+using fp6_e2x2_t = __nv_fp6x2_e2m3;
+using fp6_e2x4_t = __nv_fp6x4_e2m3;
+struct fp6_e2x8_t {
+ fp6_e2_t data[8];
+};
+struct fp6_e2x16_t {
+ fp6_e2_t data[16];
+};
+using fp6_e3_t = __nv_fp6_e3m2;
+using fp6_e3x2_t = __nv_fp6x2_e3m2;
+using fp6_e3x4_t = __nv_fp6x4_e3m2;
+struct fp6_e3x8_t {
+ fp6_e3_t data[8];
+};
+struct fp6_e3x16_t {
+ fp6_e3_t data[16];
+};
+#endif // __CUDA_ARCH__ >= 1000
+"""
+
+    if "fp4" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+#include <cuda_fp4.h>
+using fp4_e2_t = __nv_fp4_e2m1;
+using fp4_e2x2_t = __nv_fp4x2_e2m1;
+using fp4_e2x4_t = __nv_fp4x4_e2m1;
+struct fp4_e2x8_t {
+ fp4_e2_t data[8];
+};
+struct fp4_e2x16_t {
+ fp4_e2_t data[16];
+};
+#endif // __CUDA_ARCH__ >= 800
+"""
+
+    #########################################################
+    # Vector type extensions
+    #########################################################
+    if "fp16" in tags or "bf16" in tags:
+        header += R"""
+template <typename T, typename TVec2>
+struct __align__(8) half4_bfloat164 {
+  T x, y, z, w;
+  __host__ __device__ half4_bfloat164() : x(T(0)), y(T(0)), z(T(0)), w(T(0)) {}
+  __host__ __device__ half4_bfloat164(T x, T y, T z, T w) : x(x), y(y), z(z), w(w) {}
+"""
+        if "fp8" in tags:
+            header += R"""
+  __host__ __device__ explicit half4_bfloat164(const __nv_fp8x4_e4m3& fp8x4) {
+    if constexpr (std::is_same_v<T, __half>) {
+      __nv_fp8x2_e4m3 lo_part, hi_part;
+      lo_part.__x = static_cast<__nv_fp8x2_storage_t>(fp8x4.__x & 0xFFFF);
+      hi_part.__x = static_cast<__nv_fp8x2_storage_t>((fp8x4.__x >> 16) & 0xFFFF);
+      TVec2 lo_half2 = static_cast<TVec2>(lo_part);
+      TVec2 hi_half2 = static_cast<TVec2>(hi_part);
+      x = reinterpret_cast<T*>(&lo_half2)[0];
+      y = reinterpret_cast<T*>(&lo_half2)[1];
+      z = reinterpret_cast<T*>(&hi_half2)[0];
+      w = reinterpret_cast<T*>(&hi_half2)[1];
+    } else {
+      __nv_fp8_storage_t elem0_raw = static_cast<__nv_fp8_storage_t>(fp8x4.__x & 0xFF);
+      __nv_fp8_storage_t elem1_raw = static_cast<__nv_fp8_storage_t>((fp8x4.__x >> 8) & 0xFF);
+      __nv_fp8_storage_t elem2_raw = static_cast<__nv_fp8_storage_t>((fp8x4.__x >> 16) & 0xFF);
+      __nv_fp8_storage_t elem3_raw = static_cast<__nv_fp8_storage_t>((fp8x4.__x >> 24) & 0xFF);
+      __nv_fp8_e4m3 elem0, elem1, elem2, elem3;
+      elem0.__x = elem0_raw;
+      elem1.__x = elem1_raw;
+      elem2.__x = elem2_raw;
+      elem3.__x = elem3_raw;
+      x = T(elem0);
+      y = T(elem1);
+      z = T(elem2);
+      w = T(elem3);
+    }
+  }
+  __host__ __device__ explicit operator __nv_fp8x4_e4m3() const {
+    __nv_fp8x4_e4m3 result;
+    TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
+    TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
+    __nv_fp8x2_e4m3 lo_part(lo_half2), hi_part(hi_half2);
+    result.__x =
+        (static_cast<uint32_t>(lo_part.__x) | (static_cast<uint32_t>(hi_part.__x) << 16));
+    return result;
+  }
+  __host__ __device__ explicit half4_bfloat164(const __nv_fp8x4_e5m2& fp8x4) {
+      __nv_fp8x2_e5m2 lo_part, hi_part;
+      lo_part.__x = static_cast<__nv_fp8x2_storage_t>(fp8x4.__x & 0xFFFF);
+      hi_part.__x = static_cast<__nv_fp8x2_storage_t>((fp8x4.__x >> 16) & 0xFFFF);
+      TVec2 lo_half2 = static_cast<TVec2>(lo_part);
+      TVec2 hi_half2 = static_cast<TVec2>(hi_part);
+      x = reinterpret_cast<T*>(&lo_half2)[0];
+      y = reinterpret_cast<T*>(&lo_half2)[1];
+      z = reinterpret_cast<T*>(&hi_half2)[0];
+      w = reinterpret_cast<T*>(&hi_half2)[1];
+  }
+  __host__ __device__ explicit operator __nv_fp8x4_e5m2() const {
+    __nv_fp8x4_e5m2 result;
+    TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
+    TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
+    __nv_fp8x2_e5m2 lo_part(lo_half2), hi_part(hi_half2);
+    result.__x =
+        (static_cast<uint32_t>(lo_part.__x) | (static_cast<uint32_t>(hi_part.__x) << 16));
+    return result;
+  }
+  __host__ __device__ explicit half4_bfloat164(const __nv_fp8x4_e8m0& fp8x4) {
+      __nv_fp8x2_e8m0 lo_part, hi_part;
+      lo_part.__x = static_cast<__nv_fp8x2_storage_t>(fp8x4.__x & 0xFFFF);
+      hi_part.__x = static_cast<__nv_fp8x2_storage_t>((fp8x4.__x >> 16) & 0xFFFF);
+      TVec2 lo_half2 = static_cast<TVec2>(lo_part);
+      TVec2 hi_half2 = static_cast<TVec2>(hi_part);
+      x = reinterpret_cast<T*>(&lo_half2)[0];
+      y = reinterpret_cast<T*>(&lo_half2)[1];
+      z = reinterpret_cast<T*>(&hi_half2)[0];
+      w = reinterpret_cast<T*>(&hi_half2)[1];
+  }
+  __host__ __device__ explicit operator __nv_fp8x4_e8m0() const {
+    __nv_fp8x4_e8m0 result;
+    TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
+    TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
+    __nv_fp8x2_e8m0 lo_part(lo_half2), hi_part(hi_half2);
+    result.__x =
+        (static_cast<uint32_t>(lo_part.__x) | (static_cast<uint32_t>(hi_part.__x) << 16));
+    return result;
+  }
+"""
+        if "fp4" in tags:
+            header += R"""
+  __host__ __device__ explicit half4_bfloat164(const __nv_fp4x4_e2m1& fp4x4) {
+    if constexpr (std::is_same_v<T, __half>) {
+      __nv_fp4x2_storage_t lo_part = static_cast<__nv_fp4x2_storage_t>(fp4x4.__x & 0xFF);
+      __nv_fp4x2_storage_t hi_part = static_cast<__nv_fp4x2_storage_t>((fp4x4.__x >> 8) & 0xFF);
+      TVec2 lo_half2 = __half2(__nv_cvt_fp4x2_to_halfraw2(lo_part, __NV_E2M1));
+      TVec2 hi_half2 = __half2(__nv_cvt_fp4x2_to_halfraw2(hi_part, __NV_E2M1));
+      x = reinterpret_cast<T*>(&lo_half2)[0];
+      y = reinterpret_cast<T*>(&lo_half2)[1];
+      z = reinterpret_cast<T*>(&hi_half2)[0];
+      w = reinterpret_cast<T*>(&hi_half2)[1];
+    } else {
+      __nv_fp4_e2m1 elem0, elem1, elem2, elem3;
+      elem0.__x = static_cast<__nv_fp4_storage_t>(fp4x4.__x & 0xF);
+      elem1.__x = static_cast<__nv_fp4_storage_t>((fp4x4.__x >> 4) & 0xF);
+      elem2.__x = static_cast<__nv_fp4_storage_t>((fp4x4.__x >> 8) & 0xF);
+      elem3.__x = static_cast<__nv_fp4_storage_t>((fp4x4.__x >> 12) & 0xF);
+      x = T(elem0);
+      y = T(elem1);
+      z = T(elem2);
+      w = T(elem3);
+    }
+  }
+  __host__ __device__ explicit operator __nv_fp4x4_e2m1() const {
+    TVec2 lo_half2 = *reinterpret_cast<const TVec2*>(&x);
+    TVec2 hi_half2 = *reinterpret_cast<const TVec2*>(&z);
+    return __nv_fp4x4_e2m1(lo_half2, hi_half2);
+  }
+"""
+        header += R"""
+};
+"""
+    if "fp16" in tags:
+        header += R"""
+using half4 = half4_bfloat164<__half, __half2>;
+__host__ __device__ half4 make_half4(__half x, __half y, __half z, __half w) {
+    return half4(x, y, z, w);
+}
+"""
+    if "bf16" in tags:
+        header += R"""
+using nv_bfloat164 = half4_bfloat164<nv_bfloat16, nv_bfloat162>;
+__host__ __device__ nv_bfloat164 make_nv_bfloat164(nv_bfloat16 x, nv_bfloat16 y, nv_bfloat16 z, nv_bfloat16 w) {
+    return nv_bfloat164(x, y, z, w);
+}
+__host__ __device__ nv_bfloat162 make_nv_bfloat162(nv_bfloat16 x, nv_bfloat16 y) {
+    return nv_bfloat162(x, y);
+}
+"""  # noqa: E501
+        if "fp8" in tags:
+            header += R"""
+__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp8x2_e4m3& fp8x2) {
+    __nv_fp8_e4m3 elem0, elem1;
+    elem0.__x = static_cast<__nv_fp8_storage_t>(fp8x2.__x & 0xFF);
+    elem1.__x = static_cast<__nv_fp8_storage_t>((fp8x2.__x >> 8) & 0xFF);
+    nv_bfloat16 x = nv_bfloat16(elem0);
+    nv_bfloat16 y = nv_bfloat16(elem1);
+    return nv_bfloat162(x, y);
+}
+__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp8x2_e5m2& fp8x2) {
+    __nv_fp8_e5m2 elem0, elem1;
+    elem0.__x = static_cast<__nv_fp8_storage_t>(fp8x2.__x & 0xFF);
+    elem1.__x = static_cast<__nv_fp8_storage_t>((fp8x2.__x >> 8) & 0xFF);
+    nv_bfloat16 x = nv_bfloat16(elem0);
+    nv_bfloat16 y = nv_bfloat16(elem1);
+    return nv_bfloat162(x, y);
+}
+__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp8x2_e8m0& fp8x2) {
+    __nv_fp8_e8m0 elem0, elem1;
+    elem0.__x = static_cast<__nv_fp8_storage_t>(fp8x2.__x & 0xFF);
+    elem1.__x = static_cast<__nv_fp8_storage_t>((fp8x2.__x >> 8) & 0xFF);
+    nv_bfloat16 x = nv_bfloat16(elem0);
+    nv_bfloat16 y = nv_bfloat16(elem1);
+    return nv_bfloat162(x, y);
+}
+    """
+    if "fp8" in tags:
+        header += R"""
+__device__ __nv_fp8x2_e5m2 make___nv_fp8x2_e5m2(__nv_fp8_e5m2 x, __nv_fp8_e5m2 y) {
+    __nv_fp8x2_e5m2 result;
+    result.__x = (x.__x) | (y.__x << 8);
+    return result;
+}
+__device__ __nv_fp8x4_e5m2 make___nv_fp8x4_e5m2(__nv_fp8_e5m2 a, __nv_fp8_e5m2 b, __nv_fp8_e5m2 c, __nv_fp8_e5m2 d) {
+    __nv_fp8x4_e5m2 result;
+    result.__x = (a.__x) | (b.__x << 8) | (c.__x << 16) | (d.__x << 24);
+    return result;
+}
+__device__ __nv_fp8x2_e4m3 make___nv_fp8x2_e4m3(__nv_fp8_e4m3 x, __nv_fp8_e4m3 y) {
+    __nv_fp8x2_e4m3 result;
+    result.__x = (x.__x) | (y.__x << 8);
+    return result;
+}
+__device__ __nv_fp8x4_e4m3 make___nv_fp8x4_e4m3(__nv_fp8_e4m3 a, __nv_fp8_e4m3 b, __nv_fp8_e4m3 c, __nv_fp8_e4m3 d) {
+    __nv_fp8x4_e4m3 result;
+    result.__x = (a.__x) | (b.__x << 8) | (c.__x << 16) | (d.__x << 24);
+    return result;
+}
+__device__ __nv_fp8x2_e8m0 make___nv_fp8x2_e8m0(__nv_fp8_e8m0 x, __nv_fp8_e8m0 y) {
+    __nv_fp8x2_e8m0 result;
+    result.__x = (x.__x) | (y.__x << 8);
+    return result;
+}
+__device__ __nv_fp8x4_e8m0 make___nv_fp8x4_e8m0(__nv_fp8_e8m0 a, __nv_fp8_e8m0 b, __nv_fp8_e8m0 c, __nv_fp8_e8m0 d) {
+    __nv_fp8x4_e8m0 result;
+    result.__x = (a.__x) | (b.__x << 8) | (c.__x << 16) | (d.__x << 24);
+    return result;
+}
+"""  # noqa: E501
+    if "fp4" in tags:
+        header += R"""
+__host__ __device__ nv_bfloat162 cast_to_nv_bfloat162(const __nv_fp4x2_e2m1& fp4x2) {
+    __nv_fp4_e2m1 elem0, elem1;
+    elem0.__x = static_cast<__nv_fp4_storage_t>(fp4x2.__x & 0xFF);
+    elem1.__x = static_cast<__nv_fp4_storage_t>((fp4x2.__x >> 8) & 0xFF);
+    nv_bfloat16 x = nv_bfloat16(elem0);
+    nv_bfloat16 y = nv_bfloat16(elem1);
+    return nv_bfloat162(x, y);
+}
+"""
+
+    if "int8" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 610)
+#include <sm_61_intrinsics.h>
+
+#if defined(__CUDACC_RTC__)
+#define __SM_61_INTRINSICS_DECL__ __device__
+#else /* !__CUDACC_RTC__ */
+#define __SM_61_INTRINSICS_DECL__ static __device__ __inline__
+#endif /* __CUDACC_RTC__ */
+
+#ifndef __CUDA_ARCH__
+#define __DEF_IF_HOST { }
+#else  /* !__CUDA_ARCH__ */
+#define __DEF_IF_HOST ;
+#endif /* __CUDA_ARCH__ */
+
+__SM_61_INTRINSICS_DECL__ int __dp4a(unsigned int srcA, int srcB, int c) __DEF_IF_HOST
+__SM_61_INTRINSICS_DECL__ int __dp4a(int srcA, unsigned int srcB, int c) __DEF_IF_HOST
+
+#undef __DEF_IF_HOST
+
+#if !defined(__CUDACC_RTC__) && defined(__CUDA_ARCH__)
+__SM_61_INTRINSICS_DECL__ int __dp4a(unsigned int srcA, int srcB, int c) {
+    int ret;
+    asm volatile ("dp4a.u32.s32 %0, %1, %2, %3;" : "=r"(ret) : "r"(srcA), "r"(srcB), "r"(c));
+    return ret;
+}
+
+__SM_61_INTRINSICS_DECL__ int __dp4a(int srcA, unsigned int srcB, int c) {
+    int ret;
+    asm volatile ("dp4a.s32.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(srcA), "r"(srcB), "r"(c));
+    return ret;
+}
+#endif /* !__CUDACC_RTC__ && defined(__CUDA_ARCH__) */
+
+#undef __SM_61_INTRINSICS_DECL__
+
+#endif // __CUDA_ARCH__ >= 610
+"""
+    if "math_constants" in tags:
+        header += R"""
+#include <math_constants.h>
+"""
+    if "mma" in tags:
+        header += R"""
+#include <mma.h>
+"""
+
+    if "warp_shuffle" in tags:
+        header += R"""
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700)
+#define __shfl_sync(mask, var, lane, width) \
+        __shfl((var), (lane), (width))
+
+#define __shfl_down_sync(mask, var, offset, width) \
+        __shfl_down((var), (offset), (width))
+
+#define __shfl_up_sync(mask, var, offset, width) \
+        __shfl_up((var), (offset), (width))
+#endif
+"""
+
+    if "cast_smem_ptr_to_int" in tags:
+        header += R"""
+__forceinline__ __device__ unsigned int cast_smem_ptr_to_int(const void* const smem_ptr) {
+  unsigned int smem_int;
+  asm volatile ("{ .reg .u64 smem_int; cvta.to.shared.u64 smem_int, %1; cvt.u32.u64 %0, smem_int; }"
+    : "=r"(smem_int) : "l"(smem_ptr));
+  return smem_int;
+}
+"""
+    header += R"""
+#if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || \
+     (__CUDACC_VER_MAJOR__ > 11))
+#define TVM_ENABLE_L2_PREFETCH 1
+#else
+#define TVM_ENABLE_L2_PREFETCH 0
+#endif
+
+#ifdef _WIN32
+  using uint = unsigned int;
+  using uchar = unsigned char;
+  using ushort = unsigned short;
+  using int64_t = long long;
+  using uint64_t = unsigned long long;
+#else
+  #define uint unsigned int
+  #define uchar unsigned char
+  #define ushort unsigned short
+#endif
+"""
+
+    if "get_tmem_addr" in tags:
+        header += R"""
+__forceinline__ __device__ uint32_t get_tmem_addr(uint32_t idx, int row_offset, int col_offset) {
+  int col_idx = idx & 0xFFFF;
+  int row_idx = (idx >> 16) & 0xFFFF;
+  col_idx += col_offset;
+  row_idx += row_offset;
+  col_idx = col_idx & 0xFFFF;
+  row_idx = row_idx & 0xFFFF;
+
+  uint32_t new_idx = (row_idx << 16) | col_idx;
+  return new_idx;
+}
+"""
+
+    if "get_time_stamp" in tags:
+        header += R"""
+__forceinline__ __device__ uint32_t tvm_builtin_get_timestamp() {
+  volatile uint32_t ret;
+  asm volatile("mov.u32 %0, %globaltimer_lo;" : "=r"(ret));
+  return ret;
+}
+"""
+
+    if "gmma_descriptor" in tags:
+        header += R"""
+#ifndef HOST_DEVICE
+#define HOST_DEVICE __forceinline__ __host__ __device__
+#endif
+union GmmaDescriptor
+{
+  HOST_DEVICE constexpr
+  GmmaDescriptor() noexcept : desc_(0) {}
+  HOST_DEVICE constexpr
+  GmmaDescriptor(uint64_t desc) noexcept : desc_(desc) {}
+  HOST_DEVICE constexpr
+  GmmaDescriptor(GmmaDescriptor const& t) noexcept : desc_(t.desc_) {}
+  HOST_DEVICE constexpr
+  GmmaDescriptor(GmmaDescriptor && t) noexcept : desc_(t.desc_) {}
+
+  HOST_DEVICE constexpr
+  GmmaDescriptor& operator=(GmmaDescriptor const& t) noexcept {
+    desc_ = t.desc_;
+    return *this;
+  }
+
+  HOST_DEVICE constexpr
+  GmmaDescriptor& operator=(GmmaDescriptor && t) noexcept {
+    desc_ = t.desc_;
+    return *this;
+  }
+
+  uint64_t desc_;
+  uint32_t reg32_[2];
+  uint16_t reg16_[4];
+
+  // Bitfield implementation avoids the need for shifts in assignment
+  struct {
+    // start_address, bit [0,14), 4LSB not included
+    uint16_t start_address_ : 14, : 2;        // 14 bits [0,14), 2 bits unused
+    // leading dimension byte offset, bit [16,30), 4LSB not included
+    // For N: This is the stride from the first col to the second col of the 8x2 brick in INTERLEAVED
+    //   Unused for all SWIZZLE_* layouts (and assumed to be 1)
+    // For T: This is the stride from the first 8 rows to the next 8 rows.
+    uint16_t leading_byte_offset_ : 14, : 2;  // 14 bits [0,14), 2 bits unused
+    // stride dimension byte offset, bit [32,46), 4LSB not included
+    // For N: This is the stride from the first 8 rows to the next 8 rows.
+    // For T: This is the stride fro mthe first 8 cols to the next 8 cols.
+    uint16_t stride_byte_offset_ : 14, : 2;   // 14 bits [0,14), 2 bits unused
+    // base_offset, bit [49,52)
+    // Valid only for SWIZZLE_128B and SWIZZLE_64B
+    uint8_t : 1, base_offset_ : 3, : 4;       // 1 bit unused, 3 bits [1,4), 4 bits unused
+    // layout type, bit [62,64)
+    // SWIZZLE_NONE = 0, SWIZZLE_32B = 3, SWIZZLE_64B = 2, SWIZZLE_128B = 1
+    uint8_t : 6, layout_type_ : 2;            // 6 bits unused, 2 bits [6,8)
+  } bitfield;
+
+  // Decay to a uint64_t
+  HOST_DEVICE constexpr
+  operator uint64_t() const noexcept { return desc_; }
+};
+"""  # noqa: E501
+
+    if "smem_descriptor" in tags:
+        header += R"""
+#ifndef HOST_DEVICE
+#define HOST_DEVICE __forceinline__ __host__ __device__
+#endif
+union SmemDescriptor
+{
+  uint64_t desc_ = 0;
+  // Bitfield implementation avoids the need for shifts in assignment
+  struct {
+    // start_address, bit [0,14), 4LSB not included
+    uint16_t start_address_ : 14, : 2;                     // 14 bits [0,14), 2 bits unused
+    // leading dimension byte offset, bit [16,30), 4LSB not included
+    uint16_t leading_byte_offset_ : 14, : 2;               // 14 bits [0,14), 2 bits unused
+    // stride dimension byte offset, bit [32,46), 4LSB not included
+    uint16_t stride_byte_offset_ : 14, version_ : 2;       // 14 bits [0,14), 2 bits [14,16)
+    // base_offset, bit [49,52). leading_byte_offset_mode, bit [52,53).
+    uint8_t : 1, base_offset_ : 3, lbo_mode_ : 1, : 3;     // 1 bit unused, 3 bits [1,4), 1 bit [4,5), 3 bits unused
+    // layout type, bit [61,64), SWIZZLE_NONE matrix descriptor = 0, SWIZZLE_128B matrix descriptor = 2, SWIZZLE_64B descriptor = 4, SWIZZLE_32B descriptor = 6, SWIZZLE_128B_BASE32B = 1, N/A = 3, N/A = 5, N/A = 7
+    uint8_t : 5, layout_type_ : 3;                         // 6 bits unused, 3 bits [5,8)
+  };
+  // Seperate the field, as we may only update one part of desc
+  struct {
+    uint32_t lo;
+    uint32_t hi;
+  };
+
+  // Decay to a uint64_t
+  HOST_DEVICE constexpr
+  operator uint64_t() const noexcept { return desc_; }
+};
+"""  # noqa: E501
+
+    if "instr_descriptor" in tags:
+        header += R"""
+#ifndef HOST_DEVICE
+#define HOST_DEVICE __forceinline__ __host__ __device__
+#endif
+union InstrDescriptor
+{
+  uint32_t desc_;
+
+  struct {
+    // Bitfield implementation avoids the need for shifts in assignment
+    uint16_t sparse_id2_    : 2,  // bit [ 0, 2) : Sparse meta data id2
+             sparse_flag_   : 1,  // bit [ 2, 3) : 0 = dense. 1 = sparse. 1 value valid only for F32F16/S8/MXF8F6F4
+             saturate_      : 1,  // bit [ 3, 4) : 0 = no saturate. 1 = saturate. 1 value valid only for S8
+             c_format_      : 2,  // bit [ 4, 6) : 0 = F16. 1 = F32, 2 = S32
+                            : 1,  //
+             a_format_      : 3,  // bit [ 7,10) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. Boolean MMA: 0 Boolean
+             b_format_      : 3,  // bit [10,13) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. Boolean MMA: 0 Boolean
+             a_negate_      : 1,  // bit [13,14) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
+             b_negate_      : 1,  // bit [14,15) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
+             a_major_       : 1;  // bit [15,16) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
+    uint16_t b_major_       : 1,  // bit [16,17) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
+             n_dim_         : 6,  // bit [17,23) : 3 LSBs not included. Valid values range from 1 (N=8) to 32 (N=256).  All values are not valid for all instruction formats
+                            : 1,  //
+             m_dim_         : 5,  // bit [24,29) : 4 LSBs not included. Valid values are: 4 (M=64), 8 (M=128), 16 (M=256)
+                            : 1,  //
+             max_shift_     : 2;  // bit [30,32) : Maximum shift for WS instruction. Encoded as follows: 0 = no shift, 1 = maximum shift of 8, 2 = maximum shift of 16, 3 = maximum shift of 32.
+  };
+
+  // Decay to a uint32_t
+  HOST_DEVICE constexpr explicit
+  operator uint32_t() const noexcept { return desc_; }
+};
+"""  # noqa: E501
+
+    if "instr_descriptor_block_scaled" in tags:
+        header += R"""
+#ifndef HOST_DEVICE
+#define HOST_DEVICE __forceinline__ __host__ __device__
+#endif
+union InstrDescriptorBlockScaled
+{
+  uint32_t desc_;
+
+  struct {
+    // Bitfield implementation avoids the need for shifts in assignment
+    uint16_t sparse_id2_    : 2,  // bit [ 0, 2) : Sparse meta data id2
+             sparse_flag_   : 1,  // bit [ 2, 3) : 0 = dense. 1 = sparse. 1 value valid only for F32F16/S8/MXF8F6F4
+                            : 1,  //
+             b_sf_id_       : 2,  // bit [ 4, 6) : Matrix B Scale Factor ID
+                            : 1,  //
+             a_format_      : 3,  // bit [ 7, 9) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
+             b_format_      : 3,  // bit [10,12) : MXF8F6F4Format:0 = E4M3, 1 = E5M2, 3 = E2M3, 4 = E3M2, 5 = E2M1. F32F16Format: 0 = F16, 1 = BF16, 2 = TF32. S8: 0 unsigned 8 bit, 1 signed 8 bit. BMMA: 0 Boolean
+             a_negate_      : 1,  // bit [13,14) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
+             b_negate_      : 1,  // bit [14,15) : 0 = no negate. 1 = negate. 1 value valid only for F32F16Format and MXF8F6F4Format
+             a_major_       : 1;  // bit [15,16) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
+    uint16_t b_major_       : 1,  // bit [16,17) : 0 = K-major. 1 = MN-major. Major value of 1 is only valid for E4M3, E5M2, INT8 (signed and unsigned), F16, BF16 and TF32 source formats
+             n_dim_         : 6,  // bit [17,23) : 3 LSBs not included. Valid values range from 1 (N=8) to 32 (N=256).  All values are not valid for all instruction formats
+             scale_format_  : 1,  // bit [23,24) : 0=E4M3, 1=E8M0
+             m_dim_         : 5,  // bit [24,29) : 4 LSBs not included. Valid values are: 4 (M=64), 8 (M=128), 16 (M=256)
+             a_sf_id_       : 2,  // bit [29,31) : Matrix A Scale Factor ID
+                            : 1;  //
+  };
+
+  // Decay to a uint32_t
+  HOST_DEVICE constexpr
+  operator uint32_t() const noexcept { return desc_; }
+};
+"""  # noqa: E501
+
+    if "elect_one_sync" in tags:
+        header += R"""
+__forceinline__ __device__ uint32_t tvm_builtin_elect_one_sync() {{
+  uint32_t pred = 0;
+  uint32_t laneid = 0;
+  asm volatile(
+    "{\n"
+    ".reg .b32 %%rx;\n"
+    ".reg .pred %%px;\n"
+    "     elect.sync %%rx|%%px, %2;\n"
+    "@%%px mov.s32 %1, 1;\n"
+    "     mov.s32 %0, %%rx;\n"
+    "}\n"
+    : "+r"(laneid), "+r"(pred)
+    : "r"(0xFFFFFFFF));
+  return pred;
+}}
+"""
+    return header

python/tvm/backend/cuda/operator/intrinsics/math.py

@@ -0,0 +1,501 @@
+# 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
+"""Math intrinsics.
+
+PTX side:
+* ``add{.rnd}{.ftz}.f32x2`` / ``sub`` / ``mul`` / ``fma`` — packed f32x2.
+* ``ex2.approx.ftz.f32`` / ``rcp.approx.ftz.f32`` — special functions.
+* ``max.f32`` / ``min.f32`` — 3-operand reduction form.
+
+CUDA side:
+* warp / CTA reductions (templated butterfly shuffle-XOR).
+"""
+
+from tvm.backend.cuda.op import cuda_func_call
+
+from ._schema import device_intrinsic
+from .registry import register_codegen
+from .utils import parse_str, validate_power_of_two_range
+
+# =============================================================================
+# Packed f32x2 arithmetic — `add{.rnd}{.ftz}.f32x2 d, a, b ;` and friends.
+# Inputs are packed into a `.b64` register (low half = elem 0, high half =
+# elem 1); the body packs/unpacks via ``make_float2`` + ``reinterpret_cast``.
+# =============================================================================
+
+# PTX add/sub/mul/fma over (f32 | f32x2 | f64), DPS form.
+#   add{.rnd}{.ftz}{.sat}.f32     [d], a, b
+#   add{.rnd}{.ftz}.f32x2          [d], a, b      (a,b are packed-as-u64)
+#   add{.rnd}.f64                  [d], a, b
+#   (sub / mul same shape; fma adds a `c` operand)
+# Inputs a/b/c are register operands (scalar fp32 / packed u64 / scalar fp64).
+# Result is written through `d` (a pointer).
+_PACKED_ROUNDING = ("rz", "rn", "rm", "rp")
+
+
+# Per-dtype operand types and asm constraints.
+#  - c_in: C type of input register operand (matches PTX register type)
+#  - out_cast: pointer cast applied at d_addr (callers may pass float*/double*/...)
+#  - in_cstr / out_cstr: GCC asm constraint letter
+_DTYPE_INFO = {
+    "f32": {"c_in": "float", "out_cast": "float*", "in_cstr": "f", "out_cstr": "f"},
+    "f32x2": {
+        "c_in": "unsigned long long",
+        "out_cast": "uint64_t*",
+        "in_cstr": "l",
+        "out_cstr": "l",
+    },
+    "f64": {"c_in": "double", "out_cast": "double*", "in_cstr": "d", "out_cstr": "d"},
+}
+
+
+def _ptx_arith_modifier_string(dtype, rounding, ftz, sat):
+    """Build the `.rnd.ftz.sat` modifier substring + name suffix."""
+    rnd = parse_str(rounding)
+    assert rnd in _PACKED_ROUNDING, f"invalid rounding {rnd!r}, expected one of {_PACKED_ROUNDING}"
+    ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz)
+    sat_b = bool(int(sat)) if hasattr(sat, "value") else bool(sat)
+    if dtype == "f64" and (ftz_b or sat_b):
+        raise ValueError("PTX <op>.f64 does not accept .ftz or .sat")
+    if dtype == "f32x2" and sat_b:
+        raise ValueError("PTX <op>.f32x2 does not accept .sat")
+    mod = f".{rnd}"
+    if ftz_b:
+        mod += ".ftz"
+    if sat_b:
+        mod += ".sat"
+    name_suffix = f"_{rnd}"
+    if ftz_b:
+        name_suffix += "_ftz"
+    if sat_b:
+        name_suffix += "_sat"
+    return mod, name_suffix
+
+
+def _ptx_binary_arith_parts(op, dtype):
+    """Return (name_fn, sig, body_fn) for ptx_{op}_{dtype} binary form."""
+    info = _DTYPE_INFO[dtype]
+    # Destination is ``void*`` so callers can pass any element-type pointer
+    # (float* / double* / uint64_t*); body reinterpret-casts to the right type.
+    sig = f"(void* d, {info['c_in']} a, {info['c_in']} b)"
+
+    def _name(d, a, b, rounding, ftz, sat):
+        _, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
+        return f"tvm_builtin_ptx_{op}_{dtype}{suf}"
+
+    out_c = info["out_cstr"]
+    in_c = info["in_cstr"]
+    out_cast = info["out_cast"]
+
+    def _body(d, a, b, rounding, ftz, sat):
+        mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
+        return (
+            f'    asm volatile("{op}{mod}.{dtype} %0, %1, %2;"\n'
+            f'        : "={out_c}"(*reinterpret_cast<{out_cast}>(d))\n'
+            f'        : "{in_c}"(a), "{in_c}"(b));'
+        )
+
+    return _name, sig, _body
+
+
+def _ptx_fma_parts(dtype):
+    """Return (name_fn, sig, body_fn) for ptx_fma_{dtype}."""
+    info = _DTYPE_INFO[dtype]
+    sig = f"(void* d, {info['c_in']} a, {info['c_in']} b, {info['c_in']} c)"
+
+    def _name(d, a, b, c, rounding, ftz, sat):
+        _, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
+        return f"tvm_builtin_ptx_fma_{dtype}{suf}"
+
+    out_c = info["out_cstr"]
+    in_c = info["in_cstr"]
+    out_cast = info["out_cast"]
+
+    def _body(d, a, b, c, rounding, ftz, sat):
+        mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat)
+        return (
+            f'    asm volatile("fma{mod}.{dtype} %0, %1, %2, %3;"\n'
+            f'        : "={out_c}"(*reinterpret_cast<{out_cast}>(d))\n'
+            f'        : "{in_c}"(a), "{in_c}"(b), "{in_c}"(c));'
+        )
+
+    return _name, sig, _body
+
+
+# Register 12 ops: {add, sub, mul, fma} x {f32, f32x2, f64}.
+for _dtype in ("f32", "f32x2", "f64"):
+    for _op in ("add", "sub", "mul"):
+        _name_fn, _sig, _body_fn = _ptx_binary_arith_parts(_op, _dtype)
+        device_intrinsic(
+            f"ptx_{_op}_{_dtype}",
+            n_attrs=3,  # rounding, ftz, sat
+            helper_name=_name_fn,
+            c_signature=_sig,
+            body=_body_fn,
+        )
+    _name_fn, _sig, _body_fn = _ptx_fma_parts(_dtype)
+    device_intrinsic(
+        f"ptx_fma_{_dtype}",
+        n_attrs=3,
+        helper_name=_name_fn,
+        c_signature=_sig,
+        body=_body_fn,
+    )
+del _dtype, _op, _name_fn, _sig, _body_fn
+
+
+# =============================================================================
+# ex2.approx.ftz.f32 / rcp.approx.ftz.f32 — 1 form each.
+# =============================================================================
+device_intrinsic(
+    "ptx_exp2",
+    c_signature="(float x)",
+    return_type="float",
+    body=(
+        "    float result;\n"
+        '    asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));\n'
+        "    return result;"
+    ),
+)
+device_intrinsic(
+    "ptx_rcp",
+    c_signature="(float x)",
+    return_type="float",
+    body=(
+        "    float result;\n"
+        '    asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));\n'
+        "    return result;"
+    ),
+)
+
+
+# =============================================================================
+# 3-operand max.f32 / min.f32 — the f32, 3-operand form-table entry of the
+# redux/reduction-style fp32 max/min ops.
+# =============================================================================
+_ABC_SIG = "(float a, float b, float c)"
+device_intrinsic(
+    "ptx_reduce3_max_f32",
+    c_signature=_ABC_SIG,
+    return_type="float",
+    body=(
+        "    float result;\n"
+        '    asm volatile("max.f32 %0, %1, %2, %3;"\n'
+        '                 : "=f"(result) : "f"(a), "f"(b), "f"(c));\n'
+        "    return result;"
+    ),
+)
+device_intrinsic(
+    "ptx_reduce3_min_f32",
+    c_signature=_ABC_SIG,
+    return_type="float",
+    body=(
+        "    float result;\n"
+        '    asm volatile("min.f32 %0, %1, %2, %3;"\n'
+        '                 : "=f"(result) : "f"(a), "f"(b), "f"(c));\n'
+        "    return result;"
+    ),
+)
+
+
+_BINARY_F32_SIG = "(float a, float b)"
+
+
+def _ptx_max_f32_body(a, b, ftz, nan):
+    ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz)
+    nan_b = bool(int(nan)) if hasattr(nan, "value") else bool(nan)
+    ftz_suffix = ".ftz" if ftz_b else ""
+    nan_suffix = ".NaN" if nan_b else ""
+    return (
+        "    float result;\n"
+        f'    asm volatile("max{ftz_suffix}{nan_suffix}.f32 %0, %1, %2;"\n'
+        '                 : "=f"(result) : "f"(a), "f"(b));\n'
+        "    return result;"
+    )
+
+
+def _ptx_max_f32_name(a, b, ftz, nan):
+    ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz)
+    nan_b = bool(int(nan)) if hasattr(nan, "value") else bool(nan)
+    suffix = ""
+    if ftz_b:
+        suffix += "_ftz"
+    if nan_b:
+        suffix += "_nan"
+    return f"tvm_builtin_ptx_max_f32{suffix}"
+
+
+device_intrinsic(
+    "ptx_max_f32",
+    n_attrs=2,
+    helper_name=_ptx_max_f32_name,
+    c_signature=_BINARY_F32_SIG,
+    return_type="float",
+    body=_ptx_max_f32_body,
+)
+
+
+# =============================================================================
+# CUDA-side warp / CTA reductions (templated butterfly shuffle-XOR).
+# Emitted directly via ``cuda_func_call`` — the helper signature uses a
+# single template parameter ``T`` for both arg and return, which doesn't
+# match the operand-driven C signature pattern.
+# =============================================================================
+
+# (accumulation expression, identity value for cross-warp padding)
+_OP_TABLE = {
+    "sum": ("val += shuffled;", "T(0)"),
+    "max": ("val = max(val, shuffled);", "-INFINITY"),
+    "min": ("val = min(val, shuffled);", "INFINITY"),
+}
+
+
+def _validate_op(op_str, context):
+    if op_str not in _OP_TABLE:
+        raise ValueError(f"Unsupported {context} op '{op_str}', expected one of {list(_OP_TABLE)}")
+    return _OP_TABLE[op_str]
+
+
+def _warp_reduce_source(func_name, width_int, step_expr):
+    return (
+        f"\ntemplate <typename T>\n"
+        f"__forceinline__ __device__ T {func_name}(T val) {{\n"
+        f"    #pragma unroll\n"
+        f"    for (int mask = {width_int} >> 1; mask > 0; mask >>= 1) {{\n"
+        "        T shuffled = __shfl_xor_sync(0xFFFFFFFF, val, mask);\n"
+        f"        {step_expr}\n"
+        "    }\n"
+        "    return val;\n"
+        "}\n"
+    )
+
+
+@register_codegen("cuda_warp_reduce")
+def codegen_cuda_warp_reduce(value, op, width):
+    op_str = parse_str(op)
+    width_int = validate_power_of_two_range(width, 2, 32, "warp_reduce width")
+    step_expr, _ = _validate_op(op_str, "warp_reduce")
+
+    func_name = f"tvm_builtin_cuda_warp_reduce_{op_str}_{width_int}"
+    source_code = _warp_reduce_source(func_name, width_int, step_expr)
+    return cuda_func_call(func_name, value, source_code=source_code, return_type=value.ty)
+
+
+@register_codegen("cuda_cta_reduce")
+def codegen_cuda_cta_reduce(value, op, num_warps, scratch):
+    op_str = parse_str(op)
+    nw = validate_power_of_two_range(num_warps, 1, 32, "cta_reduce num_warps")
+    step_expr, identity = _validate_op(op_str, "cta_reduce")
+
+    warp_reduce_name = f"tvm_builtin_cuda_warp_reduce_{op_str}_32"
+    func_name = f"tvm_builtin_cuda_cta_reduce_{op_str}_{nw}"
+
+    cta_body = (
+        f"{_warp_reduce_source(warp_reduce_name, 32, step_expr)}"
+        "template <typename T>\n"
+        f"__forceinline__ __device__ T {func_name}(T val, void* scratch_raw) {{\n"
+        "    T* scratch = reinterpret_cast<T*>(scratch_raw);\n"
+        f"    val = {warp_reduce_name}(val);\n"
+        "    int tid = threadIdx.x + threadIdx.y * blockDim.x"
+        " + threadIdx.z * blockDim.x * blockDim.y;\n"
+        "    int warp_id = tid / 32;\n"
+        "    int lane_id = tid % 32;\n"
+        "    if (lane_id == 0) scratch[warp_id] = val;\n"
+        "    __syncthreads();\n"
+        "    if (warp_id == 0) {\n"
+        f"        T partial = (lane_id < {nw}) ? scratch[lane_id] : {identity};\n"
+        f"        partial = {warp_reduce_name}(partial);\n"
+        "        if (lane_id == 0) scratch[0] = partial;\n"
+        "    }\n"
+        "    __syncthreads();\n"
+        "    return scratch[0];\n"
+        "}\n"
+    )
+    return cuda_func_call(func_name, value, scratch, source_code=cta_body, return_type=value.ty)
+
+
+# =============================================================================
+# Additional FP8/BF16 packing, integer, and activation helpers.
+# =============================================================================
+
+# PTX integer bit-search form:
+#   fns.b32 d, mask, base, offset;
+device_intrinsic(
+    "ptx_fns_b32",
+    helper_name="tvm_builtin_ptx_fns_b32",
+    c_signature="(unsigned int mask, unsigned int base, int offset)",
+    return_type="unsigned int",
+    body=(
+        "    unsigned int ret;\n"
+        '    asm("fns.b32 %0, %1, %2, %3;" : "=r"(ret) : "r"(mask), "r"(base), "r"(offset));\n'
+        "    return ret;"
+    ),
+)
+
+device_intrinsic(
+    "cuda_ffs_u32",
+    helper_name="tvm_builtin_ffs_u32",
+    c_signature="(unsigned int value)",
+    return_type="int",
+    body="    return __ffs(value);",
+)
+
+device_intrinsic(
+    "ptx_add_rn_f32_bf16",
+    helper_name="tvm_builtin_ptx_add_rn_f32_bf16",
+    c_signature="(float acc, unsigned short x)",
+    return_type="float",
+    body=('    asm("add.rn.f32.bf16 %0, %1, %0;" : "+f"(acc) : "h"(x));\n    return acc;'),
+)
+
+
+device_intrinsic(
+    "cuda_make_float2",
+    helper_name="tvm_builtin_make_float2",
+    c_signature="(float x, float y)",
+    return_type="unsigned long long",
+    body=(
+        "    float2 value = make_float2(x, y);\n"
+        "    return *reinterpret_cast<unsigned long long*>(&value);"
+    ),
+)
+
+device_intrinsic(
+    "cuda_float2_x",
+    helper_name="tvm_builtin_float2_x",
+    c_signature="(unsigned long long packed)",
+    return_type="float",
+    body=("    float2 value = *reinterpret_cast<float2*>(&packed);\n    return value.x;"),
+)
+
+device_intrinsic(
+    "cuda_float2_y",
+    helper_name="tvm_builtin_float2_y",
+    c_signature="(unsigned long long packed)",
+    return_type="float",
+    body=("    float2 value = *reinterpret_cast<float2*>(&packed);\n    return value.y;"),
+)
+
+device_intrinsic(
+    "cuda_fmul2_rn",
+    helper_name="tvm_builtin_fmul2_rn",
+    c_signature="(unsigned long long a, unsigned long long b)",
+    return_type="unsigned long long",
+    body=(
+        "    float2 lhs = *reinterpret_cast<float2*>(&a);\n"
+        "    float2 rhs = *reinterpret_cast<float2*>(&b);\n"
+        "    float2 result = __fmul2_rn(lhs, rhs);\n"
+        "    return *reinterpret_cast<unsigned long long*>(&result);"
+    ),
+)
+
+device_intrinsic(
+    "cuda_fadd2_rn",
+    helper_name="tvm_builtin_fadd2_rn",
+    c_signature="(unsigned long long a, unsigned long long b)",
+    return_type="unsigned long long",
+    body=(
+        "    float2 lhs = *reinterpret_cast<float2*>(&a);\n"
+        "    float2 rhs = *reinterpret_cast<float2*>(&b);\n"
+        "    float2 result = __fadd2_rn(lhs, rhs);\n"
+        "    return *reinterpret_cast<unsigned long long*>(&result);"
+    ),
+)
+
+device_intrinsic(
+    "cuda_float22bfloat162_rn",
+    helper_name="tvm_builtin_float22bfloat162_rn",
+    c_signature="(float x, float y)",
+    return_type="unsigned int",
+    body=(
+        "    __nv_bfloat162 value = __float22bfloat162_rn(make_float2(x, y));\n"
+        "    return *reinterpret_cast<unsigned int*>(&value);"
+    ),
+    extra_deps=("bf16",),
+)
+
+device_intrinsic(
+    "cuda_float22bfloat162_rn_from_float2",
+    helper_name="tvm_builtin_float22bfloat162_rn_from_float2",
+    c_signature="(unsigned long long packed)",
+    return_type="unsigned int",
+    body=(
+        "    float2 value = *reinterpret_cast<float2*>(&packed);\n"
+        "    __nv_bfloat162 result = __float22bfloat162_rn(value);\n"
+        "    return *reinterpret_cast<unsigned int*>(&result);"
+    ),
+    extra_deps=("bf16",),
+)
+
+device_intrinsic(
+    "cuda_bfloat1622float2",
+    helper_name="tvm_builtin_bfloat1622float2",
+    c_signature="(unsigned int packed)",
+    return_type="unsigned long long",
+    body=(
+        "    __nv_bfloat162 value;\n"
+        "    *reinterpret_cast<unsigned int*>(&value) = packed;\n"
+        "    float2 result = __bfloat1622float2(value);\n"
+        "    return *reinterpret_cast<unsigned long long*>(&result);"
+    ),
+    extra_deps=("bf16",),
+)
+
+device_intrinsic(
+    "cuda_hmin2",
+    helper_name="tvm_builtin_hmin2",
+    c_signature="(unsigned int a, unsigned int b)",
+    return_type="unsigned int",
+    body=(
+        "    __nv_bfloat162 lhs;\n"
+        "    __nv_bfloat162 rhs;\n"
+        "    *reinterpret_cast<unsigned int*>(&lhs) = a;\n"
+        "    *reinterpret_cast<unsigned int*>(&rhs) = b;\n"
+        "    __nv_bfloat162 result = __hmin2(lhs, rhs);\n"
+        "    return *reinterpret_cast<unsigned int*>(&result);"
+    ),
+    extra_deps=("bf16",),
+)
+
+device_intrinsic(
+    "cuda_hmax2",
+    helper_name="tvm_builtin_hmax2",
+    c_signature="(unsigned int a, unsigned int b)",
+    return_type="unsigned int",
+    body=(
+        "    __nv_bfloat162 lhs;\n"
+        "    __nv_bfloat162 rhs;\n"
+        "    *reinterpret_cast<unsigned int*>(&lhs) = a;\n"
+        "    *reinterpret_cast<unsigned int*>(&rhs) = b;\n"
+        "    __nv_bfloat162 result = __hmax2(lhs, rhs);\n"
+        "    return *reinterpret_cast<unsigned int*>(&result);"
+    ),
+    extra_deps=("bf16",),
+)
+
+device_intrinsic(
+    "cuda_fp8x4_e4m3_from_float4",
+    helper_name="tvm_builtin_fp8x4_e4m3_from_float4",
+    c_signature="(float x, float y, float z, float w)",
+    return_type="unsigned int",
+    body=(
+        "    __nv_fp8x4_e4m3 result = __nv_fp8x4_e4m3(make_float4(x, y, z, w));\n"
+        "    return *reinterpret_cast<unsigned int*>(&result);"
+    ),
+    extra_deps=("fp8",),
+)

python/tvm/backend/cuda/operator/intrinsics/memory.py

@@ -0,0 +1,955 @@
+# 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
+# pylint: disable=redefined-builtin, invalid-name, too-many-arguments
+"""Memory ops (load / store / copy / atomic / address conversion / type punning).
+
+PTX side:
+* ``ld.acquire.scope{.ss}.type`` scalar load forms.
+* ``ld.volatile{.ss}.type`` scalar load forms.
+* Legacy ``ld.global.acquire.gpu`` / ``ld.global.cg`` result-argument helper.
+* ``mapa.u64`` — map a SMEM ptr to a peer CTA's SMEM in the cluster.
+
+CUDA side:
+* ``__ldg`` (cache-as-read-only load).
+* Templated ``atomicAdd`` / ``atomicCAS``.
+* half↔float type-punned conversions (single, packed, batch-of-8).
+* ``__cvta_generic_to_shared`` and ``cluster_addr → shared u32`` casts.
+"""
+
+from tvm import DataType
+from tvm.backend.cuda.op import cuda_func_call
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+from .utils import parse_str
+
+
+# =============================================================================
+# __ldg — templated read-only cached load; ``T`` resolved at call time from
+# the ``dtype`` argument. Hand-written because the helper signature uses a
+# template parameter for both arg and return.
+# =============================================================================
+@register_codegen("cuda_ldg")
+def codegen_cuda_ldg(addr, dtype):
+    dtype = DataType(parse_str(dtype))
+    func_name = "tvm_builtin_cuda_ldg"
+    source_code = f"""
+template <typename T>
+__forceinline__ __device__ T {func_name}(T* src) {{
+    return __ldg(src);
+}}
+"""
+    return cuda_func_call(func_name, addr, source_code=source_code, return_type=dtype)
+
+
+# Shared PTX scalar type metadata (ld/st/red/atom).
+_PTX_SCALAR_TYPE_INFO = {
+    "b8": ("unsigned int", "r", "uint32"),
+    "u8": ("unsigned int", "r", "uint32"),
+    "s8": ("int", "r", "int32"),
+    "b16": ("unsigned short", "h", "uint16"),
+    "u16": ("unsigned short", "h", "uint16"),
+    "s16": ("short", "h", "int16"),
+    "b32": ("unsigned int", "r", "uint32"),
+    "u32": ("unsigned int", "r", "uint32"),
+    "s32": ("int", "r", "int32"),
+    "b64": ("unsigned long long", "l", "uint64"),
+    "u64": ("unsigned long long", "l", "uint64"),
+    "s64": ("long long", "l", "int64"),
+    "f32": ("float", "f", "float32"),
+    "f64": ("double", "d", "float64"),
+}
+_PTX_LD_TYPE_RETURNS = {
+    "b32": {"uint32": "unsigned int", "int32": "int"},
+    "b64": {"uint64": "unsigned long long", "int64": "long long"},
+}
+_PTX_VEC_STORE_TYPE = {
+    16: "uint4",
+    8: "uint2",
+    4: "unsigned int",
+    2: "unsigned short",
+    1: "unsigned char",
+}
+
+
+def _safe_attr(value):
+    return parse_str(value).replace("::", "_").replace(".", "_")
+
+
+def _dot(value):
+    value = parse_str(value)
+    return f".{value}" if value else ""
+
+
+def _cache_suffix(cache):
+    return ".L2::cache_hint" if cache else ""
+
+
+def _type_info(ptx_type):
+    ptx_type = parse_str(ptx_type)
+    if ptx_type not in _PTX_SCALAR_TYPE_INFO:
+        raise ValueError(
+            f"Unsupported PTX scalar type {ptx_type!r}; expected {sorted(_PTX_SCALAR_TYPE_INFO)}"
+        )
+    return (ptx_type, *_PTX_SCALAR_TYPE_INFO[ptx_type])
+
+
+# =============================================================================
+# PTX ld forms (ISA table entries registered via ``ptx_ld`` and siblings):
+#   ld{.weak}{.ss}{.cop}{.level::cache_hint}{.level::prefetch_size}{.vec}.type  d, [a]{, cache-policy};
+#   ld{.weak}{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type  d, [a]{, cache-policy};
+#   ld.volatile{.ss}{.level::prefetch_size}{.vec}.type  d, [a];
+#   ld.relaxed.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type  d, [a]{, cache-policy};
+#   ld.acquire.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type  d, [a]{, cache-policy};
+#   ld.mmio.sem.sys{.global}.type  d, [a];
+# =============================================================================
+_PTX_LD_SCOPES = {"cta", "cluster", "gpu", "sys"}
+_PTX_LD_SPACES = {"global", "shared", "shared::cta", "shared::cluster", "local"}
+_PTX_LD_VOLATILE_SPACES = _PTX_LD_SPACES | {"const"}
+_PTX_LD_WEAK_SPACES = _PTX_LD_SPACES | {"const", "param::entry", "param::func"}
+_PTX_LD_COPS = {"", "ca", "cg", "cs", "lu", "cv"}
+_PTX_VEC = {"", "v2", "v4", "v8"}
+_PTX_L1_EVICT = {
+    "",
+    "L1::evict_normal",
+    "L1::evict_unchanged",
+    "L1::evict_first",
+    "L1::evict_last",
+    "L1::no_allocate",
+}
+_PTX_L2_EVICT = {"", "L2::evict_normal", "L2::evict_first", "L2::evict_last"}
+_PTX_PREFETCH = {"", "L2::64B", "L2::128B", "L2::256B"}
+
+
+def _bool_attr(value):
+    return bool(int(value)) if hasattr(value, "value") else bool(value)
+
+
+def _parse_ld_attrs(return_dtype, ptx_type, scope=None, space="global"):
+    return_dtype = parse_str(return_dtype)
+    ptx_type = parse_str(ptx_type)
+    scope = None if scope is None else parse_str(scope)
+    space = parse_str(space)
+    ptx_type, _ptx, constraint, default_tvm = _type_info(ptx_type)
+    if ptx_type in _PTX_LD_TYPE_RETURNS:
+        returns = _PTX_LD_TYPE_RETURNS[ptx_type]
+        if return_dtype not in returns:
+            raise ValueError(
+                f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; "
+                f"expected one of {sorted(returns)}"
+            )
+        c_type = returns[return_dtype]
+    else:
+        if return_dtype != default_tvm:
+            raise ValueError(
+                f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; "
+                f"expected {default_tvm!r}"
+            )
+        c_type = _ptx
+    if scope is not None and scope not in _PTX_LD_SCOPES:
+        raise ValueError(
+            f"Unsupported PTX ld scope {scope!r}; expected one of {sorted(_PTX_LD_SCOPES)}"
+        )
+    return return_dtype, ptx_type, scope, space, c_type, constraint
+
+
+def _validate_ld_space(space: str, allowed: set[str]) -> None:
+    if space not in allowed:
+        raise ValueError(
+            f"Unsupported PTX ld state space {space!r}; expected one of {sorted(allowed)}"
+        )
+
+
+def _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size):
+    suffix = _cache_suffix("cache" if has_cache else "")
+    l1_evict = parse_str(l1_evict)
+    l2_evict = parse_str(l2_evict)
+    prefetch_size = parse_str(prefetch_size)
+    if l1_evict:
+        suffix += f".{l1_evict}"
+    if l2_evict:
+        suffix += f".{l2_evict}"
+    if prefetch_size:
+        suffix += f".{prefetch_size}"
+    return suffix
+
+
+def _ptx_shared_addr(space, ptr_name="address"):
+    if parse_str(space).startswith("shared"):
+        return (
+            f"    unsigned int addr = (unsigned int)__cvta_generic_to_shared({ptr_name});\n",
+            '"r"(addr)',
+        )
+    return "", f'"l"({ptr_name})'
+
+
+def _ptx_ld_vec_store(num_bytes, vec_len, ptx_type):
+    if ptx_type == "u8" and vec_len == 1:
+        return "    *reinterpret_cast<unsigned char*>(dst_ptr) = static_cast<unsigned char>(r0);"
+    store_type = _PTX_VEC_STORE_TYPE[num_bytes]
+    if vec_len > 1:
+        return (
+            f"    *reinterpret_cast<{store_type}*>(dst_ptr) = "
+            + "{"
+            + ", ".join(f"r{i}" for i in range(vec_len))
+            + "};"
+        )
+    return f"    *reinterpret_cast<{store_type}*>(dst_ptr) = r0;"
+
+
+def _ptx_ld_form_parts(form, attr_args):
+    if form == "weak":
+        (
+            return_dtype,
+            weak,
+            space,
+            cop,
+            vec,
+            ptx_type,
+            has_cache_hint,
+            to_dst,
+            l1_evict,
+            l2_evict,
+            prefetch_size,
+        ) = attr_args
+        sem, scope = "", ""
+    elif form == "relaxed":
+        (
+            return_dtype,
+            scope,
+            space,
+            vec,
+            ptx_type,
+            has_cache_hint,
+            to_dst,
+            l1_evict,
+            l2_evict,
+            prefetch_size,
+        ) = attr_args
+        sem, weak, cop = "", False, ""
+    elif form == "acquire":
+        (
+            return_dtype,
+            scope,
+            space,
+            vec,
+            ptx_type,
+            has_cache_hint,
+            to_dst,
+            l1_evict,
+            l2_evict,
+            prefetch_size,
+        ) = attr_args
+        sem, weak, cop = "", False, ""
+    elif form == "volatile":
+        return_dtype, space, vec, ptx_type, to_dst, prefetch_size = attr_args
+        sem, scope, weak, cop = "", "", False, ""
+        has_cache_hint, l1_evict, l2_evict = False, "", ""
+    elif form == "mmio":
+        return_dtype, sem, scope, space, ptx_type, to_dst = attr_args
+        weak, cop, vec = False, "", ""
+        has_cache_hint, l1_evict, l2_evict, prefetch_size = False, "", "", ""
+    else:
+        raise ValueError(f"unknown ld form {form!r}")
+
+    return_dtype, ptx_type, scope, space, c_type, constraint = _parse_ld_attrs(
+        return_dtype, ptx_type, scope if form in ("relaxed", "acquire") else None, space
+    )
+    sem = parse_str(sem)
+    scope = parse_str(scope)
+    space = parse_str(space)
+    cop = parse_str(cop)
+    vec = parse_str(vec)
+    l1_evict = parse_str(l1_evict)
+    l2_evict = parse_str(l2_evict)
+    prefetch_size = parse_str(prefetch_size)
+    weak = _bool_attr(weak)
+    has_cache = _bool_attr(has_cache_hint)
+    to_dst = _bool_attr(to_dst)
+    if cop and cop not in _PTX_LD_COPS:
+        raise ValueError(f"Unsupported PTX ld cache operation {cop!r}")
+    if vec and vec not in _PTX_VEC:
+        raise ValueError(f"Unsupported PTX ld vector modifier {vec!r}")
+    if l1_evict and l1_evict not in _PTX_L1_EVICT:
+        raise ValueError(f"Unsupported PTX ld L1 eviction {l1_evict!r}")
+    if l2_evict and l2_evict not in _PTX_L2_EVICT:
+        raise ValueError(f"Unsupported PTX ld L2 eviction {l2_evict!r}")
+    if prefetch_size and prefetch_size not in _PTX_PREFETCH:
+        raise ValueError(f"Unsupported PTX ld prefetch size {prefetch_size!r}")
+    if form == "mmio":
+        if sem not in ("acquire", "relaxed") or scope != "sys" or space != "global":
+            raise ValueError("ld.mmio requires sem in {acquire, relaxed}, scope=sys, space=global")
+        prefix = f"ld.mmio.{sem}.{scope}"
+    elif form == "relaxed":
+        if not scope:
+            raise ValueError("ld.relaxed requires scope")
+        _validate_ld_space(space, _PTX_LD_SPACES)
+        prefix = f"ld.relaxed.{scope}{_dot(space)}"
+    elif form == "acquire":
+        if not scope:
+            raise ValueError("ld.acquire requires scope")
+        _validate_ld_space(space, _PTX_LD_SPACES)
+        prefix = f"ld.acquire.{scope}{_dot(space)}"
+    elif form == "volatile":
+        _validate_ld_space(space, _PTX_LD_VOLATILE_SPACES)
+        prefix = f"ld.volatile{_dot(space)}"
+    else:
+        _validate_ld_space(space, _PTX_LD_WEAK_SPACES)
+        prefix = f"ld{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}"
+    level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size)
+    vec_len = int(vec[1:]) if vec else 1
+    if vec and not to_dst:
+        raise ValueError("vector ld requires to_dst")
+    elem_bytes = (
+        8
+        if ptx_type.endswith("64")
+        else 2
+        if ptx_type in ("u16", "s16", "b16")
+        else 1
+        if ptx_type in ("u8", "s8", "b8")
+        else 4
+    )
+    num_bytes = vec_len * elem_bytes if vec else elem_bytes
+    name_parts = [
+        "tvm_builtin_ptx_ld",
+        form if form != "weak" else ("weak" if weak else "plain"),
+    ]
+    if sem:
+        name_parts.append(_safe_attr(sem))
+    if scope:
+        name_parts.append(_safe_attr(scope))
+    name_parts.extend(
+        [
+            _safe_attr(space),
+            _safe_attr(cop) if cop else "",
+            _safe_attr(vec) if vec else "",
+            ptx_type,
+            return_dtype if not to_dst else "to_dst",
+        ]
+    )
+    if has_cache:
+        name_parts.append("cache_hint")
+    if l1_evict:
+        name_parts.append(_safe_attr(l1_evict))
+    if l2_evict:
+        name_parts.append(_safe_attr(l2_evict))
+    if prefetch_size:
+        name_parts.append(_safe_attr(prefetch_size))
+    name = "_".join(p for p in name_parts if p)
+    cache_operand = ', "l"(cache_policy)' if has_cache else ""
+    addr_decl, addr_operand = _ptx_shared_addr(space, "src_ptr" if to_dst else "address")
+    if to_dst:
+        reg_decls = "".join(f"    {c_type} r{i};\n" for i in range(vec_len))
+        if vec_len > 1:
+            out_slot = "{" + ", ".join(f"%{i}" for i in range(vec_len)) + "}"
+            out_constraints = ", ".join(f'"={constraint}"(r{i})' for i in range(vec_len))
+            addr_idx = vec_len
+        else:
+            out_slot = "%0"
+            out_constraints = f'"={constraint}"(r0)'
+            addr_idx = 1
+        cache_slot = f", %{addr_idx + 1}" if has_cache else ""
+        instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
+        body = (
+            f"{addr_decl}{reg_decls}"
+            f'    asm volatile("{instr} {out_slot}, [%{addr_idx}]{cache_slot};"\n'
+            f"                 : {out_constraints}\n"
+            f"                 : {addr_operand}{cache_operand});\n"
+            f"{_ptx_ld_vec_store(num_bytes, vec_len, ptx_type)}"
+        )
+        return (
+            name,
+            "(void* dst_ptr, void* src_ptr, unsigned long long cache_policy)",
+            "void",
+            "",
+            body,
+        )
+    cache_slot = ", %2" if has_cache else ""
+    instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
+    body = (
+        f"    {c_type} ret;\n"
+        f"{addr_decl}"
+        f'    asm volatile("{instr} %0, [%1]{cache_slot};"\n'
+        f'                 : "={constraint}"(ret)\n'
+        f"                 : {addr_operand}{cache_operand});\n"
+        "    return ret;"
+    )
+    sig = (
+        "(void* address, unsigned long long cache_policy)" if form == "weak" else "(void* address)"
+    )
+    return name, sig, c_type, return_dtype, body
+
+
+def _register_ptx_ld(op_name, form, n_attrs):
+    def _parts(*args):
+        return _ptx_ld_form_parts(form, args[-n_attrs:])
+
+    device_intrinsic(
+        op_name,
+        n_attrs=n_attrs,
+        helper_name=lambda *a, _p=_parts: _p(*a)[0],
+        c_signature=lambda *a, _p=_parts: _p(*a)[1],
+        return_type=lambda *a, _p=_parts: _p(*a)[2],
+        tvm_return_type=lambda *a, _p=_parts: (
+            None if _p(*a)[2] == "void" else parse_str(a[-n_attrs])
+        ),
+        body=lambda *a, _p=_parts: _p(*a)[4],
+    )
+
+
+_register_ptx_ld("ptx_ld", "weak", 11)
+_register_ptx_ld("ptx_ld_relaxed", "relaxed", 10)
+_register_ptx_ld("ptx_ld_acquire", "acquire", 10)
+_register_ptx_ld("ptx_ld_volatile", "volatile", 6)
+_register_ptx_ld("ptx_ld_mmio", "mmio", 6)
+
+
+# =============================================================================
+# Legacy acquire-load lvalue API — compatibility wrapper over
+# ``ld.acquire.gpu.global`` / ``ld.global.cg`` forms, dispatched on dtype.
+# Wrapper picks .b32/.b64 + matching constraint by dtype.
+#
+# The body uses ``#if __CUDA_ARCH__ >= 700`` to select acquire on SM70+ and
+# fall back to .cg on older arches. This is two PTX form table entries
+# combined in one device helper for arch portability.
+# =============================================================================
+_LD_GLOBAL_ACQUIRE_DTYPES = {
+    "uint32": ("uint32_t", "b32", "r"),
+    "int32": ("int32_t", "b32", "r"),
+    "uint64": ("uint64_t", "b64", "l"),
+    "int64": ("int64_t", "b64", "l"),
+}
+
+
+def _ld_global_acquire_body(ptx_type: str, spec: str) -> str:
+    return (
+        "  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700\n"
+        f'  asm volatile ("ld.acquire.gpu.global.{ptx_type} %0, [%1];\\n"\n'
+        f'                : "={spec}"(res) : "l"(addr));\n'
+        "  #else\n"
+        f'  asm volatile ("ld.global.cg.{ptx_type} %0, [%1];\\n"\n'
+        f'                : "={spec}"(res) : "l"(addr));\n'
+        "  #endif"
+    )
+
+
+for _dtype, (_c_type, _ptx_type, _spec) in _LD_GLOBAL_ACQUIRE_DTYPES.items():
+    device_intrinsic(
+        f"ptx_ld_global_acquire_{_dtype}",
+        c_signature=f"({_c_type}& res, {_c_type}* addr)",
+        body=_ld_global_acquire_body(_ptx_type, _spec),
+    )
+del _dtype, _c_type, _ptx_type, _spec
+
+
+@register_codegen("ptx_ld_global_acquire")
+def codegen_ptx_ld_global_acquire(res, addr):
+    """Dispatch to the dtype-specific helper."""
+    dtype = str(res.ty)
+    if dtype not in _LD_GLOBAL_ACQUIRE_DTYPES:
+        raise ValueError(f"Unsupported data type for ld.global.acquire: {dtype}")
+    result = CODEGEN_REGISTRY[f"tirx.ptx_ld_global_acquire_{dtype}"]([res, addr])
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# Atomics — templated wrappers around CUDA's ``atomicAdd`` / ``atomicCAS``.
+# =============================================================================
+device_intrinsic(
+    "cuda_atomic_add",
+    helper_name="tvm_builtin_cuda_atomic_add",
+    c_signature="(T* addr, T value)",
+    body="    return atomicAdd(addr, value);",
+    return_type="T",
+    templated=True,
+    tvm_return_type=lambda _addr, value: value.ty,
+)
+device_intrinsic(
+    "cuda_atomic_cas",
+    helper_name="tvm_builtin_cuda_atomic_cas",
+    c_signature="(T* address, T compare, T val)",
+    body="    return atomicCAS(address, compare, val);",
+    return_type="T",
+    templated=True,
+    tvm_return_type=lambda _p, old, _n: old.ty,
+)
+
+
+# =============================================================================
+# half / bfloat16 ↔ float type-punned conversions.
+# =============================================================================
+device_intrinsic(
+    "cuda_half2float",
+    c_signature="(half src)",
+    body="    return __half2float(src);",
+    return_type="float",
+    tvm_return_type="float32",
+)
+device_intrinsic(
+    "cuda_bfloat162float",
+    c_signature="(nv_bfloat16 src)",
+    body="    return __bfloat162float(src);",
+    return_type="float",
+    tvm_return_type="float32",
+)
+device_intrinsic(
+    "cuda_float22half2",
+    c_signature="(void* dst, void* src)",
+    body=(
+        "    half2* dst_p = (half2*) dst;\n"
+        "    float2* src_p = (float2*) src;\n"
+        "    *dst_p = __float22half2_rn(*src_p);"
+    ),
+)
+device_intrinsic(
+    "cuda_half8tofloat8",
+    c_signature="(void* src_addr, void* dst_addr)",
+    body=(
+        "    half2* source = (half2*) src_addr;\n"
+        "    float2* dest = (float2*) dst_addr;\n"
+        "    for (int i = 0; i < 4; i++) {\n"
+        "        dest[i] = __half22float2(source[i]);\n"
+        "    }"
+    ),
+)
+device_intrinsic(
+    "cuda_float8tohalf8",
+    c_signature="(void* src_addr, void* dst_addr)",
+    body=(
+        "    float2* source = (float2*) src_addr;\n"
+        "    half2* dest = (half2*) dst_addr;\n"
+        "    for (int i = 0; i < 4; i++) {\n"
+        "        dest[i] = __float22half2_rn(source[i]);\n"
+        "    }"
+    ),
+)
+
+
+# =============================================================================
+# Address-conversion helpers used by op-wrapper-side dispatch in tvm.tirx.op.
+# Each precomputes a value that the schema's specialized op then takes as a
+# typed scalar input (instead of doing the conversion inside the asm helper).
+# =============================================================================
+device_intrinsic(
+    "cuda_cvta_generic_to_shared",
+    c_signature="(void* p)",
+    body="    return __cvta_generic_to_shared(p);",
+    return_type="unsigned int",
+    tvm_return_type="uint32",
+)
+
+device_intrinsic(
+    "cuda_smem_addr_from_uint64",
+    c_signature="(uint64_t cluster_addr)",
+    body="    return static_cast<unsigned int>(cluster_addr);",
+    return_type="unsigned int",
+    tvm_return_type="uint32",
+)
+
+# =============================================================================
+# PTX mapa form:
+#   mapa{.space}.type d, a, b;
+#   .space = {.shared::cluster}; .type = {.u32, .u64}
+# =============================================================================
+
+
+def _ptx_mapa_parts(_addr, _rank, space, ptx_type, return_dtype):
+    space = parse_str(space)
+    ptx_type = parse_str(ptx_type)
+    return_dtype = parse_str(return_dtype)
+    if space not in ("", "shared::cluster"):
+        raise ValueError(f"Unsupported mapa space {space!r}")
+    if ptx_type not in ("u32", "u64"):
+        raise ValueError(f"Unsupported mapa type {ptx_type!r}")
+    c_type = "uint32_t" if ptx_type == "u32" else "uint64_t"
+    constraint = "r" if ptx_type == "u32" else "l"
+    name = f"tvm_builtin_ptx_mapa{('_' + _safe_attr(space)) if space else ''}_{ptx_type}"
+    body = (
+        f"    {c_type} result;\n"
+        f'    asm volatile("mapa{_dot(space)}.{ptx_type} %0, %1, %2;"\n'
+        f'                 : "={constraint}"(result) : "l"(addr), "r"(rank));\n'
+        "    return result;"
+    )
+    return name, c_type, return_dtype, body
+
+
+device_intrinsic(
+    "ptx_mapa",
+    n_attrs=3,
+    helper_name=lambda *a: _ptx_mapa_parts(*a)[0],
+    c_signature="(void* addr, uint32_t rank)",
+    return_type=lambda *a: _ptx_mapa_parts(*a)[1],
+    tvm_return_type=lambda *a: _ptx_mapa_parts(*a)[2],
+    body=lambda *a: _ptx_mapa_parts(*a)[3],
+)
+
+
+# =============================================================================
+# Generic PTX memory forms. Compatibility wrappers in ``tvm.tirx.op`` bind
+# concrete sem/scope/space/op/type parameters for existing call sites.
+# =============================================================================
+
+
+# PTX red scalar form:
+#   red{.sem}{.scope}{.space}.op{.level::cache_hint}.type [a], b{, cache-policy};
+def _ptx_red_scalar_parts(*args):
+    sem, scope, space, op, ptx_type, has_cache_hint = args[-6:]
+    sem = parse_str(sem)
+    scope = parse_str(scope)
+    space = parse_str(space)
+    op = parse_str(op)
+    ptx_type, c_type, constraint, _tvm_dtype = _type_info(ptx_type)
+    has_cache = (
+        bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else bool(has_cache_hint)
+    )
+    modifiers = f"{_dot(sem)}{_dot(scope)}{_dot(space)}"
+    instr = f"red{modifiers}.{op}{_cache_suffix('cache' if has_cache else '')}.{ptx_type}"
+    name = (
+        "tvm_builtin_ptx_red_scalar"
+        f"{_dot(sem).replace('.', '_')}{_dot(scope).replace('.', '_')}"
+        f"_{_safe_attr(space)}_{op}_{ptx_type}{'_cache_hint' if has_cache else ''}"
+    )
+    cache_operand = ', "l"(cache_policy)' if has_cache else ""
+    addr_decl = ""
+    addr_operand = '"l"(address)'
+    if space.startswith("shared"):
+        addr_decl = "    unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n"
+        addr_operand = '"r"(addr)'
+    body = (
+        f"{addr_decl}"
+        f'    asm volatile("{instr} [%0], %1{", %2" if has_cache else ""};"\n'
+        "                 :\n"
+        f'                 : {addr_operand}, "{constraint}"(value)'
+        f"{cache_operand}\n"
+        '                 : "memory");'
+    )
+    return name, f"(void* address, {c_type} value, unsigned long long cache_policy)", body
+
+
+device_intrinsic(
+    "ptx_red_scalar",
+    n_attrs=6,
+    helper_name=lambda *a: _ptx_red_scalar_parts(*a)[0],
+    c_signature=lambda *a: _ptx_red_scalar_parts(*a)[1],
+    body=lambda *a: _ptx_red_scalar_parts(*a)[2],
+)
+
+
+# PTX atom scalar one-source-operand form:
+#   atom{.sem}{.scope}{.space}.op{.level::cache_hint}.type d, [a], b{, cache-policy};
+def _ptx_atom_scalar_parts(*args):
+    sem, scope, space, op, ptx_type, has_cache_hint = args[-6:]
+    sem = parse_str(sem)
+    scope = parse_str(scope)
+    space = parse_str(space)
+    op = parse_str(op)
+    ptx_type, c_type, constraint, tvm_dtype = _type_info(ptx_type)
+    has_cache = (
+        bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else bool(has_cache_hint)
+    )
+    modifiers = f"{_dot(sem)}{_dot(scope)}{_dot(space)}"
+    instr = f"atom{modifiers}.{op}{_cache_suffix('cache' if has_cache else '')}.{ptx_type}"
+    name = (
+        "tvm_builtin_ptx_atom_scalar"
+        f"{_dot(sem).replace('.', '_')}{_dot(scope).replace('.', '_')}"
+        f"_{_safe_attr(space)}_{op}_{ptx_type}{'_cache_hint' if has_cache else ''}"
+    )
+    cache_operand = ', "l"(cache_policy)' if has_cache else ""
+    addr_decl = ""
+    addr_operand = '"l"(address)'
+    if space.startswith("shared"):
+        addr_decl = "    unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n"
+        addr_operand = '"r"(addr)'
+    body = (
+        f"{addr_decl}"
+        f"    {c_type} ret;\n"
+        f'    asm volatile("{instr} %0, [%1], %2{", %3" if has_cache else ""};"\n'
+        f'                 : "={constraint}"(ret)\n'
+        f'                 : {addr_operand}, "{constraint}"(value)'
+        f"{cache_operand}\n"
+        '                 : "memory");\n'
+        "    return ret;"
+    )
+    return (
+        name,
+        f"(void* address, {c_type} value, unsigned long long cache_policy)",
+        c_type,
+        tvm_dtype,
+        body,
+    )
+
+
+device_intrinsic(
+    "ptx_atom_scalar",
+    n_attrs=6,
+    helper_name=lambda *a: _ptx_atom_scalar_parts(*a)[0],
+    c_signature=lambda *a: _ptx_atom_scalar_parts(*a)[1],
+    return_type=lambda *a: _ptx_atom_scalar_parts(*a)[2],
+    tvm_return_type=lambda *a: _ptx_atom_scalar_parts(*a)[3],
+    body=lambda *a: _ptx_atom_scalar_parts(*a)[4],
+)
+
+
+# PTX prefetch tensormap form:
+#   prefetch{.tensormap_space}.tensormap [a];
+def _prefetch_tensormap_parts(_tensor_map, tensormap_space):
+    space = parse_str(tensormap_space)
+    instr = f"prefetch{_dot(space)}.tensormap"
+    name = f"tvm_builtin_ptx_prefetch{('_' + _safe_attr(space)) if space else ''}_tensormap"
+    body = (
+        f'    asm volatile("{instr} [%0];"\n'
+        "                 :\n"
+        '                 : "l"(tensor_map_addr)\n'
+        '                 : "memory");'
+    )
+    return name, body
+
+
+device_intrinsic(
+    "ptx_prefetch_tensormap",
+    n_attrs=1,
+    helper_name=lambda *a: _prefetch_tensormap_parts(*a)[0],
+    c_signature="(unsigned long long tensor_map_addr)",
+    body=lambda *a: _prefetch_tensormap_parts(*a)[1],
+)
+
+
+# PTX st forms (ISA table entries registered via ``ptx_st`` and siblings):
+#   st{.weak}{.ss}{.cop}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
+#   st{.weak}{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
+#   st.volatile{.ss}{.vec}.type [a], b;
+#   st.relaxed.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
+#   st.release.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy};
+#   st.mmio.sem.sys{.global}.type [a], b;
+_PTX_ST_COPS = {"", "wb", "cg", "cs", "wt"}
+_PTX_ST_SPACES = {"global", "shared", "shared::cta", "shared::cluster", "local", "param::func"}
+
+
+def _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type):
+    if ptx_type == "u8" and vec_len == 1:
+        return "    unsigned int r0 = *reinterpret_cast<unsigned char*>(src_ptr);\n"
+    store_type = _PTX_VEC_STORE_TYPE[num_bytes]
+    if vec_len > 1:
+        return f"    {store_type} src_ = *reinterpret_cast<{store_type}*>(src_ptr);\n" + "".join(
+            f"    {c_type} r{i} = src_.{c};\n" for i, c in enumerate("xyzw"[:vec_len])
+        )
+    return f"    {c_type} r0 = *reinterpret_cast<{c_type}*>(src_ptr);\n"
+
+
+def _ptx_st_form_parts(form, attr_args, from_src):
+    if form == "weak":
+        weak, space, cop, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args
+        sem, scope = "", ""
+    elif form == "relaxed":
+        scope, space, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args
+        sem, weak, cop = "relaxed", False, ""
+    elif form == "release":
+        scope, space, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args
+        sem, weak, cop = "release", False, ""
+    elif form == "volatile":
+        space, vec, ptx_type = attr_args
+        sem, scope, weak, cop = "", "", False, ""
+        has_cache_hint, l1_evict, l2_evict = False, "", ""
+    elif form == "mmio":
+        sem, scope, space, ptx_type = attr_args
+        weak, cop, vec = False, "", ""
+        has_cache_hint, l1_evict, l2_evict = False, "", ""
+    else:
+        raise ValueError(f"unknown st form {form!r}")
+
+    sem = parse_str(sem)
+    scope = parse_str(scope)
+    space = parse_str(space)
+    cop = parse_str(cop)
+    vec = parse_str(vec)
+    l1_evict = parse_str(l1_evict)
+    l2_evict = parse_str(l2_evict)
+    weak = _bool_attr(weak)
+    has_cache = _bool_attr(has_cache_hint)
+    ptx_type, c_type, constraint, _tvm_dtype = _type_info(ptx_type)
+    if cop and cop not in _PTX_ST_COPS:
+        raise ValueError(f"Unsupported PTX st cache operation {cop!r}")
+    if vec and vec not in _PTX_VEC:
+        raise ValueError(f"Unsupported PTX st vector modifier {vec!r}")
+    if space not in _PTX_ST_SPACES and not (form == "mmio" and space == "global"):
+        raise ValueError(f"Unsupported PTX st state space {space!r}")
+    vec_len = int(vec[1:]) if vec else 1
+    elem_bytes = (
+        8
+        if ptx_type.endswith("64")
+        else 2
+        if ptx_type in ("u16", "s16", "b16")
+        else 1
+        if ptx_type in ("u8", "s8", "b8")
+        else 4
+    )
+    num_bytes = vec_len * elem_bytes if vec else elem_bytes
+    use_cache_policy = form in ("weak", "relaxed", "release")
+    if form == "mmio":
+        if sem not in ("acquire", "relaxed", "release") or scope != "sys" or space != "global":
+            raise ValueError("st.mmio requires sem, scope=sys, space=global")
+        prefix = f"st.mmio.{sem}.{scope}"
+    elif form == "relaxed":
+        if not scope:
+            raise ValueError("st.relaxed requires scope")
+        prefix = f"st.relaxed.{scope}{_dot(space)}"
+    elif form == "release":
+        if not scope:
+            raise ValueError("st.release requires scope")
+        prefix = f"st.release.{scope}{_dot(space)}"
+    elif form == "volatile":
+        prefix = f"st.volatile{_dot(space)}"
+    else:
+        prefix = f"st{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}"
+    level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, "")
+    instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}"
+    name_parts = ["tvm_builtin_ptx_st", form if form != "weak" else ("weak" if weak else "plain")]
+    if sem:
+        name_parts.append(_safe_attr(sem))
+    if scope:
+        name_parts.append(_safe_attr(scope))
+    name_parts.extend(
+        [
+            _safe_attr(space),
+            _safe_attr(cop) if cop else "",
+            _safe_attr(vec) if vec else "",
+            ptx_type,
+            "from_src" if from_src else "values",
+        ]
+    )
+    if has_cache:
+        name_parts.append("cache_hint")
+    name = "_".join(p for p in name_parts if p)
+    values = f"{{{', '.join(f'%{i + 1}' for i in range(vec_len))}}}" if vec_len > 1 else "%1"
+    value_constraints = "".join(f', "{constraint}"(value{i})' for i in range(vec_len))
+    cache_slot = f", %{vec_len + 1}" if has_cache else ""
+    cache_operand = ', "l"(cache_policy)' if has_cache else ""
+    addr_decl, addr_operand = _ptx_shared_addr(space, "address")
+    if from_src:
+        load_regs = _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type)
+        if vec_len > 1:
+            in_constraints = ", ".join(f'"{constraint}"(r{i})' for i in range(vec_len))
+            value_args = f", {in_constraints}"
+        else:
+            value_args = f', "{constraint}"(r0)'
+        body = (
+            f"{addr_decl}{load_regs}"
+            f'    asm volatile("{instr} [%0], {values}{cache_slot};"\n'
+            "                 :\n"
+            f"                 : {addr_operand}{value_args}{cache_operand}\n"
+            '                 : "memory");'
+        )
+        if use_cache_policy:
+            sig = "(void* address, void* src_ptr, unsigned long long cache_policy)"
+        else:
+            sig = "(void* address, void* src_ptr)"
+    else:
+        body = (
+            f"{addr_decl}"
+            f'    asm volatile("{instr} [%0], {values}{cache_slot};"\n'
+            "                 :\n"
+            f"                 : {addr_operand}{value_constraints}{cache_operand}\n"
+            '                 : "memory");'
+        )
+        if form == "mmio":
+            sig = f"(void* address, {c_type} value0)"
+        elif use_cache_policy:
+            value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len))
+            sig = f"(void* address, {value_params}, unsigned long long cache_policy)"
+        else:
+            value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len))
+            sig = f"(void* address, {value_params})"
+    return name, sig, body
+
+
+def _register_ptx_st(op_name, form, n_attrs, *, with_cache_policy=True):
+    def codegen(*args):
+        from_src = _bool_attr(args[-1])
+        st_attrs = args[-n_attrs:-1]
+        parts = _ptx_st_form_parts(form, st_attrs, from_src)
+        forward = args[:-(n_attrs)]
+        name, sig, body_str = parts
+        source_code = f"\n__forceinline__ __device__ void {name}{sig} {{\n{body_str}\n}}\n"
+        return cuda_func_call(name, *forward, source_code=source_code)
+
+    codegen.__name__ = f"codegen_{op_name}"
+    register_codegen(op_name)(codegen)
+
+
+def _register_ptx_st_mmio(op_name, form, n_attrs):
+    def codegen(*args):
+        parts = _ptx_st_form_parts(form, args[-n_attrs:], False)
+        forward = args[:-n_attrs]
+        name, sig, body_str = parts
+        source_code = f"\n__forceinline__ __device__ void {name}{sig} {{\n{body_str}\n}}\n"
+        return cuda_func_call(name, *forward, source_code=source_code)
+
+    codegen.__name__ = f"codegen_{op_name}"
+    register_codegen(op_name)(codegen)
+
+
+_register_ptx_st("ptx_st", "weak", 9)
+_register_ptx_st("ptx_st_relaxed", "relaxed", 8)
+_register_ptx_st("ptx_st_release", "release", 8)
+_register_ptx_st("ptx_st_volatile", "volatile", 4)
+_register_ptx_st_mmio("ptx_st_mmio", "mmio", 4)
+
+
+# PTX st.bulk form:
+#   st.bulk{.weak}{.shared::cta} [a], size, initval;
+# ``initval`` is an immediate operand whose only legal value is 0.
+def _ptx_st_bulk_parts(_ptr, _num_bytes, weak, space):
+    weak = bool(int(weak)) if hasattr(weak, "value") else bool(weak)
+    space = parse_str(space)
+    instr = f"st.bulk{'.weak' if weak else ''}{_dot(space)}"
+    name = f"tvm_builtin_ptx_st_bulk{'_weak' if weak else ''}{('_' + _safe_attr(space)) if space else ''}"
+    addr_arg = (
+        '"r"((unsigned int)__cvta_generic_to_shared(ptr))' if space == "shared::cta" else '"l"(ptr)'
+    )
+    body = (
+        f'    asm volatile("{instr} [%0], %1, 0;"\n'
+        "                 :\n"
+        f"                 : {addr_arg}, "
+        '"l"(static_cast<uint64_t>(num_bytes))\n'
+        '                 : "memory");'
+    )
+    return name, body
+
+
+device_intrinsic(
+    "ptx_st_bulk",
+    n_attrs=2,
+    helper_name=lambda *a: _ptx_st_bulk_parts(*a)[0],
+    c_signature="(void* ptr, unsigned int num_bytes)",
+    body=lambda *a: _ptx_st_bulk_parts(*a)[1],
+)
+
+device_intrinsic(
+    "cuda_uint_as_float",
+    helper_name="tvm_builtin_uint_as_float",
+    c_signature="(unsigned int bits)",
+    return_type="float",
+    body="    return __uint_as_float(bits);",
+)
+device_intrinsic(
+    "cuda_float_as_uint",
+    helper_name="tvm_builtin_float_as_uint",
+    c_signature="(float x)",
+    return_type="unsigned int",
+    body="    return __float_as_uint(x);",
+)

python/tvm/backend/cuda/operator/intrinsics/misc.py

@@ -0,0 +1,253 @@
+# 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
+# pylint: disable=redefined-builtin, invalid-name
+"""Miscellaneous device helpers.
+
+Catch-all for ops that don't fit the (sync / mma / cp_async / memory / math /
+nvshmem) feature buckets:
+
+* PTX register-allocation control: ``setmaxnreg`` / ``mov`` from special reg.
+* Per-thread queries / scheduling hints: ``thread_rank`` / ``nano_sleep``.
+* Profiler timer hooks (``timer_init/start/end/finalize``).
+* Debug helpers: ``printf`` / ``trap`` on assert failure.
+"""
+
+import hashlib
+import json
+
+import tvm
+from tvm.backend.cuda.op import cuda_func_call
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+from .utils import parse_str
+
+# =============================================================================
+# setmaxnreg.{inc,dec}.sync.aligned.u32 — 1 PTX form (.action picks inc/dec).
+# =============================================================================
+
+
+def _ptx_setmaxnreg(inc, nreg):
+    inc = bool(int(inc)) if hasattr(inc, "value") else bool(inc)
+    nreg = int(nreg)
+    action = "inc" if inc else "dec"
+    return (
+        f"tvm_builtin_ptx_setmaxnreg_{action}_{nreg}",
+        f'    asm volatile("setmaxnreg.{action}.sync.aligned.u32 {nreg};");',
+    )
+
+
+device_intrinsic(
+    "ptx_setmaxnreg",
+    n_attrs=2,
+    helper_name=lambda inc, nreg: _ptx_setmaxnreg(inc, nreg)[0],
+    body=lambda inc, nreg: _ptx_setmaxnreg(inc, nreg)[1],
+)
+
+
+# =============================================================================
+# mov.u32/u64 from special register — 1 PTX form (Form 2 of mov.type d, sreg).
+# Each (bits, reg) emits a distinct helper because the special reg name is
+# baked into the PTX text.
+# =============================================================================
+
+
+def _ptx_fetch_register_body(bits):
+    spec = "l" if bits == 64 else "r"
+
+    def _body(reg):
+        reg = parse_str(reg)
+        return (
+            f"    uint{bits}_t x;\n"
+            f'    asm volatile("mov.u{bits} %0, %{reg};" : "={spec}"(x));\n'
+            f"    return (int{bits}_t)x;"
+        )
+
+    return _body
+
+
+for _bits in (32, 64):
+    device_intrinsic(
+        f"ptx_fetch_register_{_bits}",
+        n_attrs=1,
+        helper_name=(
+            lambda *a, bits=_bits: (
+                f"tvm_builtin_ptx_fetch_register_"
+                f"{parse_str(a[-1]).replace('::', '_').replace('.', '_')}"
+            )
+        ),
+        return_type=f"int{_bits}_t",
+        body=_ptx_fetch_register_body(_bits),
+    )
+del _bits
+
+
+@register_codegen("ptx_fetch_register")
+def codegen_ptx_fetch_register(bits, reg):
+    bits = int(bits)
+    reg = parse_str(reg)
+    if bits not in (32, 64):
+        raise ValueError(f"Only support 32/64 bits for ptx_fetch_register, but got {bits}.")
+    result = CODEGEN_REGISTRY[f"tirx.ptx_fetch_register_{bits}"]([reg])
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# Per-thread queries / scheduling hints.
+# =============================================================================
+device_intrinsic(
+    "cuda_thread_rank",
+    body=(
+        "    namespace cg = cooperative_groups;\n    return cg::this_thread_block().thread_rank();"
+    ),
+    return_type="int",
+    tvm_return_type="int32",
+    extra_deps=("cooperative_groups",),
+)
+device_intrinsic("cuda_nano_sleep", c_signature="(uint64_t time)", body="    __nanosleep(time);")
+
+
+# =============================================================================
+# Profiler timer hooks.
+# =============================================================================
+_COMMON_PARAMS = (
+    "uint64_t* profiler_buffer, uint64_t* profiler_tag, "
+    "uint32_t* profiler_write_offset, int profiler_write_stride, bool leader_cond"
+)
+_EVENT_PARAMS = f"int event_type, {_COMMON_PARAMS}"
+
+
+def _write_event(event_bits: str) -> str:
+    return (
+        "profiler_buffer[profiler_write_offset[0]] = "
+        "((uint64_t)tvm_builtin_get_timestamp() << 32) | "
+        f"(profiler_tag[0] | {event_bits});\n"
+        "        profiler_write_offset[0] += profiler_write_stride;"
+    )
+
+
+device_intrinsic(
+    "timer_init_cuda",
+    c_signature=(
+        "(uint64_t* profiler_buffer, uint64_t* profiler_tag, "
+        "uint32_t* profiler_write_offset, int num_groups, int group_id)"
+    ),
+    body=(
+        "    const uint32_t NBLOCKS = (uint32_t)(gridDim.x * gridDim.y * gridDim.z);\n"
+        "    const uint32_t BLOCK_IDX = (uint32_t)("
+        "(blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x);\n"
+        "    const uint32_t NGROUPS = num_groups;\n"
+        "    const uint32_t GROUP_ID = group_id;\n"
+        "    const uint32_t BLOCK_GROUP_IDX = BLOCK_IDX * NGROUPS + GROUP_ID;\n"
+        "    if ((blockIdx.x == 0) && (blockIdx.y == 0) && "
+        "(blockIdx.z == 0) && (threadIdx.x == 0)) {\n"
+        "        profiler_buffer[0] = ((uint64_t)NGROUPS << 32) | NBLOCKS;\n"
+        "    }\n"
+        "    profiler_write_offset[0] = 1 + BLOCK_GROUP_IDX;\n"
+        "    profiler_tag[0] = (uint64_t)BLOCK_GROUP_IDX << 12;"
+    ),
+)
+
+device_intrinsic(
+    "timer_start_cuda",
+    c_signature=f"({_EVENT_PARAMS})",
+    body=(
+        f"    if (leader_cond) {{\n        {_write_event('(uint32_t)event_type << 2 | 0x0')}\n    }}\n"
+        "    __threadfence_block();"
+    ),
+    extra_deps=("get_time_stamp",),
+)
+
+device_intrinsic(
+    "timer_end_cuda",
+    c_signature=f"({_EVENT_PARAMS})",
+    body=(
+        "    __threadfence_block();\n"
+        f"    if (leader_cond) {{\n        {_write_event('(uint32_t)event_type << 2 | 0x1')}\n    }}"
+    ),
+    extra_deps=("get_time_stamp",),
+)
+
+device_intrinsic(
+    "timer_finalize_cuda",
+    c_signature=f"({_COMMON_PARAMS})",
+    body=(
+        f"    __threadfence_block();\n    if (leader_cond) {{\n        {_write_event('0x3')}\n    }}"
+    ),
+    extra_deps=("get_time_stamp",),
+)
+
+
+# =============================================================================
+# Debug helpers — ``printf`` (variadic templated) and ``trap`` on assert.
+# =============================================================================
+device_intrinsic(
+    "cuda_trap_when_assert_failed",
+    c_signature="(bool cond)",
+    body='    do {\n        if (not (cond))\n            asm("trap;");\n    } while (0);',
+)
+
+
+@register_codegen("cuda_printf")
+def codegen_cuda_printf(fmt, *args):
+    if isinstance(fmt, tvm.tirx.StringImm):
+        fmt = fmt.value
+    if not isinstance(fmt, str):
+        raise ValueError("T.cuda.printf format must be a string literal")
+    fmt_literal = json.dumps(fmt)
+    arg_dtypes = [str(arg.ty) for arg in args]
+    signature = "|".join([fmt, *arg_dtypes])
+    digest = hashlib.sha1(signature.encode("utf-8")).hexdigest()
+    func_name = f"tvm_builtin_cuda_printf_{len(args)}_{digest}"
+
+    def c_type(dtype: str) -> str:
+        if dtype == "float32":
+            return "float"
+        if dtype == "float64":
+            return "double"
+        if dtype in {"int8", "int16", "int32"}:
+            return "int"
+        if dtype == "int64":
+            return "long long"
+        if dtype in {"uint8", "uint16", "uint32"}:
+            return "unsigned int"
+        if dtype == "uint64":
+            return "unsigned long long"
+        if dtype == "bool":
+            return "int"
+        if dtype == "handle":
+            return "void*"
+        raise ValueError(f"Unsupported T.cuda.printf argument dtype: {dtype}")
+
+    params = ", ".join(f"{c_type(dtype)} arg{i}" for i, dtype in enumerate(arg_dtypes))
+    call_args = ", ".join(f"arg{i}" for i in range(len(args)))
+    comma_call_args = f", {call_args}" if call_args else ""
+    source_code = f"""
+__noinline__ __device__ void {func_name}({params}) {{
+    printf({fmt_literal}{comma_call_args});
+}}
+"""
+    return cuda_func_call(func_name, *args, source_code=source_code)
+
+
+device_intrinsic(
+    "cuda_clock64",
+    helper_name="tvm_builtin_clock64",
+    return_type="unsigned long long",
+    body="    return clock64();",
+)

python/tvm/backend/cuda/operator/intrinsics/mma.py

@@ -0,0 +1,482 @@
+# 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, too-many-arguments, too-many-locals, too-many-positional-arguments
+"""PTX MMA / ldmatrix / stmatrix intrinsics.
+
+mma.sync.aligned has 7 form_kinds per the PTX docs (f16 / tf32 / bf16 / fp64
+/ int8 / fp8 / subbyte). Each form_kind is one ``device_intrinsic`` registration;
+the (shape, layouts, dtypes) modifier slots are attrs. Body computes the per-
+fragment register counts at codegen time from M*N*bits/threads/frag_size and
+hand-builds the asm constraint list.
+
+ldmatrix / stmatrix each have a single PTX form (the .m8n8 .b16/.b8 variant
+that TIRx uses); ``num`` and ``trans`` are modifier attrs.
+"""
+
+import re
+from dataclasses import dataclass
+
+from tvm import DataType
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+from .types import PTXDataType
+from .utils import parse_str
+
+
+@dataclass
+class FragAttrs:
+    reg_type: str  # asm constraint letter (r / f / d)
+    size: int  # bit width per register slot (32 or 64)
+    ptr_type: str  # C type for the cast
+
+
+_FRAG_ATTRS_MAP = {
+    PTXDataType.BIT1: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.INT4: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.UINT4: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.INT8: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.UINT8: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.FLOAT8_E4M3FN: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.FLOAT8_E5M2: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.BIT16: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.FLOAT16: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.BFLOAT16: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.TENSOR_FLOAT32: FragAttrs("r", 32, "uint32_t"),
+    PTXDataType.INT32: FragAttrs("r", 32, "int32_t"),
+    PTXDataType.FLOAT32: FragAttrs("f", 32, "float"),
+    PTXDataType.FLOAT64: FragAttrs("d", 64, "double"),
+}
+
+
+def _parse_mma_shape(shape_str):
+    match = re.search(r"m(\d+)n(\d+)k(\d+)", shape_str)
+    if not match:
+        raise ValueError(f"Cannot parse MMA shape: {shape_str!r}")
+    return tuple(map(int, match.groups()))
+
+
+def _classify_mma_form(d_type, a_type, b_type):
+    """Map (d, a, b) dtype triple to one of the 7 PTX form_kind tags."""
+    fp16 = {"float16", "fp16"}
+    tf32 = {"tensor_float32", "tf32"}
+    bf16 = {"bfloat16", "bf16"}
+    fp64 = {"float64", "fp64"}
+    int_a = {"int8", "uint8", "s8", "u8"}
+    fp8 = {"e4m3", "e5m2", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2"}
+    subbyte = {"int4", "uint4", "bit1", "s4", "u4", "b1", "int1", "uint1"}
+    if a_type in fp16 and b_type in fp16:
+        return "f16"
+    if a_type in tf32 and b_type in tf32:
+        return "tf32"
+    if a_type in bf16 and b_type in bf16:
+        return "bf16"
+    if a_type in fp64 and b_type in fp64:
+        return "fp64"
+    if a_type in int_a and b_type in int_a:
+        return "int8"
+    if a_type in fp8 and b_type in fp8:
+        return "fp8"
+    if a_type in subbyte and b_type in subbyte:
+        return "subbyte"
+    raise ValueError(
+        f"Unknown ptx.mma form for d_type={d_type!r}, a_type={a_type!r}, b_type={b_type!r}"
+    )
+
+
+def _frag(dtype_str):
+    return _FRAG_ATTRS_MAP[PTXDataType.from_string(dtype_str)]
+
+
+def _mma_threads(shape, a_type):
+    """Special case: m8n8k4 with f16 a/b uses 8 threads per fragment."""
+    m, n, k = _parse_mma_shape(shape)
+    if m == 8 and n == 8 and k == 4 and a_type == "float16":
+        return 8
+    return 32
+
+
+# PTX dtype abbreviation -> element bit width. Used by _frag_count so that
+# callers passing the PTX abbreviation (e.g. "fp32") don't blow up in
+# ``DataType("fp32")``.
+_PTX_BITS = {
+    "fp16": 16,
+    "fp32": 32,
+    "fp64": 64,
+    "bf16": 16,
+    "tf32": 32,  # tensor-float32 packs 19 significant bits into a 32-bit slot
+    "s8": 8,
+    "u8": 8,
+    "s32": 32,
+    "s4": 4,
+    "u4": 4,
+    "b1": 1,
+    "b16": 16,
+    "e4m3": 8,
+    "e5m2": 8,
+}
+
+
+def _frag_count(dtype, dim_a, dim_b, threads):
+    if dtype in _PTX_BITS:
+        bits = _PTX_BITS[dtype]
+    else:
+        bits = DataType(dtype).bits
+    size = _frag(dtype).size
+    return dim_a * dim_b * bits // threads // size
+
+
+# =============================================================================
+# Shared helpers for the 7 mma form_kinds.
+# Args layout for each form:
+#   (d_ptr_in, a_ptr_in, b_ptr_in [, c_ptr_in], shape, a_layout, b_layout,
+#    d_type, a_type, b_type, c_type, no_c_ptr [, saturate or bit_op])
+# n_attrs = 8 for f16/tf32/bf16/fp64/fp8 (last 8 = shape, layouts, 4 dtypes, no_c_ptr)
+# n_attrs = 9 for int8 (+ saturate) and subbyte (+ bit_op)
+# =============================================================================
+
+
+def _mma_form_parts(args, *, has_saturate=False, has_bit_op=False):
+    """Compute (helper_name, c_signature, body) for one mma form invocation.
+
+    ``args`` is the full positional arg tuple as received by codegen.
+    The trailing ``n_attrs`` (8 or 9) entries are attrs.
+    """
+    n_extra = (1 if has_saturate else 0) + (1 if has_bit_op else 0)
+    n_attrs = 8 + n_extra
+    # Split off attr args from the tail (operand args are ahead).
+    attrs = args[-n_attrs:]
+    shape = parse_str(attrs[0])
+    a_layout = parse_str(attrs[1])
+    b_layout = parse_str(attrs[2])
+    d_type = parse_str(attrs[3])
+    a_type = parse_str(attrs[4])
+    b_type = parse_str(attrs[5])
+    c_type = parse_str(attrs[6])
+    no_c_ptr_raw = attrs[7]
+    no_c_ptr = bool(int(no_c_ptr_raw)) if hasattr(no_c_ptr_raw, "value") else bool(no_c_ptr_raw)
+    saturate = False
+    bit_op = ""
+    if has_saturate:
+        s = attrs[8]
+        saturate = bool(int(s)) if hasattr(s, "value") else bool(s)
+    if has_bit_op:
+        bit_op = parse_str(attrs[8])
+
+    # Fragment counts (same derivation as the contiguous form).
+    m, n, k = _parse_mma_shape(shape)
+    threads = _mma_threads(shape, a_type)
+    d_cnt = _frag_count(d_type, m, n, threads)
+    a_cnt = _frag_count(a_type, m, k, threads)
+    b_cnt = _frag_count(b_type, k, n, threads)
+    c_cnt = _frag_count(c_type, m, n, threads)
+
+    # C signature: one void* per register, ordered D regs, A regs, B regs,
+    # then C regs (only when the accumulator is used).
+    sig_parts = (
+        [f"void* d_ptr{i}" for i in range(d_cnt)]
+        + [f"void* a_ptr{i}" for i in range(a_cnt)]
+        + [f"void* b_ptr{i}" for i in range(b_cnt)]
+    )
+    if not no_c_ptr:
+        sig_parts += [f"void* c_ptr{i}" for i in range(c_cnt)]
+    sig = "(" + ", ".join(sig_parts) + ")"
+
+    def _safe(s):
+        return s.replace("::", "_").replace(".", "_")
+
+    name = (
+        f"ptx_mma_{shape}_{a_layout}_{b_layout}"
+        f"_{_safe(d_type)}_{_safe(a_type)}_{_safe(b_type)}_{_safe(c_type)}"
+        f"{'_no_c_ptr' if no_c_ptr else ''}"
+        f"{'_saturate' if saturate else ''}"
+    )
+
+    d_frag = _frag(d_type)
+    a_frag = _frag(a_type)
+    b_frag = _frag(b_type)
+    c_frag = _frag(c_type)
+
+    saturate_inst = ".satfinite" if saturate else ""
+    # PTX b1 mma requires a `.popc` suffix after the bit op (e.g. `.xor.popc`).
+    bit_op_inst = f".{bit_op}.popc" if bit_op else ""
+
+    d_type_inst = PTXDataType.from_string(d_type).to_string()
+    c_type_inst = PTXDataType.from_string(c_type).to_string()
+    a_type_inst = PTXDataType.from_string(a_type).to_string()
+    b_type_inst = PTXDataType.from_string(b_type).to_string()
+
+    def _slot_arr(start, cnt):
+        return "{" + ", ".join(f"%{start + i}" for i in range(cnt)) + "}"
+
+    args_template = (
+        f"{_slot_arr(0, d_cnt)}, {_slot_arr(d_cnt, a_cnt)}, "
+        f"{_slot_arr(d_cnt + a_cnt, b_cnt)}, {_slot_arr(d_cnt + a_cnt + b_cnt, c_cnt)}"
+    )
+
+    # Each register binds to its OWN pointer via *(T*)X_ptrN (scatter).
+    d_outs = ", ".join(
+        f'"=r"(*({d_frag.ptr_type}*)d_ptr{i})'
+        if d_frag.reg_type == "r"
+        else f'"={d_frag.reg_type}"(*({d_frag.ptr_type}*)d_ptr{i})'
+        for i in range(d_cnt)
+    )
+    a_inputs = ", ".join(
+        f'"{a_frag.reg_type}"(*({a_frag.ptr_type}*)a_ptr{i})' for i in range(a_cnt)
+    )
+    b_inputs = ", ".join(
+        f'"{b_frag.reg_type}"(*({b_frag.ptr_type}*)b_ptr{i})' for i in range(b_cnt)
+    )
+    if no_c_ptr:
+        c_value = "0.f" if c_frag.reg_type == "f" else "0"
+        c_inputs = ", ".join(f'"{c_frag.reg_type}"({c_value})' for _ in range(c_cnt))
+    else:
+        c_inputs = ", ".join(
+            f'"{c_frag.reg_type}"(*({c_frag.ptr_type}*)c_ptr{i})' for i in range(c_cnt)
+        )
+
+    body = (
+        "    asm volatile(\n"
+        f'        "mma.sync.aligned.{shape}.{a_layout}.{b_layout}{saturate_inst}'
+        f'{d_type_inst}{a_type_inst}{b_type_inst}{c_type_inst}{bit_op_inst} "\n'
+        f'        "{args_template};\\n"\n'
+        f"        : {d_outs}\n"
+        f"        : {a_inputs}, {b_inputs}, {c_inputs}\n"
+        "    );"
+    )
+    return name, sig, body
+
+
+def _register_mma_form(form_kind, *, has_saturate=False, has_bit_op=False):
+    n_attrs = 8 + (1 if has_saturate else 0) + (1 if has_bit_op else 0)
+
+    def _parts(*args, hs=has_saturate, hb=has_bit_op):
+        return _mma_form_parts(args, has_saturate=hs, has_bit_op=hb)
+
+    device_intrinsic(
+        f"_ptx_mma_{form_kind}",
+        n_attrs=n_attrs,
+        helper_name=lambda *a: _parts(*a)[0],
+        c_signature=lambda *a: _parts(*a)[1],
+        body=lambda *a: _parts(*a)[2],
+    )
+
+
+# Form 1 — f16. Form 2 — tf32. Form 3 — bf16. Form 4 — fp64. Form 6 — fp8.
+# All share the same 8-attr layout (no saturate / bit_op).
+for _kind in ("f16", "tf32", "bf16", "fp64", "fp8"):
+    _register_mma_form(_kind)
+del _kind
+
+# Form 5 — int8 (+ saturate).
+_register_mma_form("int8", has_saturate=True)
+
+# Form 7 — subbyte (+ bit_op for b1).
+_register_mma_form("subbyte", has_bit_op=True)
+
+
+@register_codegen("ptx_mma")
+def codegen_ptx_mma(
+    shape,
+    a_layout,
+    b_layout,
+    d_type,
+    a_type,
+    b_type,
+    c_type,
+    d_cnt,
+    a_cnt,
+    b_cnt,
+    c_cnt,
+    no_c_ptr,
+    *rest,
+):
+    """Classify (d, a, b) dtype triple to one of 7 form_kinds and forward.
+
+    ``rest`` = flattened per-register pointers (d_cnt + a_cnt + b_cnt + c_cnt of
+    them) followed by ``saturate`` and optionally ``bit_op``.
+    """
+    shape = parse_str(shape)
+    a_layout = parse_str(a_layout)
+    b_layout = parse_str(b_layout)
+    d_type = parse_str(d_type)
+    a_type = parse_str(a_type)
+    b_type = parse_str(b_type)
+    c_type = parse_str(c_type)
+    d_cnt = int(d_cnt)
+    a_cnt = int(a_cnt)
+    b_cnt = int(b_cnt)
+    c_cnt = int(c_cnt)
+    no_c_ptr = bool(int(no_c_ptr)) if hasattr(no_c_ptr, "value") else bool(no_c_ptr)
+
+    n_ptrs = d_cnt + a_cnt + b_cnt + (0 if no_c_ptr else c_cnt)
+    ptrs = list(rest[:n_ptrs])
+    trailing = list(rest[n_ptrs:])
+    saturate = bool(trailing[0]) if trailing else False
+    bit_op_v = ""
+    if len(trailing) >= 2:
+        bo = trailing[1]
+        bit_op_v = parse_str(bo) if isinstance(bo, str) else (bo if bo is not None else "")
+
+    kind = _classify_mma_form(d_type, a_type, b_type)
+
+    # op_args are the flattened per-register pointers (already in PTX order:
+    # D regs, A regs, B regs, then C regs unless no_c_ptr).
+    op_args = ptrs
+
+    attr_args = [shape, a_layout, b_layout, d_type, a_type, b_type, c_type, no_c_ptr]
+    if kind == "int8":
+        attr_args.append(saturate)
+    elif kind == "subbyte":
+        attr_args.append(bit_op_v)
+
+    result = CODEGEN_REGISTRY[f"tirx._ptx_mma_{kind}"](op_args + attr_args)
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# ldmatrix / stmatrix — m8n8 fragment load/store. PTX docs lists 3 ldmatrix
+# forms (m8n8 + m8n16 + m16n16); TIRx uses only the m8n8 form. 1
+# device_intrinsic each. ``num`` (.x1/.x2/.x4) and ``trans`` are modifier
+# attrs; the asm body loops over per-register constraints based on
+# (num, dtype).
+# =============================================================================
+
+
+def _ldmatrix_parts(*args):
+    # args = (smem_ptr, dst0, dst1, ..., dst{N-1}, num, dtype, trans)
+    # The last 3 entries are the codegen attrs (n_attrs=3).
+    num = int(args[-3])
+    dtype = parse_str(args[-2])
+    trans_b = bool(int(args[-1])) if hasattr(args[-1], "value") else bool(args[-1])
+    if num not in (1, 2, 4):
+        raise ValueError(f"ldmatrix .num must be one of {{1, 2, 4}}, got {num}")
+    if dtype not in ("b16", "b8"):
+        raise ValueError(f"ldmatrix dtype must be 'b16' or 'b8', got {dtype!r}")
+    n_regs = num if dtype == "b16" else num // 2
+    trans_inst = ".trans" if trans_b else ""
+    slot_list = "{" + ", ".join(f"%{i}" for i in range(n_regs)) + "}"
+    reg_decls = ", ".join(f"r{i}" for i in range(n_regs))
+    out_constraints = ", ".join(f'"=r"(r{i})' for i in range(n_regs))
+    dst_assigns = "\n".join(f"    *(uint32_t*)dst{i} = r{i};" for i in range(n_regs))
+    name = f"ptx_ldmatrix_{num}_{dtype.replace('::', '_').replace('.', '_')}_{1 if trans_b else 0}"
+    sig = "(void* smem_ptr, " + ", ".join(f"void* dst{i}" for i in range(n_regs)) + ")"
+    body = (
+        f"    uint32_t {reg_decls};\n"
+        "    unsigned int addr = __cvta_generic_to_shared(smem_ptr);\n"
+        "    asm volatile(\n"
+        f'        "ldmatrix.sync.aligned.m8n8.x{num}{trans_inst}.shared.{dtype} '
+        f'{slot_list}, [%{n_regs}];"\n'
+        f"        : {out_constraints}\n"
+        f'        : "r"(addr));\n'
+        f"{dst_assigns}"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "_ptx_ldmatrix_impl",
+    n_attrs=3,
+    c_signature=lambda *a: _ldmatrix_parts(*a)[1],
+    helper_name=lambda *a: _ldmatrix_parts(*a)[0],
+    body=lambda *a: _ldmatrix_parts(*a)[2],
+)
+
+
+@register_codegen("ptx_ldmatrix")
+def codegen_ptx_ldmatrix(trans, num, dtype, smem_ptr, *dst_handles):
+    trans = bool(trans)
+    num = int(num)
+    dtype = parse_str(dtype)
+    if dtype.startswith("."):
+        dtype = dtype[1:]
+    n_regs = num if dtype == "b16" else num // 2
+    if len(dst_handles) != n_regs:
+        raise ValueError(
+            f"ldmatrix .x{num}.{dtype} codegen expects {n_regs} dst handles, got {len(dst_handles)}"
+        )
+    result = CODEGEN_REGISTRY["tirx._ptx_ldmatrix_impl"](
+        [smem_ptr, *dst_handles, num, dtype, trans]
+    )
+    return result[0] if isinstance(result, tuple) else result
+
+
+def _stmatrix_parts(*args):
+    # args = (smem_ptr, src0, src1, ..., src{N-1}, trans, num, dtype, shape, space)
+    # The last 5 entries are codegen attrs (n_attrs=5).
+    trans_arg, num_arg, dtype_arg, shape_arg, space_arg = args[-5:]
+    n_regs = int(num_arg)
+    trans_b = bool(int(trans_arg)) if hasattr(trans_arg, "value") else bool(trans_arg)
+    dtype = parse_str(dtype_arg)
+    shape = parse_str(shape_arg)
+    space = parse_str(space_arg)
+    if dtype.startswith("."):
+        dtype = dtype[1:]
+    if n_regs not in (1, 2, 4):
+        raise ValueError(f"stmatrix .num must be one of {{1, 2, 4}}, got {n_regs}")
+    if dtype not in ("b16", "b8"):
+        raise ValueError(f"stmatrix .type must be b16 or b8, got {dtype!r}")
+    if shape not in ("m8n8", "m16n8"):
+        raise ValueError(f"stmatrix .shape must be m8n8 or m16n8, got {shape!r}")
+    if space not in ("shared", "shared::cta"):
+        raise ValueError(f"stmatrix state space must be shared or shared::cta, got {space!r}")
+    if shape == "m16n8" and not trans_b:
+        raise ValueError("stmatrix .m16n8 requires .trans")
+    trans_inst = ".trans" if trans_b else ""
+    slot_list = "{" + ", ".join(f"%{i}" for i in range(n_regs)) + "}"
+    src_loads = "\n".join(f"    uint32_t r{i} = *(uint32_t*)src{i};" for i in range(n_regs))
+    in_constraints = ", ".join(f'"r"(r{i})' for i in range(n_regs))
+    name = f"ptx_stmatrix_{shape}_{n_regs}_{1 if trans_b else 0}_{space.replace('::', '_')}_{dtype}"
+    sig = "(void* smem_ptr, " + ", ".join(f"void* src{i}" for i in range(n_regs)) + ")"
+    body = (
+        f"{src_loads}\n"
+        "    unsigned int addr = __cvta_generic_to_shared(smem_ptr);\n"
+        "    asm volatile(\n"
+        f'        "stmatrix.sync.aligned.{shape}.x{n_regs}{trans_inst}.{space}.{dtype} '
+        f'[%{n_regs}], {slot_list};"\n'
+        "        :\n"
+        f'        : {in_constraints}, "r"(addr));'
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "_ptx_stmatrix_impl",
+    n_attrs=5,
+    c_signature=lambda *a: _stmatrix_parts(*a)[1],
+    helper_name=lambda *a: _stmatrix_parts(*a)[0],
+    body=lambda *a: _stmatrix_parts(*a)[2],
+)
+
+
+@register_codegen("ptx_stmatrix")
+def codegen_ptx_stmatrix(trans, num, dtype, shape, space, smem_ptr, *src_handles):
+    trans = bool(trans)
+    num = int(num)
+    dtype_str = parse_str(dtype)
+    if dtype_str.startswith("."):
+        dtype_str = dtype_str[1:]
+    n_regs = num
+    if len(src_handles) != n_regs:
+        raise ValueError(
+            f"stmatrix .x{num}.{dtype_str} codegen expects {n_regs} src handles, "
+            f"got {len(src_handles)}"
+        )
+    result = CODEGEN_REGISTRY["tirx._ptx_stmatrix_impl"](
+        [smem_ptr, *src_handles, trans, num, dtype, shape, space]
+    )
+    return result[0] if isinstance(result, tuple) else result

python/tvm/backend/cuda/operator/intrinsics/nvshmem.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=redefined-builtin, invalid-name
+"""NVSHMEM intrinsics. Each backend call is one ``device_intrinsic(...)``."""
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+
+_NVSHMEM = ("nvshmem",)
+
+# =============================================================================
+# No-arg helpers: PE queries, quiet, fence, barrier_all.
+# =============================================================================
+for _op, _call, _ret, _tvm_ret in [
+    ("nvshmem_my_pe", "nvshmem_my_pe", "int32_t", "int32"),
+    ("nvshmem_n_pes", "nvshmem_n_pes", "int32_t", "int32"),
+    ("nvshmem_quiet", "nvshmem_quiet", "void", None),
+    ("nvshmem_fence", "nvshmem_fence", "void", None),
+    ("nvshmem_barrier_all", "nvshmem_barrier_all", "void", None),
+]:
+    device_intrinsic(
+        _op,
+        body=("    " + (f"return {_call}();" if _ret != "void" else f"{_call}();")),
+        return_type=_ret,
+        tvm_return_type=_tvm_ret,
+        extra_deps=_NVSHMEM,
+    )
+del _op, _call, _ret, _tvm_ret
+
+
+# =============================================================================
+# RMA get/put (thread/warp/block).
+# =============================================================================
+_RMA_SIG = "(void *dest, const void *source, size_t nelems, int pe)"
+for _op, _backend_call in [
+    ("nvshmem_getmem_nbi", "nvshmem_getmem_nbi"),
+    ("nvshmem_putmem_nbi", "nvshmem_putmem_nbi"),
+    ("nvshmem_getmem_nbi_warp", "nvshmemx_getmem_nbi_warp"),
+    ("nvshmem_putmem_nbi_warp", "nvshmemx_putmem_nbi_warp"),
+    ("nvshmem_getmem_nbi_block", "nvshmemx_getmem_nbi_block"),
+    ("nvshmem_putmem_nbi_block", "nvshmemx_putmem_nbi_block"),
+]:
+    device_intrinsic(
+        _op,
+        c_signature=_RMA_SIG,
+        body=f"    {_backend_call}(dest, source, nelems, pe);",
+        extra_deps=_NVSHMEM,
+    )
+del _op, _backend_call
+
+
+# =============================================================================
+# Signal / wait_until — each backend call is one device_intrinsic. String
+# attrs (sig_op, cmp) are mapped to NVSHMEM integer constants in the
+# user-facing dispatcher below.
+# =============================================================================
+
+_SIG_OP_VAL = {"set": 0, "add": 1}
+_CMP_VAL = {"eq": 0, "ne": 1, "gt": 2, "ge": 3, "lt": 4, "le": 5}
+
+
+def _resolve_attr(value, table, label):
+    s = value if isinstance(value, str) else value.value
+    if s not in table:
+        raise ValueError(f"Unsupported {label}: {s}")
+    return table[s]
+
+
+device_intrinsic(
+    "_nvshmem_signal_op_impl",
+    helper_name="tvm_builtin_nvshmem_signal_op",
+    c_signature="(uint64_t* sig_addr, uint64_t signal, int sig_op, int pe)",
+    body="    nvshmemx_signal_op(sig_addr, signal, sig_op, pe);",
+    extra_deps=_NVSHMEM,
+)
+
+
+@register_codegen("nvshmem_signal_op")
+def codegen_nvshmem_signal_op(sig_addr, signal, sig_op, pe):
+    """Map ``sig_op`` (string) to its NVSHMEM int constant, then forward."""
+    sig_op_int = _resolve_attr(sig_op, _SIG_OP_VAL, "signal op")
+    result = CODEGEN_REGISTRY["tirx._nvshmem_signal_op_impl"]([sig_addr, signal, sig_op_int, pe])
+    return result
+
+
+# nvshmem_<type>_wait_until — one device_intrinsic per supported type.
+_WAIT_UNTIL_TYPES = {"uint64_t": "uint64", "uint64": "uint64"}
+
+for _c_type, _suffix in [("uint64_t", "uint64")]:
+    device_intrinsic(
+        f"_nvshmem_{_suffix}_wait_until_impl",
+        helper_name=f"tvm_builtin_nvshmem_{_suffix}_wait_until",
+        c_signature=f"({_c_type}* ivar, int cmp, {_c_type} cmp_value)",
+        body=f"    nvshmem_{_suffix}_wait_until(ivar, cmp, cmp_value);",
+        extra_deps=_NVSHMEM,
+    )
+del _c_type, _suffix
+
+
+@register_codegen("nvshmem_wait_until")
+def codegen_nvshmem_wait_until(ivar, cmp, cmp_value, type):
+    """Dispatch to the type-specific wait_until helper after mapping ``cmp``
+    (string) to its NVSHMEM int constant."""
+    type_str = type if isinstance(type, str) else type.value
+    if type_str not in _WAIT_UNTIL_TYPES:
+        raise ValueError(f"Unsupported type for nvshmem_wait_until: {type_str}")
+    suffix = _WAIT_UNTIL_TYPES[type_str]
+    cmp_int = _resolve_attr(cmp, _CMP_VAL, "cmp operation")
+    result = CODEGEN_REGISTRY[f"tirx._nvshmem_{suffix}_wait_until_impl"]([ivar, cmp_int, cmp_value])
+    return result
+
+
+# putmem_signal_nbi (thread / warp / block) — three scope-specific helpers.
+_PUTMEM_SIG_SIG = (
+    "(void* dest, const void* source, size_t nelems, "
+    "uint64_t* sig_addr, uint64_t signal, int sig_op, int pe)"
+)
+for _scope_suffix, _backend_call in [
+    ("", "nvshmem_putmem_signal_nbi"),
+    ("_warp", "nvshmemx_putmem_signal_nbi_warp"),
+    ("_block", "nvshmemx_putmem_signal_nbi_block"),
+]:
+    device_intrinsic(
+        f"_nvshmem_putmem_signal_nbi{_scope_suffix}_impl",
+        helper_name=f"tvm_builtin_nvshmem_putmem_signal_nbi{_scope_suffix}",
+        c_signature=_PUTMEM_SIG_SIG,
+        body=f"    {_backend_call}(dest, source, nelems, sig_addr, signal, sig_op, pe);",
+        extra_deps=_NVSHMEM,
+    )
+del _scope_suffix, _backend_call
+
+
+def _make_putmem_signal_dispatcher(scope_suffix):
+    @register_codegen(f"nvshmem_putmem_signal_nbi{scope_suffix}")
+    def _codegen(dest, source, nelems, sig_addr, signal, sig_op, pe):
+        sig_op_int = _resolve_attr(sig_op, _SIG_OP_VAL, "signal op")
+        result = CODEGEN_REGISTRY[f"tirx._nvshmem_putmem_signal_nbi{scope_suffix}_impl"](
+            [dest, source, nelems, sig_addr, signal, sig_op_int, pe]
+        )
+        return result
+
+    return _codegen
+
+
+for _suffix in ("", "_warp", "_block"):
+    _make_putmem_signal_dispatcher(_suffix)
+del _suffix

python/tvm/backend/cuda/operator/intrinsics/registry.py

@@ -0,0 +1,77 @@
+# 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.
+"""Codegen registry for CUDA HW ops.
+
+User-facing Python wrappers are hand-written in :mod:`tvm.tirx.op` so that
+editors / static analyzers (Cursor, Pyright) can see their signatures. This
+module only handles the backend codegen side.
+"""
+
+import functools
+
+import tvm_ffi
+
+CODEGEN_REGISTRY = {}
+
+
+def _canonical_device_intrin_name(op_name: str) -> str:
+    if not op_name.startswith("tirx."):
+        return op_name
+    basename = op_name[len("tirx.") :]
+    if "." in basename:
+        return op_name
+    for prefix, namespace in (
+        ("cuda_", "cuda"),
+        ("ptx_", "ptx"),
+        ("nvshmem_", "nvshmem"),
+        ("nki_", "nki"),
+    ):
+        if basename.startswith(prefix):
+            return f"tirx.{namespace}.{basename[len(prefix) :]}"
+    return op_name
+
+
+@tvm_ffi.register_global_func("tirx.intrinsics.cuda.get_codegen")
+def get_codegen(op):
+    """get the codegen function for a given op"""
+    return CODEGEN_REGISTRY.get(op, None)
+
+
+def register_codegen(op, backend="cuda"):
+    """Register a codegen function for a given op.
+
+    The codegen function should return a ``cuda_func_call`` statement, and
+    optionally a list of tags that the codegen function needs.
+    """
+
+    def decorator(func):
+        full_op_name = "tirx." + op
+        canonical_op_name = _canonical_device_intrin_name(full_op_name)
+        op_names = {full_op_name, canonical_op_name}
+
+        @functools.wraps(func)
+        def wrapper(arg_list):
+            res = func(*arg_list)  # pylint: disable=not-callable
+            if isinstance(res, tuple):
+                return res[0], res[1]
+            return res, list()
+
+        for op_name in op_names:
+            CODEGEN_REGISTRY[op_name] = wrapper
+        return wrapper
+
+    return decorator

python/tvm/backend/cuda/operator/intrinsics/sync.py

@@ -0,0 +1,570 @@
+# 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
+"""Synchronization primitives.
+
+PTX side:
+* ``bar.arrive`` / ``bar.sync`` — named-barrier alias of ``barrier.arrive/sync``
+* ``fence{.sem}.scope`` / ``fence.proxy.async`` / ``fence.mbarrier_init``
+* ``barrier.cluster.arrive`` / ``barrier.cluster.wait``
+* ``mbarrier.init`` / ``mbarrier.arrive[.expect_tx]`` (local + remote) / ``mbarrier.try_wait``
+* ``elect.sync``  — warp leader election
+* warp-vote ``__any_sync``
+
+CUDA-side helpers:
+* ``__threadfence`` / ``__syncwarp`` / ``__syncthreads`` / ``__syncthreads_and|or``
+* cooperative-groups grid sync
+* cluster sync (open-coded ``barrier.cluster.arrive/wait`` pair)
+* warpgroup sync (``bar.sync``)
+"""
+
+from tvm.tirx.operator.intrinsics._common import (
+    CLUSTER_BARRIER_SEM,
+    FENCE_PROXY_ASYNC_SPACE,
+    FENCE_SCOPE,
+    FENCE_SEM,
+)
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+from .utils import parse_str
+
+# =============================================================================
+# bar.arrive / bar.sync — alias of barrier.arrive/sync. 1 form each.
+#   bar.sync   a, b ;
+#   bar.arrive a, b ;
+# =============================================================================
+device_intrinsic(
+    "ptx_bar_arrive",
+    c_signature="(int name_bar_id, int thread_count)",
+    body=(
+        '    asm volatile("bar.arrive %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory");'
+    ),
+)
+device_intrinsic(
+    "ptx_bar_sync",
+    c_signature="(int name_bar_id, int thread_count)",
+    body=(
+        '    asm volatile("bar.sync %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory");'
+    ),
+)
+
+
+# =============================================================================
+# fence{.sem}.scope — 1 form (sem/scope are modifier values).
+# =============================================================================
+def _ptx_fence(sem, scope):
+    sem, scope = parse_str(sem), parse_str(scope)
+    assert sem in FENCE_SEM, f"invalid fence sem {sem!r}, expected one of {FENCE_SEM}"
+    assert scope in FENCE_SCOPE, f"invalid fence scope {scope!r}, expected one of {FENCE_SCOPE}"
+    return (
+        f"tvm_builtin_ptx_fence_{sem}_{scope}",
+        f'    asm volatile("fence.{sem}.{scope};" ::: "memory");',
+    )
+
+
+device_intrinsic(
+    "ptx_fence",
+    n_attrs=2,
+    helper_name=lambda sem, scope: _ptx_fence(sem, scope)[0],
+    body=lambda sem, scope: _ptx_fence(sem, scope)[1],
+)
+
+
+# =============================================================================
+# fence.proxy.async{.<space>} — 1 form, optional .space modifier.
+# =============================================================================
+def _ptx_fence_proxy_async(space):
+    space = parse_str(space)
+    assert space in FENCE_PROXY_ASYNC_SPACE, (
+        f"invalid fence.proxy.async space {space!r}, expected one of {FENCE_PROXY_ASYNC_SPACE}"
+    )
+    suffix = f".{space}" if space else ""
+    name_safe = "_" + space.replace("::", "_").replace(".", "_") if space else ""
+    return (
+        f"tvm_builtin_ptx_fence_proxy_async{name_safe}",
+        f'    asm volatile("fence.proxy.async{suffix};" ::: "memory");',
+    )
+
+
+device_intrinsic(
+    "ptx_fence_proxy_async",
+    n_attrs=1,
+    helper_name=lambda space: _ptx_fence_proxy_async(space)[0],
+    body=lambda space: _ptx_fence_proxy_async(space)[1],
+)
+
+
+# =============================================================================
+# fence.mbarrier_init.release.cluster — 1 form, no operands.
+# =============================================================================
+device_intrinsic(
+    "ptx_fence_mbarrier_init",
+    body='    asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory");',
+)
+
+
+# =============================================================================
+# barrier.cluster.arrive{.sem}{.aligned} — 1 form.
+# =============================================================================
+def _ptx_barrier_cluster_arrive(sem, aligned):
+    sem = parse_str(sem)
+    aligned = bool(int(aligned)) if hasattr(aligned, "value") else bool(aligned)
+    assert sem in CLUSTER_BARRIER_SEM, (
+        f"invalid cluster.arrive sem {sem!r}, expected one of {CLUSTER_BARRIER_SEM}"
+    )
+    sem_suffix = f".{sem}" if sem else ""
+    aligned_suffix = ".aligned" if aligned else ""
+    name_sem = "_" + sem.replace("::", "_").replace(".", "_") if sem else ""
+    name_aligned = "_aligned" if aligned else ""
+    return (
+        f"tvm_builtin_ptx_barrier_cluster_arrive{name_sem}{name_aligned}",
+        f'    asm volatile("barrier.cluster.arrive{sem_suffix}{aligned_suffix};" ::: "memory");',
+    )
+
+
+device_intrinsic(
+    "ptx_barrier_cluster_arrive",
+    n_attrs=2,
+    helper_name=lambda sem, aligned: _ptx_barrier_cluster_arrive(sem, aligned)[0],
+    body=lambda sem, aligned: _ptx_barrier_cluster_arrive(sem, aligned)[1],
+)
+
+
+# =============================================================================
+# barrier.cluster.wait{.acquire}{.aligned} — 1 form.
+# =============================================================================
+def _ptx_barrier_cluster_wait(acquire, aligned):
+    acquire = bool(int(acquire)) if hasattr(acquire, "value") else bool(acquire)
+    aligned = bool(int(aligned)) if hasattr(aligned, "value") else bool(aligned)
+    acq_suffix = ".acquire" if acquire else ""
+    aligned_suffix = ".aligned" if aligned else ""
+    return (
+        f"tvm_builtin_ptx_barrier_cluster_wait"
+        f"{'_acquire' if acquire else ''}{'_aligned' if aligned else ''}",
+        f'    asm volatile("barrier.cluster.wait{acq_suffix}{aligned_suffix};" ::: "memory");',
+    )
+
+
+device_intrinsic(
+    "ptx_barrier_cluster_wait",
+    n_attrs=2,
+    helper_name=lambda acquire, aligned: _ptx_barrier_cluster_wait(acquire, aligned)[0],
+    body=lambda acquire, aligned: _ptx_barrier_cluster_wait(acquire, aligned)[1],
+)
+
+
+# =============================================================================
+# clusterlaunchcontrol.try_cancel / query_cancel — Blackwell Cluster Launch
+# Control (CLC) work-stealing, written from the PTX ISA spec (section
+# "clusterlaunchcontrol", PTX ISA 8.6). try_cancel async-requests cancelling the
+# next cluster's launch, writing a 16B response to smem + signalling mbar. query
+# decodes the response: on success it extracts the cancelled cluster's first
+# ctaid.x (via the get_first_ctaid::x form); a single uint32 is returned, with
+# 0xFFFFFFFF as the "no work stolen" sentinel (a device helper returns one scalar).
+# =============================================================================
+device_intrinsic(
+    "ptx_clc_try_cancel",
+    c_signature="(void* handle, void* mbar)",
+    body=(
+        "    unsigned int addr = (unsigned int)__cvta_generic_to_shared(handle);\n"
+        "    unsigned int bar = (unsigned int)__cvta_generic_to_shared(mbar);\n"
+        "    asm volatile(\n"
+        '        "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes"\n'
+        '        ".multicast::cluster::all.b128 [%0], [%1];\\n"\n'
+        '        :: "r"(addr), "r"(bar) : "memory");'
+    ),
+)
+
+
+device_intrinsic(
+    "ptx_clc_query_cancel",
+    c_signature="(void* handle)",
+    return_type="uint32_t",
+    tvm_return_type="uint32",
+    body=(
+        "    unsigned int addr = (unsigned int)__cvta_generic_to_shared(handle);\n"
+        "    unsigned int first_ctaid_x;\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred canceled;\\n"\n'
+        '        ".reg .b128 response;\\n"\n'
+        '        "ld.shared.b128 response, [%1];\\n"\n'
+        '        "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 canceled, response;\\n"\n'
+        '        "mov.u32 %0, 0xffffffff;\\n"\n'
+        '        "@canceled clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128"\n'
+        '        " %0, response;\\n"\n'
+        '        "}\\n"\n'
+        '        : "=r"(first_ctaid_x) : "r"(addr) : "memory");\n'
+        '    asm volatile("fence.proxy.async.shared::cta;\\n" ::: "memory");\n'
+        "    return first_ctaid_x;"
+    ),
+)
+
+
+# =============================================================================
+# mbarrier.init.shared.b64 [addr], count ; — 1 form.
+# =============================================================================
+device_intrinsic(
+    "ptx_mbarrier_init",
+    c_signature="(void* barrier, int thread_count)",
+    body=(
+        "    unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
+        '    asm volatile("mbarrier.init.shared.b64 [%0], %1;"'
+        ' : : "r"(barrier_addr), "r"(thread_count) : "memory");'
+    ),
+)
+
+
+# =============================================================================
+# mbarrier.arrive — local + remote (cluster-mapped) forms. 2 PTX forms.
+#   Form local:  mbarrier.arrive.shared.b64 _, [bar];
+#   Form remote: { setp+@p mapa.shared::cluster.u32 + @p mbarrier.arrive.shared::cluster.b64 }
+# Dispatcher picks by arg count (1 vs 3).
+# =============================================================================
+device_intrinsic(
+    "_ptx_mbarrier_arrive_local",
+    helper_name="tvm_builtin_ptx_mbarrier_arrive",
+    c_signature="(void* barrier)",
+    body=(
+        "    unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
+        '    asm volatile("mbarrier.arrive.shared.b64 _, [%0];"\n'
+        '                 :: "r"(barrier_addr) : "memory");'
+    ),
+)
+device_intrinsic(
+    "_ptx_mbarrier_arrive_remote",
+    helper_name="tvm_builtin_ptx_mbarrier_arrive_remote",
+    c_signature="(void* barrier, int cta_id, int pred)",
+    body=(
+        "    unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred p;\\n"\n'
+        '        ".reg .b32 remAddr32;\\n"\n'
+        '        "setp.ne.s32 p, %2, 0;\\n"\n'
+        '        "@p mapa.shared::cluster.u32  remAddr32, %0, %1;\\n"\n'
+        '        "@p mbarrier.arrive.shared::cluster.b64  _, [remAddr32];\\n"\n'
+        '        "}\\n"\n'
+        '        :: "r"(barrier_addr), "r"(cta_id), "r"(pred) : "memory");'
+    ),
+)
+
+
+# Same cross-CTA arrive, but with an explicit arrival-count operand
+# (``..., [remAddr32], count``). Matches the ``tma::cluster::arrive`` spelling.
+device_intrinsic(
+    "_ptx_mbarrier_arrive_remote_count",
+    helper_name="tvm_builtin_ptx_mbarrier_arrive_remote_count",
+    c_signature="(void* barrier, int cta_id, int pred, int count)",
+    body=(
+        "    unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred p;\\n"\n'
+        '        ".reg .b32 remAddr32;\\n"\n'
+        '        "setp.ne.s32 p, %2, 0;\\n"\n'
+        '        "@p mapa.shared::cluster.u32  remAddr32, %0, %1;\\n"\n'
+        '        "@p mbarrier.arrive.shared::cluster.b64  _, [remAddr32], %3;\\n"\n'
+        '        "}\\n"\n'
+        '        :: "r"(barrier_addr), "r"(cta_id), "r"(pred), "r"(count) : "memory");'
+    ),
+)
+
+
+@register_codegen("ptx_mbarrier_arrive")
+def _codegen_mbarrier_arrive(*args):
+    """Dispatch by arg count: 1 -> local, 3 -> remote, 4 -> remote+count."""
+    if len(args) == 1:
+        result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_local"](list(args))
+    elif len(args) == 3:
+        result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_remote"](list(args))
+    elif len(args) == 4:
+        result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_remote_count"](list(args))
+    else:
+        raise ValueError(f"ptx_mbarrier_arrive expects 1, 3, or 4 args, got {len(args)}")
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# mbarrier.arrive.expect_tx — local + remote (cluster-mapped) forms.
+# =============================================================================
+device_intrinsic(
+    "_ptx_mbarrier_arrive_expect_tx_local",
+    helper_name="tvm_builtin_ptx_mbarrier_arrive_expect_tx",
+    c_signature="(void* barrier, int byte_count)",
+    body=(
+        "    unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
+        '    asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"\n'
+        '                 :: "r"(barrier_addr), "r"(byte_count) : "memory");'
+    ),
+)
+device_intrinsic(
+    "_ptx_mbarrier_arrive_expect_tx_remote",
+    helper_name="tvm_builtin_ptx_mbarrier_arrive_expect_tx_remote",
+    c_signature="(void* barrier, int cta_id, int pred, int byte_count)",
+    body=(
+        "    unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred p;\\n"\n'
+        '        ".reg .b32 remAddr32;\\n"\n'
+        '        "setp.ne.s32 p, %2, 0;\\n"\n'
+        '        "@p mapa.shared::cluster.u32  remAddr32, %0, %1;\\n"\n'
+        '        "@p mbarrier.arrive.expect_tx.shared::cluster.b64  _, [remAddr32], %3;\\n"\n'
+        '        "}\\n"\n'
+        '        :: "r"(barrier_addr), "r"(cta_id), "r"(pred), "r"(byte_count) : "memory");'
+    ),
+)
+
+
+@register_codegen("ptx_mbarrier_arrive_expect_tx")
+def _codegen_mbarrier_arrive_expect_tx(*args):
+    """Dispatch by arg count: 2 -> local, 4 -> remote. Remote arg order from
+    the user is (bar, byte_count, cta_id, pred); reorder to match the helper
+    signature (bar, cta_id, pred, byte_count)."""
+    if len(args) == 2:
+        result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_expect_tx_local"](list(args))
+    elif len(args) == 4:
+        bar, byte_count, cta_id, pred = args
+        result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_expect_tx_remote"](
+            [bar, cta_id, pred, byte_count]
+        )
+    else:
+        raise ValueError(f"ptx_mbarrier_arrive_expect_tx expects 2 or 4 args, got {len(args)}")
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# mbarrier.try_wait.parity.shared::cta.b64 — 1 form. Body wraps the asm in a
+# label loop (TIRx convention; the magic ``ticks = 0x989680`` is the timeout
+# hint in ns).
+# =============================================================================
+device_intrinsic(
+    "ptx_mbarrier_try_wait",
+    c_signature="(void* barrier, int phase)",
+    body=(
+        "    unsigned int barrier_addr_int = __cvta_generic_to_shared(barrier);\n"
+        "    unsigned int ticks = 0x989680;\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred                P1;\\n"\n'
+        '        "LAB_WAIT:\\n"\n'
+        '        "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2;\\n"\n'
+        '        "@P1                       bra.uni DONE;\\n"\n'
+        '        "bra.uni                   LAB_WAIT;\\n"\n'
+        '        "DONE:\\n"\n'
+        '        "}\\n"\n'
+        '        :: "r"(barrier_addr_int), "r"(phase), "r"(ticks) : "memory");'
+    ),
+)
+
+
+# mbarrier.try_wait.parity.acquire.cluster — cluster-scope acquire wait used for
+# cross-CTA barrier handshakes (e.g. the tmem-finished handoff).
+device_intrinsic(
+    "ptx_mbarrier_try_wait_acquire_cluster",
+    c_signature="(void* barrier, int phase)",
+    body=(
+        "    unsigned int barrier_addr_int = __cvta_generic_to_shared(barrier);\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred                P1;\\n"\n'
+        '        "LAB_WAIT_AC:\\n"\n'
+        '        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\\n"\n'
+        '        "@P1                       bra.uni DONE_AC;\\n"\n'
+        '        "bra.uni                   LAB_WAIT_AC;\\n"\n'
+        '        "DONE_AC:\\n"\n'
+        '        "}\\n"\n'
+        '        :: "r"(barrier_addr_int), "r"(phase) : "memory");'
+    ),
+)
+
+
+# =============================================================================
+# mbarrier.try_wait.parity — ONE-SHOT non-blocking variant. Returns true
+# if the requested parity has already been reached, false otherwise.
+# The TIRx-standard ``ptx_mbarrier_try_wait`` above wraps this in a
+# label loop that retries until success; this one-shot form is the
+# building block for bounded-retry debug waits (Nymph's
+# ``debug_bounded_wait`` lowering mode wraps it in a Python-counted
+# loop so the kernel cannot hang forever at a mis-protocoled wait).
+# =============================================================================
+device_intrinsic(
+    "ptx_mbarrier_try_wait_once",
+    c_signature="(void* barrier, int phase, int ticks)",
+    return_type="uint32_t",
+    body=(
+        "    unsigned int barrier_addr_int = __cvta_generic_to_shared(barrier);\n"
+        "    unsigned int ticks_u = (unsigned int)ticks;\n"
+        "    unsigned int result;\n"
+        "    asm volatile(\n"
+        '        "{\\n"\n'
+        '        ".reg .pred                P1;\\n"\n'
+        '        "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2, %3;\\n"\n'
+        '        "selp.u32                  %0, 1, 0, P1;\\n"\n'
+        '        "}\\n"\n'
+        '        : "=r"(result) : "r"(barrier_addr_int), "r"(phase), "r"(ticks_u) : "memory");\n'
+        "    return result;"
+    ),
+)
+
+
+# =============================================================================
+# elect.sync — TIRx uses the CUDA builtin ``tvm_builtin_elect_one_sync()``
+# helper (declared in the CUDA header tags), not direct PTX.
+# =============================================================================
+device_intrinsic(
+    "ptx_elect_sync",
+    helper_name="tvm_builtin_elect_one_sync_op",
+    return_type="uint32_t",
+    body="    return tvm_builtin_elect_one_sync();",
+    extra_deps=("elect_one_sync",),
+)
+
+
+# =============================================================================
+# __any_sync — warp-vote (pure CUDA helper).
+# =============================================================================
+device_intrinsic(
+    "ptx_any_sync",
+    c_signature="(unsigned mask, int pred)",
+    body="    return __any_sync(mask, pred);",
+    return_type="int",
+    tvm_return_type="int32",
+)
+
+
+# =============================================================================
+# CUDA-side sync helpers (zero-arg void unless noted).
+# =============================================================================
+device_intrinsic("cuda_thread_fence", body="    __threadfence();")
+device_intrinsic("cuda_warp_sync", body="    __syncwarp();")
+device_intrinsic("cuda_cta_sync", body="    __syncthreads();")
+device_intrinsic(
+    "cuda_grid_sync",
+    body="    namespace cg = cooperative_groups;\n    cg::this_grid().sync();",
+    extra_deps=("cooperative_groups",),
+)
+device_intrinsic(
+    "cuda_cluster_sync",
+    body=('    asm("barrier.cluster.arrive.aligned;");\n    asm("barrier.cluster.wait.aligned;");'),
+)
+device_intrinsic(
+    "cuda_warpgroup_sync",
+    c_signature="(int name_bar_id)",
+    body='    asm volatile("bar.sync %0, 128;" : : "r"(name_bar_id));',
+)
+device_intrinsic(
+    "cuda_syncthreads_and",
+    c_signature="(int predicate)",
+    body="    return __syncthreads_and(predicate);",
+    return_type="int",
+    tvm_return_type="int32",
+)
+device_intrinsic(
+    "cuda_syncthreads_or",
+    c_signature="(int predicate)",
+    body="    return __syncthreads_or(predicate);",
+    return_type="int",
+    tvm_return_type="int32",
+)
+
+
+# =============================================================================
+# Additional mbarrier, grid-sync, and warp collective helpers.
+# =============================================================================
+
+
+# PTX mbarrier parity wait form:
+#   mbarrier.test_wait.parity{.sem.scope}{.shared{::cta}}.b64 waitComplete, [addr], phaseParity;
+def _mbarrier_test_wait_parity_parts(_barrier, _phase, sem, scope, space):
+    sem = parse_str(sem)
+    scope = parse_str(scope)
+    space = parse_str(space)
+    if sem and sem not in ("acquire", "relaxed"):
+        raise ValueError(f"Unsupported mbarrier.test_wait.parity sem {sem!r}")
+    if scope and scope not in ("cta", "cluster"):
+        raise ValueError(f"Unsupported mbarrier.test_wait.parity scope {scope!r}")
+    if space not in ("shared", "shared::cta"):
+        raise ValueError(f"Unsupported mbarrier.test_wait.parity space {space!r}")
+    sem_scope = f".{sem}.{scope}" if sem else ""
+    name = (
+        "tvm_builtin_ptx_mbarrier_test_wait_parity"
+        f"{('_' + sem + '_' + scope) if sem else ''}_{space.replace('::', '_')}_b64"
+    )
+    body = (
+        "    unsigned int ready = 0;\n"
+        "    asm volatile(\n"
+        '        "{\\n\\t"\n'
+        '        ".reg .pred P1; \\n\\t"\n'
+        f'        "mbarrier.test_wait.parity{sem_scope}.{space}.b64 P1, [%1], %2; \\n\\t"\n'
+        '        "selp.b32 %0, 1, 0, P1; \\n\\t"\n'
+        '        "}" : "=r"(ready) : "r"((unsigned int)__cvta_generic_to_shared(barrier)), '
+        '"r"(phase) : "memory");\n'
+        "    return ready;"
+    )
+    return name, body
+
+
+device_intrinsic(
+    "ptx_mbarrier_test_wait_parity",
+    n_attrs=3,
+    helper_name=lambda *a: _mbarrier_test_wait_parity_parts(*a)[0],
+    c_signature="(void* barrier, int phase)",
+    return_type="unsigned int",
+    tvm_return_type="uint32",
+    body=lambda *a: _mbarrier_test_wait_parity_parts(*a)[1],
+)
+
+device_intrinsic(
+    "cuda_ballot_sync",
+    helper_name="tvm_builtin_ballot_sync",
+    c_signature="(unsigned int mask, int pred)",
+    return_type="unsigned int",
+    body="    return __ballot_sync(mask, pred);",
+)
+device_intrinsic(
+    "cuda_reduce_add_sync_u32",
+    helper_name="tvm_builtin_reduce_add_sync_u32",
+    c_signature="(unsigned int mask, unsigned int value)",
+    return_type="unsigned int",
+    body="    return __reduce_add_sync(mask, value);",
+)
+device_intrinsic(
+    "cuda_reduce_min_sync_u32",
+    helper_name="tvm_builtin_reduce_min_sync_u32",
+    c_signature="(unsigned int mask, unsigned int value)",
+    return_type="unsigned int",
+    body="    return __reduce_min_sync(mask, value);",
+)
+
+
+# =============================================================================
+# griddepcontrol.wait / griddepcontrol.launch_dependents (sm_90+)
+# Programmatic Dependent Launch (PDL) synchronization. Both carry memory
+# clobber to prevent CSE / cross-barrier reordering.
+# =============================================================================
+device_intrinsic(
+    "ptx_griddepcontrol_wait",
+    body='    asm volatile("griddepcontrol.wait;" ::: "memory");',
+)
+
+device_intrinsic(
+    "ptx_griddepcontrol_launch_dependents",
+    body='    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");',
+)

python/tvm/backend/cuda/operator/intrinsics/types.py

@@ -0,0 +1,71 @@
+# 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.
+"""PTX data types for CUDA codegen."""
+
+import enum
+
+import tvm_ffi
+
+from_string_func = tvm_ffi.get_global_func("tirx.intrinsics.cuda.PTXDTypeFromString")
+to_string_func = tvm_ffi.get_global_func("tirx.intrinsics.cuda.PTXDTypeToString")
+
+
+class PTXDataType(enum.Enum):
+    """
+    A Python equivalent of the provided C++ DataType enum class.
+
+    Inherits from IntEnum so that members behave both as enum members
+    and as integers, mirroring the C++ behavior.
+
+    see also src/target/source/ptx.cc
+    """
+
+    INT4 = 0
+    UINT4 = 1
+    INT8 = 2
+    UINT8 = 3
+    INT16 = 4
+    UINT16 = 5
+    INT32 = 6
+    UINT32 = 7
+    INT64 = 8
+    UINT64 = 9
+    FLOAT4_E2M1FN = 10
+    FLOAT6_E2M3FN = 11
+    FLOAT6_E3M2FN = 12
+    FLOAT8_E4M3FN = 13
+    FLOAT8_E4M3FNUZ = 14
+    FLOAT8_E5M2 = 15
+    FLOAT8_E8M0FNU = 16
+    FLOAT16 = 17
+    BFLOAT16 = 18
+    FLOAT16X2 = 19
+    FLOAT32 = 20
+    TENSOR_FLOAT32 = 21
+    FLOAT64 = 22
+    BIT1 = 23
+    BIT8 = 24
+    BIT16 = 25
+    BIT32 = 26
+    BIT64 = 27
+
+    @classmethod
+    def from_string(cls, s_type: str) -> "PTXDataType":
+        return PTXDataType(from_string_func(s_type))
+
+    def to_string(self) -> str:
+        return to_string_func(self.value)

python/tvm/backend/cuda/operator/intrinsics/utils.py

@@ -0,0 +1,82 @@
+# 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 utility functions for CUDA op codegen."""
+
+
+def parse_str(arg) -> str:
+    """Parse TIR StringImm or Python str to a plain str.
+
+    TIR StringImm values stringify to quoted strings, e.g., ``'"float16"'``;
+    Python strs do not. Idempotent — passing an already-parsed str returns it
+    unchanged, so dispatchers that parse once before forwarding to inner
+    codegens won't double-strip the value.
+    """
+    s = str(arg)
+    if len(s) >= 2 and s[0] == '"' and s[-1] == '"':
+        return s[1:-1]
+    return s
+
+
+def is_power_of_two(n: int) -> bool:
+    """Check if n is a power of two."""
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def validate_cta_group(cta_group, context: str = "") -> int:
+    """Validate that cta_group is 1 or 2 and return it as int.
+
+    Args:
+        cta_group: The cta_group value (can be int or TIR IntImm)
+        context: Optional context string for error message (e.g., "allocating Tensor Memory")
+
+    Returns:
+        The validated cta_group as int
+
+    Raises:
+        ValueError: If cta_group is not 1 or 2
+    """
+    cta_group = int(cta_group)
+    if cta_group not in [1, 2]:
+        ctx = f" involved in {context}" if context else ""
+        raise ValueError(
+            f"The number of cta_group{ctx} is incorrect, expected 1 or 2, got {cta_group}"
+        )
+    return cta_group
+
+
+def validate_power_of_two_range(value, min_val: int, max_val: int, name: str) -> int:
+    """Validate that value is within range and is a power of two.
+
+    Args:
+        value: The value to validate
+        min_val: Minimum allowed value (inclusive)
+        max_val: Maximum allowed value (inclusive)
+        name: Name of the parameter for error messages
+
+    Returns:
+        The validated value as int
+
+    Raises:
+        ValueError: If value is out of range or not a power of two
+    """
+    value = int(value)
+    if not (min_val <= value <= max_val and is_power_of_two(value)):
+        raise ValueError(
+            f"The {name} is invalid, expect a value within range [{min_val}, {max_val}] "
+            f"and be a power of 2, got {value}"
+        )
+    return value

python/tvm/backend/cuda/operator/intrinsics/wgmma.py

@@ -0,0 +1,403 @@
+# 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, too-many-arguments, too-many-locals, too-many-positional-arguments
+"""PTX WGMMA operations (Hopper warpgroup MMA).
+
+One ``device_intrinsic`` registration per PTX form table entry. Bodies are
+hand-written ``asm volatile(...)`` strings. Variable-arity register vectors
+(``mma_async`` accumulators / A fragments) materialize via the same
+device_intrinsic with an attr-driven ``c_signature`` callable.
+"""
+
+import tvm
+
+from ._schema import device_intrinsic
+from .registry import CODEGEN_REGISTRY, register_codegen
+from .types import PTXDataType
+from .utils import parse_str
+
+# =============================================================================
+# wgmma.fence / commit_group / wait_group — one PTX form each.
+# =============================================================================
+device_intrinsic(
+    "ptx_wgmma_fence",
+    helper_name="ptx_wgmma_fence",
+    body='    asm volatile("wgmma.fence.sync.aligned;" ::: "memory");',
+)
+device_intrinsic(
+    "ptx_wgmma_commit_group",
+    helper_name="ptx_wgmma_commit_group",
+    body='    asm volatile("wgmma.commit_group.sync.aligned;" ::: "memory");',
+)
+device_intrinsic(
+    "ptx_wgmma_wait_group",
+    n_attrs=1,
+    helper_name=lambda n: f"ptx_wgmma_wait_group_{int(n)}",
+    body=lambda n: f'    asm volatile("wgmma.wait_group.sync.aligned {int(n)};" ::: "memory");',
+)
+
+
+# =============================================================================
+# wgmma_encode_matrix_descriptor — pure-C bitfield struct fill (no asm).
+# =============================================================================
+device_intrinsic(
+    "ptx_wgmma_encode_matrix_descriptor",
+    helper_name="ptx_wgmma_encode_matrix_descriptor",
+    c_signature="(uint64_t* desc, void* addr, int ldo, int sdo, int swizzle)",
+    body=(
+        "  GmmaDescriptor _desc{};  // value-init: reading uncovered pad bits is UB\n"
+        "\n"
+        "  switch (swizzle) {\n"
+        "    case 0: _desc.bitfield.layout_type_ = uint8_t(0); break; // No swizzle\n"
+        "    case 1: _desc.bitfield.layout_type_ = uint8_t(3); break; // 32B swizzle\n"
+        "    case 2: _desc.bitfield.layout_type_ = uint8_t(2); break; // 64B swizzle\n"
+        "    case 3: _desc.bitfield.layout_type_ = uint8_t(1); break; // 128B swizzle\n"
+        "  }\n"
+        "\n"
+        "  uint32_t start_address = __cvta_generic_to_shared(addr);\n"
+        "  _desc.bitfield.start_address_ = static_cast<uint16_t>(start_address >> 4);\n"
+        "\n"
+        "  constexpr uint8_t base_offset = 0;\n"
+        "  _desc.bitfield.base_offset_ = base_offset;\n"
+        "\n"
+        "  _desc.bitfield.stride_byte_offset_  = static_cast<uint32_t>(sdo);\n"
+        "  _desc.bitfield.leading_byte_offset_ = static_cast<uint32_t>(ldo);\n"
+        "\n"
+        "  *desc = (uint64_t)_desc;"
+    ),
+    extra_deps=("gmma_descriptor",),
+)
+
+
+# =============================================================================
+# wgmma_noop_barrier — empty asm with one inout register operand. Two
+# device_intrinsic calls, one per supported dtype; dispatcher picks the form
+# based on the operand's runtime dtype.
+# =============================================================================
+device_intrinsic(
+    "ptx_wgmma_noop_barrier_uint32",
+    helper_name="ptx_wgmma_fence_uint32_t",
+    c_signature="(uint32_t reg)",
+    body='    asm volatile("" : "+r"(reg) :: "memory");',
+)
+device_intrinsic(
+    "ptx_wgmma_noop_barrier_float32",
+    helper_name="ptx_wgmma_fence_float",
+    c_signature="(float reg)",
+    body='    asm volatile("" : "+f"(reg) :: "memory");',
+)
+
+
+@register_codegen("ptx_wgmma_noop_barrier")
+def codegen_ptx_wgmma_noop_barrier(reg):
+    dtype = str(reg.dtype)
+    dtype_enum = PTXDataType.from_string(dtype)
+    if dtype_enum == PTXDataType.UINT32:
+        op_name = "tirx.ptx_wgmma_noop_barrier_uint32"
+    elif dtype_enum == PTXDataType.FLOAT32:
+        op_name = "tirx.ptx_wgmma_noop_barrier_float32"
+    else:
+        raise ValueError(f"Only support uint32/float32 for wgmma_fence, but got {dtype}.")
+    result = CODEGEN_REGISTRY[op_name]([reg])
+    return result[0] if isinstance(result, tuple) else result
+
+
+# =============================================================================
+# wgmma.mma_async ss / rs — 2 PTX form table entries. Accumulator count and
+# A-register count vary with (M, N, K, in_dtype) but are fully determined by
+# attrs at codegen time.
+#
+# Args layout for ss form (forwarded operand args first, then 9 attr args):
+#   *p_acc[0..num_accums-1], p_descA, p_descB, p_scaleD,
+#   M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
+#
+# Args layout for rs form:
+#   *p_acc[0..num_accums-1], *p_A[0..num_A_regs-1], p_descB, p_scaleD,
+#   M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
+# =============================================================================
+
+
+def _coerce_wgmma_attrs(attrs):
+    """Decode the trailing 9 attrs (M, N, K, in_dtype, out_dtype, transA,
+    transB, scaleA, scaleB) into native Python types."""
+    M, N, K = int(attrs[0]), int(attrs[1]), int(attrs[2])
+    in_dtype = parse_str(attrs[3])
+    out_dtype = parse_str(attrs[4])
+    transA = bool(int(attrs[5])) if hasattr(attrs[5], "value") else bool(attrs[5])
+    transB = bool(int(attrs[6])) if hasattr(attrs[6], "value") else bool(attrs[6])
+    scaleA = bool(int(float(attrs[7])))
+    scaleB = bool(int(float(attrs[8])))
+    if out_dtype != "float32":
+        raise ValueError("WGMMA codegen only supports float32 as output dtype.")
+    allow_transpose = in_dtype in {"float16", "bfloat16"}
+    if not allow_transpose and (transA or transB):
+        raise ValueError("Transpose is only supported for .f16/.bf16 types in WGMMA.")
+    return M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, allow_transpose
+
+
+def _safe(s):
+    return s.replace("::", "_").replace(".", "_")
+
+
+def _wgmma_helper_name(prefix, M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB):
+    return (
+        f"{prefix}_{M}x{N}x{K}_{_safe(in_dtype)}_{_safe(out_dtype)}"
+        f"_{1 if scaleA else 0}_{1 if scaleB else 0}"
+        f"_{1 if transA else 0}_{1 if transB else 0}"
+    )
+
+
+def _wgmma_in_bits(in_dtype):
+    return tvm.runtime.DataType(in_dtype).bits
+
+
+def _wgmma_ss_parts(*args):
+    M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, allow_transpose = (
+        _coerce_wgmma_attrs(args[-9:])
+    )
+    num_accums = M * N // 128
+
+    name = _wgmma_helper_name(
+        "ptx_wgmma_mma_async_ss", M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
+    )
+    sig = (
+        "("
+        + ", ".join(
+            [f"float& p_acc{i}" for i in range(num_accums)]
+            + ["uint64_t p_descA", "uint64_t p_descB", "int p_scaleD"]
+        )
+        + ")"
+    )
+    descA_idx = num_accums
+    descB_idx = num_accums + 1
+    scaleD_idx = num_accums + 2
+    scaleA_idx = num_accums + 3
+    scaleB_idx = num_accums + 4
+    transA_idx = num_accums + 5
+    transB_idx = num_accums + 6
+    accum_r_list = ", ".join(f"%{i}" for i in range(num_accums))
+    accum_constraints = ", ".join(f'"+f"(p_acc{i})' for i in range(num_accums))
+    itype = PTXDataType.from_string(in_dtype)
+    otype = PTXDataType.from_string(out_dtype)
+    if allow_transpose:
+        transpose_r_code = f", %{transA_idx}, %{transB_idx}"
+        transpose_constraints = f', "n"({1 if transA else 0}), "n"({1 if transB else 0})'
+    else:
+        transpose_r_code = ""
+        transpose_constraints = ""
+    instr = (
+        f"wgmma.mma_async.sync.aligned.m{M}n{N}k{K}"
+        f"{otype.to_string()}{itype.to_string()}{itype.to_string()}"
+    )
+    asm_inputs = (
+        f'"l"(p_descA), "l"(p_descB), "r"(p_scaleD),'
+        f' "n"({1 if scaleA else 0}), "n"({1 if scaleB else 0})'
+        f"{transpose_constraints}"
+    )
+    body = (
+        "    asm volatile(\n"
+        '      "{    \\n"\n'
+        '      ".reg .pred p;\\n"\n'
+        f'      "setp.ne.b32 p, %{scaleD_idx}, 0;\\n"\n'
+        f'      "{instr} "\n'
+        f'      "{{{accum_r_list}}},"\n'
+        f'      "%{descA_idx}, %{descB_idx},"\n'
+        f'      "p, %{scaleA_idx}, %{scaleB_idx}{transpose_r_code};\\n"\n'
+        '      "}\\n"\n'
+        f"      : {accum_constraints}\n"
+        f"      : {asm_inputs}\n"
+        "    );"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "_ptx_wgmma_mma_async_ss_impl",
+    n_attrs=9,
+    helper_name=lambda *a: _wgmma_ss_parts(*a)[0],
+    c_signature=lambda *a: _wgmma_ss_parts(*a)[1],
+    body=lambda *a: _wgmma_ss_parts(*a)[2],
+)
+
+
+def _wgmma_rs_parts(*args):
+    M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, allow_transpose = (
+        _coerce_wgmma_attrs(args[-9:])
+    )
+    num_accums = M * N // 128
+    in_bits = _wgmma_in_bits(in_dtype)
+    num_A_regs = M * K // 128 // (32 // in_bits)
+
+    name = _wgmma_helper_name(
+        "ptx_wgmma_mma_async_rs", M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB
+    )
+    sig = (
+        "("
+        + ", ".join(
+            [f"float& p_acc{i}" for i in range(num_accums)]
+            + [f"uint32_t& p_A{i}" for i in range(num_A_regs)]
+            + ["uint64_t p_descB", "int p_scaleD"]
+        )
+        + ")"
+    )
+
+    accum_r_list = ", ".join(f"%{i}" for i in range(num_accums))
+    A_reg_r_list = ", ".join(f"%{num_accums + i}" for i in range(num_A_regs))
+    base_idx = num_accums + num_A_regs
+    descB_idx = base_idx
+    scaleD_idx = base_idx + 1
+    scaleA_idx = base_idx + 2
+    scaleB_idx = base_idx + 3
+    transB_idx = base_idx + 4
+    accum_constraints = ", ".join(f'"+f"(p_acc{i})' for i in range(num_accums))
+    A_reg_constraints = ", ".join(f'"r"(p_A{i})' for i in range(num_A_regs))
+    itype = PTXDataType.from_string(in_dtype)
+    otype = PTXDataType.from_string(out_dtype)
+    if allow_transpose:
+        transpose_r_code = f", %{transB_idx}"
+        transpose_constraints = f', "n"({1 if transB else 0})'
+    else:
+        transpose_r_code, transpose_constraints = "", ""
+    instr = (
+        f"wgmma.mma_async.sync.aligned.m{M}n{N}k{K}"
+        f"{otype.to_string()}{itype.to_string()}{itype.to_string()}"
+    )
+    asm_inputs = (
+        f'{A_reg_constraints}, "l"(p_descB), "r"(p_scaleD),'
+        f' "n"({1 if scaleA else 0}), "n"({1 if scaleB else 0})'
+        f"{transpose_constraints}"
+    )
+    body = (
+        "    asm volatile(\n"
+        '      "{    \\n"\n'
+        '      ".reg .pred p;\\n"\n'
+        f'      "setp.ne.b32 p, %{scaleD_idx}, 0;\\n"\n'
+        f'      "{instr} "\n'
+        f'      "{{{accum_r_list}}},"\n'
+        f'      "{{{A_reg_r_list}}}, %{descB_idx},"\n'
+        f'      "p, %{scaleA_idx}, %{scaleB_idx}{transpose_r_code};\\n"\n'
+        '      "}\\n"\n'
+        f"      : {accum_constraints}\n"
+        f"      : {asm_inputs}\n"
+        "    );"
+    )
+    return name, sig, body
+
+
+device_intrinsic(
+    "_ptx_wgmma_mma_async_rs_impl",
+    n_attrs=9,
+    helper_name=lambda *a: _wgmma_rs_parts(*a)[0],
+    c_signature=lambda *a: _wgmma_rs_parts(*a)[1],
+    body=lambda *a: _wgmma_rs_parts(*a)[2],
+)
+
+
+# User-facing wrappers: just normalise types + reorder positional args to
+# put operands first, then attrs, matching the schema convention.
+
+
+def _wgmma_user_wrapper_ss(*args):
+    M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, scaleD, descA, descB, *accums = (
+        args
+    )
+    M = int(M)
+    N = int(N)
+    K = int(K)
+    in_dtype = parse_str(in_dtype)
+    out_dtype = parse_str(out_dtype)
+    transA = bool(transA)
+    transB = bool(transB)
+    scaleA = bool(int(float(scaleA)))
+    scaleB = bool(int(float(scaleB)))
+    expected = M * N // 128
+    if len(accums) != expected:
+        raise ValueError(
+            "The number of arguments is incorrect. Expected "
+            f"{12 + expected} total args (meaning {expected} accumulator args), "
+            f"but got {len(accums)}."
+        )
+    return [
+        *accums,
+        descA,
+        descB,
+        scaleD,
+        M,
+        N,
+        K,
+        in_dtype,
+        out_dtype,
+        transA,
+        transB,
+        scaleA,
+        scaleB,
+    ]
+
+
+@register_codegen("ptx_wgmma_mma_async_ss")
+def codegen_ptx_wgmma_mma_async_ss(*args):
+    forwarded = _wgmma_user_wrapper_ss(*args)
+    result = CODEGEN_REGISTRY["tirx._ptx_wgmma_mma_async_ss_impl"](forwarded)
+    return result[0] if isinstance(result, tuple) else result
+
+
+def _wgmma_user_wrapper_rs(*args):
+    M, N, K, in_dtype, out_dtype, transA, transB, scaleA, scaleB, scaleD, descB, *reg_list = args
+    M = int(M)
+    N = int(N)
+    K = int(K)
+    in_dtype = parse_str(in_dtype)
+    out_dtype = parse_str(out_dtype)
+    transA = bool(transA)
+    transB = bool(transB)
+    scaleA = bool(int(float(scaleA)))
+    scaleB = bool(int(float(scaleB)))
+    if out_dtype != "float32":
+        raise ValueError("This generator only supports float32 as the output dtype for WGMMA.")
+    in_dtype_bits = tvm.runtime.DataType(in_dtype).bits
+    if in_dtype_bits is None:
+        raise ValueError(f"Bit width not defined for input dtype: {in_dtype}")
+    expected_A_cnt = M * K // 128 // (32 // in_dtype_bits)
+    expected_accm_cnt = M * N // 128
+    if len(reg_list) != expected_A_cnt + expected_accm_cnt:
+        raise ValueError(
+            f"Incorrect number of A registers. Expected {expected_A_cnt}, got {len(reg_list)}"
+        )
+    A_regs = reg_list[:expected_A_cnt]
+    accums = reg_list[expected_A_cnt:]
+    return [
+        *accums,
+        *A_regs,
+        descB,
+        scaleD,
+        M,
+        N,
+        K,
+        in_dtype,
+        out_dtype,
+        transA,
+        transB,
+        scaleA,
+        scaleB,
+    ]
+
+
+@register_codegen("ptx_wgmma_mma_async_rs")
+def codegen_ptx_wgmma_mma_async_rs(*args):
+    forwarded = _wgmma_user_wrapper_rs(*args)
+    result = CODEGEN_REGISTRY["tirx._ptx_wgmma_mma_async_rs_impl"](forwarded)
+    return result[0] if isinstance(result, tuple) else result

python/tvm/backend/cuda/operator/tile_primitive/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+
+from .copy import *
+from .copy_async import *
+from .elementwise import *
+from .gemm import *
+from .gemm_async import *
+from .permute_layout import *
+from .reduction import *

python/tvm/backend/cuda/operator/tile_primitive/common.py

@@ -0,0 +1,283 @@
+# 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 utilities for CUDA operator scheduling (basic helpers and copy ops)."""
+
+import functools
+import operator
+import re
+from enum import Enum
+
+from tvm.arith.analyzer import Analyzer
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, BufferRegion, PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def next_power_of_2(x: int) -> int:
+    """Return the smallest power of 2 greater than or equal to x."""
+    if x <= 1:
+        return 1
+    return 1 << (x - 1).bit_length()
+
+
+def get_st_extent(buffer_region: BufferRegion):
+    """Get the start and extent of a buffer region."""
+    region = buffer_region.region
+    return [r.min for r in region], [r.extent for r in region]
+
+
+def get_indices(nth, start, extent):
+    """Convert a fused index into multi-dimensional indices."""
+    assert len(start) == len(extent)
+    if len(start) == 1:
+        return [start[0] + nth]
+    relative = []
+    for e in reversed(extent):
+        relative.append(nth % e)
+        nth //= e
+    return [r + s for r, s in zip(reversed(relative), start)]
+
+
+def smem_desc_add_16B_offset(desc_val, offset):
+    """Add a 16B-aligned byte offset to the lower 32 bits of a SMEM descriptor.
+
+    Uses the SmemDescriptor union defined in the CUDA header (header.py).
+    All callers must share a single implementation to avoid codegen conflicts.
+    """
+    func_name = "tvm_builtin_smem_desc_add_16B_offset"
+    source_code = f"""
+__forceinline__ __device__ uint64_t {func_name}(uint64_t desc_base, int32_t offset) {{
+    SmemDescriptor desc;
+    desc.desc_ = desc_base;
+    desc.lo += static_cast<uint32_t>(offset);
+    return desc.desc_;
+}}
+"""
+    return T.cuda.func_call(
+        func_name, desc_val, offset, source_code=source_code, return_type="uint64"
+    )
+
+
+class CopyInstType(Enum):
+    """Enumeration of instruction types for memory operations."""
+
+    NORMAL = 0
+    CP_ASYNC = 1
+
+
+def validate_copy_op(
+    op_call: TilePrimitiveCall,
+    sctx: DispatchContext,  # pylint: disable=unused-argument
+) -> bool:
+    """Sanity check for copy op"""
+    dst_buffer_region, src_buffer_region = op_call.args[:2]
+    src: Buffer = src_buffer_region.buffer
+    dst: Buffer = dst_buffer_region.buffer
+    if not (src.layout and dst.layout and src.dtype == dst.dtype):
+        return False
+    # Extract regions and validate dimensions
+    analyzer = Analyzer()
+    src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
+    # Extract extents and validate non-unit dimensions match
+    src_extent_ = [r.extent for r in src_region if r.extent != 1]
+    dst_extent_ = [r.extent for r in dst_region if r.extent != 1]
+    if len(src_extent_) != len(dst_extent_) or not all(
+        analyzer.can_prove_equal(s, d) for s, d in zip(src_extent_, dst_extent_)
+    ):
+        return False
+    return True
+
+
+def get_vec_len(
+    dst_buffer_region: BufferRegion,
+    src_buffer_region: BufferRegion,
+    vec_candidates: list[int],
+    thread_cnt=1,
+) -> int | None:
+    """Get the vector length for the copy operation."""
+
+    dst: Buffer = dst_buffer_region.buffer
+    src: Buffer = src_buffer_region.buffer
+    # layout=None (flat local buffer) is treated as trivial for vectorization purposes
+    if not (
+        (dst.layout is None or dst.layout.is_trivial())
+        and (src.layout is None or src.layout.is_trivial())
+    ):
+        return None
+
+    # Extract regions and validate dimensions
+    analyzer = Analyzer()
+    src_st, src_extent = get_st_extent(src_buffer_region)
+    dst_st, dst_extent = get_st_extent(dst_buffer_region)
+
+    # Thread and vectorization setup
+    DataType(src.dtype).bits  # in bits
+    n_elements = functools.reduce(operator.mul, src_extent, 1)
+    if n_elements % thread_cnt != 0:
+        return None
+
+    # Find valid vector length
+    for vec_len in vec_candidates:
+        if vec_len > 0 and all(
+            analyzer.can_prove_equal(x % vec_len, 0)
+            for x in [
+                src_st[-1],
+                dst_st[-1],
+                src.shape[-1] if len(src.shape) > 1 else 0,
+                dst.shape[-1] if len(dst.shape) > 1 else 0,
+                src_extent[-1],
+                dst_extent[-1],
+                n_elements // thread_cnt,
+            ]
+        ):
+            return vec_len
+    else:
+        return None
+
+
+def copy_vec_load_impl(
+    op_call: TilePrimitiveCall, sctx: DispatchContext, inst_type: CopyInstType
+) -> PrimFunc | None:
+    """Schedule copy operation between global and local/shared memory on CUDA across a CTA/thread.
+    The implementation tries to vectorize the copy operation and parallelize over
+    threads in a CTA/using a single thread.
+    """
+    dst_buffer_region, src_buffer_region = op_call.args[:2]
+    src: Buffer = src_buffer_region.buffer
+    dst: Buffer = dst_buffer_region.buffer
+    if not (
+        (src.scope() == "global" and dst.scope().startswith("shared"))
+        or (src.scope().startswith("shared") and dst.scope() == "global")
+        or (src.scope() == "global" and dst.scope() == "local")
+        or (src.scope() == "local" and dst.scope() == "global")
+        or (src.scope().startswith("shared") and dst.scope() == "local")
+        or (dst.scope().startswith("shared") and src.scope() == "local")
+    ):
+        fail(f"unsupported memory scopes src={src.scope()} dst={dst.scope()}")
+
+    # Thread and vectorization setup
+    if sctx.is_cta:
+        tx = sctx.launch_params["threadIdx.x"].dom.extent
+        assert "threadIdx.y" not in sctx.launch_params and "threadIdx.z" not in sctx.launch_params
+    elif sctx.is_thread:
+        tx = 1
+    else:
+        fail(f"unsupported exec_scope {sctx.scope_kind}")
+
+    elem_size = DataType(src.dtype).bits  # in bits
+    vec_len = op_call.config.get("vec_len", None)
+    if vec_len is None:
+        vec_len = get_vec_len(
+            dst_buffer_region,
+            src_buffer_region,
+            [128 // elem_size, 64 // elem_size, 32 // elem_size, 1],
+            thread_cnt=tx,
+        )
+    if vec_len is None:
+        fail("no valid vector length; check alignment/extents/thread-count")
+
+    # cp-size (the size of data in bytes) can only be 4, 8 and 16 for cp.async
+    if inst_type == CopyInstType.CP_ASYNC:
+        cp_size = vec_len * elem_size // 8  # in bytes
+        if cp_size not in [4, 8, 16]:
+            fail("invalid cp.async cp_size; expected 4, 8 or 16 bytes")
+
+    src_st, src_extent = get_st_extent(src_buffer_region)
+    dst_st, dst_extent = get_st_extent(dst_buffer_region)
+    n_elements = functools.reduce(operator.mul, src_extent, 1)
+
+    if sctx.is_cta:
+        # fmt: off
+        @T.prim_func
+        def impl():
+            """Implement copy operation with vectorized loads/stores."""
+            for s in T.serial(0, n_elements // (tx * vec_len)):
+                for tid_x in T.thread_binding(tx, "threadIdx.x"):
+                    if inst_type == CopyInstType.NORMAL:
+                        for vec in T.vectorized(vec_len):
+                            fused = T.meta_var((s * tx + tid_x) * vec_len + vec)
+                            dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
+                            src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
+                            dst[tuple(dst_indices)] = src[tuple(src_indices)]
+                    elif inst_type == CopyInstType.CP_ASYNC:
+                        fused = T.meta_var((s * tx + tid_x) * vec_len)
+                        dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
+                        src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
+                        T.evaluate(T.ptx.cp_async(dst.ptr_to(dst_indices), src.ptr_to(src_indices), cp_size))  # noqa: E501
+            if dst.scope().startswith("shared") and inst_type == CopyInstType.NORMAL:
+                T.tvm_storage_sync("shared")
+        # fmt: on
+    elif sctx.is_thread:
+        # fmt: off
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            for s in T.serial(0, n_elements // (vec_len)):
+                if inst_type == CopyInstType.NORMAL:
+                    for vec in T.vectorized(vec_len):
+                        fused = T.meta_var(s * vec_len + vec)
+                        dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
+                        src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
+                        dst[tuple(dst_indices)] = src[tuple(src_indices)]
+                elif inst_type == CopyInstType.CP_ASYNC:
+                    fused = T.meta_var(s * vec_len)
+                    dst_indices = T.meta_var(get_indices(fused, dst_st, dst_extent))
+                    src_indices = T.meta_var(get_indices(fused, src_st, src_extent))
+                    T.evaluate(T.ptx.cp_async(dst.ptr_to(dst_indices), src.ptr_to(src_indices), cp_size))  # noqa: E501
+        # fmt: on
+    else:
+        fail(f"unsupported exec_scope {sctx.scope_kind}")
+    return impl
+
+
+def match_scope(scope: str | None, pattern: str) -> bool:
+    """Glob-lite scope matching: 'shared*' => prefix match; otherwise exact.
+
+    Returns True when scope is None (meaning "any scope is fine").
+    """
+    if scope is None:
+        return True
+    if pattern.endswith("*"):
+        return scope.startswith(pattern[:-1])
+    return scope == pattern
+
+
+def get_thread_cnt(sctx: DispatchContext) -> int | None:
+    """Get thread count for the current execution scope."""
+    scope_name = sctx.scope_kind
+    if scope_name == "cta":
+        return sctx.launch_params["threadIdx.x"].dom.extent
+    if scope_name == "warpgroup":
+        return 128
+    if scope_name == "warp":
+        return 32
+    if scope_name == "thread":
+        return 1
+    return None
+
+
+def sm_version_ok(
+    op: TilePrimitiveCall, sctx: DispatchContext, min_version: int
+) -> tuple[bool, str | None]:
+    """Check if SM version >= min_version. Usable as a dispatch predicate."""
+    target_arch = sctx.target.arch if hasattr(sctx.target, "arch") else ""
+    sm_match = re.match(r"sm_(\d+)", target_arch)
+    sm_version = int(sm_match.group(1)) if sm_match else 0
+    ok = sm_version >= min_version
+    return (ok, None if ok else f"sm_version {sm_version} < {min_version}")

python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+
+from .fallback import *
+from .gmem_smem import *
+from .ld_stmatrix import *
+from .reg import *
+from .utils import (
+    _is_valid_copy,
+    _is_valid_smem_tmem_copy,
+    _scope_allowed,
+    _single_thread_exec,
+)

python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py

@@ -0,0 +1,556 @@
+# 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.
+
+"""Shared partition / layout algorithm for synthesized-partition copy
+dispatches (currently ``gmem_smem`` and ``ldgsts``).
+
+``gmem_smem`` (sync ``Tx.copy`` global ↔ shared) and ``ldgsts`` (async
+``Tx.copy_async`` global → shared via cp.async / SASS LDGSTS) share the
+same algorithm to pick a vec-isolating + thread-distributing layout for
+``G ↔ S`` copies. Only emit-time details differ (which copy instruction
+to call, allowed vec widths). All the layout/partition logic lives here.
+"""
+
+from tvm import arith
+from tvm.tirx.layout import ComposeLayout, Iter, S, SwizzleLayout, TileLayout
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+
+
+def _alignment_ok(vec_len: int, terms) -> bool:
+    """Every term must be a multiple of ``vec_len``. Constants checked
+    directly; Expr / symbolic terms checked via ``arith.Analyzer``.
+
+    ``vec_len=1`` always passes (the scalar fallback). When a symbolic
+    term can't be proved divisible, returns ``False`` conservatively —
+    the candidate loop will then try a smaller ``vec_len``.
+    """
+    if vec_len <= 1:
+        return True
+    analyzer = arith.Analyzer()
+    for t in terms:
+        if isinstance(t, int):
+            if t % vec_len != 0:
+                return False
+        else:
+            if not analyzer.can_prove_equal(t % vec_len, 0):
+                return False
+    return True
+
+
+# scope_kind → name of the scope_id that decomposes the scope into per-thread.
+_TID_AXIS_FOR_SCOPE = {
+    "warp": "laneid",
+    "warpgroup": "tid_in_wg",
+    "cta": "tx",
+}
+
+
+def _thread_cnt(sctx: DispatchContext) -> int:
+    """Total threads active in the current scope = ∏ intra-axis extents.
+
+    For thread scope ``sctx.intra`` is empty → returns 1.
+    """
+    n = 1
+    for ext, _off in sctx.intra.values():
+        n *= int(ext)
+    return n
+
+
+# -----------------------------------------------------------------------------
+# Layout primitives
+# -----------------------------------------------------------------------------
+
+
+def _contig_group(iters: list) -> list[int]:
+    """Indices (in iters) of the maximal physical-contiguous chain starting
+    at the stride=1 iter, ordered stride-ascending.
+
+    Returns [] if no stride=1 iter exists.
+    """
+    one_idx = next(
+        (i for i, it in enumerate(iters) if int(it.stride) == 1),
+        None,
+    )
+    if one_idx is None:
+        return []
+    chain = [one_idx]
+    acc = int(iters[one_idx].extent)
+    used = {one_idx}
+    while True:
+        nxt = next(
+            (i for i, it in enumerate(iters) if i not in used and int(it.stride) == acc),
+            None,
+        )
+        if nxt is None:
+            break
+        chain.append(nxt)
+        acc *= int(iters[nxt].extent)
+        used.add(nxt)
+    return chain
+
+
+def _try_split_vec(iters: list, vec_len: int):
+    """Try to walk ``vec_len`` consecutive elements along the contig chain.
+
+    Returns ``(new_iters, selected_positions)`` on success, ``None`` on
+    failure. ``new_iters`` may contain a freshly-split iter (replacing one
+    entry with its "outer" half, with the "inner" half appended at the end);
+    ``selected_positions`` are positions in ``new_iters`` that together
+    cover the ``vec_len`` contig elements.
+    """
+    chain = _contig_group(iters)
+    if not chain:
+        return None
+    rem = vec_len
+    new_iters = list(iters)
+    selected: list[int] = []
+    for orig_idx in chain:
+        if rem == 0:
+            break
+        it = new_iters[orig_idx]
+        ext = int(it.extent)
+        if ext <= rem:
+            if rem % ext != 0:
+                return None
+            selected.append(orig_idx)
+            rem //= ext
+        else:
+            if ext % rem != 0:
+                return None
+            stride = int(it.stride)
+            outer = Iter(ext // rem, stride * rem, it.axis)
+            inner = Iter(rem, stride, it.axis)
+            new_iters[orig_idx] = outer
+            new_iters.append(inner)
+            selected.append(len(new_iters) - 1)
+            rem = 0
+            break
+    if rem != 0:
+        return None
+    return new_iters, selected
+
+
+def _isolated_shape(iters: list, selected: list[int]) -> tuple[list[int], list[tuple[int, int]]]:
+    """Build the isolated shape: each selected iter is its own segment;
+    adjacent unselected iters are merged into a single segment.
+
+    Returns ``(shape, segments)`` where ``segments[i] = (start, end)`` is
+    the half-open range in ``iters`` covered by shape entry ``i``.
+    """
+    sel_set = set(selected)
+    shape: list[int] = []
+    segments: list[tuple[int, int]] = []
+    cur_start = None
+    cur_ext = 1
+    for i, it in enumerate(iters):
+        if i in sel_set:
+            if cur_start is not None:
+                shape.append(cur_ext)
+                segments.append((cur_start, i))
+                cur_start = None
+                cur_ext = 1
+            shape.append(int(it.extent))
+            segments.append((i, i + 1))
+        else:
+            if cur_start is None:
+                cur_start = i
+            cur_ext *= int(it.extent)
+    if cur_start is not None:
+        shape.append(cur_ext)
+        segments.append((cur_start, len(iters)))
+    return shape, segments
+
+
+def _vec_perm(iters: list, selected: list[int]) -> list[int]:
+    """Reorder ``iters`` into ``[outer, vec]``, both ordered by stride
+    descending so the stride=1 iter ends up at the very last position."""
+    sel_set = set(selected)
+    unsel_sorted = sorted(
+        (i for i in range(len(iters)) if i not in sel_set),
+        key=lambda i: -int(iters[i].stride),
+    )
+    sel_sorted = sorted(selected, key=lambda i: -int(iters[i].stride))
+    return list(unsel_sorted) + sel_sorted
+
+
+def _try_split_thread(iters: list, vec_selected: list[int], thread_cnt: int):
+    """After ``_try_split_vec``, carve ``thread_cnt`` from the OUTER tail
+    (smallest-stride outer iter, then towards bigger stride if needed).
+
+    Unlike vec split, this doesn't require physical contiguity — T
+    consecutive fused indices map to per-thread offsets via the layout's
+    stride (which may be > 1).
+
+    Returns ``(new_iters, thread_selected_positions)`` on success, ``None``
+    on failure (outer doesn't divide T cleanly, or no outer iters left).
+    """
+    if thread_cnt == 1:
+        return list(iters), []
+    vec_set = set(vec_selected)
+    outer = [i for i in range(len(iters)) if i not in vec_set]
+    if not outer:
+        return None
+    outer_by_stride_desc = sorted(outer, key=lambda i: -int(iters[i].stride))
+    rem = thread_cnt
+    new_iters = list(iters)
+    thread_selected: list[int] = []
+    for orig_idx in reversed(outer_by_stride_desc):
+        if rem == 0:
+            break
+        it = new_iters[orig_idx]
+        ext = int(it.extent)
+        if ext <= rem:
+            if rem % ext != 0:
+                return None
+            thread_selected.append(orig_idx)
+            rem //= ext
+        else:
+            if ext % rem != 0:
+                return None
+            stride = int(it.stride)
+            new_iters[orig_idx] = Iter(ext // rem, stride * rem, it.axis)
+            new_iters.append(Iter(rem, stride, it.axis))
+            thread_selected.append(len(new_iters) - 1)
+            rem = 0
+            break
+    if rem != 0:
+        return None
+    return new_iters, thread_selected
+
+
+def _three_segment_perm(iters: list, t_selected: list[int], vec_selected: list[int]) -> list[int]:
+    """Reorder ``iters`` into ``[outer, T, vec]`` segments. Within each
+    segment, stride descending so stride=1 sits at the very end."""
+    t_set = set(t_selected)
+    vec_set = set(vec_selected)
+    outer = sorted(
+        (i for i in range(len(iters)) if i not in t_set and i not in vec_set),
+        key=lambda i: -int(iters[i].stride),
+    )
+    t_sorted = sorted(t_selected, key=lambda i: -int(iters[i].stride))
+    vec_sorted = sorted(vec_selected, key=lambda i: -int(iters[i].stride))
+    return list(outer) + t_sorted + vec_sorted
+
+
+def _shape_perm_for_isolated(
+    shape_segments: list[tuple[int, int]], iter_perm: list[int]
+) -> list[int]:
+    """Given segments (one per shape entry, each = (start, end) in
+    pre-perm iter positions) and the iter permutation, compute the
+    corresponding shape permutation."""
+    seg_of = [0] * sum(end - start for start, end in shape_segments)
+    for seg_idx, (start, end) in enumerate(shape_segments):
+        for k in range(start, end):
+            seg_of[k] = seg_idx
+    seen: set[int] = set()
+    perm: list[int] = []
+    for orig_idx in iter_perm:
+        seg_idx = seg_of[orig_idx]
+        if seg_idx not in seen:
+            seen.add(seg_idx)
+            perm.append(seg_idx)
+    return perm
+
+
+def _verify_s_tail_contig(s_p: TileLayout, vec_len: int) -> bool:
+    """Check the last iters of ``s_p`` form a stride=1 contig chain whose
+    extent product equals ``vec_len``."""
+    iters = list(s_p.shard)
+    if not iters:
+        return vec_len == 1
+    last = iters[-1]
+    if int(last.stride) != 1:
+        return False
+    acc = int(last.extent)
+    if acc == vec_len:
+        return True
+    for k in range(len(iters) - 2, -1, -1):
+        it = iters[k]
+        if int(it.stride) != acc:
+            break
+        acc *= int(it.extent)
+        if acc == vec_len:
+            return True
+        if acc > vec_len:
+            return False
+    return acc >= vec_len and acc % vec_len == 0
+
+
+_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8)
+
+
+def _vec_len_candidates(elem_bits: int, allowed_bits: tuple | None = None) -> list[int]:
+    """Vec-length candidates (in elements) for the given element width.
+
+    ``allowed_bits`` optionally filters the per-instruction allowed widths
+    (e.g. cp.async only accepts {128, 64, 32} bits = 16/8/4 bytes).
+    Defaults to ``_VEC_BITS_CANDIDATES`` if not specified.
+    """
+    bits_tuple = allowed_bits if allowed_bits is not None else _VEC_BITS_CANDIDATES
+    out: list[int] = []
+    for vb in bits_tuple:
+        if vb < elem_bits or vb % elem_bits != 0:
+            continue
+        n = vb // elem_bits
+        if n not in out:
+            out.append(n)
+    if 1 not in out and allowed_bits is None:
+        # Scalar fallback is only added for the unrestricted candidate set;
+        # an instruction-specific list (cp.async etc.) keeps its strictness.
+        out.append(1)
+    return out
+
+
+def _extract_tile(layout, region):
+    """Strip swizzle so we can perm/group as a TileLayout."""
+    if isinstance(layout, ComposeLayout):
+        return layout.tile_layout
+    if isinstance(layout, SwizzleLayout):
+        extents = [int(end - start) for (start, end) in region]
+        return TileLayout(S[tuple(extents)])
+    return layout
+
+
+def _sort_by_stride_desc(layout: TileLayout) -> TileLayout:
+    """Reorder shard so list order = traversal order (outer first, stride=1
+    last). Required before canonicalize() can fuse non-adjacent-but-contig
+    iters."""
+    iters = list(layout.shard)
+    perm = sorted(range(len(iters)), key=lambda i: -int(iters[i].stride))
+    if perm == list(range(len(iters))):
+        return layout
+    return layout.permute_dims(perm)
+
+
+def _carve_tail(iters: list, chunk: int):
+    """Carve ``chunk`` elements off the tail. Walk back across multiple
+    iters as needed; at most one iter is split.
+
+    Per iter (from last to first), let ``ext`` = iter extent and ``rem``
+    = remaining chunk to fill:
+
+    * ``ext == rem``: eat this iter whole, done.
+    * ``ext < rem``: must divide ``rem``; eat whole, ``rem //= ext``.
+    * ``ext > rem``: must divide ``ext``; split into
+      ``(ext/rem, stride*rem) + (rem, stride)``, take the inner. Done.
+
+    Returns the new iter list on success, ``None`` on failure.
+    """
+    if not iters or chunk <= 0:
+        return None
+    rem = chunk
+    work = list(iters)
+    for idx in range(len(work) - 1, -1, -1):
+        it = work[idx]
+        ext = int(it.extent)
+        if ext == rem:
+            return work
+        if ext < rem:
+            if rem % ext != 0:
+                return None
+            rem //= ext
+            continue
+        if ext % rem != 0:
+            return None
+        stride = int(it.stride)
+        work[idx] = Iter(ext // rem, stride * rem, it.axis)
+        work.insert(idx + 1, Iter(rem, stride, it.axis))
+        return work
+    return None
+
+
+def align_layouts_gs(
+    g_layout,
+    g_shape,
+    g_region,
+    s_layout,
+    s_shape,
+    s_region,
+    elem_bits,
+    thread_cnt: int,
+    vec_bits_candidates: tuple | None = None,
+    debug: bool = False,
+):
+    """Align G and S layouts for a synthesized G↔S copy.
+
+    Algorithm:
+      1. Sort G iters by stride desc + canonicalize (fuses anything physically
+         contig). Same for S.
+      2. Group S by G's iter shape (per-iter shape list) + permute_by_groups
+         with identity so S's groups line up with G's iters.
+      3. Carve ``vec_len`` off G's tail (a single iter split if needed).
+      4. Carve ``T = thread_cnt`` off the iter before vec (multi-iter walk,
+         single split).
+      5. Re-group S (already permuted) by the new finer per-iter shape. No
+         further permute needed because steps 3-4 only refine the tail.
+      6. Verify S's tail (vec segment) is physically contig.
+
+    ``vec_bits_candidates`` optionally restricts the per-instruction allowed
+    widths (e.g. cp.async only accepts {128, 64, 32} bits). When ``None``
+    (default), uses the full {128, 64, 32, 16, 8} set plus a scalar (1)
+    fallback.
+
+    Returns ``(g_p, s_p, vec_len)``. ``g_p.shard`` ends as
+    ``[outer iters..., T iter, vec iter]``; ``s_p.shard`` has the same iter
+    count and matching iter-by-iter extents.
+    """
+    g = g_layout.slice(list(g_shape), g_region)
+    s = s_layout.slice(list(s_shape), s_region)
+    # Detect a SwizzleLayout on the S side BEFORE _extract_tile strips it.
+    # vec_len must fit inside one swizzle chunk (C = 2^per_element elements);
+    # otherwise the vec ld/st crosses a swizzle XOR boundary and hits the
+    # wrong physical bytes mid-vec.
+    s_swizzle_chunk_elems = None
+    if isinstance(s_layout, ComposeLayout):
+        s_swizzle_chunk_elems = 1 << int(s_layout.swizzle.per_element)
+    elif isinstance(s_layout, SwizzleLayout):
+        s_swizzle_chunk_elems = 1 << int(s_layout.per_element)
+    g = _extract_tile(g, g_region)
+    s = _extract_tile(s, s_region)
+
+    # Only G drives the canonical form. S's iter order is derived from G's
+    # pre-sort iter extents (used as the grouping shape) and then permuted
+    # by G's stride-desc permutation. Independently sorting/canonicalizing
+    # S would fuse iters whose strides happen to chain into one contiguous
+    # range — that loses the layout's logical `(i, j) → addr` mapping for
+    # layout-permuting copies (e.g. row-major GMEM → K-tiled SMEM, where
+    # both layouts cover the same byte range but with different coord maps).
+    g_pre_sort_extents = [int(it.extent) for it in g.shard]
+    g_perm = sorted(range(len(g.shard)), key=lambda i: -int(g.shard[i].stride))
+    try:
+        s_grp1, seps1 = s.group(g_pre_sort_extents)
+    except Exception as e:
+        if debug:
+            print(f"    Step-2 S.group({g_pre_sort_extents}) failed: {e}")
+        return _sort_by_stride_desc(g).canonicalize(), s, 1
+    # S iters keep their original strides; only the *group order* is
+    # rearranged to follow G's sorted iter order. Canonicalize after the
+    # permute is safe — it only fuses iters whose strides genuinely chain
+    # in the post-permute order, which preserves the logical structure.
+    s_aligned = s_grp1.permute_by_groups(list(seps1), g_perm).canonicalize()
+    g = _sort_by_stride_desc(g).canonicalize()
+    if debug:
+        print(f"    g (sort+canon): shard={[(int(it.extent), int(it.stride)) for it in g.shard]}")
+        print(
+            f"    s (grouped+permuted by G shape {g_pre_sort_extents}): shard="
+            f"{[(int(it.extent), int(it.stride)) for it in s_aligned.shard]}"
+        )
+        print(f"    thread_cnt: {thread_cnt}")
+
+    dummy_axis = g.shard[-1].axis if g.shard else None
+
+    for vec_len in _vec_len_candidates(elem_bits, vec_bits_candidates):
+        if vec_len == 1:
+            g_after_vec = [*list(g.shard), Iter(1, 1, dummy_axis)]
+        else:
+            # Swizzle chunk-size cap: vec must fit in one swizzle chunk.
+            if s_swizzle_chunk_elems is not None and vec_len > s_swizzle_chunk_elems:
+                if debug:
+                    print(
+                        f"    vec_len={vec_len}: exceeds swizzle chunk "
+                        f"({s_swizzle_chunk_elems} elements)"
+                    )
+                continue
+            g_after_vec = _carve_tail(list(g.shard), vec_len)
+            if g_after_vec is None:
+                if debug:
+                    print(f"    vec_len={vec_len}: G vec carve failed")
+                continue
+        outer_for_t = g_after_vec[:-1]
+        outer_after_t = _carve_tail(outer_for_t, thread_cnt) if thread_cnt > 1 else outer_for_t
+        if outer_after_t is None:
+            if debug:
+                print(f"    vec_len={vec_len}: T={thread_cnt} carve failed")
+            continue
+        g_final_iters = [*outer_after_t, g_after_vec[-1]]
+        new_shape = [int(it.extent) for it in g_final_iters]
+        try:
+            s_final_grp, _ = s_aligned.group(new_shape)
+        except Exception as e:
+            if debug:
+                print(f"    vec_len={vec_len}: S.group({new_shape}) failed: {e}")
+            continue
+        g_p = TileLayout.from_iters(g_final_iters, list(g.replica), dict(g.offset))
+        s_p = s_final_grp
+
+        ok = _verify_s_tail_contig(s_p, vec_len)
+        if debug:
+            print(
+                f"    vec_len={vec_len}: shape={new_shape}, "
+                f"g_p.shard={[(int(it.extent), int(it.stride)) for it in g_p.shard]}, "
+                f"s_p.shard={[(int(it.extent), int(it.stride)) for it in s_p.shard]}, "
+                f"s_tail_contig={ok}"
+            )
+        if not ok:
+            continue
+        # Alignment: per-thread starting addr = base_ptr + sizeof(elem) *
+        # (region_base + tid*t_stride + outer_iter_strides). For the
+        # vec_bits/8 byte vector op to be naturally aligned, every one of
+        # those element-count terms must be a multiple of vec_len.
+        align_terms = []
+        for it in s_p.shard[:-1]:
+            align_terms.append(int(it.stride))
+        for it in g_p.shard[:-1]:
+            align_terms.append(int(it.stride))
+        align_terms.extend(s_p.offset.values())
+        align_terms.extend(g_p.offset.values())
+        if not _alignment_ok(vec_len, align_terms):
+            if debug:
+                print(f"    vec_len={vec_len}: alignment check failed")
+            continue
+        return g_p, s_p, vec_len
+
+    return g, s_aligned, 1
+
+
+def _flat_outer_coords(outer_exts: list[int], flat_idx: int) -> list[int]:
+    """Decode a row-major flat index into per-iter coords. ``outer_exts`` is
+    in stride-desc order (outermost first), so the first coord changes
+    slowest."""
+    coords: list[int] = []
+    rem = flat_idx
+    for ext in reversed(outer_exts):
+        coords.append(rem % ext)
+        rem //= ext
+    coords.reverse()
+    return coords
+
+
+def _outer_offsets(outer_iters_s, outer_iters_g, flat_idx):
+    """Returns ``(ds, dg)``: constant offsets on S and G sides for the
+    given flat outer-loop iteration."""
+    outer_exts = [int(it.extent) for it in outer_iters_s]
+    coords = _flat_outer_coords(outer_exts, flat_idx)
+    ds = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_s))
+    dg = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_g))
+    return ds, dg
+
+
+def copy_ptx_form(num_bytes: int) -> tuple[str, str]:
+    """Map copy width (bytes) to PTX ``(vec, ptx_type)`` for ``T.ptx.ld`` / ``T.ptx.st``."""
+    return {
+        16: ("v4", "u32"),
+        8: ("v2", "u32"),
+        4: ("", "u32"),
+        2: ("", "u16"),
+        1: ("", "u8"),
+    }[num_bytes]
+
+
+def copy_ptx_ld_return_type(ptx_type: str) -> str:
+    """TVM dtype string for ``T.ptx.ld``'s ``return_type`` argument."""
+    return {"u32": "uint32", "u16": "uint16", "u8": "uint32"}[ptx_type]

python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py

@@ -0,0 +1,406 @@
+# 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.
+
+"""Generic swizzle-aware iter pattern for CUDA copy dispatches.
+
+When the per-thread outer-iter loop satisfies (C1)+(C2) below for a
+``SwizzleLayout(per_element=p, swizzle_len=sw, atom_len=at,
+swizzle_inner=True)`` on the SMEM side, the swizzled physical address
+at unrolled iter ``k`` reduces to
+
+    addr(k) = base_off + sum_{j : bit_j(k)=1} signed_strides[j]
+
+where ``base_off`` and the ``signed_strides[j]`` are per-thread runtime
+constants set once at thread setup. Per-iter cost is then ``popcount(k)``
+register adds instead of a full ``swizzle.apply(...)`` per iter.
+
+Notation. Each binary outer iter has element-stride ``2^(bj + p)`` for
+some chunk bit position ``bj >= 0`` (so ``stride / C = 2^bj`` where
+``C = 1 << p``). The chunk index ``q(M0) = M0 // C`` partitions into
+four bit ranges by where ``bj`` lands:
+
+  * ``[0, sw)``        — "inner" (Case 1.A in the proof)
+  * ``[sw, at)``       — "mid"   (Case 1.B)
+  * ``[at, at + sw)``  — "outer" (Case 1.C; the bit overlaps the swizzle
+                         outer mask, so its addition produces a
+                         secondary contribution at ``bj - at``)
+  * ``[at + sw, ∞)``   — "above" (Case 1.D)
+
+Conditions for the linear-combination fast path:
+
+  (C1) bit-clear no-carry: ``bit_bj(q(M0)) = 0`` for every binary iter.
+  (C2) support disjointness: no inner-outer pair ``(bj_A, bj_C)`` with
+       ``bj_C in [at, at+sw)`` and ``bj_A = bj_C - at`` both present.
+
+  (distinctness) The ``bj`` values across all binary iters must be
+       distinct — two iters at the same ``bj`` collapse into bit
+       ``bj + 1`` whose case behavior may differ.
+
+Under (C1)+(C2)+(distinctness), for each binary iter at position ``bj``:
+
+    T(bj)        = 2^(bj + p)            # element stride
+    sigma_b(M0)  = 1 - 2 * bit_b(q(M0))  # ∈ {+1, -1}
+
+    signed_strides[j] = sigma_(at + bj)(M0) * T(bj)        bj in [0, sw)
+                     = T(bj)                               bj in [sw, at)
+                     = T(bj) + sigma_(bj - at)(M0) * T(bj - at)
+                                                           bj in [at, at + sw)
+                     = T(bj)                               bj >= at + sw
+
+The ``swizzle_inner=False`` mode swaps the inner/outer roles and is not
+yet covered; ``try_recognize`` gates on this.
+"""
+
+from dataclasses import dataclass
+
+import tvm
+from tvm import arith
+from tvm.script import tirx as T
+from tvm.tirx.expr import IntImm as _IntImm
+from tvm.tirx.layout import ComposeLayout, SwizzleLayout
+
+
+@dataclass
+class _BitIter:
+    """Pow2-extent outer iter, binary-split into ``n_bits`` chunk-bit flips.
+
+    ``slot_start..slot_start + n_bits`` is this iter's range in the global
+    ``bit_positions`` / ``iter_strides_elems`` / ``signed_strides`` arrays.
+    Slot ``slot_start + b`` corresponds to bit position ``n_bits - 1 - b``
+    of this iter's per-iter coord (outermost binary bit first).
+    """
+
+    ext: int
+    n_bits: int
+    slot_start: int
+
+
+@dataclass
+class _LinearIter:
+    """Outer iter contributing ``c * stride`` to the offset (no bit decomp).
+
+    Used when ``stride`` is a multiple of ``2^(p + at + sw)`` (pure Case 1.D
+    regime: swizzle XOR has no effect on bits the iter flips). ``ext`` does
+    not need to be a power of two.
+    """
+
+    ext: int
+    stride: int
+
+
+@dataclass
+class SwizzlePattern:
+    """A recognized swizzle iter pattern.
+
+    ``bit_positions[j]`` and ``iter_strides_elems[j]`` collect the binary
+    sub-iters from every BitIter in outer-iter order (outermost first).
+    ``outer_iters`` lists every outer iter (BitIter or LinearIter) in
+    outermost-first order; ``emit_iter_offset`` walks this list to
+    decompose ``mm`` per-iter. Empty lists = trivially recognized
+    degenerate case (no outer iter, just base_off).
+    """
+
+    swizzle: SwizzleLayout
+    bit_positions: list[int]
+    iter_strides_elems: list[int]
+    outer_iters: "list[_BitIter | _LinearIter]"
+
+    @property
+    def n_binary_iters(self) -> int:
+        return len(self.bit_positions)
+
+
+def get_swizzle(layout) -> SwizzleLayout | None:
+    """Return the SwizzleLayout from ``layout`` if present, else ``None``.
+
+    Accepts ``ComposeLayout(SwizzleLayout, TileLayout)`` (the common case
+    when a TileLayout is wrapped by a swizzle), or a bare ``SwizzleLayout``.
+    """
+    if isinstance(layout, ComposeLayout):
+        return layout.swizzle
+    if isinstance(layout, SwizzleLayout):
+        return layout
+    return None
+
+
+def _is_pow2(n: int) -> bool:
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def try_recognize(
+    swizzle: SwizzleLayout,
+    iter_extents: list[int],
+    iter_strides: list[int],
+    s_off_template,
+    var_bounds: dict | None = None,
+) -> SwizzlePattern | None:
+    """Return a ``SwizzlePattern`` if (C1)+(C2)+(distinctness) hold, else ``None``.
+
+    ``iter_extents`` / ``iter_strides``: the outer-iter list on the S side
+    (excluding T iter and vec iter), in outermost-first order matching
+    ``s_p.shard[:-2]`` (or the atom-derived analog in ``reg.py``).
+    Strides are in element units.
+
+    Each outer iter with ``extent=2^k`` and ``stride=s`` is conceptually
+    split into ``k`` binary iters of strides ``2^(k-1)*s, ..., 2*s, s``
+    (outermost first within the split — this matches ``_flat_outer_coords``
+    semantics, since the highest-stride iter must change slowest in the
+    flat-index decomposition).
+
+    ``s_off_template`` is the per-thread linear base offset expression
+    (with a placeholder var for the thread-id contribution). It is used
+    only to check condition (C1) symbolically via ``arith.Analyzer``;
+    ``emit_init`` takes the resolved form separately.
+
+    ``var_bounds`` is an optional ``{Var: tvm.ir.Range}`` map of placeholder
+    bounds to ``analyzer.bind`` before the (C1) check. Without bounds,
+    structurally-OK forms like ``(lane // 8) * 8 + (lane % 8) * Q`` where
+    ``lane < 32`` make ``(... // (C·2^bj)) % 2 == 0`` unprovable — the
+    bit is in fact always 0 but the analyzer can't conclude it universally.
+    Pass ``{lane_ph: Range(0, 32), warp_ph: Range(0, n_warps)}`` (or the
+    scope's equivalents) to let the (C1) check fire on these templates.
+    """
+    # swizzle_inner=False swaps the inner/outer xor direction — Cases 1.A
+    # and 1.C roles flip. Not derived/tested yet; reject for safety.
+    if not swizzle.swizzle_inner:
+        return None
+
+    p = swizzle.per_element
+    sw = swizzle.swizzle_len
+    at = swizzle.atom_len
+    C = 1 << p
+    # Pure Case 1.D threshold: stride a multiple of this means every chunk-bit
+    # the iter flips is at position >= at + sw (above the swizzle XOR region),
+    # so swizzle has no effect and the contribution is purely linear in the
+    # iter coord — no power-of-2 ext requirement.
+    pure_1d = 1 << (p + at + sw)
+
+    bit_positions: list[int] = []
+    iter_strides_elems: list[int] = []
+    outer_iters: list = []
+
+    for ext, stride in zip(iter_extents, iter_strides):
+        # Zero-stride iters degrade dq=0 → log2 undefined. Explicit guard.
+        if stride == 0 or stride % C != 0:
+            return None
+        if ext <= 0:
+            return None
+        if ext == 1:
+            # Trivial iter contributes nothing; skip without forcing pow2.
+            continue
+        if not _is_pow2(ext):
+            # Non-pow2 ext can only be handled by the linear path. That in turn
+            # requires the iter to be in pure Case 1.D (stride a multiple of
+            # the swizzle period) so the swizzle does not interact with the
+            # per-coord contribution.
+            if stride % pure_1d != 0:
+                return None
+            outer_iters.append(_LinearIter(ext=ext, stride=stride))
+            continue
+        # pow2 ext: binary split (existing path).
+        k = ext.bit_length() - 1  # log2(ext)
+        slot_start = len(bit_positions)
+        # Split into k binary iters; the outermost (within this split) carries
+        # the largest stride so that flat-index bit decomp matches our
+        # outer-iter list ordering.
+        for j in range(k - 1, -1, -1):
+            substride = stride * (1 << j)
+            dq = substride // C
+            # dq must be a single bit set (so this binary iter flips exactly
+            # one bit of the chunk index). _is_pow2 also rejects dq=0.
+            if not _is_pow2(dq):
+                return None
+            bj = dq.bit_length() - 1
+            # All bj >= 0 accepted; case branching happens in emit_init.
+            bit_positions.append(bj)
+            iter_strides_elems.append(substride)
+        outer_iters.append(_BitIter(ext=ext, n_bits=k, slot_start=slot_start))
+
+    # Distinctness: two binary iters at the same bj collapse to bj+1, whose
+    # case behavior may differ from bj. See module docstring NB.
+    if len(set(bit_positions)) != len(bit_positions):
+        return None
+
+    bj_set = set(bit_positions)
+
+    # (C2) support disjointness: the only possible collision is between a
+    # Case-1.A iter at bj_A and a Case-1.C iter at bj_A + at. Checking the
+    # 1.C direction alone is symmetric and complete.
+    for bj in bj_set:
+        if at <= bj < at + sw and (bj - at) in bj_set:
+            return None  # inner-outer pair collision
+
+    # (C1) per-iter no-carry on q(M0). Must hold *symbolically over all*
+    # free lane / warp placeholders in s_off_template — ``can_prove_equal``
+    # returns False if the analyzer can't discharge the equality
+    # universally, conservatively forcing a fallback.
+    analyzer = arith.Analyzer()
+    if var_bounds:
+        for var, rng in var_bounds.items():
+            analyzer.bind(var, rng)
+    for bj in bj_set:
+        divisor = C * (1 << bj)
+        check = tvm.tirx.floormod(
+            tvm.tirx.floordiv(s_off_template, _IntImm("int32", divisor)),
+            _IntImm("int32", 2),
+        )
+        if not analyzer.can_prove_equal(check, _IntImm("int32", 0)):
+            return None
+
+    return SwizzlePattern(
+        swizzle=swizzle,
+        bit_positions=bit_positions,
+        iter_strides_elems=iter_strides_elems,
+        outer_iters=outer_iters,
+    )
+
+
+def emit_init(pattern: SwizzlePattern, s_off_resolved):
+    """Emit at thread setup (call from inside the @T.prim_func body):
+
+      1. ``base_off = swizzle.apply(s_off_resolved)`` — runtime, per-thread,
+         computed once.
+      2. ``signed_strides[j]`` for each binary iter j, written into a local
+         buffer using the sigma formula above.
+
+    Returns ``(signed_strides_buffer_or_None, base_off_primexpr)``. The
+    buffer is ``None`` when ``pattern.n_binary_iters == 0`` (no outer
+    iter, no signed_strides needed).
+
+    ``s_off_resolved`` is the per-thread offset with the real tid Var
+    substituted in (not the placeholder).
+    """
+    swizzle = pattern.swizzle
+    p = swizzle.per_element
+    sw = swizzle.swizzle_len
+    at = swizzle.atom_len
+    C = 1 << p
+
+    base_off = swizzle.apply(s_off_resolved)["m"]
+
+    n = pattern.n_binary_iters
+    if n == 0:
+        return None, base_off
+
+    signed_strides = T.alloc_buffer([n], "int32", scope="local")
+    q = tvm.tirx.floordiv(s_off_resolved, C)
+
+    def _sigma_bit(bit_pos: int):
+        # 1 - 2 * bit_(bit_pos)(q); ∈ {+1, -1}.
+        row_bit = tvm.tirx.bitwise_and(
+            tvm.tirx.shift_right(q, _IntImm("int32", bit_pos)),
+            _IntImm("int32", 1),
+        )
+        return _IntImm("int32", 1) - row_bit * _IntImm("int32", 2)
+
+    for j, (bj, stride) in enumerate(zip(pattern.bit_positions, pattern.iter_strides_elems)):
+        stride_pow = stride  # = 2^(bj + p) elements
+        if 0 <= bj < sw:
+            # Case 1.A (inner): signed_stride = sigma_(at + bj) · T.
+            value = _sigma_bit(at + bj) * _IntImm("int32", stride_pow)
+        elif sw <= bj < at:
+            # Case 1.B (mid): signed_stride = +T.
+            value = _IntImm("int32", stride_pow)
+        elif at <= bj < at + sw:
+            # Case 1.C (outer): signed_stride = T + sigma_(bj - at) · T_sec.
+            # Invariant: bj >= at, so T_sec = T >> at = 2^(bj - at + p)
+            # = T(bj - at) is well-defined (no underflow).
+            stride_sec = stride_pow >> at
+            value = _IntImm("int32", stride_pow) + _sigma_bit(bj - at) * _IntImm(
+                "int32", stride_sec
+            )
+        else:  # bj >= at + sw, Case 1.D (above)
+            # No swizzle effect at this bit; signed_stride = +T.
+            value = _IntImm("int32", stride_pow)
+        # NB: Buffer.__setitem__ syntax (``signed_strides[j] = value``) is
+        # intercepted by the TIRx script parser but not by raw Python when
+        # this function is called from outside an @T.inline body. Use the
+        # low-level buffer_store builder instead.
+        T.buffer_store(signed_strides, value, [_IntImm("int32", j)])
+
+    return signed_strides, base_off
+
+
+def emit_iter_offset(pattern: SwizzlePattern, signed_strides, base_off, k):
+    """Compute the per-mm physical S offset = ``base_off`` + sum of per-iter
+    contributions.
+
+    ``k`` is the flat outer iter index ∈ ``[0, prod(it.ext for it in outer_iters))``.
+    Decomposed innermost-first across ``pattern.outer_iters`` into per-iter
+    coords ``c_i``. Each iter contributes:
+
+      * ``_BitIter``: ``sum_b bit_(n_bits-1-b)(c_i) * signed_strides[slot_start + b]``,
+        i.e. each binary bit of ``c_i`` selects its precomputed sigma-stride.
+        The slot order (outermost-first within the iter) means the highest
+        bit of ``c_i`` indexes the slot at ``slot_start``.
+      * ``_LinearIter``: ``c_i * stride`` (no bit decomposition; used when
+        ``stride`` is a multiple of ``2^(p + at + sw)`` so swizzle has no
+        XOR effect and ``ext`` need not be pow2).
+
+    Two paths per iter:
+      * Python int ``k`` — coords and bits known at parse time; emits only
+        the necessary adds, no runtime shift/mask.
+      * TIRx Var ``k`` — emits floormod/floordiv + bit-and/shift; relies on
+        downstream unroll + constant-fold.
+    """
+    if not pattern.outer_iters:
+        return base_off
+
+    off = base_off
+    remaining = k
+    is_const = isinstance(k, int)
+    for it in reversed(pattern.outer_iters):  # innermost first
+        ext = it.ext
+        if is_const:
+            c = remaining % ext
+            remaining = remaining // ext
+        else:
+            c = tvm.tirx.floormod(remaining, _IntImm("int32", ext))
+            remaining = tvm.tirx.floordiv(remaining, _IntImm("int32", ext))
+        if isinstance(it, _LinearIter):
+            if is_const:
+                if c != 0:
+                    off = off + c * it.stride
+            else:
+                off = off + c * _IntImm("int32", it.stride)
+            continue
+        # _BitIter
+        for b in range(it.n_bits):
+            bit_pos = it.n_bits - 1 - b
+            slot = it.slot_start + b
+            if is_const:
+                if (c >> bit_pos) & 1:
+                    off = off + signed_strides[slot]
+            else:
+                bit = tvm.tirx.bitwise_and(
+                    tvm.tirx.shift_right(c, _IntImm("int32", bit_pos)),
+                    _IntImm("int32", 1),
+                )
+                off = off + bit * signed_strides[slot]
+    return off
+
+
+def emit_fallback_offset(swizzle: SwizzleLayout, s_off_resolved, ds_k):
+    """Slow but always-correct path: full ``swizzle.apply(s_off + ds_k)``
+    per iter. Use when ``try_recognize`` returns ``None``.
+
+    ``ds_k`` is the outer-iter delta for unrolled iter k — typically a
+    Expr (a function of the unroll var that simplifies to a constant
+    after unrolling) or a Python int. ``s_off_resolved`` is the per-thread
+    base linear offset with the real tid Var substituted.
+    """
+    return swizzle.apply(s_off_resolved + ds_k)["m"]

python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py

@@ -0,0 +1,116 @@
+# 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.
+
+"""Scalar single-thread copy fallback (priority=0)."""
+
+import warnings
+
+import tvm
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx.operator.tile_primitive.dispatcher import (
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ._common import _TID_AXIS_FOR_SCOPE
+from .reg import _axis_decl
+from .utils import _is_valid_copy
+
+
+def _region_st_extent(buffer_region):
+    region = buffer_region.region
+    return [r.min for r in region], [r.extent for r in region]
+
+
+def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src: Buffer = op_call.src.buffer
+    dst: Buffer = op_call.dst.buffer
+    src_st, src_extent = _region_st_extent(op_call.src)
+    dst_st, dst_extent = _region_st_extent(op_call.dst)
+
+    warnings.warn(
+        f"copy/fallback (scalar single-thread) picked for {src.scope()} -> "
+        f"{dst.scope()} at scope_kind={sctx.scope_kind}; all faster variants "
+        f"rejected.",
+        stacklevel=2,
+    )
+
+    def _copy_body(dst_buf, src_buf):
+        dst_indices = [i for i in range(len(dst_buf.shape)) if dst_extent[i] != 1]
+        src_indices = [i for i in range(len(src_buf.shape)) if src_extent[i] != 1]
+        assert len(dst_indices) == len(src_indices)
+        copy_extents = [dst_extent[i] for i in dst_indices]
+
+        def _dst_coord(lvs):
+            if isinstance(lvs, tvm.tirx.Var):
+                lvs = [lvs]
+            coord = list(dst_st)
+            for k, lv in enumerate(lvs):
+                coord[dst_indices[k]] += lv
+            return coord
+
+        def _src_coord(lvs):
+            if isinstance(lvs, tvm.tirx.Var):
+                lvs = [lvs]
+            coord = list(src_st)
+            for k, lv in enumerate(lvs):
+                coord[src_indices[k]] += lv
+            return coord
+
+        with T.grid(*copy_extents) as lvs:
+            T.buffer_store(dst_buf, src_buf[tuple(_src_coord(lvs))], _dst_coord(lvs))
+
+    scope_kind = sctx.scope_kind
+
+    if scope_kind == "thread":
+
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            _copy_body(dst, src)
+
+        return impl
+
+    tid_axis_name = _TID_AXIS_FOR_SCOPE[scope_kind]
+    # first-active tid = composition of per-axis offsets (radix-32, since a warp is 32 lanes)
+    first_tid = int(sctx.intra["laneid"][1])
+    if scope_kind == "warpgroup":
+        first_tid += 32 * int(sctx.intra["wid_in_wg"][1])
+    elif scope_kind == "cta":
+        first_tid += 32 * int(sctx.intra["warpid"][1])
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        tid = _axis_decl(tid_axis_name, sctx)
+        if tid == first_tid:
+            _copy_body(dst, src)
+
+    return impl
+
+
+@register_dispatch(
+    "copy",
+    "cuda",
+    variant="fallback",
+    priority=0,
+    when=[predicate("validate_copy_op", _is_valid_copy)],
+)
+def copy_schedule_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return _emit_fallback(op_call, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py

@@ -0,0 +1,328 @@
+# 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.
+
+"""Copy dispatch for ``global ↔ shared`` (no register side).
+
+There's no per-thread register side to inherit a partition from — both sides
+are cross-thread storage. The partition is synthesized from the surrounding
+scope context (warp / warpgroup / cta / thread): ``thread_cnt`` is derived
+from ``sctx.intra`` and each thread takes ``n_elements / thread_cnt``
+consecutive fused-index slots. Layout / partition algorithm lives in
+``_common.py`` and is shared with ``ldgsts.py``.
+"""
+
+import tvm
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx import Var as _TirVar
+from tvm.tirx.expr import IntImm as _IntImm
+from tvm.tirx.operator.tile_primitive.dispatcher import (
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ._common import (
+    _TID_AXIS_FOR_SCOPE,
+    _thread_cnt,
+    align_layouts_gs,
+    copy_ptx_form,
+    copy_ptx_ld_return_type,
+)
+from ._swizzle_iter import (
+    emit_init,
+    emit_iter_offset,
+    get_swizzle,
+    try_recognize,
+)
+from .reg import _all_threads_active, _axis_decl, _ptr_off
+from .utils import _is_valid_copy, _scope_allowed
+
+_GMEM_SMEM_PAIRS = [
+    ("global", "shared*"),
+    ("shared*", "global"),
+]
+
+
+def _divides_thread_cnt(
+    op_call: TilePrimitiveCall, sctx: DispatchContext
+) -> tuple[bool, str | None]:
+    """Reject copies whose region element count does not divide ``thread_cnt``.
+
+    Without this guard the emit's ``[outer, T, vec]`` partition has no
+    integer solution: either every thread gets fractional work, or
+    ``thread_cnt=0`` (degenerate scope) hits a modulo-by-zero. Both cases
+    indicate a poorly-shaped copy (e.g. 1024-thread CTA writing a 64-elem
+    tail) that this dispatch refuses to paper over with a slow scalar emit.
+    """
+    op_call = TilePrimitiveCall.downcast(op_call)
+    thread_cnt = _thread_cnt(sctx)
+    if thread_cnt <= 0:
+        return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
+    g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst
+    n_elements = 1
+    for r in g_br.region:
+        ext = r.extent
+        try:
+            n_elements *= int(ext)
+        except (TypeError, ValueError):
+            return False, f"non-constant region extent {ext}"
+    if n_elements % thread_cnt != 0:
+        return False, (f"region size {n_elements} not divisible by thread_cnt={thread_cnt}")
+    return True, None
+
+
+def _is_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+    if not sctx.is_target("cuda"):
+        return False, "non-cuda target"
+    if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
+        return False, f"unsupported exec_scope {sctx.scope_kind}"
+    for check in (
+        lambda: _all_threads_active(sctx),
+        lambda: _is_valid_copy(op_call, sctx),
+        lambda: _scope_allowed(op_call, sctx, allowed_pairs=_GMEM_SMEM_PAIRS),
+        lambda: _divides_thread_cnt(op_call, sctx),
+    ):
+        ok, msg = check()
+        if not ok:
+            return False, msg
+    return True, None
+
+
+def _emit_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src: Buffer = op_call.src.buffer
+    dst: Buffer = op_call.dst.buffer
+    if src.scope() == "global":
+        g_buf, g_br, s_buf, s_br = src, op_call.src, dst, op_call.dst
+        g_is_src = True
+    else:
+        g_buf, g_br, s_buf, s_br = dst, op_call.dst, src, op_call.src
+        g_is_src = False
+
+    g_region = [(r.min, r.min + r.extent) for r in g_br.region]
+    s_region = [(r.min, r.min + r.extent) for r in s_br.region]
+
+    elem_bits = DataType(src.dtype).bits
+    thread_cnt = _thread_cnt(sctx)
+
+    with sctx.target:
+        g_p, s_p, vec_len = align_layouts_gs(
+            g_buf.layout,
+            g_buf.shape,
+            g_region,
+            s_buf.layout,
+            s_buf.shape,
+            s_region,
+            elem_bits,
+            thread_cnt,
+        )
+
+    # vec_len=1 is the scalar fallback — uses the same unified
+    # [outer x thread x vec] coord scheme below.
+
+    vec_bits = vec_len * elem_bits
+    num_bytes = vec_bits // 8
+    vec, ptx_type = copy_ptx_form(num_bytes)
+
+    # Partition guarantees ``prod(s_p.shard.extents) == prod(g_p.shard.extents)
+    # == n_elements`` (the total transfer count). Express the per-thread
+    # per-round address as a 3D coord ``(f, tid, 0)`` against shape
+    # ``[total_outer, thread_cnt, vec_len]``, and let ``layout.apply`` flatten
+    # it through whatever multi-iter T / outer-iter structure ``align_layouts_gs``
+    # picked. This makes the emit oblivious to how many iters the partition
+    # split T or outer across.
+    n_elements = 1
+    for it in s_p.shard:
+        n_elements *= int(it.extent)
+    assert n_elements % (thread_cnt * vec_len) == 0, (
+        f"partition produced {n_elements} elements but thread_cnt({thread_cnt}) * "
+        f"vec_len({vec_len}) = {thread_cnt * vec_len} doesn't divide it"
+    )
+    total_outer = n_elements // (thread_cnt * vec_len)
+    apply_shape = [
+        _IntImm("int32", total_outer),
+        _IntImm("int32", thread_cnt),
+        _IntImm("int32", vec_len),
+    ]
+
+    s_zero = [0] * len(s_buf.shape)
+    g_zero = [0] * len(g_buf.shape)
+
+    tid_axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind] if thread_cnt > 1 else None
+
+    # Walk shard from the vec iter backward to find the prefix that covers
+    # the T region exactly (∏ext == thread_cnt). The iters consumed are T
+    # iters; the leading prefix is the outer iter list — handed to
+    # ``try_recognize`` so the swizzle fast path can decide whether the
+    # outer iter strides match a pattern it can lower to signed_strides.
+    if thread_cnt > 1:
+        acc, _i = 1, len(s_p.shard) - 2
+        while _i >= 0 and acc < thread_cnt:
+            _ext = int(s_p.shard[_i].extent)
+            if acc * _ext > thread_cnt:
+                break
+            acc *= _ext
+            _i -= 1
+        outer_iters_s = list(s_p.shard[: _i + 1]) if acc == thread_cnt else []
+    else:
+        outer_iters_s = list(s_p.shard[:-1])
+
+    # SwizzleLayout on s_buf: try the closed-form signed-strides pattern
+    # (precomputed once per thread, then per-iter is a sum of register
+    # adds); fall back to per-iter ``swizzle.apply`` (one full XOR +
+    # decompose per iter). Closure picked at parse time so the TIRx parser
+    # doesn't AST-evaluate a "dead" ternary branch.
+    swizzle = get_swizzle(s_buf.layout)
+    swizzle_pattern = None
+    if swizzle is not None and outer_iters_s:
+        if tid_axis_name is not None:
+            _tid_placeholder = _TirVar(tid_axis_name, "int32")
+        else:
+            _tid_placeholder = _IntImm("int32", 0)
+        s_off_template = s_p.apply(
+            _IntImm("int32", 0),
+            _tid_placeholder,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )["m"]
+        # Bind the tid placeholder's range so the (C1) analyzer check can
+        # discharge ``bit_bj(s_off // C) == 0`` for high bj's. Outer iter
+        # stride here is ``thread_cnt * vec_len`` ⇒ bj ∈ [log2(thread_cnt),
+        # ...]; without bounds the analyzer can't prove the lane's high bits
+        # are 0 and rejects.
+        var_bounds = {}
+        if tid_axis_name is not None:
+            var_bounds[_tid_placeholder] = tvm.ir.Range.from_min_extent(0, thread_cnt)
+        swizzle_pattern = try_recognize(
+            swizzle,
+            [int(it.extent) for it in outer_iters_s],
+            [int(it.stride) for it in outer_iters_s],
+            s_off_template,
+            var_bounds=var_bounds or None,
+        )
+
+    class _SwizzleState:
+        def __init__(self):
+            self.signed_strides = None
+            self.base_off = None
+
+    state = _SwizzleState()
+
+    def _decl_tid():
+        if tid_axis_name is not None:
+            return _axis_decl(tid_axis_name, sctx)
+        return _IntImm("int32", 0)
+
+    def _setup_swizzle(tid):
+        if swizzle_pattern is None:
+            return
+        s_off_resolved = s_p.apply(
+            _IntImm("int32", 0),
+            tid,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )["m"]
+        state.signed_strides, state.base_off = emit_init(
+            swizzle_pattern,
+            s_off_resolved,
+        )
+
+    if swizzle_pattern is not None:
+
+        def _s_off(f, s_lin):
+            return emit_iter_offset(
+                swizzle_pattern,
+                state.signed_strides,
+                state.base_off,
+                f,
+            )
+    elif swizzle is not None:
+        _sw = swizzle
+
+        def _s_off(f, s_lin):
+            return _sw.apply(s_lin)["m"]
+    else:
+
+        def _s_off(f, s_lin):
+            return s_lin
+
+    v0 = _IntImm("int32", 0)
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        tid = _decl_tid()
+        _setup_swizzle(tid)
+        tmp = T.alloc_local((vec_len,), src.dtype)
+        tmp_ptr = tmp.ptr_to([0])
+        # NB: pass typed ptr_to(...) directly to _ptr_off; caching in a
+        # local var turns it into void* + offset = byte arithmetic →
+        # misaligned vector ops.
+        #
+        # Use a serial TIR loop and let ptxas unroll downstream. Mirrors
+        # the reg.py rationale in commit ac7ecf70f0: explicit ``T.unroll``
+        # materializes the per-iter scratch (s_lin/g_lin/s_off/s_ptr/g_ptr)
+        # as N copies of each ``alignas(64)`` declaration. For large
+        # ``total_outer`` (e.g. thread-scope fp32 swizzled copies of 32x256
+        # at vec=4 ⇒ 2048 iters; ldgsts test4 ⇒ ~4k iters once both
+        # g2s/s2g sites add up) this floods the kernel and nvcc times out.
+        for f in range(total_outer):
+            s_lin = s_p.apply(f, tid, v0, shape=apply_shape)["m"]
+            g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
+            s_off = _s_off(f, s_lin)
+            s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
+            g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
+            if g_is_src:
+                T.ptx.ld(
+                    g_ptr,
+                    copy_ptx_ld_return_type(ptx_type),
+                    ptx_type,
+                    dst=tmp_ptr,
+                    space="global",
+                    vec=vec,
+                )
+                T.ptx.st(
+                    s_ptr, src=tmp_ptr, space="shared", vec=vec, ptx_type=ptx_type
+                )
+            else:
+                T.ptx.ld(
+                    s_ptr,
+                    copy_ptx_ld_return_type(ptx_type),
+                    ptx_type,
+                    dst=tmp_ptr,
+                    space="shared",
+                    vec=vec,
+                )
+                T.ptx.st(
+                    g_ptr, src=tmp_ptr, space="global", vec=vec, ptx_type=ptx_type
+                )
+    # fmt: on
+    return impl
+
+
+@register_dispatch(
+    "copy",
+    "cuda",
+    variant="gmem_smem",
+    priority=10,
+    when=[predicate("gmem_smem_applicable", _is_gmem_smem)],
+)
+def copy_schedule_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return _emit_gmem_smem(op_call, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py

@@ -0,0 +1,450 @@
+# 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.
+
+"""copy dispatch variant: ldmatrix / stmatrix (TBD algorithm).
+
+Handles register ↔ shared copies on CUDA via PTX ``ldmatrix`` / ``stmatrix``.
+Direction (ld vs st) and exec scope (warp / warpgroup) are decided inside
+``_emit`` from the src/dst scopes and ``sctx.scope_kind``.
+"""
+
+from math import prod
+
+import tvm
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+from tvm.tirx import Var as _TirVar
+from tvm.tirx.expr import IntImm as _IntImm
+from tvm.tirx.layout import S, TileLayout
+from tvm.tirx.operator.tile_primitive.dispatcher import fail, predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ._common import (  # noqa: F401  (_carve_tail reserved for future variants)
+    _carve_tail,
+    _extract_tile,
+)
+from ._swizzle_iter import emit_init, emit_iter_offset, get_swizzle, try_recognize
+from .reg import _all_threads_active, _ptr_off
+from .utils import _is_valid_copy, _scope_allowed
+
+_REG_SMEM_PAIRS = [
+    ("local", "shared*"),
+    ("shared*", "local"),
+]
+
+_VALID_R_LANE_AXES = {"laneid", "tid_in_wg", "tx"}
+
+
+def _compute_r_perm(r):
+    """Permutation: thread iters first (stride-desc), then memory iters (stride-desc)."""
+
+    def key(p):
+        it = p[1]
+        return (0 if it.axis.is_thread() else 1, -int(it.stride))
+
+    return [i for i, _ in sorted(enumerate(r.shard), key=key)]
+
+
+def _is_ldstmatrix(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+    if not sctx.is_target("cuda"):
+        return False, "non-cuda target"
+    if sctx.scope_kind not in ("warp", "warpgroup", "cta"):
+        return False, f"unsupported exec_scope {sctx.scope_kind} (need warp, warpgroup, or cta)"
+    for check in (
+        lambda: _all_threads_active(sctx),
+        lambda: _is_valid_copy(op_call, sctx),
+        lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_SMEM_PAIRS),
+    ):
+        ok, msg = check()
+        if not ok:
+            return False, msg
+    return True, None
+
+
+def _emit(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    op_call = TilePrimitiveCall.downcast(op_call)
+
+    # Step 1: identify reg / smem sides and pull their tensor shape + layout.
+    src_br = op_call.src
+    dst_br = op_call.dst
+    if src_br.buffer.scope() == "local":
+        r_br, s_br = src_br, dst_br
+        direction = "st"  # reg -> smem (stmatrix)
+    else:
+        r_br, s_br = dst_br, src_br
+        direction = "ld"  # smem -> reg (ldmatrix)
+    r_buf = r_br.buffer
+    s_buf = s_br.buffer
+    r_shape = list(r_buf.shape)
+    r_layout = r_buf.layout
+    s_shape = list(s_buf.shape)
+    s_layout = s_buf.layout
+
+    r_region = [(r.min, r.min + r.extent) for r in r_br.region]
+    s_region = [(r.min, r.min + r.extent) for r in s_br.region]
+    with sctx.target:
+        r_sliced = r_layout.slice(r_shape, r_region)
+        s_sliced = s_layout.slice(s_shape, s_region)
+        r = r_sliced.canonicalize()
+        s = s_sliced.canonicalize()
+
+    # Step 2.5: peel any S-side swizzle wrapper to expose the underlying
+    # TileLayout. The ComposeLayout doesn't have ``.replica`` / ``.shard``,
+    # so we must peel *before* the structural checks below. Capture the
+    # swizzle separately for use at emit time.
+    # NB: read swizzle from the *buffer* layout, not the post-canon ``s``.
+    # When the underlying tile is trivial, ``ComposeLayoutNode::Canonicalize``
+    # returns a bare ``SwizzleLayout``; isinstance(s, ComposeLayout) is then
+    # False and we'd miss the swizzle here.
+    s_swizzle = get_swizzle(s_buf.layout)
+    if s_swizzle is not None and s_swizzle.per_element < 3:
+        # ldmatrix/stmatrix .b16 reads/writes 8 fp16 = 128b per lane in one
+        # contiguous chunk. The swizzle preserves the lowest ``per_element``
+        # bits of the address (in-chunk offset). For the per-lane 128b unit
+        # to stay contiguous post-swizzle, ``2^per_element >= 8`` ⇒ p >= 3.
+        fail(
+            f"swizzle per_element={s_swizzle.per_element} < 3 incompatible "
+            f"with .b16 ldmatrix/stmatrix (need 8-fp16 chunk integrity)"
+        )
+    s = _extract_tile(s, s_region)
+
+    # Step 3: ldstmatrix doesn't broadcast — require zero replica on both sides.
+    if len(r.replica) != 0:
+        fail(f"R layout has replica {list(r.replica)}; ldstmatrix requires no replica")
+    if len(s.replica) != 0:
+        fail(f"S layout has replica {list(s.replica)}; ldstmatrix requires no replica")
+
+    # Step 4: R must have exactly one kind of lane axis from the valid set.
+    r_thread_axes = {it.axis.name for it in r.shard if it.axis.is_thread()}
+    if len(r_thread_axes) != 1:
+        fail(f"R must have exactly one thread axis name; got {sorted(r_thread_axes)}")
+    r_lane_axis = next(iter(r_thread_axes))
+    if r_lane_axis not in _VALID_R_LANE_AXES:
+        fail(f"R thread axis {r_lane_axis!r} not in {sorted(_VALID_R_LANE_AXES)}")
+
+    # Step 5: group S by R's iter extents (one S group per R iter, outer→inner).
+    r_group_shape = [int(it.extent) for it in r.shard]
+    s_grp, s_seps = s.group(r_group_shape)
+
+    # Step 6: permute R so thread iters come first (stride-desc), then memory
+    # iters (stride-desc).
+    r_perm = _compute_r_perm(r)
+    r = r.permute_dims(r_perm)
+
+    # Step 7: apply R's perm to S in group units (1-to-1 with R's iters), and
+    # rebuild s_seps to track group boundaries in the new order.
+    s = s_grp.permute_by_groups(list(s_seps), r_perm)
+    old_sizes = [s_seps[i + 1] - s_seps[i] for i in range(len(s_seps) - 1)]
+    s_seps = [0]
+    for pi in r_perm:
+        s_seps.append(s_seps[-1] + old_sizes[pi])
+
+    # Step 7.5: canonicalize both R and S after permute. Fuses adjacent
+    # contig iters — keeps step 8's group input clean. Push target so
+    # scope-aware fusers run (laneid+wid_in_wg → tid_in_wg, etc.).
+    with sctx.target:
+        r = r.canonicalize()
+        s = s.canonicalize()
+
+    t_total = prod(int(it.extent) for it in r.shard if it.axis.is_thread())
+    m_total = prod(int(it.extent) for it in r.shard if not it.axis.is_thread())
+    if t_total % 32 != 0:
+        fail(f"R thread section total {t_total} not divisible by 32")
+
+    def _strs(lay, seps):
+        # Atoms 8 / 4 / 2 (segs 1, 2, 5) must be single iters — their strides
+        # feed downstream stride checks (lane partition + fragment 2-fp16
+        # contig). The num atom (seg 4) may be MULTI-ITER: we return its iter
+        # list and let layout.apply handle the decomposition at emit time.
+        fixed_segs = [list(lay.shard[seps[i] : seps[i + 1]]) for i in (1, 2, 5)]
+        if not all(len(g) == 1 for g in fixed_segs):
+            return None
+        num_iters = list(lay.shard[seps[4] : seps[5]])
+        return (
+            int(fixed_segs[0][0].stride),  # 8 atom stride
+            int(fixed_segs[1][0].stride),  # 4 atom stride
+            num_iters,  # num atom iter list (multi-iter OK)
+            int(fixed_segs[2][0].stride),  # 2 atom stride
+        )
+
+    def _try_num(r_in, s_in, num):
+        """Try grouping (r_in, s_in) with [T/32, 8, 4, M/(2num), num, 2].
+
+        Returns (rg, rsep, sg, ssep, trans, p, num) if structural checks pass,
+        else None. ``trans`` is the ldmatrix .trans flag; ``p`` is the
+        per-tile-row S stride used at emit.
+        """
+        gs = [t_total // 32, 8, 4, m_total // (num * 2), num, 2]
+        try:
+            rg, rsep = r_in.group(gs)
+            sg, ssep = s_in.group(gs)
+        except Exception:
+            return None
+        # R seg 0 (T/32 outer): require single iter with stride 32. When
+        # T/32 == 1 the segment is trivial — skip.
+        if t_total > 32:
+            seg0 = list(rg.shard[rsep[0] : rsep[1]])
+            if len(seg0) != 1 or int(seg0[0].stride) != 32:
+                return None
+        rs, ss = _strs(rg, rsep), _strs(sg, ssep)
+        if rs is None or ss is None:
+            return None
+        r8, r4, _r_num_iters, r2 = rs
+        s8, s4, s_num_iters, s2 = ss
+        if (r8, r4, r2) != (4, 1, 1):
+            return None
+        # S num atom: every iter must have stride > 0 and multiple of 8 (the
+        # per-tile spacing geometry of ldmatrix m8n8; 8 fp16 = 16 bytes = one
+        # tile column dimension).
+        if num > 1 and not all(
+            int(it.stride) > 0 and int(it.stride) % 8 == 0 for it in s_num_iters
+        ):
+            return None
+        # m_outer (seg 3) iters: each per-mm advance must keep the per-lane
+        # SMEM address 16-byte aligned (ldmatrix .b16 reads 8 fp16 = 16 bytes
+        # per lane), so the m_outer S-stride must also be a multiple of 8.
+        # Without this, mm > 0 iterations land at unaligned addresses and
+        # silently read garbage even though the layout group succeeds.
+        # Skip extent-1 trivial iters — they contribute no per-mm advance,
+        # so their (placeholder) stride is irrelevant.
+        m_outer_iters = list(sg.shard[ssep[3] : ssep[4]])
+        if not all(int(it.extent) == 1 or int(it.stride) % 8 == 0 for it in m_outer_iters):
+            return None
+        if (s4, s2) == (2, 1) and s8 > 0 and s8 % 8 == 0:
+            return (rg, rsep, sg, ssep, False, s8, num)
+        if s8 == 1 and s2 > 0 and s2 % 8 == 0 and s4 == 2 * s2:
+            return (rg, rsep, sg, ssep, True, s2, num)
+        return None
+
+    # Try the **sorted** variant: 5D-group, sub-group R's M/2 by S's M/2
+    # extents, sort the sub-groups by descending S-stride, rebuild. This
+    # makes the m_outer iter list carry the largest S-strides on top, which
+    # maximizes the §2 swizzle fast-path applicability later. If anything
+    # in the rebuild raises (e.g. M/2 can't be sub-grouped by S's extents),
+    # we silently fall back to the no-sort path below.
+    r_sort = s_sort = None
+    try:
+        gs5 = [t_total // 32, 8, 4, m_total // 2, 2]
+        rg5, rsep5 = r.group(gs5)
+        sg5, ssep5 = s.group(gs5)
+        r_m_iters = list(rg5.shard[rsep5[3] : rsep5[4]])
+        s_m_iters = list(sg5.shard[ssep5[3] : ssep5[4]])
+        s_m_extents = [int(it.extent) for it in s_m_iters]
+        # Sub-group R's M/2 iters by S's M/2 iter extents. This 1-to-1's
+        # the R sub-groups with the S iters so we can permute them together.
+        r_m_sub = TileLayout.from_iters(r_m_iters)
+        r_m_grouped, r_m_seps = r_m_sub.group(s_m_extents)
+        # Sort S iters by S-stride descending; permute R sub-groups in lockstep.
+        perm = sorted(range(len(s_m_iters)), key=lambda i: -int(s_m_iters[i].stride))
+        if perm != list(range(len(perm))):
+            r_m_permuted = r_m_grouped.permute_by_groups(list(r_m_seps), perm)
+            s_m_permuted = [s_m_iters[i] for i in perm]
+            r_sort = TileLayout.from_iters(
+                list(rg5.shard[: rsep5[3]])
+                + list(r_m_permuted.shard)
+                + list(rg5.shard[rsep5[4] :]),
+                offset=dict(rg5.offset),
+            )
+            s_sort = TileLayout.from_iters(
+                list(sg5.shard[: ssep5[3]]) + list(s_m_permuted) + list(sg5.shard[ssep5[4] :]),
+                offset=dict(sg5.offset),
+            )
+        # If perm is identity, sorted == unsorted; no need to build duplicate layouts.
+    except Exception:
+        r_sort = s_sort = None
+
+    # Enumerate num largest-first; for each num try sorted then unsorted.
+    chosen = None
+    for num in (4, 2, 1):
+        if m_total % (num * 2):
+            continue
+        if r_sort is not None:
+            res = _try_num(r_sort, s_sort, num)
+            if res is not None:
+                chosen = res
+                break
+        res = _try_num(r, s, num)
+        if res is not None:
+            chosen = res
+            break
+
+    if chosen is None:
+        fail("ldstmatrix layout doesn't fit any num ∈ {4,2,1}")
+    r, r_seps, s, s_seps, trans, p, num = chosen
+
+    # Step 10: emit one ldmatrix/stmatrix per mm, per warp.
+
+    def _get_warp_idx_in_T():
+        # T.warp_id_in_wg() / T.warp_id() must be called from inside a
+        # @T.prim_func body — wrap so the prim_func parser calls us at parse
+        # time (Python `if` here is plain control flow, not TIR-intercepted).
+        if r_lane_axis == "laneid":
+            return 0
+        if r_lane_axis == "tid_in_wg":
+            return T.warp_id_in_wg()
+        return T.warp_id()  # "tx"
+
+    def _seg4_coord(laneid_expr):
+        # num=1: seg 4 trivially extent-1, pass 0. num>1: use lane//8 (tile
+        # index in ldmatrix lane convention); layout.apply decomposes through
+        # the seg's iter structure (single or multi-iter).
+        if num > 1:
+            return laneid_expr // 8
+        return 0
+
+    apply_shape = [t_total // 32, 8, 4, m_total // (num * 2), num, 2]
+    r_mem_axis = r.shard[r_seps[5]].axis.name
+    s_mem_axis = s.shard[s_seps[5]].axis.name
+    m_outer = m_total // (num * 2)
+    s_zero = [0] * len(s_buf.shape)
+
+    # Swizzle fast-path setup. When S is swizzled, the per-mm `tile_off +
+    # row_off` is a logical offset; the physical SMEM address is
+    # `swizzle.apply(logical)`. The slow path computes that per iter; the
+    # fast path (§2.E of the swizzle-iter plan) reduces it to
+    # `base_off + sum_j bit_j(mm) · signed_strides[j]` where base_off and
+    # signed_strides are per-thread constants set once. We try to
+    # recognize the m_outer iter list as such a pattern; if it fails (e.g.
+    # the analyzer can't discharge condition C1 over the lane/warp
+    # placeholders) we silently fall through to the slow path.
+    swizzle_pattern = None
+    s_off_template = None
+    lane_ph = warp_ph = None
+    if s_swizzle is not None:
+        m_outer_iters = list(s.shard[s_seps[3] : s_seps[4]])
+        iter_extents = [int(it.extent) for it in m_outer_iters]
+        iter_strides = [int(it.stride) for it in m_outer_iters]
+        # Build s_off at mm=0 with placeholder vars for lane and warp_idx.
+        lane_ph = _TirVar("lane_ph", "int32")
+        seg4_ph = (lane_ph // 8) if num > 1 else _IntImm("int32", 0)
+        if r_lane_axis == "laneid":
+            warp_ph_expr = _IntImm("int32", 0)
+        else:
+            warp_ph = _TirVar("warp_ph", "int32")
+            warp_ph_expr = warp_ph
+        s_off_template = s.apply(
+            warp_ph_expr,
+            _IntImm("int32", 0),
+            _IntImm("int32", 0),
+            _IntImm("int32", 0),
+            seg4_ph,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )[s_mem_axis] + (lane_ph % 8) * _IntImm("int32", p)
+        # Bind lane / warp placeholder bounds for the (C1) analyzer. ``lane_ph``
+        # is the per-warp lane id ∈ [0, 32); ``warp_ph`` (when present) is the
+        # warp index inside the scope: warpgroup ⇒ [0, 4), cta ⇒ [0, t_total/32).
+        var_bounds = {lane_ph: tvm.ir.Range.from_min_extent(0, 32)}
+        if warp_ph is not None:
+            var_bounds[warp_ph] = tvm.ir.Range.from_min_extent(0, t_total // 32)
+        swizzle_pattern = try_recognize(
+            s_swizzle,
+            iter_extents,
+            iter_strides,
+            s_off_template,
+            var_bounds=var_bounds,
+        )
+
+    class _SwizzleState:
+        def __init__(self):
+            self.signed_strides = None
+            self.base_off = None
+
+    state = _SwizzleState()
+
+    def _resolve_s_off(laneid_var, warp_var):
+        # Build the placeholder→runtime-var map and substitute. Keep this in a
+        # regular Python helper — the @T.prim_func parser intercepts dict
+        # literals when written directly in the body.
+        vmap = {lane_ph: laneid_var}
+        if warp_ph is not None:
+            vmap[warp_ph] = warp_var
+        return tvm.tirx.stmt_functor.substitute(s_off_template, vmap)
+
+    def _setup_swizzle(s_off_resolved):
+        if swizzle_pattern is None:
+            return
+        state.signed_strides, state.base_off = emit_init(
+            swizzle_pattern,
+            s_off_resolved,
+        )
+
+    def _smem_off(mm_idx, logical_off):
+        # Three paths:
+        #   * pattern matched: physical off = base_off + Σ bit_j(mm)·ss[j].
+        #   * swizzle present, pattern missed: per-iter swizzle.apply(logical).
+        #   * no swizzle: identity.
+        if swizzle_pattern is not None:
+            return emit_iter_offset(
+                swizzle_pattern,
+                state.signed_strides,
+                state.base_off,
+                mm_idx,
+            )
+        if s_swizzle is not None:
+            return s_swizzle.apply(logical_off)["m"]
+        return logical_off
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        r_local = r_buf.local(m_total, layout=TileLayout(S[(m_total,)]))
+        laneid = T.lane_id()
+        warp_idx_in_T = _get_warp_idx_in_T()
+        # Resolve s_off_template by substituting placeholders → actual
+        # scope-id vars (via _resolve_s_off helper to keep the dict literal
+        # out of the parser's view). Only the swizzle fast path needs this;
+        # without swizzle we keep using the per-iter s.apply directly.
+        if swizzle_pattern is not None:
+            _setup_swizzle(_resolve_s_off(laneid, warp_idx_in_T))
+        for mm in T.unroll(m_outer):
+            tile_off = s.apply(
+                warp_idx_in_T, 0, 0, mm, _seg4_coord(laneid), 0, shape=apply_shape,
+            )[s_mem_axis]
+            row_off = (laneid % 8) * p
+            logical_off = tile_off + row_off
+            smem_ptr = _ptr_off(s_buf.ptr_to(s_zero), _smem_off(mm, logical_off))
+            handles = [
+                r_local.ptr_to([
+                    r.apply(0, 0, 0, mm, i, 0, shape=apply_shape)[r_mem_axis]
+                ])
+                for i in range(num)
+            ]
+            if direction == "ld":
+                T.ptx.ldmatrix(trans, num, ".b16", smem_ptr, *handles)
+            else:
+                T.ptx.stmatrix(
+                    trans, num, ".b16", smem_ptr, *handles,
+                    shape="m8n8", space="shared",
+                )
+    # fmt: on
+    return impl
+
+
+@register_dispatch(
+    "copy",
+    "cuda",
+    variant="ldstmatrix",
+    priority=10,
+    when=[predicate("ldstmatrix_applicable", _is_ldstmatrix)],
+)
+def copy_schedule_ldstmatrix(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return _emit(op_call, sctx)
+
+
+__all__ = ["copy_schedule_ldstmatrix"]

python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py

@@ -0,0 +1,596 @@
+# 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.
+
+"""Non-ldmatrix copy dispatch for register ↔ memory.
+
+This file owns every copy where one side is per-thread local (``R`` =
+register). That R side carries the partition: its ``TileLayout`` ``shard``
+has thread-axis iters telling us which thread owns which logical coordinate.
+The other side (``S``) can be ``shared*`` or ``global`` — the algorithm is
+identical either way.
+
+Slice/canonicalize both sides, align via perm+group, then emit a per-thread
+vectorized copy loop. Direction-symmetric: covers R2S / S2R / R2G / G2R.
+"""
+
+import tvm
+from tvm.arith import Analyzer
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx import Var as _TirVar
+from tvm.tirx.expr import IntImm as _IntImm
+from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
+from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_return_type
+from ._swizzle_iter import (
+    emit_fallback_offset,
+    emit_init,
+    emit_iter_offset,
+    get_swizzle,
+    try_recognize,
+)
+from .utils import _is_valid_copy, _scope_allowed
+
+
+def _extract_tile(layout, region):
+    """Strip swizzle off ``layout`` so we can perm/group it as a TileLayout.
+
+    ``region`` is the per-axis ``(start, end)`` pair list — we only consume
+    its extents when ``layout`` is a bare ``SwizzleLayout`` (rebuilding a
+    trivial TileLayout for it). Plain ``TileLayout`` / ``ComposeLayout``
+    don't need the extent, so symbolic regions are fine for them.
+    """
+    if isinstance(layout, ComposeLayout):
+        return layout.tile_layout
+    if isinstance(layout, SwizzleLayout):
+        # TODO: keep swizzle info around for later (addressing in emit).
+        extents = [int(end - start) for (start, end) in region]
+        return TileLayout(S[tuple(extents)])
+    return layout
+
+
+_REG_PAIRS = [
+    ("local", "shared*"),
+    ("shared*", "local"),
+    ("local", "global"),
+    ("global", "local"),
+]
+_SCOPE_RANK = {"thread": 0, "warp": 1, "warpgroup": 2, "cta": 3}
+_VALID_R_SUBSCOPES = {"thread", "warp", "warpgroup"}
+
+
+def _all_threads_active(sctx: DispatchContext) -> tuple[bool, str | None]:
+    if sctx.scope_kind == "thread":
+        return True, None
+    required: dict[str, int] = {}
+    if sctx.scope_kind in ("warp", "warpgroup", "cta"):
+        required["laneid"] = 32
+    if sctx.scope_kind == "warpgroup":
+        required["wid_in_wg"] = 4
+    if sctx.scope_kind == "cta":
+        tx_iv = sctx.launch_params.get("threadIdx.x")
+        if tx_iv is None:
+            return False, "cta scope missing threadIdx.x launch_params"
+        try:
+            required["warpid"] = int(tx_iv.dom.extent) // 32
+        except (TypeError, ValueError):
+            return False, f"non-static threadIdx.x extent: {tx_iv.dom.extent}"
+    for axis_name, expected in required.items():
+        if axis_name not in sctx.intra:
+            return False, f"sctx.intra missing {axis_name!r}"
+        ext_raw, off_raw = sctx.intra[axis_name]
+        try:
+            ext, off = int(ext_raw), int(off_raw)
+        except (TypeError, ValueError):
+            return False, f"non-static range for {axis_name}: ({ext_raw}, {off_raw})"
+        if ext != expected or off != 0:
+            return False, f"{axis_name} narrowed to [{off}, {off + ext}) vs full [0, {expected})"
+    return True, None
+
+
+def _r_side_layout_valid(
+    op_call: TilePrimitiveCall, sctx: DispatchContext
+) -> tuple[bool, str | None]:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src: Buffer = op_call.src.buffer
+    dst: Buffer = op_call.dst.buffer
+    r_buf = src if src.scope() == "local" else dst
+    layout = r_buf.layout
+    if layout is None:
+        return False, "R has no layout"
+    if layout.is_swizzle():
+        return False, "R layout is swizzle"
+    if not isinstance(layout, TileLayout):
+        return False, f"R layout is {type(layout).__name__}, not TileLayout"
+
+    scope_rank = _SCOPE_RANK[sctx.scope_kind]
+    for it in layout.shard:
+        ax = it.axis
+        if not ax.is_thread():
+            continue
+        ax_scope = ax.get_scope()
+        ax_sub = ax.get_subscope()
+        if ax_scope is None or ax_sub is None:
+            return False, f"R thread axis {ax.name!r} missing scope/subscope"
+        if ax_sub.name not in _VALID_R_SUBSCOPES:
+            return False, f"R thread axis {ax.name!r} subscope={ax_sub.name!r} (not register-level)"
+        if ax_scope.name not in _SCOPE_RANK or _SCOPE_RANK[ax_scope.name] > scope_rank:
+            return (
+                False,
+                f"R thread axis {ax.name!r} scope={ax_scope.name!r} > exec {sctx.scope_kind!r}",
+            )
+    r_br = op_call.src if src.scope() == "local" else op_call.dst
+    region = [(r.min, r.min + r.extent) for r in r_br.region]
+    sliced = layout.slice(list(r_buf.shape), region)
+    if sliced is None:
+        return False, "R layout slice failed"
+    analyzer = Analyzer()
+    for axis, off in sliced.offset.items():
+        if axis.is_thread() and not analyzer.can_prove_equal(off, 0):
+            return False, f"R sliced offset on thread axis {axis.name!r} = {off}"
+    return True, None
+
+
+def _s_side_slice_ok(op_call: TilePrimitiveCall) -> tuple[bool, str | None]:
+    """S is the non-local side (shared* or global). Slice must succeed."""
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src_br = op_call.src
+    dst_br = op_call.dst
+    s_br = dst_br if src_br.buffer.scope() == "local" else src_br
+    s_buf: Buffer = s_br.buffer
+    layout = s_buf.layout
+    if layout is None:
+        return False, "S has no layout"
+    region = [(r.min, r.min + r.extent) for r in s_br.region]
+    if layout.slice(list(s_buf.shape), region) is None:
+        return False, "S layout slice failed"
+    return True, None
+
+
+def _is_reg_copy(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+    if not sctx.is_target("cuda"):
+        return False, "non-cuda target"
+    if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
+        return False, f"unsupported exec_scope {sctx.scope_kind}"
+    for check in (
+        lambda: _all_threads_active(sctx),
+        lambda: _is_valid_copy(op_call, sctx),
+        lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_PAIRS),
+        lambda: _r_side_layout_valid(op_call, sctx),
+        lambda: _s_side_slice_ok(op_call),
+    ):
+        ok, msg = check()
+        if not ok:
+            return False, msg
+    return True, None
+
+
+def _compute_perm_r(r):
+    # thread axes first, then by stride descending
+    def key(p):
+        it = p[1]
+        return (0 if it.axis.is_thread() else 1, -int(it.stride))
+
+    return [i for i, _ in sorted(enumerate(r.shard), key=key)]
+
+
+def align_layouts_raw(r_sliced, s_sliced, s_region):
+    """Returns (r_p, s_p, s_seps, r_perm).
+
+    ``r_p`` for ``_s_thread_offset``; ``r_perm`` for ``_split_thread_loop`` pairing.
+    """
+    r = r_sliced.canonicalize()
+    s = s_sliced.canonicalize()
+    s = _extract_tile(s, s_region)
+    perm = _compute_perm_r(r)
+    r_shape_for_group = [int(it.extent) for it in r.shard]
+    s_grp, seps = s.group(r_shape_for_group)
+    s_p = s_grp.permute_by_groups(list(seps), perm)
+    r_perm = r.permute_dims(perm)
+    r_p = r_perm.canonicalize()
+    sizes = [seps[i + 1] - seps[i] for i in range(len(seps) - 1)]
+    s_seps = [0]
+    for p in perm:
+        s_seps.append(s_seps[-1] + sizes[p])
+    return r_p, s_p, s_seps, r_perm
+
+
+def _split_thread_loop(r_perm, s_p, s_seps):
+    """Drop R's thread-axis positions and return per-R-position bundles:
+    (r_iters, s_groups) — same length lists; s_groups[k] is the list of S
+    iters belonging to the k-th kept R position.
+
+    ``r_perm`` is the permuted-but-uncanonicalized R layout: its iter list lines
+    up one-to-one with ``s_seps`` (which is built from the same ``perm``), so a
+    multi-group memory axis keeps each of its groups paired with the matching S
+    group instead of collapsing into one (see ``align_layouts_raw``)."""
+    r_iters = []
+    s_groups = []
+    for k, r_it in enumerate(r_perm.shard):
+        if r_it.axis.is_thread():
+            continue
+        r_iters.append(r_it)
+        s_groups.append(list(s_p.shard[s_seps[k] : s_seps[k + 1]]))
+    return r_iters, s_groups
+
+
+def _build_atoms(r_iters, s_groups):
+    """One atom per (R position, intra-group S iter): (extent, s_stride, r_mul).
+    r_mul = R_stride_at_position * (product of S group extents to the right
+    of this intra-position index) — i.e. how much R address advances per unit
+    of this iter's loop input."""
+    atoms = []
+    for r_it, s_group in zip(r_iters, s_groups, strict=True):
+        rs = int(r_it.stride)
+        extents = [int(it.extent) for it in s_group]
+        for j, s_it in enumerate(s_group):
+            inner_prod = 1
+            for e in extents[j + 1 :]:
+                inner_prod *= e
+            atoms.append((int(s_it.extent), int(s_it.stride), rs * inner_prod))
+    return atoms
+
+
+def _atoms_contiguous_tail_extent(atoms) -> int:
+    """Like _contiguous_tail_extent but on atoms (uses s_stride for chaining)."""
+    if not atoms or atoms[-1][1] != 1:
+        return 0
+    acc = atoms[-1][0]
+    for k in range(len(atoms) - 2, -1, -1):
+        if atoms[k][1] == acc:
+            acc *= atoms[k][0]
+        else:
+            break
+    return acc
+
+
+def _split_atoms_for_vec(atoms, vec_len):
+    """Returns outer atoms (the inner vec_len-element tail is consumed by one
+    vec ld/st and dropped). Splits the boundary atom if needed."""
+    outer = list(atoms)
+    acc = 1
+    while outer:
+        ext, ss, rm = outer[-1]
+        new_acc = acc * ext
+        if new_acc == vec_len:
+            outer.pop()
+            return outer
+        if new_acc > vec_len:
+            inner_factor = vec_len // acc
+            outer[-1] = (ext // inner_factor, ss * inner_factor, rm * inner_factor)
+            return outer
+        acc = new_acc
+        outer.pop()
+    raise ValueError(f"tail too short for vec_len {vec_len}")
+
+
+def _align_layouts(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src_br = op_call.src
+    dst_br = op_call.dst
+    if src_br.buffer.scope() == "local":
+        r_br, s_br = src_br, dst_br
+    else:
+        r_br, s_br = dst_br, src_br
+    r_buf = r_br.buffer
+    s_buf = s_br.buffer
+    r_region = [(r.min, r.min + r.extent) for r in r_br.region]
+    s_region = [(r.min, r.min + r.extent) for r in s_br.region]
+    with sctx.target:
+        r_sliced = r_buf.layout.slice(list(r_buf.shape), r_region)
+        s_sliced = s_buf.layout.slice(list(s_buf.shape), s_region)
+        return align_layouts_raw(r_sliced, s_sliced, s_region)
+
+
+def _make_thread_placeholders(r_p) -> dict[str, _TirVar]:
+    placeholders: dict[str, _TirVar] = {}
+    for it in r_p.shard:
+        name = it.axis.name
+        if it.axis.is_thread() and name not in placeholders:
+            placeholders[name] = _TirVar(name, "int32")
+    return placeholders
+
+
+def _s_thread_offset(r_p, s_p, placeholders: dict[str, _TirVar]):
+    """Per-thread S base offset. Coord per R position is placeholder (thread
+    axis) or 0 (memory axis); apply_to_shape decomposes across s_p iters.
+    Includes layout-level offsets (e.g. from slicing a non-zero S region)."""
+    coord = [
+        placeholders[it.axis.name] if it.axis.is_thread() else _IntImm("int32", 0)
+        for it in r_p.shard
+    ]
+    input_shape = [int(it.extent) for it in r_p.shard]
+    per_iter = s_p.apply_to_shape(coord, input_shape)
+    off = _IntImm("int32", 0)
+    for c, it in zip(per_iter, s_p.shard, strict=True):
+        off = off + c * it.stride
+    for _axis, val in s_p.offset.items():
+        off = off + val
+    return off
+
+
+_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8)
+
+
+def _vec_len_candidates(elem_bits: int) -> list[int]:
+    """Widest-first element counts to try; always ends with scalar (1)."""
+    out: list[int] = []
+    for vb in _VEC_BITS_CANDIDATES:
+        if vb < elem_bits or vb % elem_bits != 0:
+            continue
+        n = vb // elem_bits
+        if n not in out:
+            out.append(n)
+    if 1 not in out:
+        out.append(1)
+    return out
+
+
+def _choose_vec_len(elem_bits: int, atoms, r_p, s_p) -> int:
+    """Widest candidate that:
+      1. divides the atom contiguous-tail extent (so vec_len consecutive
+         R-side regs map to vec_len contiguous S-side elements), AND
+      2. keeps every per-thread / per-round address-offset term a
+         multiple of vec_len, so the resulting vec ld/st pointer is
+         naturally aligned to vec_bits/8 bytes.
+
+    Only **mem-axis** strides contribute to physical address. Thread-axis
+    iter strides live in partition-coord space (which thread owns which
+    logical position), not in the buffer's storage space — they're
+    redistributed through ``apply_to_shape`` into the mem iters and don't
+    appear directly in the per-thread address. So neither r-side nor
+    s-side thread-axis strides belong in the alignment check.
+
+    The contig-tail atoms (whose extents the vec ld/st consumes) have
+    stride 1 by definition; they live entirely inside the vec and
+    contribute nothing to the per-round address delta. Only the
+    **post-vec-split** outer atom strides matter for the per-round delta.
+    """
+    t = _atoms_contiguous_tail_extent(atoms)
+    # Region-base offsets are real address contributions. Thread-iter
+    # strides on either side are partition-virtual, not storage-physical,
+    # so they don't enter the per-thread address — exclude them.
+    shared_terms = list(s_p.offset.values()) + list(r_p.offset.values())
+    for n in _vec_len_candidates(elem_bits):
+        if n == 1:
+            return n
+        if t % n != 0 or t < n:
+            continue
+        # Post-vec-split outer atoms: these are the strides that contribute
+        # to per-round address deltas after the vec consumes the inner tail.
+        outer = _split_atoms_for_vec(atoms, n)
+        outer_atom_terms = [a[1] for a in outer] + [a[2] for a in outer]
+        if not _alignment_ok(n, outer_atom_terms + shared_terms):
+            continue
+        return n
+    return 1
+
+
+def _axis_decl(axis_name: str, sctx: DispatchContext):
+    """Declare the runtime Var for one thread axis (called inside impl body).
+
+    Each scope_id declarator emits a ``ScopeIdDef`` stmt at the current
+    builder frame. ``TilePrimitiveDispatch`` re-gathers + resolves all
+    ScopeIdDefs after dispatch (see ``ResolveAllScopeBinds`` in
+    ``tile_primitive_dispatch.cc``), so dispatch-introduced vars are bound
+    alongside kernel-declared ones.
+
+    Extents are deferred: the kernel header is expected to declare the full
+    scope-id chain (``cta_id`` / ``warpgroup_id`` / ``warp_id_in_wg`` /
+    ``lane_id`` / ``thread_id`` / ``thread_id_in_wg``) — the verifier then
+    fills our deferred defs from those siblings.
+    """
+    if axis_name == "tx":
+        return sctx.launch_params["threadIdx.x"].var
+    if axis_name == "laneid":
+        return T.lane_id()
+    if axis_name == "wid_in_wg":
+        return T.warp_id_in_wg()
+    if axis_name == "tid_in_wg":
+        return T.thread_id_in_wg()
+    if axis_name == "warpid":
+        return T.warp_id()
+    if axis_name == "wgid":
+        return T.warpgroup_id()
+    raise ValueError(f"unsupported thread axis {axis_name}")
+
+
+def _s_thread_offset_with_vars(r_p, s_p, axis_var_map: dict):
+    coord = [
+        axis_var_map[it.axis.name] if it.axis.is_thread() else _IntImm("int32", 0)
+        for it in r_p.shard
+    ]
+    input_shape = [int(it.extent) for it in r_p.shard]
+    per_iter = s_p.apply_to_shape(coord, input_shape)
+    off = _IntImm("int32", 0)
+    for c, it in zip(per_iter, s_p.shard, strict=True):
+        off = off + c * it.stride
+    for _ax, val in s_p.offset.items():
+        off = off + val
+    return off
+
+
+def _substitute_axes(s_off_template, placeholders: dict[str, _TirVar], sctx: DispatchContext):
+    """Inside an impl body: declare real scope_ids and rewrite the
+    placeholder-built ``s_off_template`` to use them."""
+    vmap = {placeholders[name]: _axis_decl(name, sctx) for name in placeholders}
+    return tvm.tirx.stmt_functor.substitute(s_off_template, vmap)
+
+
+def _flat_coords(outer_atoms, flat_idx: int) -> list[int]:
+    coords = []
+    rem = flat_idx
+    for a in reversed(outer_atoms):
+        coords.append(rem % a[0])
+        rem //= a[0]
+    coords.reverse()
+    return coords
+
+
+_POINTER_OFFSET_SRC = (
+    "\ntemplate <typename T>\n"
+    "__forceinline__ __device__ T* tvm_builtin_pointer_offset(T* ptr, int offset) {\n"
+    "    return ptr + offset;\n"
+    "}\n"
+)
+
+
+def _ptr_off(base_ptr, off):
+    return T.cuda.func_call(
+        "tvm_builtin_pointer_offset",
+        base_ptr,
+        off,
+        source_code=_POINTER_OFFSET_SRC,
+        return_type=base_ptr.ty,
+    )
+
+
+def _outer_const_offsets(outer_atoms, flat_idx: int) -> tuple[int, int]:
+    """Returns (s_offset_const, r_offset_const) for one outer-loop flat index."""
+    coords = _flat_coords(outer_atoms, flat_idx)
+    ds = sum(c * a[1] for c, a in zip(coords, outer_atoms))
+    dr = sum(c * a[2] for c, a in zip(coords, outer_atoms))
+    return ds, dr
+
+
+def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src: Buffer = op_call.src.buffer
+    dst: Buffer = op_call.dst.buffer
+    if src.scope() == "local":
+        r_buf, s_buf, r_is_src = src, dst, True
+    else:
+        r_buf, s_buf, r_is_src = dst, src, False
+
+    with sctx.target:
+        r_p, s_p, s_seps, r_perm = _align_layouts(op_call, sctx)
+    r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps)
+    atoms = _build_atoms(r_iters, s_groups)
+    elem_bits = DataType(src.dtype).bits
+    vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p)
+    vec_bits = vec_len * elem_bits
+    outer = _split_atoms_for_vec(atoms, vec_len)
+    per_thread_r_total = 1
+    for it in r_iters:
+        per_thread_r_total *= int(it.extent)
+    per_thread_r_shape = [per_thread_r_total or 1]
+
+    # Build the per-thread S offset OUTSIDE the impl using placeholder Vars
+    # (one per thread axis). Inside the impl we'll declare the real scope_ids
+    # via T.lane_id/T.thread_id_in_wg/... and substitute them in.
+    placeholders = _make_thread_placeholders(r_p)
+    s_off_template = _s_thread_offset(r_p, s_p, placeholders)
+
+    # R-side base offset from slicing (e.g. ``R[i*8:i*8+8]`` ⇒ ``i*8``). The
+    # canonicalize() result lives in ``r_p.offset``; sum across axes (memory
+    # or thread — irrelevant once it's all on R's local stride-1 storage).
+    r_off_base = _IntImm("int32", 0)
+    for _ax, val in r_p.offset.items():
+        r_off_base = r_off_base + val
+
+    s_is_shared = s_buf.scope().startswith("shared")
+    num_bytes = vec_bits // 8
+    vec, ptx_type = copy_ptx_form(num_bytes)
+    space = "shared" if s_is_shared else "global"
+
+    total_outer = 1
+    for a in outer:
+        total_outer *= a[0]
+
+    # Swizzle handling: recognize the iter-pattern on S side from the atom
+    # extents/strides (atom = (extent, s_stride, r_mul); a[1] is the S-side
+    # stride per outer round, equivalent to outer_iter strides in gmem_smem).
+    swizzle = get_swizzle(s_buf.layout)
+    swizzle_pattern = None
+    if swizzle is not None:
+        swizzle_pattern = try_recognize(
+            swizzle,
+            [a[0] for a in outer],
+            [a[1] for a in outer],
+            s_off_template,
+        )
+
+    class _SwizzleState:
+        def __init__(self):
+            self.signed_strides = None
+            self.base_off = None
+
+    state = _SwizzleState()
+
+    def _setup_swizzle(s_off):
+        if swizzle_pattern is None:
+            return
+        state.signed_strides, state.base_off = emit_init(swizzle_pattern, s_off)
+
+    def _s_iter_off(f, ds, s_off):
+        if swizzle_pattern is not None:
+            return emit_iter_offset(swizzle_pattern, state.signed_strides, state.base_off, f)
+        if swizzle is not None:
+            return emit_fallback_offset(swizzle, s_off, ds)
+        return s_off + ds
+
+    # fmt: off
+    s_zero_indices = [0] * len(s_buf.shape)
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        s_off = _substitute_axes(s_off_template, placeholders, sctx)
+        _setup_swizzle(s_off)
+        r_local = r_buf.local(*per_thread_r_shape)
+        # Keep as a serial TIR loop and let ptxas unroll downstream. An
+        # explicit ``T.unroll`` materializes the per-iter scratch
+        # (ds/dr/s_ptr/r_ptr, swizzle ``v_<n>[]`` signed-strides) as N
+        # copies of each buffer declaration; on kernels with many R↔S copy
+        # sites and large ``total_outer`` (FA4 writeback) this floods the
+        # function with ``alignas(64) int`` arrays and pressures registers.
+        for f in range(total_outer):
+            ds, dr = _outer_const_offsets(outer, f)
+            s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), _s_iter_off(f, ds, s_off))
+            r_ptr = _ptr_off(r_local.ptr_to([0]), r_off_base + dr)
+            if r_is_src:
+                T.ptx.st(s_ptr, src=r_ptr, space=space, vec=vec, ptx_type=ptx_type)
+            else:
+                T.ptx.ld(
+                    s_ptr,
+                    copy_ptx_ld_return_type(ptx_type),
+                    ptx_type,
+                    dst=r_ptr,
+                    space=space,
+                    vec=vec,
+                )
+    # fmt: on
+    import os
+
+    if os.environ.get("R2S_DUMP"):
+        print("=== emitted impl ===")
+        print(impl.script())
+    return impl
+
+
+@register_dispatch(
+    "copy",
+    "cuda",
+    variant="reg",
+    priority=10,
+    when=[predicate("reg_applicable", _is_reg_copy)],
+)
+def copy_schedule_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return _emit_reg(op_call, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py

@@ -0,0 +1,101 @@
+# 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.
+"""Shared helpers for copy operator dispatches on CUDA targets."""
+
+from collections.abc import Iterable
+
+from tvm.tirx import Buffer
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import match_scope, validate_copy_op
+
+
+def _is_valid_smem_tmem_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    """Validate smem->tmem copy operation.
+
+    The new tcgen05.cp.32x128b.warpx4 dispatch requires the destination tmem
+    buffer to declare warpx4 broadcast as ``R[4 : 32@TLane]``. The legacy
+    128-row dispatch path (no replica) goes through a separate code path and
+    is not handled here.
+    """
+    dst_region, src_region = op_call.args[:2]
+    src: Buffer = src_region.buffer
+    dst: Buffer = dst_region.buffer
+    if not (src.scope().startswith("shared") and dst.scope() == "tmem"):
+        return (False, f"expected shared->tmem, got {src.scope()}->{dst.scope()}")
+    if not (src.layout and dst.layout):
+        return (False, "both buffers must have layouts")
+    if dst.allocated_addr is None:
+        return (False, "tmem buffer must have allocated_addr")
+    # Require warpx4 router on TMEM side so this dispatch only handles the
+    # 32x128b.warpx4 case; other shapes (128x256b/128x128b etc.) fall back
+    # to the legacy dispatch.
+    rep = dst.layout.replica
+    if not (
+        len(rep) == 1
+        and int(rep[0].extent) == 4
+        and int(rep[0].stride) == 32
+        and "TLane" in str(rep[0].axis)
+    ):
+        return (False, f"requires R[4:32@TLane] on tmem, got replica={list(rep)}")
+    return (True, None)
+
+
+def _single_thread_exec(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    """Predicate: exec scope must be single-thread."""
+    exec_scope = sctx.scope_kind
+    ok = exec_scope == "thread"
+    return (ok, None if ok else f"expected thread exec_scope, got {exec_scope}")
+
+
+DEFAULT_ALLOWED_PAIRS: tuple[tuple[str, str], ...] = (
+    ("global", "shared*"),
+    ("shared*", "global"),
+    ("global", "local"),
+    ("local", "global"),
+    ("shared*", "local"),
+    ("local", "shared*"),
+)
+
+
+def _scope_allowed(
+    op_call: TilePrimitiveCall,
+    sctx: DispatchContext,
+    allowed_pairs: Iterable[tuple[str, str]] = DEFAULT_ALLOWED_PAIRS,
+):
+    op_call = TilePrimitiveCall.downcast(op_call)
+    dst_buffer_region, src_buffer_region = (op_call.dst, op_call.src)
+    src_scope = src_buffer_region.buffer.scope()
+    dst_scope = dst_buffer_region.buffer.scope()
+    ok = any(
+        (
+            match_scope(src_scope, src_pat) and match_scope(dst_scope, dst_pat)
+            for src_pat, dst_pat in allowed_pairs
+        )
+    )
+    if not ok:
+        allowed_str = ", ".join((f"{a}->{b}" for a, b in allowed_pairs))
+        return (
+            False,
+            f"unsupported memory scopes src={src_scope} dst={dst_scope}; allowed: {allowed_str}",
+        )
+    return (True, None)
+
+
+def _is_valid_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    return (validate_copy_op(op_call, sctx), "validate_copy_op failed")

python/tvm/backend/cuda/operator/tile_primitive/copy_async/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+
+"""Implementation of copy_async operator dispatches for CUDA targets.
+
+Registered op: copy_async (4 variants).
+See the @register_dispatch blocks in each submodule for detailed documentation
+with before/after IR examples.
+"""
+
+from .dsmem import *
+from .ldgsts import *
+from .tcgen05_cp import *
+from .tcgen05_ldst import *
+from .tma import *

python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py

@@ -0,0 +1,226 @@
+# 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.
+
+"""copy_async dispatch variant: dsmem (shared::cta -> shared::cluster)."""
+
+import functools
+import operator
+
+import tvm
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx.operator.tile_primitive import (
+    DispatchContext,
+    fail,
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import validate_copy_op
+from ..exec_scope_utils import single_thread
+from .utils import find_contiguous_region, to_tile_layout
+
+
+def _is_shared_to_shared(op_call: TilePrimitiveCall) -> bool:
+    """Check if both src and dst are in shared memory."""
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src_scope = op_call.src.buffer.scope()
+    dst_scope = op_call.dst.buffer.scope()
+    return src_scope.startswith("shared") and dst_scope.startswith("shared")
+
+
+def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    """Implement shared-to-shared cross-CTA copy using cp.async.bulk.
+
+    Uses cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes
+    to copy data from the executing CTA's shared memory to a remote CTA's shared
+    memory within the same cluster.
+
+    The copy region is decomposed into contiguous byte chunks based on layout
+    analysis of both src and dst buffers. Non-contiguous dimensions are iterated
+    over, emitting one cp.async.bulk instruction per contiguous chunk.
+    """
+    op_call = TilePrimitiveCall.downcast(op_call)
+
+    # Extract config
+    remote_cta_id = op_call.config.get("remote_cta_id", None)
+    if remote_cta_id is None:
+        fail("remote_cta_id not set in config")
+    mbar = op_call.config.get("mbar", None)
+    if mbar is None:
+        fail("mbar not set in config")
+
+    # Extract buffer regions
+    dst_buffer_region = op_call.dst
+    src_buffer_region = op_call.src
+    src_buf: Buffer = src_buffer_region.buffer
+    dst_buf: Buffer = dst_buffer_region.buffer
+
+    src_st = [r.min for r in src_buffer_region.region]
+    src_ext = [r.extent for r in src_buffer_region.region]
+    dst_st = [r.min for r in dst_buffer_region.region]
+    dst_ext = [r.extent for r in dst_buffer_region.region]
+
+    dtype_bytes = tvm.DataType(src_buf.dtype).bits // 8
+
+    # Get tile layouts for both buffers
+    src_tile_layout = to_tile_layout(src_buf.layout, src_buf.shape)
+    dst_tile_layout = to_tile_layout(dst_buf.layout, dst_buf.shape)
+
+    # Slice layouts to copy region
+    src_region_tuples = [(src_st[i], src_st[i] + src_ext[i]) for i in range(len(src_st))]
+    sliced_src = src_tile_layout.slice([s for s in src_buf.shape], src_region_tuples)
+    if sliced_src is None:
+        fail("Cannot slice src layout for DSMEM copy")
+
+    dst_region_tuples = [(dst_st[i], dst_st[i] + dst_ext[i]) for i in range(len(dst_st))]
+    sliced_dst = dst_tile_layout.slice([s for s in dst_buf.shape], dst_region_tuples)
+    if sliced_dst is None:
+        fail("Cannot slice dst layout for DSMEM copy")
+
+    # Group src layout by region extents, then group dst by src's shard extents
+    # This creates 1:1 shard correspondence between the two layouts
+    grouped_src, src_seps = sliced_src.canonicalize().group(src_ext)
+    src_shard_extents = [s.extent for s in grouped_src.shard]
+    grouped_dst, dst_seps = sliced_dst.canonicalize().group(src_shard_extents)
+
+    # Find contiguous regions in both layouts
+    src_contig_indices, _ = find_contiguous_region(grouped_src)
+    dst_contig_indices, _ = find_contiguous_region(grouped_dst)
+
+    # Intersect: walk from innermost outward, include only matching shard indices
+    shared_contig_indices = []
+    for s_idx, d_idx in zip(src_contig_indices, dst_contig_indices):
+        if s_idx != d_idx:
+            break
+        shared_contig_indices.append(s_idx)
+
+    # Compute chunk size
+    if shared_contig_indices:
+        chunk_elements = functools.reduce(
+            operator.mul, [grouped_src.shard[i].extent for i in shared_contig_indices], 1
+        )
+    else:
+        chunk_elements = 1
+
+    chunk_bytes = chunk_elements * dtype_bytes
+    if chunk_bytes < 16 or chunk_bytes % 16 != 0:
+        fail(
+            f"Layouts not compatible for bulk DSMEM copy: "
+            f"chunk_bytes={chunk_bytes} (need >= 16 and multiple of 16)"
+        )
+
+    # Build iteration space over non-contiguous (outer) shards
+    shared_contig_set = set(shared_contig_indices)
+    outer_shard_indices = [i for i in range(len(grouped_src.shard)) if i not in shared_contig_set]
+    outer_extents = [grouped_src.shard[i].extent for i in outer_shard_indices]
+    outer_src_strides = [grouped_src.shard[i].stride for i in outer_shard_indices]
+    outer_dst_strides = [grouped_dst.shard[i].stride for i in outer_shard_indices]
+
+    # Helper to compute element offsets from loop variables (called via T.meta_var)
+    def compute_offsets(loop_vars):
+        if len(outer_extents) == 1:
+            lvs = [loop_vars]
+        else:
+            lvs = list(loop_vars)
+        src_off = 0
+        dst_off = 0
+        for j, v in enumerate(lvs):
+            src_off = src_off + v * outer_src_strides[j]
+            dst_off = dst_off + v * outer_dst_strides[j]
+        return src_off, dst_off
+
+    src_tile = to_tile_layout(src_buf.layout, src_buf.shape)
+    dst_tile = to_tile_layout(dst_buf.layout, dst_buf.shape)
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        # Map mbar to remote CTA (complete_tx targets the destination's mbar)
+        remote_mbar = T.ptx.map_shared_rank(mbar, remote_cta_id)
+
+        if not outer_extents:
+            # Single contiguous chunk — no iteration needed
+            src_ptr = src_buf.ptr_to(src_st)
+            cluster_dst = T.ptx.map_shared_rank(dst_buf.ptr_to(dst_st), remote_cta_id)
+            T.ptx.cp_async.bulk.s2c(cluster_dst, src_ptr, chunk_bytes, remote_mbar)
+        else:
+            for loop_vars in T.grid(*outer_extents):
+                src_elem_offset, dst_elem_offset = T.meta_var(compute_offsets(loop_vars))
+
+                src_buf_w = T.decl_buffer(
+                    src_buf.shape, src_buf.dtype, src_buf.data,
+                    elem_offset=src_buf.elem_offset + src_elem_offset,
+                    scope=src_buf.scope(),
+                    layout=src_tile,
+                )
+                dst_buf_w = T.decl_buffer(
+                    dst_buf.shape, dst_buf.dtype, dst_buf.data,
+                    elem_offset=dst_buf.elem_offset + dst_elem_offset,
+                    scope=dst_buf.scope(),
+                    layout=dst_tile,
+                )
+
+                src_ptr = src_buf_w.ptr_to(src_st)
+                cluster_dst = T.ptx.map_shared_rank(dst_buf_w.ptr_to(dst_st), remote_cta_id)
+                T.ptx.cp_async.bulk.s2c(cluster_dst, src_ptr, chunk_bytes, remote_mbar)
+    # fmt: on
+
+    return impl
+
+
+# === Variant: copy_async/dsmem (priority=10) ===
+#
+# When: valid async copy at single-thread scope where both src and dst are in
+# shared memory. Used for intra-cluster DSMEM copies (shared::cta -> shared::cluster).
+#
+# Before (TilePrimitiveCall):
+#     Tx.copy_async(
+#         dst_smem[0:128, 0:64],
+#         src_smem[0:128, 0:64],
+#         config={"mbar": mbar, "remote_cta_id": cta_id}
+#     )
+#
+# After (emits cp.async.bulk.shared::cluster.shared::cta):
+#   cluster_dst = mapa(dst_smem.ptr, cta_id)
+#   cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes
+#       [cluster_dst], [src_smem.ptr], size, [mbar]
+@register_dispatch(
+    "copy_async",
+    "cuda",
+    variant="dsmem",
+    priority=10,
+    when=[
+        predicate(
+            "validate_copy_op", lambda op, sctx: (validate_copy_op(op, sctx), "not a valid copy op")
+        ),
+        predicate(
+            "single_thread",
+            lambda op, sctx: (
+                single_thread(op, sctx),
+                f"unsupported exec_scope {sctx.exec_scope}, expected single thread",
+            ),
+        ),
+        predicate(
+            "is_shared_to_shared",
+            lambda op, sctx: (_is_shared_to_shared(op), "not a shared-to-shared copy"),
+        ),
+    ],
+)
+def copy_async_dispatch_dsmem(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return copy_dsmem_impl(op, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py

@@ -0,0 +1,275 @@
+# 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.
+
+"""``copy_async`` dispatch for ``global → shared`` via ``cp.async``
+(SASS: ``LDGSTS``).
+
+Shares the partition / layout-alignment algorithm with
+``cuda/copy/gmem_smem.py`` (sync ``T.copy`` global ↔ shared); differs at
+emit time only:
+
+* direction: ``cp.async`` is global → shared only (hardware restriction).
+* cp_size: PTX ``cp.async`` only accepts 4 / 8 / 16 bytes, so the vec-width
+  candidate set is restricted to ``{32, 64, 128}`` bits.
+* emit: ``T.evaluate(T.ptx.cp_async(dst, src, cp_size))`` instead of the
+  synchronous ``T.cuda.copy_{vec_bits}b(dst, src)``.
+
+Note: ``cp.async`` does **not** sync at emit time — caller is responsible
+for ``commit_group`` / ``wait_group`` / ``cta_sync`` plumbing around the
+async pipeline.
+"""
+
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx import Var as _TirVar
+from tvm.tirx.expr import IntImm as _IntImm
+from tvm.tirx.operator.tile_primitive.dispatcher import (
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..copy._common import (
+    _TID_AXIS_FOR_SCOPE,
+    _thread_cnt,
+    align_layouts_gs,
+)
+from ..copy._swizzle_iter import (
+    emit_init,
+    emit_iter_offset,
+    get_swizzle,
+    try_recognize,
+)
+from ..copy.reg import _all_threads_active, _axis_decl, _ptr_off
+from ..copy.utils import _is_valid_copy, _scope_allowed
+
+# cp.async is unidirectional: global → shared.
+_LDGSTS_PAIRS = [("global", "shared*")]
+# cp.async cp_size ∈ {4, 8, 16} bytes ⇒ vec_bits ∈ {32, 64, 128}.
+_LDGSTS_VEC_BITS = (128, 64, 32)
+
+
+def _divides_thread_cnt_ldgsts(
+    op_call: TilePrimitiveCall, sctx: DispatchContext
+) -> tuple[bool, str | None]:
+    """Mirror of ``gmem_smem._divides_thread_cnt``: reject copies whose
+    region element count doesn't divide ``thread_cnt`` (and reject
+    ``thread_cnt=0`` scopes outright). See that docstring for rationale."""
+    op_call = TilePrimitiveCall.downcast(op_call)
+    thread_cnt = _thread_cnt(sctx)
+    if thread_cnt <= 0:
+        return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)"
+    g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst
+    n_elements = 1
+    for r in g_br.region:
+        ext = r.extent
+        try:
+            n_elements *= int(ext)
+        except (TypeError, ValueError):
+            return False, f"non-constant region extent {ext}"
+    if n_elements % thread_cnt != 0:
+        return False, (f"region size {n_elements} not divisible by thread_cnt={thread_cnt}")
+    return True, None
+
+
+def _is_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+    if not sctx.is_target("cuda"):
+        return False, "non-cuda target"
+    if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
+        return False, f"unsupported exec_scope {sctx.scope_kind}"
+    for check in (
+        lambda: _all_threads_active(sctx),
+        lambda: _is_valid_copy(op_call, sctx),
+        lambda: _scope_allowed(op_call, sctx, allowed_pairs=_LDGSTS_PAIRS),
+        lambda: _divides_thread_cnt_ldgsts(op_call, sctx),
+    ):
+        ok, msg = check()
+        if not ok:
+            return False, msg
+    return True, None
+
+
+def _emit_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src: Buffer = op_call.src.buffer
+    dst: Buffer = op_call.dst.buffer
+    # Predicate above guarantees src is global, dst is shared.
+    g_buf, g_br = src, op_call.src
+    s_buf, s_br = dst, op_call.dst
+
+    g_region = [(r.min, r.min + r.extent) for r in g_br.region]
+    s_region = [(r.min, r.min + r.extent) for r in s_br.region]
+
+    elem_bits = DataType(src.dtype).bits
+    thread_cnt = _thread_cnt(sctx)
+
+    with sctx.target:
+        g_p, s_p, vec_len = align_layouts_gs(
+            g_buf.layout,
+            g_buf.shape,
+            g_region,
+            s_buf.layout,
+            s_buf.shape,
+            s_region,
+            elem_bits,
+            thread_cnt,
+            vec_bits_candidates=_LDGSTS_VEC_BITS,
+        )
+
+    vec_bits = vec_len * elem_bits
+    cp_size = vec_bits // 8  # cp.async cp_size is in bytes
+    if cp_size not in (4, 8, 16):
+        # align_layouts_gs already restricted candidates to _LDGSTS_VEC_BITS,
+        # so reaching here means no candidate worked at all.
+        from tvm.tirx.operator.tile_primitive.dispatcher import fail
+
+        fail(f"ldgsts: cannot find a cp.async-compatible vec_len for elem_bits={elem_bits}")
+
+    # Mirror gmem_smem.py: build 3D `(f, tid, 0)` against
+    # `[total_outer, thread_cnt, vec_len]` and let `s_p.apply(coord, shape)`
+    # flatten + resplit into whatever multi-iter T / outer-iter structure
+    # `align_layouts_gs` picked. Emit is oblivious to how many shard iters
+    # cover T.
+    n_elements = 1
+    for it in s_p.shard:
+        n_elements *= int(it.extent)
+    assert n_elements % (thread_cnt * vec_len) == 0, (
+        f"partition produced {n_elements} elements but thread_cnt({thread_cnt}) * "
+        f"vec_len({vec_len}) = {thread_cnt * vec_len} doesn't divide it"
+    )
+    total_outer = n_elements // (thread_cnt * vec_len)
+    apply_shape = [
+        _IntImm("int32", total_outer),
+        _IntImm("int32", thread_cnt),
+        _IntImm("int32", vec_len),
+    ]
+
+    s_zero = [0] * len(s_buf.shape)
+    g_zero = [0] * len(g_buf.shape)
+
+    tid_axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind] if thread_cnt > 1 else None
+
+    # T-iters-walk-back to recover outer_iters_s for the fast-path
+    # recognizer. Same trick as gmem_smem.py.
+    if thread_cnt > 1:
+        acc, _i = 1, len(s_p.shard) - 2
+        while _i >= 0 and acc < thread_cnt:
+            _ext = int(s_p.shard[_i].extent)
+            if acc * _ext > thread_cnt:
+                break
+            acc *= _ext
+            _i -= 1
+        outer_iters_s = list(s_p.shard[: _i + 1]) if acc == thread_cnt else []
+    else:
+        outer_iters_s = list(s_p.shard[:-1])
+
+    swizzle = get_swizzle(s_buf.layout)
+    swizzle_pattern = None
+    if swizzle is not None and outer_iters_s:
+        if tid_axis_name is not None:
+            _tid_placeholder = _TirVar(tid_axis_name, "int32")
+        else:
+            _tid_placeholder = _IntImm("int32", 0)
+        s_off_template = s_p.apply(
+            _IntImm("int32", 0),
+            _tid_placeholder,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )["m"]
+        swizzle_pattern = try_recognize(
+            swizzle,
+            [int(it.extent) for it in outer_iters_s],
+            [int(it.stride) for it in outer_iters_s],
+            s_off_template,
+        )
+
+    class _SwizzleState:
+        def __init__(self):
+            self.signed_strides = None
+            self.base_off = None
+
+    state = _SwizzleState()
+
+    def _decl_tid():
+        if tid_axis_name is not None:
+            return _axis_decl(tid_axis_name, sctx)
+        return _IntImm("int32", 0)
+
+    def _setup_swizzle(tid):
+        if swizzle_pattern is None:
+            return
+        s_off_resolved = s_p.apply(
+            _IntImm("int32", 0),
+            tid,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )["m"]
+        state.signed_strides, state.base_off = emit_init(
+            swizzle_pattern,
+            s_off_resolved,
+        )
+
+    if swizzle_pattern is not None:
+
+        def _s_off(f, s_lin):
+            return emit_iter_offset(
+                swizzle_pattern,
+                state.signed_strides,
+                state.base_off,
+                f,
+            )
+    elif swizzle is not None:
+        _sw = swizzle
+
+        def _s_off(f, s_lin):
+            return _sw.apply(s_lin)["m"]
+    else:
+
+        def _s_off(f, s_lin):
+            return s_lin
+
+    v0 = _IntImm("int32", 0)
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        tid = _decl_tid()
+        _setup_swizzle(tid)
+        for f in T.unroll(total_outer):
+            s_lin = s_p.apply(f, tid, v0, shape=apply_shape)["m"]
+            g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"]
+            s_off = _s_off(f, s_lin)
+            s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
+            g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
+            T.evaluate(T.ptx.cp_async(s_ptr, g_ptr, cp_size))
+        # cp.async is caller-synced — no cta_sync here (commit_group /
+        # wait_group / cta_sync are the caller's responsibility).
+    # fmt: on
+    return impl
+
+
+@register_dispatch(
+    "copy_async",
+    "cuda",
+    variant="ldgsts",
+    priority=20,
+    when=[predicate("ldgsts_applicable", _is_ldgsts)],
+)
+def copy_schedule_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return _emit_ldgsts(op_call, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py

@@ -0,0 +1,485 @@
+# 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.
+
+"""smem->tmem dispatch via tcgen05.cp.32x128b.warpx4.
+
+``tcgen05.cp`` is inherently async; this dispatch emits the cp loop only and
+leaves completion signaling (``tcgen05.commit`` against a barrier) to the
+caller. Callers who want sync semantics should issue ``tcgen05.commit``
+themselves after the copy.
+
+Algorithm
+---------
+Given ``Tx.copy_async(t_region, s_region)`` where t is in tmem (with
+R[4:32@TLane] indicating warpx4 broadcast), and s is in shared memory:
+
+A. Slice + canonicalize both layouts at the given regions.
+B. Verify ``t.replica == [4:32@TLane]`` (warpx4 router).
+C. Compute permutation that puts TLane first, then TCol stride-descending;
+   apply to t.permute_dims and to s via group + permute_by_groups.
+D. Canonicalize again.
+E. Isolate broadcast: split-by-stride-zero on both t and s; their split
+   sequences must match (same distinct prefix prods + broadcast extents).
+   Drop stride-0 iters → ``t_iso`` and ``s_iso``.
+F. Group both into ``(32, middle, elem_per_128b)``. Validate:
+   - t_lane = (32, 1@TLane)
+   - t_col = (elem_per_128b, 1@TCol)
+   - s_col = (elem_per_128b, 1)
+   - s_lane refines into (4, 8) on m axis with strides (SDO_stride, atom_K_stride)
+   - atom_K_byte ∈ {16, 32, 64, 128} → swizzle_mode 0..3
+   - swizzle_mode matches s_buf.layout's SwizzleLayout (if any)
+G. Alignment checks:
+   - t_iso TCol offset ≡ 0 (mod 32-bit)
+   - s_iso m offset ≡ 0 (mod 16B for sw=0; mod atom_size for sw>0)
+   - middle iter strides 16B-aligned
+H. middle 1-1 correspondence (simple-mode): t_middle and s_middle have same
+   iter count and matching extents per position.
+I. Emit:
+   - SmemDescriptor encoded once at SMEM base (hoisted via post_buffer_def_stmt).
+   - Loop over middle iters; each cp uses ``desc.add_16B_offset(init + loop)``
+     and writes to ``tmem_addr + t_col0 + Σ i_j * t_step_j``.
+"""
+
+import functools
+import operator
+
+import tvm
+from tvm.arith import Analyzer
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx.layout import ComposeLayout, SwizzleLayout, TCol, TileLayout, TLane
+from tvm.tirx.layout import m as m_axis
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt, TilePrimitiveCall
+
+from ..copy import _is_valid_smem_tmem_copy, _single_thread_exec
+
+
+# -----------------------------------------------------------------------------
+# Helpers
+# -----------------------------------------------------------------------------
+def _compute_perm(t):
+    def key(p):
+        it = p[1]
+        return (0 if it.axis == TLane else 1, -int(it.stride))
+
+    return [i for i, _ in sorted(enumerate(t.shard), key=key)]
+
+
+def _split_by_zero(lay):
+    """Split lay.shard into segments at stride==0 positions.
+    Returns (split_seq, kept_iters_with_nonzero_stride)."""
+    new_seq = []
+    keep = []
+    cur = 1
+    for it in lay.shard:
+        e, st = int(it.extent), int(it.stride)
+        if st == 0:
+            if cur > 1:
+                new_seq.append(cur)
+            new_seq.append(e)
+            cur = 1
+        else:
+            cur *= e
+            keep.append(it)
+    if cur > 1:
+        new_seq.append(cur)
+    return new_seq, keep
+
+
+def _align_middles(t_middle, s_middle):
+    """Sub-group both middles by union-of-boundaries so they become 1-1.
+
+    Both inputs must be post-canonicalize iter lists with equal extent products.
+    The shape is the consecutive ratios of sorted(B_t U B_s) where B_x is the
+    set of cumulative extent boundaries of x_middle. Each segment then contains
+    at most one iter per side (whole or sub-divided), so trivially single iter.
+
+    Returns (new_t_middle, new_s_middle) with len() == len() == k segments,
+    each segment a single Iter on each side.
+    """
+
+    def cum_bounds(iters):
+        b, p = [], 1
+        for it in iters:
+            p *= int(it.extent)
+            b.append(p)
+        return b
+
+    t_bounds = cum_bounds(t_middle)
+    s_bounds = cum_bounds(s_middle)
+    if not t_bounds and not s_bounds:
+        return t_middle, s_middle
+    N = t_bounds[-1] if t_bounds else s_bounds[-1]
+    if (s_bounds and s_bounds[-1] != N) or (t_bounds and t_bounds[-1] != N):
+        raise ValueError(f"middle extent mismatch: t={N} s={s_bounds[-1] if s_bounds else 0}")
+
+    cuts = sorted(set(t_bounds) | set(s_bounds))
+    shape, prev = [], 1
+    for c in cuts:
+        if c % prev != 0:
+            raise ValueError(
+                f"middle align failed: cut {c} not divisible by prev cut {prev} "
+                f"(t_bounds={t_bounds}, s_bounds={s_bounds})"
+            )
+        shape.append(c // prev)
+        prev = c
+
+    def subgroup(iters):
+        if len(iters) == 1 and shape == [int(iters[0].extent)]:
+            return iters
+        lay, _seps = TileLayout.from_iters(iters, [], {}).group(shape)
+        seps = list(_seps)
+        out = []
+        for i in range(len(shape)):
+            seg = list(lay.shard[seps[i] : seps[i + 1]])
+            seg_canon = list(TileLayout.from_iters(seg, [], {}).canonicalize().shard)
+            if len(seg_canon) != 1:
+                raise ValueError(
+                    f"middle sub-group seg[{i}] not single iter after canon: {seg_canon}"
+                )
+            out.append(seg_canon[0])
+        return out
+
+    return subgroup(t_middle), subgroup(s_middle)
+
+
+# -----------------------------------------------------------------------------
+# Plan (state object)
+# -----------------------------------------------------------------------------
+def _build_plan(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    """Run A..H and return a dispatch plan.
+
+    Plan fields:
+      - s_buf, t_buf
+      - dtype, dtype_bits
+      - elem_per_128b, elem_per_32b
+      - SmemSwizzleMode (int)
+      - SDO_field, atom_K_byte
+      - middle_iters: list of (extent, s_step_16B, t_step_32bcol)
+      - init_off_16B (Expr)
+      - t_col0 (Expr, TMEM 32-bit col offset for cp's first call)
+    """
+    op_call = TilePrimitiveCall.downcast(op_call)
+    dst_region, src_region = op_call.args[:2]
+    s_buf: Buffer = src_region.buffer
+    t_buf: Buffer = dst_region.buffer
+    dtype = s_buf.dtype
+    dtype_bits = DataType(dtype).bits
+    elem_per_128b = 128 // dtype_bits
+    elem_per_32b = 32 // dtype_bits
+
+    # C: slice + canonicalize.
+    s_region = [(r.min, r.min + r.extent) for r in src_region.region]
+    t_region = [(r.min, r.min + r.extent) for r in dst_region.region]
+    s = s_buf.layout.slice(list(s_buf.shape), s_region).canonicalize()
+    t = t_buf.layout.slice(list(t_buf.shape), t_region).canonicalize()
+
+    # If s is ComposeLayout (SwizzleLayout∘TileLayout), peel off the swizzle
+    # for stride analysis; record swizzle_len for cross-check.
+    s_swizzle_mode_from_layout = 0
+    if isinstance(s, ComposeLayout):
+        s_swizzle_mode_from_layout = int(s.swizzle.swizzle_len)
+        s = s.tile_layout
+    elif isinstance(s, SwizzleLayout):
+        raise ValueError("s slice produced bare SwizzleLayout (unexpected)")
+
+    # B: warpx4 router check.
+    rep = t.replica
+    if not (
+        len(rep) == 1
+        and int(rep[0].extent) == 4
+        and int(rep[0].stride) == 32
+        and rep[0].axis == TLane
+    ):
+        raise ValueError(
+            f"warpx4 router fail: t.replica = "
+            f"{[(int(r.extent), int(r.stride), str(r.axis)) for r in rep]}"
+        )
+
+    # C: permute (TLane first, TCol stride desc).
+    perm = _compute_perm(t)
+    t_shape_for_group = [int(it.extent) for it in t.shard]
+    s_grp, seps = s.group(t_shape_for_group)
+    s_p = s_grp.permute_by_groups(list(seps), perm).canonicalize()
+    t_p = t.permute_dims(perm).canonicalize()
+
+    # E: isolate broadcast.
+    seq_t, keep_t = _split_by_zero(t_p)
+    seq_s, keep_s = _split_by_zero(s_p)
+    if seq_t != seq_s:
+        raise ValueError(f"isolate split mismatch: t={seq_t} s={seq_s}")
+    s_iso = TileLayout.from_iters(keep_s, list(s_p.replica), dict(s_p.offset))
+    t_iso = TileLayout.from_iters(keep_t, list(t_p.replica), dict(t_p.offset))
+
+    # F: group into (32, middle, elem_per_128b).
+    def shard_prod(lay):
+        return functools.reduce(operator.mul, [int(it.extent) for it in lay.shard], 1)
+
+    n_lane, n_col = 32, elem_per_128b
+    n_mid_t = shard_prod(t_iso) // (n_lane * n_col)
+    n_mid_s = shard_prod(s_iso) // (n_lane * n_col)
+    t_grp, t_seps = t_iso.group([n_lane, n_mid_t, n_col])
+    s_grp2, s_seps = s_iso.group([n_lane, n_mid_s, n_col])
+    t_seps = list(t_seps)
+    s_seps = list(s_seps)
+
+    def _canon_segment(iters):
+        return TileLayout.from_iters(iters, [], {}).canonicalize().shard
+
+    t_lane = list(_canon_segment(list(t_grp.shard[t_seps[0] : t_seps[1]])))
+    t_middle = list(_canon_segment(list(t_grp.shard[t_seps[1] : t_seps[2]])))
+    t_col = list(_canon_segment(list(t_grp.shard[t_seps[2] : t_seps[3]])))
+    s_lane = list(s_grp2.shard[s_seps[0] : s_seps[1]])
+    s_middle = list(_canon_segment(list(s_grp2.shard[s_seps[1] : s_seps[2]])))
+    s_col = list(_canon_segment(list(s_grp2.shard[s_seps[2] : s_seps[3]])))
+
+    # F.5: align middles via union-cut sub-grouping. Both t_middle and s_middle
+    # are post-canonicalize. To make their structure 1-1 we sub-group both by
+    # the union of their internal cumulative-extent boundaries.
+    t_middle, s_middle = _align_middles(t_middle, s_middle)
+
+    # F.1: lane / col validation.
+    if len(t_lane) != 1:
+        raise ValueError(f"t_lane must canonicalize to single iter, got {t_lane}")
+    if len(t_col) != 1:
+        raise ValueError(f"t_col must canonicalize to single iter, got {t_col}")
+    if len(s_col) != 1:
+        raise ValueError(f"s_col must canonicalize to single iter, got {s_col}")
+    li = t_lane[0]
+    if not (int(li.extent) == 32 and int(li.stride) == 1 and li.axis == TLane):
+        raise ValueError(f"t_lane must be (32, 1@TLane), got {li}")
+    ci = t_col[0]
+    if not (int(ci.extent) == elem_per_128b and int(ci.stride) == 1 and ci.axis == TCol):
+        raise ValueError(f"t_col must be ({elem_per_128b}, 1@TCol), got {ci}")
+    sci = s_col[0]
+    if not (int(sci.extent) == elem_per_128b and int(sci.stride) == 1):
+        raise ValueError(f"s_col must be ({elem_per_128b}, 1, m), got {sci}")
+
+    # F.2: s_lane → group (4, 8) → (SDO_stride, atom_K_stride)
+    s_lane_layout = TileLayout.from_iters(s_lane, [], {})
+    s_lane_grp, s_lane_seps = s_lane_layout.group([4, 8])
+    s_lane_seps = list(s_lane_seps)
+    blk_4 = list(s_lane_grp.shard[s_lane_seps[0] : s_lane_seps[1]])
+    blk_8 = list(s_lane_grp.shard[s_lane_seps[1] : s_lane_seps[2]])
+    if len(blk_4) != 1 or len(blk_8) != 1:
+        raise ValueError(
+            f"s_lane must group into single iter per block: blk_4={blk_4}, blk_8={blk_8}"
+        )
+    SDO_byte = int(blk_4[0].stride) * dtype_bits // 8
+    atom_K_byte = int(blk_8[0].stride) * dtype_bits // 8
+    sw_candidates = {16: 0, 32: 1, 64: 2, 128: 3}
+    if atom_K_byte not in sw_candidates:
+        raise ValueError(f"atom_K_byte {atom_K_byte} not in {{16,32,64,128}}")
+    derived_sw = sw_candidates[atom_K_byte]
+    if s_swizzle_mode_from_layout != derived_sw:
+        raise ValueError(
+            f"swizzle mode mismatch: s_layout swizzle_len="
+            f"{s_swizzle_mode_from_layout} but atom_K_byte={atom_K_byte} "
+            f"implies sw={derived_sw}"
+        )
+
+    analyzer = Analyzer()
+
+    # G: alignments.
+    # G.1: t_iso TCol offset ≡ 0 (mod 32-bit element count).
+    t_col_offset_expr = 0
+    for ax, val in t_iso.offset.items():
+        if ax == TCol:
+            t_col_offset_expr = val
+            break
+    if not analyzer.can_prove_equal(t_col_offset_expr % elem_per_32b, 0):
+        raise ValueError(f"t TCol offset {t_col_offset_expr} not provably 32b-aligned")
+
+    # G.2: s_iso m offset alignment.
+    s_m_offset_expr = 0
+    for ax, val in s_iso.offset.items():
+        if ax == m_axis:
+            s_m_offset_expr = val
+            break
+    elem_per_16B = 16 * 8 // dtype_bits
+    if derived_sw == 0:
+        align_elem = elem_per_16B
+        align_label = "16B"
+    else:
+        atom_size_byte = 8 * atom_K_byte
+        align_elem = atom_size_byte * 8 // dtype_bits
+        align_label = f"atom={atom_size_byte}B"
+    if not analyzer.can_prove_equal(s_m_offset_expr % align_elem, 0):
+        raise ValueError(
+            f"s offset {s_m_offset_expr} not provably aligned to {align_label} "
+            f"({align_elem} {dtype} elements)"
+        )
+
+    # H: middle 1-1 correspondence.
+    if len(t_middle) != len(s_middle):
+        raise ValueError(
+            f"t_middle iter count {len(t_middle)} != s_middle {len(s_middle)} "
+            "(simple-mode requires 1-1)"
+        )
+    middle_iters = []
+    for i, (ti, si) in enumerate(zip(t_middle, s_middle)):
+        if int(ti.extent) != int(si.extent):
+            raise ValueError(f"middle[{i}] extent: t={int(ti.extent)} s={int(si.extent)}")
+        n = int(ti.extent)
+        if n == 1:
+            continue
+        if ti.axis != TCol:
+            raise ValueError(f"middle[{i}] t axis must be TCol, got {ti.axis}")
+        s_stride_byte = int(si.stride) * dtype_bits // 8
+        if s_stride_byte % 16 != 0:
+            raise ValueError(f"s_middle[{i}] stride {s_stride_byte}B not 16B-aligned")
+        middle_iters.append((n, s_stride_byte // 16, int(ti.stride) // elem_per_32b))
+
+    SDO_field = SDO_byte // 16
+    init_off_16B = s_m_offset_expr * dtype_bits // 8 // 16
+    t_col0 = t_col_offset_expr // elem_per_32b
+
+    return {
+        "s_buf": s_buf,
+        "t_buf": t_buf,
+        "dtype": dtype,
+        "dtype_bits": dtype_bits,
+        "elem_per_128b": elem_per_128b,
+        "elem_per_32b": elem_per_32b,
+        "swizzle_mode": derived_sw,
+        "SDO_field": SDO_field,
+        "atom_K_byte": atom_K_byte,
+        "middle_iters": middle_iters,
+        "init_off_16B": init_off_16B,
+        "t_col0": t_col0,
+    }
+
+
+# -----------------------------------------------------------------------------
+# Descriptor caching: one (smem_buf, ldo, sdo, swizzle) → one desc_buf,
+# encoded once at SMEM base, hoisted to right after SMEM alloc via
+# add_post_buffer_def_stmt.
+# -----------------------------------------------------------------------------
+def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle):
+    # Cache descriptor template at SMEM 0; patch addr per cp.
+    cache_key = f"smem_tmem_desc:{int(ldo)}:{int(sdo)}:{int(swizzle)}"
+    cached = sctx.cache_get(cache_key)
+    if cached is not None:
+        return cached
+
+    desc_buf = tvm.tirx.decl_buffer((1,), "uint64", name="cp_desc", scope="local")
+    encode_call = T.ptx.tcgen05.encode_matrix_descriptor(
+        desc_buf.data, T.reinterpret("handle", T.uint64(0)), ldo, sdo, swizzle
+    )
+    wrap = SeqStmt([AllocBuffer(desc_buf), Evaluate(encode_call)])
+    sctx.add_post_buffer_def_stmt(s_buf, wrap)
+    sctx.cache_set(cache_key, desc_buf)
+    return desc_buf
+
+
+def _desc_set_addr(desc_val, addr_ptr):
+    """Patch a SMEM matrix descriptor's 14-bit address field with cvta(addr)>>4 —
+    matches the hand-rolled ``replace_smem_desc_addr`` (descriptor encoded at 0)."""
+    start_addr = T.cast(
+        T.bitwise_and(
+            T.shift_right(T.cuda.cvta_generic_to_shared(addr_ptr), T.uint32(4)),
+            T.uint32(0x3FFF),
+        ),
+        "uint64",
+    )
+    return T.bitwise_or(T.bitwise_and(desc_val, T.bitwise_not(T.uint64(0x3FFF))), start_addr)
+
+
+# -----------------------------------------------------------------------------
+# Core impl: emits the cp loop given a plan + cp config. Async only — caller
+# is responsible for issuing ``tcgen05.commit`` against a barrier if they
+# need synchronization.
+# -----------------------------------------------------------------------------
+def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    plan = _build_plan(op_call, sctx)
+    s_buf = plan["s_buf"]
+    t_buf = plan["t_buf"]
+    SDO_field = plan["SDO_field"]
+    sw = plan["swizzle_mode"]
+    middle_iters = plan["middle_iters"]
+    init_off_16B = plan["init_off_16B"]
+    t_col0 = plan["t_col0"]
+
+    # cp ignores LDO for data; non-zero LDO bloats the addr-patch codegen.
+    LDO_field = 0
+
+    cta_group = op_call.config.get("cta_group", 1)
+
+    desc_buf = _get_or_create_desc(sctx, s_buf, LDO_field, SDO_field, sw)
+    t_addr = t_buf.allocated_addr
+    s_rank = len(s_buf.shape)
+
+    def _cp_desc(off_16B):
+        addr = T.ptr_byte_offset(s_buf.ptr_to([0] * s_rank), off_16B * 16, s_buf.dtype)
+        return _desc_set_addr(desc_buf[0], addr)
+
+    # Flatten the N-D middle iteration into a single T.unroll. Each iteration's
+    # per-dim index is (flat // stride) % extent, summed into the t/s offsets.
+    # Works uniformly for n_mid ∈ {0, 1, 2, ...}; total == 1 (no middle dims) is
+    # special-cased to avoid a degenerate T.unroll(1).
+    total = functools.reduce(operator.mul, [n for n, _, _ in middle_iters], 1)
+
+    # fmt: off
+    if total == 1:
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            T.ptx.tcgen05.cp(
+                t_addr[0] + t_col0,
+                _cp_desc(init_off_16B),
+                shape="32x128b", cta_group=cta_group, multicast="warpx4",
+            )
+    else:
+        def compute_offsets(flat):
+            t_off = 0
+            s_off = 0
+            div = 1
+            for n, s_step, t_step in middle_iters:
+                idx = (flat // div) % n
+                div = div * n
+                t_off = t_off + idx * t_step
+                s_off = s_off + idx * s_step
+            return t_off, s_off
+
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            for flat in T.unroll(total):
+                t_off, s_off = T.meta_var(compute_offsets(flat))
+                T.ptx.tcgen05.cp(
+                    t_addr[0] + t_col0 + t_off,
+                    _cp_desc(init_off_16B + s_off),
+                    shape="32x128b", cta_group=cta_group, multicast="warpx4",
+                )
+    # fmt: on
+
+    return impl
+
+
+# === Variant: copy_async/smem->tmem (priority=10) ===
+@register_dispatch(
+    "copy_async",
+    "cuda",
+    variant="smem->tmem",
+    priority=10,
+    when=[
+        predicate("validate_smem_tmem_copy", _is_valid_smem_tmem_copy),
+        predicate("exec_scope", _single_thread_exec),
+    ],
+)
+def copy_async_schedule_smem_tmem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return copy_smem_tmem_impl(op_call, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py

@@ -0,0 +1,458 @@
+# 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.
+
+"""copy_async dispatch: ``tcgen05.ld`` / ``tcgen05.st`` (tmem <-> local registers).
+
+Both are inherently async; this dispatch emits the PTX instruction only and
+leaves completion (``tcgen05.wait.ld`` / ``tcgen05.wait.st``) to the caller.
+Callers that want sync semantics should issue the matching wait after the copy.
+"""
+
+import tvm
+from tvm.arith import Analyzer
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx.layout import (
+    S,
+    TCol,
+    TileLayout,
+    TLane,
+    tcgen05_atom_layout,
+    tid_in_wg,
+    tmem_datapath_layout,
+)
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import get_st_extent
+from ..copy import _is_valid_copy, _scope_allowed
+from ..exec_scope_utils import exec_scope_ok
+
+# Per-warp fp32-column factor for each instr_shape (mirrors
+# ``_TCGEN05_COL_FACTOR_FP32`` in ``tvm.tirx.layout``; .16x64b → 2,
+# .16x128b → 4, .16x256b → 8). Source: PTX ISA Table 49.
+_TCGEN05_COL_FACTOR_FP32 = {"16x64b": 2, "16x128b": 4, "16x256b": 8}
+
+
+def _match_tcgen05_atom_layout(buf):
+    """Return ``(instr_shape, rep, frag_rows)`` if ``buf.layout`` matches a
+    tcgen05 ``.16x*b`` atom layout for some supported ``instr_shape``.
+
+    The local buffer shape ``(frag_rows, K)`` (``frag_rows`` ∈ {64, 128})
+    together with the dtype determines the candidate ``rep`` for each
+    ``instr_shape``; we just probe the three shapes x two frag_rows and
+    structurally compare. ``None`` if no atom layout matches.
+    """
+    if len(buf.shape) != 2:
+        return None
+    rows, cols = int(buf.shape[0]), int(buf.shape[1])
+    if rows not in (64, 128):
+        return None
+    dtype = buf.dtype
+    layout_c = buf.layout.canonicalize()
+    for shape in _TCGEN05_COL_FACTOR_FP32:
+        try:
+            cand = tcgen05_atom_layout(shape, (rows, cols), dtype).canonicalize()
+        except ValueError:
+            continue
+        try:
+            tvm.ir.assert_structural_equal(layout_c, cand)
+        except (AssertionError, ValueError):
+            continue
+        # Recover rep from cols (same arithmetic the factory uses).
+        elem_per_32b = 32 // DataType(dtype).bits
+        rep = cols // (_TCGEN05_COL_FACTOR_FP32[shape] * elem_per_32b)
+        return shape, rep, rows
+    return None
+
+
+def _classify_tmem_datapath(tmem_buf):
+    """Return ``"D"`` / ``"F"`` if ``tmem_buf.layout`` matches a known tcgen05
+    datapath (PTX ISA §9.7.16.10.5), else ``None``.
+
+    Layout D (M=128, identity row→lane) is the default returned by
+    ``_default_tmem_layout``. Layout F (M=64 non-``.ws``, scattered) is the
+    explicit opt-in produced by ``tmem_pool.alloc(..., datapath="F")``.
+    The dispatch uses this to pair each ``.16x*b`` / ``.32x32b`` atom with a
+    compatible layout — see ``_check_tmem_layout_for_atom``.
+    """
+    if tmem_buf.layout is None:
+        return None
+    buf_layout = tmem_buf.layout.canonicalize()
+    rows = int(tmem_buf.shape[0])
+    if rows == 128:
+        cand = tmem_datapath_layout("D", 128, tmem_buf.shape[1]).canonicalize()
+        try:
+            tvm.ir.assert_structural_equal(buf_layout, cand)
+            return "D"
+        except (AssertionError, ValueError):
+            return None
+    if rows == 64:
+        cand = tmem_datapath_layout("F", 64, tmem_buf.shape[1]).canonicalize()
+        try:
+            tvm.ir.assert_structural_equal(buf_layout, cand)
+            return "F"
+        except (AssertionError, ValueError):
+            return None
+    return None
+
+
+# Compatibility matrix between the TMEM buffer's datapath layout and the
+# tcgen05 ld/st atom requested by ``T.copy_async``:
+#
+#   datapath x atom              | accepted? | rationale
+#   ---------------------------- | --------- | --------------------------------
+#   D (M=128 full)  x .32x32b    | yes       | full 128 lanes, all 32 per warp
+#   D (M=128 full)  x .16x*b M=64| yes       | reads first half-slab (lanes
+#                                |           |   0..15 of each warp partition)
+#                                |           |   — the rest of acc is wasted
+#                                |           |   for this atom but valid data
+#   D (M=128 full)  x .16x*b M=128| yes      | reads all 128 lanes via row=0
+#                                |           |   and row=16 PTX issues
+#   F (M=64 scatter)x .16x*b M=64| yes       | canonical pairing - F's row
+#                                |           |   indexing matches the atom's
+#                                |           |   scatter access
+#   F (M=64 scatter)x .16x*b M=128| no       | F only writes the low slab; the
+#                                |           |   high slab (row=16) is garbage
+#   F (M=64 scatter)x .32x32b    | no       | F only utilizes 16 of each
+#                                |           |   warp's 32 lanes
+_TMEM_ATOM_COMPAT = {
+    ("D", "32x32b", 128): True,
+    ("D", "16x*b", 64): True,
+    ("D", "16x*b", 128): True,
+    ("F", "32x32b", 128): False,
+    ("F", "16x*b", 64): True,
+    ("F", "16x*b", 128): False,
+}
+
+
+def _check_tmem_layout_for_atom(tmem_buf, atom_kind, frag_rows):
+    """Raise ``ValueError`` if the TMEM buffer's datapath layout is
+    incompatible with the requested ``tcgen05`` atom.
+
+    ``atom_kind`` is ``"32x32b"`` or ``"16x*b"``; ``frag_rows`` is the
+    register-side fragment row count (128 for ``.32x32b`` and ``.16x*b``
+    M=128 variants, 64 for ``.16x*b`` M=64). If the buffer's layout is
+    unrecognized (i.e. it isn't Layout D or Layout F), the dispatch falls
+    back to the structural assertions below.
+    """
+    datapath = _classify_tmem_datapath(tmem_buf)
+    if datapath is None:
+        return None
+    allowed = _TMEM_ATOM_COMPAT.get((datapath, atom_kind, frag_rows), False)
+    if not allowed:
+        raise ValueError(
+            f"tcgen05 dispatch: TMEM buffer with datapath={datapath!r} is "
+            f"incompatible with atom={atom_kind!r} (frag_rows={frag_rows}). "
+            f"See PTX ISA §9.7.16.10.5 for datapath/atom pairings; the "
+            f"buffer was allocated via tmem_pool.alloc(..., "
+            f"datapath={datapath!r})."
+        )
+    return datapath
+
+
+def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    dst_buffer_region, src_buffer_region = op_call.dst, op_call.src
+    dst: Buffer = dst_buffer_region.buffer
+    src: Buffer = src_buffer_region.buffer
+
+    if src.scope() == "tmem" and dst.scope() == "local":
+        direction = "tmem2local"
+        tmem_region, local_region = src_buffer_region, dst_buffer_region
+    elif src.scope() == "local" and dst.scope() == "tmem":
+        direction = "local2tmem"
+        local_region, tmem_region = src_buffer_region, dst_buffer_region
+    else:
+        raise ValueError(f"Unsupported src scope {src.scope()} and dst scope {dst.scope()}")
+
+    tmem_buf, local_buf = tmem_region.buffer, local_region.buffer
+
+    assert tmem_buf.layout is not None
+    assert local_buf.layout is not None
+    assert tmem_buf.dtype == local_buf.dtype
+    assert tmem_buf.allocated_addr is not None
+
+    analyzer = Analyzer()
+    elem_size = DataType(local_buf.dtype).bits
+    elem_per_32b = 32 // elem_size
+    assert len(local_buf.shape) == len(tmem_buf.shape) == 2
+
+    # Try the .16x* (M=64) path first by structural-matching the register-side
+    # layout against ``tcgen05_atom_layout(instr_shape, (64, K), dtype)``. The
+    # TMEM-side layout is the standard (128, W):(1@TLane, 1@TCol); the M=64
+    # fragment lives at lanes 0..15 of each warp's accessible slab (per PTX
+    # 9.7.16.8.1), so each warp issues with row_offset=0 and collectively the
+    # 4 warps cover all 64 rows.
+    atom_match = _match_tcgen05_atom_layout(local_buf)
+
+    if atom_match is not None:
+        shape, num, frag_rows = atom_match
+        return _emit_16xnb_path(
+            shape=shape,
+            num=num,
+            frag_rows=frag_rows,
+            direction=direction,
+            tmem_buf=tmem_buf,
+            local_buf=local_buf,
+            tmem_region=tmem_region,
+            local_region=local_region,
+            elem_per_32b=elem_per_32b,
+            analyzer=analyzer,
+        )
+
+    # Fall through to the existing .32x32b (M=128) path.
+    return _emit_32x32b_path(
+        direction=direction,
+        tmem_buf=tmem_buf,
+        local_buf=local_buf,
+        tmem_region=tmem_region,
+        local_region=local_region,
+        elem_per_32b=elem_per_32b,
+        analyzer=analyzer,
+    )
+
+
+def _emit_32x32b_path(
+    *, direction, tmem_buf, local_buf, tmem_region, local_region, elem_per_32b, analyzer
+) -> PrimFunc:
+    """Original M=128 fragment path using ``tcgen05.{ld,st}.32x32b.xN``."""
+    # local: 128xWIDTH <-> tmem: 128xSHAPE[1]
+    # ``.32x32b`` accesses 32 lanes per warp — the full warp partition — so
+    # the TMEM buffer must be Layout D (M=128 full datapath). Reject Layout F.
+    _check_tmem_layout_for_atom(tmem_buf, "32x32b", 128)
+    assert analyzer.can_prove_equal(local_buf.shape[0], 128)
+    assert analyzer.can_prove_equal(tmem_buf.shape[0], 128)
+
+    # Check width is valid for 32x32b, and determine num
+    width = local_region.region[1].extent
+    candidates = [1, 2, 4, 8, 16, 32, 64, 128]
+
+    if not analyzer.can_prove_equal(tvm.tirx.floormod(width, elem_per_32b), 0):
+        raise ValueError(f"Width {width} is not valid for tcgen05.ld/st with shape 32x32b")
+
+    num = None
+    for n in candidates:
+        if analyzer.can_prove_equal(tvm.tirx.floordiv(width, elem_per_32b), n):
+            num = n
+            break
+    else:
+        raise ValueError(f"Width {width} is not valid for tcgen05.ld/st with shape 32x32b")
+
+    tmem_st, tmem_extent = get_st_extent(tmem_region)
+    local_st, local_extent = get_st_extent(local_region)
+    # tmem layout (128, WIDTH):(1@TLane, 1@TCol)
+    tmem_layout = TileLayout(S[(128, tmem_buf.shape[1]) : (1 @ TLane, 1 @ TCol)]).canonicalize()
+    # local layout
+    TileLayout(S[(128, width) : (1 @ tid_in_wg, 1)]).canonicalize()
+
+    tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), tmem_layout)
+    # local: [0:128, 0:WIDTH] <-> tmem: [0:128, st:st+WIDTH]
+    assert analyzer.can_prove_equal(tmem_st[0], 0)
+    assert analyzer.can_prove_equal(tmem_extent[0], 128)
+
+    assert analyzer.can_prove_equal(local_st[0], 0)
+    assert analyzer.can_prove_equal(local_extent[0], 128)
+
+    offset = tmem_st[1]
+    assert analyzer.can_prove_equal(tvm.tirx.floormod(offset, elem_per_32b), 0)
+    offset_32b = tvm.tirx.floordiv(offset, elem_per_32b)
+    assert analyzer.can_prove_equal(tmem_extent[1], width), (
+        f"tmem_extent[1]: {tmem_extent[1]}, width: {width}"
+    )
+
+    # assert analyzer.can_prove_equal(local_st[1], 0)
+    assert analyzer.can_prove_equal(local_extent[1], width)
+
+    op = T.ptx.tcgen05.ld if direction == "tmem2local" else T.ptx.tcgen05.st
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        local_storage = local_buf.view(local_buf.shape[1] * elem_per_32b, layout=TileLayout(S[num * elem_per_32b]))  # noqa: E501
+        local_32b = local_storage.view("uint32")
+        op(tmem_buf.allocated_addr[0], *[local_32b[local_st[1] // elem_per_32b+i] for i in range(num)], shape="32x32b", num=num, row=0, col=offset_32b)  # noqa: E501
+    # fmt: on
+    return impl
+
+
+def _emit_16xnb_path(
+    *,
+    shape,
+    num,
+    frag_rows,
+    direction,
+    tmem_buf,
+    local_buf,
+    tmem_region,
+    local_region,
+    elem_per_32b,
+    analyzer,
+) -> PrimFunc:
+    """``.16x*b`` fragment path using ``tcgen05.{ld,st}.<shape>.x<num>`` (one
+    of ``.16x64b``, ``.16x128b``, ``.16x256b``).
+
+    Each of the warpgroup's 4 warps issues the atom with ``row_offset=0`` to
+    cover lanes 0..15 of its 32-lane TMEM partition (one 16-row slab); the
+    four warps collectively span M=64 rows. When ``frag_rows == 128`` the
+    dispatch emits a second issue with ``row_offset=16`` to also cover lanes
+    16..31 of each warp's partition, doubling the fragment's row coverage to
+    M=128. The two atoms share the same column footprint; the layout factory
+    surfaces the combined per-thread register vector with the second slab's
+    regs in the high half of the m-axis (so the dispatch can split regs
+    contiguously between the two PTX calls).
+    """
+    # Per-atom column footprint in fp32 columns:
+    #   .16x64b  → 2N    .16x128b → 4N    .16x256b → 8N
+    col_factor_fp32 = {"16x64b": 2, "16x128b": 4, "16x256b": 8}[shape]
+    # Per-thread register count per 16-row slab (in 32-bit units):
+    #   .16x64b.xN  → N        .16x128b.xN → 2N      .16x256b.xN → 4N
+    regs_per_thread_per_slab = {"16x64b": num, "16x128b": 2 * num, "16x256b": 4 * num}[shape]
+    n_slabs = frag_rows // 64  # 1 for M=64, 2 for M=128
+    assert n_slabs in (1, 2)
+    regs_per_thread = regs_per_thread_per_slab * n_slabs
+    # Logical column width that the local buffer view exposes (in element units).
+    width_elems = col_factor_fp32 * num * elem_per_32b
+    # Per-thread storage in element units (same total bits as the register vector).
+    per_thread_elems = regs_per_thread * elem_per_32b
+
+    # Local-side: shape (frag_rows, K_cols)
+    assert analyzer.can_prove_equal(local_buf.shape[0], frag_rows), (
+        f".16x*b path expects local_buf rows={frag_rows}, got {local_buf.shape[0]}"
+    )
+    assert analyzer.can_prove_equal(local_buf.shape[1], width_elems), (
+        f".16x*b path expects local_buf cols={width_elems}, got {local_buf.shape[1]}"
+    )
+
+    # TMEM-side: structurally classify the buffer's datapath (D or F) and
+    # reject incompatible pairings. The PTX is identical in either case (the
+    # warp partition rule and the atom's lane access pattern are baked into
+    # the hardware); the layout classification just keeps the buffer's
+    # logical row indexing in sync with the physical TMEM occupation.
+    datapath = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows)
+
+    if datapath == "F":
+        # Layout F: buffer shape (64, W), scattered row→lane.
+        assert analyzer.can_prove_equal(tmem_buf.shape[0], 64), (
+            f".16x*b Layout F expects tmem_buf rows=64, got {tmem_buf.shape[0]}"
+        )
+        tmem_rows = 64
+    else:
+        # Layout D (or untagged legacy buffers): shape (128, W), identity.
+        # The legacy structural check below still fires for untagged buffers
+        # so we don't silently accept arbitrary layouts.
+        assert analyzer.can_prove_equal(tmem_buf.shape[0], 128), (
+            f".16x*b path expects tmem_buf rows=128, got {tmem_buf.shape[0]}"
+        )
+        if datapath is None:
+            tmem_layout = TileLayout(
+                S[(128, tmem_buf.shape[1]) : (1 @ TLane, 1 @ TCol)]
+            ).canonicalize()
+            tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), tmem_layout)
+        tmem_rows = 128
+
+    tmem_st, tmem_extent = get_st_extent(tmem_region)
+    local_st, local_extent = get_st_extent(local_region)
+
+    # Rows must span the full frag. The COLUMN extent may be a sub-multiple of
+    # the atom's full width ``width_elems`` — i.e. a per-chunk column slice of a
+    # wider frag (e.g. an epilogue that loads one big (128, MMA_N) frag in
+    # EPI_TILE-wide chunks). The atom layout maps consecutive columns to
+    # consecutive registers within each slab, so a column slice occupies a
+    # contiguous register window; we emit ``num_eff`` (the slice's atom rep) at
+    # the slab base + the column's register offset. When the slice IS the full
+    # atom (the common case), num_eff == num and reg offset == 0 (no change).
+    assert analyzer.can_prove_equal(local_st[0], 0)
+    assert analyzer.can_prove_equal(local_extent[0], frag_rows)
+    assert analyzer.can_prove_equal(tmem_st[0], 0)
+    assert analyzer.can_prove_equal(tmem_extent[0], frag_rows)
+    # local and tmem column slices must match and divide the atom's full width.
+    assert analyzer.can_prove_equal(local_extent[1], tmem_extent[1])
+    slice_w = int(local_extent[1])
+    assert width_elems % slice_w == 0, f"slice width {slice_w} must divide atom width {width_elems}"
+    num_eff = num * slice_w // width_elems
+    regs_eff = regs_per_thread_per_slab * slice_w // width_elems
+    del tmem_rows  # only used for the structural check above
+
+    col_off = tmem_st[1]
+    assert analyzer.can_prove_equal(tvm.tirx.floormod(col_off, elem_per_32b), 0)
+    col_off_32b = tvm.tirx.floordiv(col_off, elem_per_32b)
+    local_col_off = local_st[1]
+    assert analyzer.can_prove_equal(tvm.tirx.floormod(local_col_off, elem_per_32b), 0)
+    local_col_off_elems = local_col_off
+
+    is_load = direction == "tmem2local"
+    op = T.ptx.tcgen05.ld if is_load else T.ptx.tcgen05.st
+    # We intentionally do *not* emit ``.pack::16b`` / ``.unpack::16b`` for
+    # 16-bit dtypes. That qualifier would store one 16-bit element per 32-bit
+    # TMEM cell (LOW half only, HIGH half wasted) — fine for some CUTLASS
+    # epilogues but a 2x TMEM waste vs. the existing ``.32x32b`` convention,
+    # which packs two 16-bit elements per cell. By using the plain ``.b32``
+    # form we keep TMEM dense (2 elements per 32-bit cell); the per-thread
+    # register file holds two packed 16-bit values per 32-bit register, and
+    # the layout factory's iters describe that packing.
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        # Per-thread 1-D flat view of the local storage, then a uint32 view
+        # for the register-pointer arguments of the PTX builtin.
+        local_storage = local_buf.view(per_thread_elems, layout=TileLayout(S[per_thread_elems]))
+        local_32b = local_storage.view("uint32")
+        # Register offset of the column slice within each slab. The old
+        # ``local_col_off // elem_per_32b`` is only correct when the slice IS the
+        # full atom; in general consecutive columns advance registers at the rate
+        # (regs_per_thread_per_slab / width_elems). For a full-atom load the
+        # offset is 0 either way, so existing callers are unaffected.
+        local_reg_base = local_col_off_elems * regs_per_thread_per_slab // width_elems
+        for slab in range(n_slabs):
+            reg_base = slab * regs_per_thread_per_slab
+            op(
+                tmem_buf.allocated_addr[0],
+                *[local_32b[local_reg_base + reg_base + i] for i in range(regs_eff)],
+                shape=shape, num=num_eff, row=slab * 16, col=col_off_32b,
+            )
+    # fmt: on
+    return impl
+
+
+# === Variant: copy_async/tmem<->local (priority=10) ===
+#
+# When: one buffer is in tmem (tensor memory, Blackwell SM100+) and the other
+# is in local scope, at warpgroup exec scope.
+#
+# Emits: T.ptx.tcgen05.ld / T.ptx.tcgen05.st (async). The caller is
+# responsible for issuing the matching ``T.ptx.tcgen05.wait.ld`` /
+# ``T.ptx.tcgen05.wait.st`` when synchronization is required.
+@register_dispatch(
+    "copy_async",
+    "cuda",
+    variant="tmem<->local",
+    priority=10,
+    when=[
+        predicate("validate_copy_op", _is_valid_copy),
+        predicate("exec_scope", exec_scope_ok, expected_scopes=["warpgroup"]),
+        predicate(
+            "storage_scope", _scope_allowed, allowed_pairs=[("tmem", "local"), ("local", "tmem")]
+        ),
+    ],
+)
+def copy_async_schedule_tmem_local_async(
+    op_call: TilePrimitiveCall, sctx: DispatchContext
+) -> PrimFunc:
+    return copy_tmem_local_impl(op_call, sctx)

python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py

@@ -0,0 +1,78 @@
+# 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.
+
+"""Shared helpers for copy_async operator dispatch variants.
+
+The TMA-specific lowering moved to ``tma.py``. What remains here are the tiny
+layout helpers other variants (e.g. ``dsmem.py``) still import.
+"""
+
+from tvm.arith import Analyzer
+from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout
+
+
+def find_contiguous_region(layout: TileLayout) -> tuple:
+    """Return the maximal stride-1 contiguous memory-shard chain.
+
+    Starts from stride==1 and repeatedly picks the shard whose stride equals
+    the running product of extents, stopping when no shard matches. Returns
+    the maximal chain; callers that need a shorter prefix should take one
+    themselves (e.g. to satisfy TMA's rank<=5 or a per-path reduction step).
+    Stride/extent comparisons go through an ``Analyzer`` so symbolic strides
+    work.
+    """
+
+    analyzer = Analyzer()
+    memory_shards = [
+        (i, s)
+        for i, s in enumerate(layout.shard)
+        if s.axis.is_memory() and not analyzer.can_prove_equal(s.extent, 1)
+    ]
+    if not memory_shards:
+        return [], 1
+
+    contiguous_indices: list[int] = []
+    contiguous_extent = 1
+    expected_stride = 1
+    consumed: set[int] = set()
+
+    while True:
+        for idx, shard in memory_shards:
+            if idx in consumed:
+                continue
+            if analyzer.can_prove_equal(shard.stride, expected_stride):
+                consumed.add(idx)
+                contiguous_indices.append(idx)
+                contiguous_extent *= shard.extent
+                expected_stride = contiguous_extent
+                break
+        else:
+            break
+
+    if not contiguous_indices:
+        return [], 0
+    return contiguous_indices, contiguous_extent
+
+
+def to_tile_layout(layout: Layout, shape: list[int]) -> TileLayout:
+    """Normalize any layout kind to a TileLayout for pointer arithmetic."""
+
+    if isinstance(layout, ComposeLayout):
+        return layout.tile_layout
+    if isinstance(layout, SwizzleLayout):
+        return TileLayout(S[tuple(shape)])
+    return layout

python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py

@@ -0,0 +1,38 @@
+# 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.
+
+"""CUDA elementwise dispatch.
+
+Split by storage scope to mirror ``cuda/copy/``:
+
+  reg.py  — operands all in ``local`` (registers)   → induced partition
+  smem.py — operands all in ``shared*``             → synthesized partition
+
+Each op in ``ops.ALL_OPS`` is registered under both variants. Per-op packed
+PTX/CUDA intrinsics live in ``vec_emit/`` (``binary_f32x2`` / ``cast_vec2``
+/ ``fma_f32x2``) and are attached to the relevant ``OpSpec.vec_impls``.
+"""
+
+from .register import *
+
+# Suppress submodule-attribute leakage. Without an explicit ``__all__`` here,
+# ``from tvm.backend.cuda.operator.tile_primitive.elementwise import *`` (run by
+# tile_primitive/__init__.py) re-exports the implicit submodule attributes
+# (``ops``, ``reg``, ``smem``, ``vec_emit``) — and ``ops`` in particular
+# shadows the top-level ``tile_primitive/ops.py`` (BinaryReduce / UnaryReduce
+# / ...) when downstream code does ``from tile_primitive import ops``.
+__all__: list[str] = []

python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py

@@ -0,0 +1,412 @@
+# 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.
+
+"""Shared layout / vec-selection / emit helpers for ``reg.py`` and ``smem.py``.
+
+Borrows directly from ``cuda/copy/reg.py`` (induced partition) and
+``cuda/copy/_common.py`` (synthesized partition), extended to N operands.
+
+The dispatch split mirrors copy:
+    reg.py   — all operands in ``local``  → partition induced by anchor's layout
+    smem.py  — all operands in ``shared*`` → partition synthesized from ``sctx.intra``
+"""
+
+from __future__ import annotations
+
+import functools
+import operator
+
+from tvm.arith.analyzer import Analyzer
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion
+from tvm.tirx.layout import Axis, Iter, TileLayout
+
+from ..common import get_indices, get_st_extent
+
+# Re-use copy's primitives (PR-640) — same algorithm, same scope_id machinery.
+from ..copy._common import _TID_AXIS_FOR_SCOPE, _extract_tile, _thread_cnt
+from ..copy.reg import _all_threads_active, _axis_decl, _compute_perm_r
+
+
+# -----------------------------------------------------------------------------
+# Plan helpers
+# -----------------------------------------------------------------------------
+def buffer_regions(plan) -> list[BufferRegion]:
+    """All BufferRegion args (dst + buffer-region srcs), in plan order."""
+    out: list[BufferRegion] = [plan.dst]
+    for s in plan.srcs:
+        if s.buf_region is not None:
+            out.append(s.buf_region)
+    return out
+
+
+def dtype_name(dtype) -> str:
+    dtype_obj = getattr(dtype, "dtype", None)
+    if dtype_obj is not None:
+        return str(dtype_obj)
+    return str(dtype)
+
+
+def dtype_bits(dtype) -> int:
+    return DataType(dtype_name(dtype)).bits
+
+
+def compute_dtype_of(plan) -> str:
+    """Widest dtype in bits across dst + buffer/scalar srcs (dst breaks ties)."""
+    candidates = [dtype_name(plan.dst.buffer.dtype)]
+    for s in plan.srcs:
+        if s.buf_region is not None:
+            candidates.append(dtype_name(s.buf_region.buffer.dtype))
+        elif s.scalar is not None:
+            candidates.append(scalar_dtype(s.scalar))
+    widest = candidates[0]
+    widest_bits = DataType(widest).bits
+    for d in candidates[1:]:
+        b = DataType(d).bits
+        if b > widest_bits:
+            widest, widest_bits = d, b
+    return widest
+
+
+def scalar_dtype(scalar) -> str:
+    dtype = getattr(scalar, "dtype", None)
+    if dtype is not None:
+        return str(dtype)
+    ty = getattr(scalar, "ty", None)
+    if ty is None and hasattr(scalar, "expr_ty"):
+        ty = scalar.expr_ty()
+    dtype = getattr(ty, "dtype", None)
+    if dtype is None:
+        raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing PrimType")
+    return str(dtype)
+
+
+def n_elements(buf_region: BufferRegion) -> int:
+    _, ext = get_st_extent(buf_region)
+    return functools.reduce(operator.mul, ext, 1)
+
+
+# -----------------------------------------------------------------------------
+# Anchor selection (reg.py)
+# -----------------------------------------------------------------------------
+def pick_anchor(plan) -> BufferRegion:
+    """Anchor is always ``plan.dst`` — every operand must have a layout
+    (enforced by predicate); dst's layout drives iteration. No choice to make.
+    """
+    return plan.dst
+
+
+# -----------------------------------------------------------------------------
+# Broadcast support (NumPy-style right-aligned, anchor = result shape)
+# -----------------------------------------------------------------------------
+def _tensor_shape_of(region) -> tuple[int, ...]:
+    """Per-dim region extent (post-slice tensor shape, NOT layout shape).
+
+    Accepts either ``[(start, end), ...]`` pairs (as built locally from a
+    ``BufferRegion``) or the ``BufferRegion.region`` sequence of ``Range``
+    objects directly. ``Range.extent`` is already simplified by the
+    front-end, so we avoid computing ``end - start`` on raw Expr (which
+    yields an un-simplified ``Sub`` and breaks ``int(...)``).
+    """
+    out = []
+    a = Analyzer()
+    for r in region:
+        if hasattr(r, "extent"):
+            ext = r.extent
+        else:
+            start, end = r
+            ext = a.simplify(end - start)
+        out.append(int(ext))
+    return tuple(out)
+
+
+def shape_broadcast_compat(op_shape, anchor_shape) -> tuple[bool, str | None]:
+    """NumPy-style: right-align op against anchor; per-dim extent must equal
+    anchor's or be 1. anchor is the result shape; op broadcasts TO anchor.
+    """
+    pad = len(anchor_shape) - len(op_shape)
+    if pad < 0:
+        return False, f"op rank {len(op_shape)} > anchor rank {len(anchor_shape)}"
+    for d in range(len(op_shape)):
+        e_op = int(op_shape[d])
+        e_a = int(anchor_shape[pad + d])
+        if e_op != e_a and e_op != 1:
+            return False, f"dim {d}: op extent {e_op} vs anchor {e_a} (need equal or 1)"
+    return True, None
+
+
+def _broadcast_lift(op_layout, op_tensor_shape, anchor_tensor_shape):
+    """Lift ``op_layout`` to ``anchor_tensor_shape`` by inserting stride-0
+    iters for padded leading dims and replacing extent-1 buckets with a
+    single stride-0 iter of anchor's extent. Offset and replica list are
+    preserved untouched.
+
+    The lift preserves the physical-address function: new iters have
+    stride 0 so they contribute ``coord * 0 = 0`` to the address
+    regardless of which virtual index is supplied, and dropped extent-1
+    iters contributed ``0 * stride = 0`` already.
+    """
+    pad = len(anchor_tensor_shape) - len(op_tensor_shape)
+    assert pad >= 0, "shape_broadcast_compat should have rejected this"
+
+    try:
+        grouped, seps = op_layout.group(list(op_tensor_shape))
+    except Exception as e:  # pylint: disable=broad-except
+        raise ValueError(
+            f"op layout {op_layout} not groupable by tensor shape {op_tensor_shape}: {e}"
+        ) from e
+
+    new_shard: list = []
+    # (1) Padded leading dims — one stride-0 iter each.
+    for d in range(pad):
+        new_shard.append(Iter(int(anchor_tensor_shape[d]), 0, Axis.get("m")))
+    # (2) Aligned dims.
+    for d_op in range(len(op_tensor_shape)):
+        e_op = int(op_tensor_shape[d_op])
+        e_a = int(anchor_tensor_shape[pad + d_op])
+        bucket = list(grouped.shard[seps[d_op] : seps[d_op + 1]])
+        if e_op == e_a:
+            new_shard.extend(bucket)
+        elif e_op == 1:
+            new_shard.append(Iter(e_a, 0, Axis.get("m")))
+        else:
+            raise ValueError(
+                f"dim {d_op}: op extent {e_op} vs anchor {e_a}"
+                " (shape_broadcast_compat should have rejected)"
+            )
+
+    return TileLayout.from_iters(new_shard, grouped.replica, grouped.offset)
+
+
+# -----------------------------------------------------------------------------
+# Shared preprocess: slice each operand by its region, broadcast-lift to
+# anchor's tensor shape. Output: every operand has a layout whose logical
+# shape equals ``anchor_tensor_shape``. Reg / smem diverge from here.
+# -----------------------------------------------------------------------------
+def preprocess_operand(op_br, anchor_tshape):
+    """Slice ``op_br``'s buffer layout by region (region offset absorbed into
+    layout.offset), then broadcast-lift to ``anchor_tshape`` if shapes differ.
+
+    Raises ``ValueError`` if the lift is not broadcast-compatible (caller
+    should have verified via ``shape_broadcast_compat`` in the predicate).
+    """
+    op_layout = op_br.buffer.layout
+    op_shape = op_br.buffer.shape
+    op_region = [(r.min, r.min + r.extent) for r in op_br.region]
+    sliced = op_layout.slice(list(op_shape), op_region).canonicalize()
+    sliced = _extract_tile(sliced, op_region)
+    op_tshape = _tensor_shape_of(op_br.region)
+    if op_tshape != tuple(anchor_tshape):
+        sliced = _broadcast_lift(sliced, op_tshape, anchor_tshape)
+    return sliced
+
+
+def preprocess_operands(plan):
+    """Shared entry for reg.py and smem.py: returns
+    ``(anchor_tensor_shape, {op_br: sliced_lifted_layout})``.
+
+    Every output layout has logical shape == ``anchor_tensor_shape``.
+    Broadcast iters carry stride 0 with the default mem axis. Reg's induced
+    partition (`align_operands_to_anchor`) and smem's synthesized partition
+    both build on this output.
+    """
+    anchor_tshape = _tensor_shape_of(plan.dst.region)
+    out: dict = {}
+    for br in buffer_regions(plan):
+        out[br] = preprocess_operand(br, anchor_tshape)
+    return anchor_tshape, out
+
+
+# -----------------------------------------------------------------------------
+# Multi-operand layout alignment for reg.py (induced)
+# -----------------------------------------------------------------------------
+def _align_layouts_no_post_canon(r_layout, r_shape, r_region, s_layout, s_shape, s_region):
+    """Variant of copy ``reg.py:align_layouts_raw`` that omits the final
+    ``canonicalize()`` on ``r_p``.
+
+    Copy's version returns ``r_p = r.permute_dims(perm).canonicalize()`` —
+    that post-permute canonicalize can fuse adjacent iters (e.g. wgmma
+    layout's 5 iters collapse to 2), but ``s_seps`` is built from
+    ``perm`` of length ``len(r.shard) pre-canon``. The two lengths then
+    disagree and ``s_p.shard[s_seps[k]:s_seps[k+1]]`` indexes into the
+    wrong sub-range.
+
+    Dropping the final canonicalize keeps ``r_p.shard`` and ``s_seps`` in
+    1-to-1 correspondence. Copy's tests don't hit this because R is
+    typically 1D and doesn't fuse further after permute.
+    """
+    r = r_layout.slice(list(r_shape), r_region).canonicalize()
+    s = s_layout.slice(list(s_shape), s_region).canonicalize()
+    s = _extract_tile(s, s_region)
+    # Broadcast lift: when op's post-slice tensor shape != anchor's, expand
+    # s via stride-0 iters so group() below can partition along anchor's
+    # iter structure. Legality must be enforced upstream by the predicate.
+    r_tshape = _tensor_shape_of(r_region)
+    s_tshape = _tensor_shape_of(s_region)
+    if s_tshape != r_tshape:
+        s = _broadcast_lift(s, s_tshape, r_tshape)
+    perm = _compute_perm_r(r)
+    r_shape_for_group = [int(it.extent) for it in r.shard]
+    s_grp, seps = s.group(r_shape_for_group)
+    s_p = s_grp.permute_by_groups(list(seps), perm)
+    r_p = r.permute_dims(perm)  # NO post-canonicalize
+    sizes = [seps[i + 1] - seps[i] for i in range(len(seps) - 1)]
+    s_seps = [0]
+    for p in perm:
+        s_seps.append(s_seps[-1] + sizes[p])
+    return r_p, s_p, s_seps
+
+
+def align_operands_to_anchor(anchor_br, layout_others_br):
+    """Align every layout-bearing non-anchor operand to ``anchor_br``.
+
+    Returns ``(anchor_p, per_op_aligned)`` where ``per_op_aligned[op_br] =
+    (op_p, op_seps)``. Trivial-layout operands are NOT included here —
+    caller indexes them directly via their region. Scalar srcs likewise
+    live outside this map.
+
+    Caller must enter ``with sctx.target:`` so ``canonicalize()`` runs the
+    scope-aware fusers (e.g. laneid+wid_in_wg → tid_in_wg).
+
+    Uses ``_align_layouts_no_post_canon`` (not copy's ``align_layouts_raw``
+    directly) so ``anchor_p.shard`` length matches ``op_seps`` groupings.
+    """
+    anchor_layout = anchor_br.buffer.layout
+    anchor_shape = anchor_br.buffer.shape
+    anchor_region = [(r.min, r.min + r.extent) for r in anchor_br.region]
+
+    per_op_aligned: dict = {}
+    anchor_p = None
+    if not layout_others_br:
+        # Just slice + permute anchor alone (no post-canon — keep iters
+        # 1-to-1 with how they'd appear with srcs).
+        r = anchor_layout.slice(list(anchor_shape), anchor_region).canonicalize()
+        perm = _compute_perm_r(r)
+        anchor_p = r.permute_dims(perm)
+        return anchor_p, per_op_aligned
+
+    for op_br in layout_others_br:
+        op_layout = op_br.buffer.layout
+        op_shape = op_br.buffer.shape
+        op_region = [(r.min, r.min + r.extent) for r in op_br.region]
+        r_p, op_p, op_seps = _align_layouts_no_post_canon(
+            anchor_layout,
+            anchor_shape,
+            anchor_region,
+            op_layout,
+            op_shape,
+            op_region,
+        )
+        if anchor_p is None:
+            anchor_p = r_p
+        per_op_aligned[op_br] = (op_p, op_seps)
+    return anchor_p, per_op_aligned
+
+
+# -----------------------------------------------------------------------------
+# vec_chunk selection
+# -----------------------------------------------------------------------------
+def pick_vec_chunk(spec, op_call, sctx, plan, max_layout_vec_len: int):
+    """Pick widest ``(vec_chunk, vec_impl)`` such that:
+      - ``vec_impl.vec_len`` divides ``max_layout_vec_len`` AND ``vec_impl.applies(...)``
+      - Or no vec_impl matches → scalar fallback ``(max_layout_vec_len, None)``
+
+    ``spec.vec_impls`` is assumed pre-sorted widest-first.
+    """
+    if max_layout_vec_len <= 0:
+        return 1, None
+    for impl in getattr(spec, "vec_impls", []):
+        if impl.vec_len > max_layout_vec_len:
+            continue
+        if max_layout_vec_len % impl.vec_len != 0:
+            continue
+        ok, _ = impl.applies(op_call, sctx, plan)
+        if ok:
+            return impl.vec_len, impl
+    return max_layout_vec_len, None
+
+
+# -----------------------------------------------------------------------------
+# Emit-time helpers
+# -----------------------------------------------------------------------------
+def _broadcast_indices(dst_indices, dst_start, dst_extent, op_start, op_ext):
+    """NumPy-style right-aligned broadcast: derive op's per-dim indices from
+    dst's. For matching extents copies dst's coord (rebased); for op_ext[d]
+    == 1 returns the constant start (the only valid index for that dim).
+    """
+    pad = len(dst_extent) - len(op_ext)
+    return [
+        (dst_indices[i + pad] - dst_start[i + pad]) + op_start[i]
+        if int(op_ext[i]) != 1
+        else op_start[i]
+        for i in range(len(op_ext))
+    ]
+
+
+def fetch_src_value(src, fused, dst_indices, dst_start, dst_extent):
+    """Per-element load Expr for one src. Handles buffer / scalar / broadcast srcs."""
+    if src.is_scalar:
+        return src.scalar
+    region = src.buf_region
+    src_st, src_ext = get_st_extent(region)
+    if src.index_fn is not None:
+        idx = src.index_fn(dst_indices, dst_start, dst_extent, src_st, src_ext)
+    elif tuple(int(e) for e in src_ext) != tuple(int(e) for e in dst_extent):
+        # Broadcast — derive src indices from dst's via right-aligned compat.
+        idx = _broadcast_indices(dst_indices, dst_start, dst_extent, src_st, src_ext)
+    else:
+        idx = get_indices(fused, src_st, src_ext)
+    return region.buffer[tuple(idx)]
+
+
+def emit_scope_sync(scope_kind: str):
+    """Returns an ``@T.inline`` sync helper matched to the exec scope."""
+
+    @T.inline
+    def sync():
+        if scope_kind == "cta":
+            T.cuda.cta_sync()
+        elif scope_kind == "warpgroup":
+            T.cuda.warpgroup_sync(8)
+        elif scope_kind == "warp":
+            T.cuda.warp_sync()
+
+    return sync
+
+
+__all__ = [
+    "_TID_AXIS_FOR_SCOPE",
+    "_all_threads_active",
+    "_axis_decl",
+    "_broadcast_indices",
+    "_tensor_shape_of",
+    "_thread_cnt",
+    "align_operands_to_anchor",
+    "buffer_regions",
+    "compute_dtype_of",
+    "dtype_bits",
+    "dtype_name",
+    "emit_scope_sync",
+    "fetch_src_value",
+    "n_elements",
+    "pick_anchor",
+    "pick_vec_chunk",
+    "preprocess_operand",
+    "preprocess_operands",
+    "shape_broadcast_compat",
+]

python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/__init__.py

@@ -0,0 +1,121 @@
+# 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.
+
+"""Per-op data model + ALL_OPS registry.
+
+``OpSpec`` describes one elementwise op. ``VecImpl`` describes one packed-PTX
+or CUDA-intrinsic emit available for that op (e.g. ``add_f32x2``); a list of
+these (widest-first) lets ``reg.py``/``smem.py`` pick the widest matching
+both the layout and the op's available intrinsics, like copy picks
+``copy_{128,64,32,16,8}b`` based on bit-width and tail contiguity.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from tvm.ir.expr import Expr
+from tvm.tirx import BufferRegion, TilePrimitiveCall
+
+
+@dataclass
+class SrcSpec:
+    """One operand of an elementwise op.
+
+    Either a ``BufferRegion`` (per-element load) or a scalar ``Expr``.
+    ``index_fn``, if given, derives per-element indices for broadcasting srcs:
+        ``index_fn(dst_indices, dst_start, dst_extent, src_start, src_extent) -> list[Expr]``
+    Default is the standard ``get_indices`` over the src's own region.
+    """
+
+    buf_region: BufferRegion | None = None
+    scalar: Expr | None = None
+    index_fn: Callable | None = None
+
+    @property
+    def is_scalar(self) -> bool:
+        return self.scalar is not None
+
+    @property
+    def buffer(self):
+        return self.buf_region.buffer if self.buf_region is not None else None
+
+
+@dataclass
+class Plan:
+    """Parsed elementwise op ready for a schedule to consume."""
+
+    dst: BufferRegion
+    srcs: list[SrcSpec]
+    extras: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class VecImpl:
+    """One packed-vector implementation registered for an op.
+
+    Mirrors the ``copy_{Nb}`` menu in copy: each entry says "I can process
+    ``vec_len`` consecutive elements per call". The schedule picks the widest
+    one whose ``vec_len`` divides the layout's contig tail AND whose
+    ``applies()`` returns ``True``.
+    """
+
+    vec_len: int  # elements per packed call
+    applies: Callable[[TilePrimitiveCall, Any, Plan], tuple[bool, str | None]]
+    # emit(dst_ptr, src_ptrs, extras) -> Stmt
+    #   dst_ptr: typed ptr to ``vec_len`` consecutive dst elements
+    #   src_ptrs[i]: typed ptr to ``vec_len`` consecutive src[i] elements,
+    #                OR a scalar Expr if src[i].is_scalar.
+    # Runs in Python at @T.prim_func build time -- branching on src kind is a
+    # normal Python ``if``, not a TVMScript shape limitation. This is what
+    # collapses the old 4x2 shape-explosion in schema.py's factories.
+    emit: Callable
+
+
+@dataclass
+class OpSpec:
+    """Metadata for an elementwise op."""
+
+    name: str
+    # parse(op_call) -> (Plan, msg|None); msg explains why parse failed.
+    parse: Callable[[TilePrimitiveCall], tuple[Plan | None, str | None]]
+    # Scalar compute used by the fallback emit path (wrapped in Tx.vectorized).
+    # compute_scalar(src_vals_at_one_idx, extras, dst_dtype) -> Expr
+    compute_scalar: Callable[[list, dict, str], Any]
+    # Optional dtype check on plan.extras (e.g. unary bias/scale dtype agreement).
+    check_extras: Callable | None = None
+    # Widest-first vec impls. Schedule picks first matching layout+applies.
+    vec_impls: list[VecImpl] = field(default_factory=list)
+
+
+def _build_all_ops() -> dict[str, OpSpec]:
+    """Aggregate per-family op specs. Deferred imports avoid cycles
+    (vec_emit/* imports VecImpl from this module)."""
+    from .binary import BINARY_OPS
+    from .cast import CAST_OPS
+    from .fma import FMA_OPS
+    from .unary import UNARY_OPS
+
+    return {**UNARY_OPS, **BINARY_OPS, **CAST_OPS, **FMA_OPS}
+
+
+ALL_OPS: dict[str, OpSpec] = _build_all_ops()
+
+
+__all__ = ["ALL_OPS", "OpSpec", "Plan", "SrcSpec", "VecImpl"]

python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py

@@ -0,0 +1,139 @@
+# 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.
+
+"""Binary elementwise ops: add / sub / mul / fdiv / maximum.
+
+Includes constant-lhs commute logic. Broadcasting (extent=1 dims) is
+handled at the layout level in dispatch's ``_broadcast_lift``, not here —
+parser just records each src as-is.
+
+``add``/``sub``/``mul`` attach a ``VecImpl`` for sm_100+ packed f32x2;
+``fdiv``/``maximum`` have no packed PTX (scalar fallback only — ``max`` lowers
+to a single ``FMNMX``/``max.f32``, which is exact, so there is no rounding/ftz
+variant to pack).
+"""
+
+from __future__ import annotations
+
+import functools
+import operator
+from typing import Any
+
+from tvm.script import tirx as Tx
+from tvm.tirx import BufferRegion, TilePrimitiveCall
+
+from ..vec_emit.binary_f32x2 import BINARY_F32X2_IMPLS
+from . import OpSpec, Plan, SrcSpec
+
+_COMMUTATIVE = frozenset({"add", "mul", "maximum"})
+
+
+def _parse_binary_for(op_name: str):
+    """Build a ``parse(op_call) -> (Plan, msg)`` for a specific binary op."""
+
+    def parse(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
+        _dst: BufferRegion = op.args[0]
+        _src1 = op.args[1]
+        _src2 = op.args[2]
+
+        s1_scalar = not isinstance(_src1, BufferRegion)
+        s2_scalar = not isinstance(_src2, BufferRegion)
+        if s1_scalar and s2_scalar:
+            return None, "both inputs are constants"
+
+        # Move constant to rhs (commute if allowed; else reject).
+        if s1_scalar:
+            if op_name not in _COMMUTATIVE:
+                return None, f"non-commutative op {op_name} cannot have constant lhs"
+            _src1, _src2 = _src2, _src1
+            s2_scalar = True
+
+        # If rhs is a smaller buffer (broadcast), swap if commutative so the
+        # bigger one is in src1 — keeps src1 == dst convention.
+        if not s2_scalar:
+            s1_n = functools.reduce(operator.mul, [r.extent for r in _src1.region], 1)
+            s2_n = functools.reduce(operator.mul, [r.extent for r in _src2.region], 1)
+            if s1_n < s2_n:
+                if op_name not in _COMMUTATIVE:
+                    return None, f"non-commutative op {op_name} cannot swap to broadcast"
+                _src1, _src2 = _src2, _src1
+
+        srcs: list[SrcSpec] = [SrcSpec(buf_region=_src1)]
+        if s2_scalar:
+            srcs.append(SrcSpec(scalar=_src2))
+        else:
+            srcs.append(SrcSpec(buf_region=_src2))
+
+        extras: dict[str, Any] = {}
+        rm = op.config.get("rounding_mode", None)
+        if rm is not None:
+            extras["rounding_mode"] = rm
+        return Plan(dst=_dst, srcs=srcs, extras=extras), None
+
+    return parse
+
+
+def _compute_add(src_vals, extras, dt):
+    return src_vals[0] + src_vals[1]
+
+
+def _compute_sub(src_vals, extras, dt):
+    return src_vals[0] - src_vals[1]
+
+
+def _compute_mul(src_vals, extras, dt):
+    return src_vals[0] * src_vals[1]
+
+
+def _compute_fdiv(src_vals, extras, dt):
+    return src_vals[0] / src_vals[1]
+
+
+def _compute_maximum(src_vals, extras, dt):
+    return Tx.max(src_vals[0], src_vals[1])
+
+
+BINARY_OPS: dict[str, OpSpec] = {
+    "add": OpSpec(
+        "add",
+        _parse_binary_for("add"),
+        _compute_add,
+        vec_impls=[BINARY_F32X2_IMPLS["add"]],
+    ),
+    "sub": OpSpec(
+        "sub",
+        _parse_binary_for("sub"),
+        _compute_sub,
+        vec_impls=[BINARY_F32X2_IMPLS["sub"]],
+    ),
+    "mul": OpSpec(
+        "mul",
+        _parse_binary_for("mul"),
+        _compute_mul,
+        vec_impls=[BINARY_F32X2_IMPLS["mul"]],
+    ),
+    "fdiv": OpSpec(
+        "fdiv",
+        _parse_binary_for("fdiv"),
+        _compute_fdiv,
+    ),
+    "maximum": OpSpec(
+        "maximum",
+        _parse_binary_for("maximum"),
+        _compute_maximum,
+    ),
+}

python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/cast.py

@@ -0,0 +1,45 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Cast op: ``Tx.cast(dst, src)``. Outer ``Tx.cast(..., dst.dtype)`` in the
+schedule handles the scalar conversion; the vec-impl packs pairs via
+CUDA intrinsics like ``__float22half2_rn``."""
+
+from __future__ import annotations
+
+from tvm.tirx import BufferRegion, TilePrimitiveCall
+
+from ..vec_emit.cast_vec2 import CAST_VEC2_IMPL
+from . import OpSpec, Plan, SrcSpec
+
+
+def _parse_cast(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
+    _dst: BufferRegion = op.args[0]
+    _src = op.args[1]
+    if not isinstance(_src, BufferRegion):
+        return None, "cast src must be a buffer region"
+    return Plan(dst=_dst, srcs=[SrcSpec(buf_region=_src)], extras={}), None
+
+
+def _compute_cast(src_vals, extras, dt):
+    # Schedule wraps with Tx.cast(..., dst.dtype) — just pass through.
+    return src_vals[0]
+
+
+CAST_OPS: dict[str, OpSpec] = {
+    "cast": OpSpec("cast", _parse_cast, _compute_cast, vec_impls=[CAST_VEC2_IMPL]),
+}

python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/fma.py

@@ -0,0 +1,50 @@
+# 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.
+
+"""FMA op: ``Tx.fma(dst, a, b, c)`` → ``dst = a*b + c``.
+
+Attaches ``fma_f32x2`` VecImpl for sm_100+ f32; falls back to scalar
+``a*b + c`` otherwise.
+"""
+
+from __future__ import annotations
+
+from tvm.tirx import BufferRegion, TilePrimitiveCall
+
+from ..vec_emit.fma_f32x2 import FMA_F32X2_IMPL
+from . import OpSpec, Plan, SrcSpec
+
+
+def _parse_fma(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
+    _dst: BufferRegion = op.args[0]
+    args = op.args[1:4]
+    srcs: list[SrcSpec] = []
+    for a in args:
+        if isinstance(a, BufferRegion):
+            srcs.append(SrcSpec(buf_region=a))
+        else:
+            srcs.append(SrcSpec(scalar=a))
+    return Plan(dst=_dst, srcs=srcs, extras={}), None
+
+
+def _compute_fma(src_vals, extras, dt):
+    return src_vals[0] * src_vals[1] + src_vals[2]
+
+
+FMA_OPS: dict[str, OpSpec] = {
+    "fma": OpSpec("fma", _parse_fma, _compute_fma, vec_impls=[FMA_F32X2_IMPL]),
+}

python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py

@@ -0,0 +1,121 @@
+# 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.
+
+"""Unary elementwise ops: zero / fill / reciprocal / sqrt / exp / exp2 / silu.
+
+All carry the same ``T.<unary>(dst, src[, bias, scale])`` shape (bias / scale
+optional; ``silu`` ignores bias/scale to preserve legacy behavior).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from tvm.ir import is_prim_expr
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, TilePrimitiveCall
+from tvm.tirx.expr import FloatImm
+
+from .._common import scalar_dtype
+from . import OpSpec, Plan, SrcSpec
+
+
+def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
+    """T.<unary>(dst, src[, bias, scale]) → Plan."""
+    _dst: BufferRegion = op.args[0]
+    _src = op.args[1]
+    _bias = op.args[2] if len(op.args) > 2 else None
+    _scale = op.args[3] if len(op.args) > 2 else None
+
+    srcs: list[SrcSpec] = []
+    if isinstance(_src, BufferRegion):
+        srcs.append(SrcSpec(buf_region=_src))
+    elif is_prim_expr(_src):
+        srcs.append(SrcSpec(scalar=_src))
+    else:
+        return None, f"unsupported src type {type(_src).__name__}"
+
+    extras: dict[str, Any] = {
+        "scale": _scale,
+        "bias_const": _bias if isinstance(_bias, FloatImm) else None,
+    }
+    if isinstance(_bias, BufferRegion):
+        srcs.append(SrcSpec(buf_region=_bias))
+        extras["has_bias_buf"] = True
+    else:
+        extras["has_bias_buf"] = False
+    return Plan(dst=_dst, srcs=srcs, extras=extras), None
+
+
+def _check_unary_extras(extras: dict, compute_dtype: str) -> tuple[bool, str | None]:
+    scale = extras.get("scale")
+    if scale is not None and scalar_dtype(scale) != compute_dtype:
+        return False, f"scale dtype {scalar_dtype(scale)} != compute dtype {compute_dtype}"
+    bias_const = extras.get("bias_const")
+    if bias_const is not None and scalar_dtype(bias_const) != compute_dtype:
+        return (
+            False,
+            f"bias_const dtype {scalar_dtype(bias_const)} != compute dtype {compute_dtype}",
+        )
+    return True, None
+
+
+def _with_bias_scale(raw_op):
+    """Wrap ``raw_op`` (e.g. ``T.exp``) into a compute that applies bias/scale first."""
+
+    def compute(src_vals, extras, dt):
+        x = src_vals[0]
+        scale = extras.get("scale")
+        if scale is not None:
+            x = x * scale
+        if extras.get("has_bias_buf"):
+            x = x + src_vals[1]
+        elif extras.get("bias_const") is not None:
+            x = x + extras["bias_const"]
+        return raw_op(x)
+
+    return compute
+
+
+def _compute_zero(src_vals, extras, dt):
+    return 0.0
+
+
+def _compute_fill(src_vals, extras, dt):
+    return src_vals[0]
+
+
+def _compute_reciprocal(src_vals, extras, dt):
+    x = src_vals[0]
+    return T.FloatImm(x.ty, 1.0) / x
+
+
+def _compute_silu(src_vals, extras, dt):
+    # Legacy: silu doesn't apply bias/scale.
+    x = src_vals[0]
+    return x / (T.FloatImm(x.ty, 1.0) + T.exp(T.FloatImm(x.ty, 0.0) - x))
+
+
+UNARY_OPS: dict[str, OpSpec] = {
+    "zero": OpSpec("zero", _parse_unary, _compute_zero, _check_unary_extras),
+    "fill": OpSpec("fill", _parse_unary, _compute_fill, _check_unary_extras),
+    "reciprocal": OpSpec("reciprocal", _parse_unary, _compute_reciprocal, _check_unary_extras),
+    "sqrt": OpSpec("sqrt", _parse_unary, _with_bias_scale(T.sqrt), _check_unary_extras),
+    "exp": OpSpec("exp", _parse_unary, _with_bias_scale(T.exp), _check_unary_extras),
+    "exp2": OpSpec("exp2", _parse_unary, _with_bias_scale(T.exp2), _check_unary_extras),
+    "silu": OpSpec("silu", _parse_unary, _compute_silu, _check_unary_extras),
+}

python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py

@@ -0,0 +1,429 @@
+# 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.
+
+"""Elementwise dispatch when all operands live in ``local`` (registers).
+
+Mirrors ``cuda/copy/reg.py``: the partition is *induced* by the layout that
+carries thread-axis info (the "anchor" operand). The region slice is absorbed
+into the sliced layout up front via ``align_operands_to_anchor`` — emit
+operates on a flat 1D per-thread view and indexes it with a scalar offset, so
+codegen never sees multi-dim ``get_indices`` inside ``T.vectorized``.
+
+Two paths inside emit:
+  * induced (anchor exists)   — atom-based, exactly mirrors copy reg.py
+  * trivial (no anchor)       — flat full region, every thread runs the full
+                                loop on its private storage
+"""
+
+from __future__ import annotations
+
+import functools
+import operator
+
+from tvm.arith import Analyzer
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc, TilePrimitiveCall
+from tvm.tirx.layout import TileLayout
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.operator.tile_primitive.dispatcher import fail
+
+from ..common import get_st_extent
+from ..copy._common import _carve_tail, _verify_s_tail_contig
+from ..layout_utils import get_sublayout_from_region, layout_signature
+from ._common import (
+    _TID_AXIS_FOR_SCOPE,
+    _all_threads_active,
+    _tensor_shape_of,
+    _thread_cnt,
+    align_operands_to_anchor,
+    buffer_regions,
+    compute_dtype_of,
+    pick_anchor,
+    shape_broadcast_compat,
+)
+
+
+# -----------------------------------------------------------------------------
+# Predicate
+# -----------------------------------------------------------------------------
+def _validate_anchor_layout(anchor_br) -> tuple[bool, str | None]:
+    layout = anchor_br.buffer.layout
+    if layout.is_swizzle():
+        return False, "anchor layout is swizzle"
+    if not isinstance(layout, TileLayout):
+        return False, f"anchor layout is {type(layout).__name__}, not TileLayout"
+    return True, None
+
+
+def _validate_scope_level_anchor(anchor_br, sctx: DispatchContext) -> tuple[bool, str | None]:
+    """For warp/warpgroup/cta scope, require dst to be scope-level: after
+    canonicalizing with the target its thread axes are the scope's intra-thread
+    axis (laneid/tid_in_wg/tx) and, sorted by stride, tile a complete ``T:1``
+    chain over all ``T`` threads of the scope. Rejects thread-local ``.local()``
+    views; thread scope is exempt.
+    """
+    scope = sctx.scope_kind
+    if scope == "thread":
+        return True, None
+    expected_axis = _TID_AXIS_FOR_SCOPE.get(scope)
+    if expected_axis is None:
+        return True, None
+    expected_cnt = _thread_cnt(sctx)
+
+    # Canonicalize the sliced anchor with the target so warp/lane axes fuse.
+    st, ext = get_st_extent(anchor_br)
+    sliced = get_sublayout_from_region(anchor_br.buffer.layout, anchor_br.buffer.shape, st, ext)
+    with sctx.target:
+        canon = sliced.canonicalize()
+    shard = getattr(canon, "shard", None)
+    if shard is None:
+        return False, f"{scope}-scope op operand layout is not a TileLayout after slicing"
+
+    thread_iters = [it for it in shard if it.axis.is_thread()]
+    if not thread_iters:
+        return (
+            False,
+            f"{scope}-scope op needs a {scope}-level operand whose layout carries "
+            f"thread axes ({expected_axis} composing to {expected_cnt}:1); got a "
+            f"thread-local view with no thread axes — pass the {scope}-level tensor, "
+            f"not its `.local()` (per-thread) view",
+        )
+    bad = sorted({it.axis.name for it in thread_iters if it.axis.name != expected_axis})
+    if bad:
+        return (
+            False,
+            f"{scope}-scope op operand carries thread axes {bad}; after "
+            f"canonicalization a {scope}-level layout must use only {expected_axis!r}",
+        )
+    # Sorted by stride the thread iters must tile a complete chain 1, e0,
+    # e0*e1, ... up to the scope thread count — i.e. cover all T threads with
+    # no gap or overlap (extents alone would miss gaps/overlaps).
+    running = 1
+    for it in sorted(thread_iters, key=lambda i: int(i.stride)):
+        stride, extent = int(it.stride), int(it.extent)
+        if stride != running:
+            return (
+                False,
+                f"{scope}-scope op operand thread axes do not tile a complete "
+                f"{expected_cnt}:1 (sorted by stride: expected {running}, got {stride})",
+            )
+        running *= extent
+    if running != expected_cnt:
+        return (
+            False,
+            f"{scope}-scope op operand thread axes span {running} threads, not the "
+            f"full {expected_cnt} of the {scope}",
+        )
+    return True, None
+
+
+def _check_layout_operands_agree(plan, sctx) -> tuple[bool, str | None]:
+    """Replica sigs must match across non-trivial-layout operands.
+
+    ``align_operands_to_anchor`` normalizes thread + local parts via
+    permute/group, but the replica part isn't touched by alignment — if
+    operands disagree there, alignment can't fix it and emit will be wrong.
+    Thread / local mismatches that alignment can't resolve will raise
+    cleanly at align time, so we don't pre-check them.
+    """
+    # All operands have a layout (predicate already enforced this); just
+    # iterate them all.
+    layout_brs = list(buffer_regions(plan))
+    if len(layout_brs) < 2:
+        return True, None
+    analyzer = Analyzer()
+    replica_sigs = []
+    for br in layout_brs:
+        st, ext = get_st_extent(br)
+        with sctx.target:
+            sliced = get_sublayout_from_region(br.buffer.layout, br.buffer.shape, st, ext)
+            canon = sliced.canonicalize()
+        sig = layout_signature(canon)
+        if sig is None:
+            return False, "layout has no signature (not a TileLayout?)"
+        # layout_signature returns (thread_sig, local_sig, replica_sig)
+        replica_sigs.append(sig[2])
+    for s in replica_sigs[1:]:
+        # Compare replica entries (axis_key, extent, stride) element-wise.
+        if len(s) != len(replica_sigs[0]):
+            return False, "replica sig mismatch (different number of replica iters)"
+        for (k_a, e_a, st_a), (k_b, e_b, st_b) in zip(replica_sigs[0], s):
+            if k_a != k_b:
+                return False, "replica sig mismatch (axis key)"
+            if not analyzer.can_prove_equal(e_a, e_b):
+                return False, "replica sig mismatch (extent)"
+            if not analyzer.can_prove_equal(st_a, st_b):
+                return False, "replica sig mismatch (stride)"
+    return True, None
+
+
+def is_reg_ewise(spec):
+    """Predicate factory: dispatch accepted iff all operands in ``local`` scope."""
+
+    def check(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+        if not sctx.is_target("cuda"):
+            return False, "non-cuda target"
+        if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
+            return False, f"unsupported scope {sctx.scope_kind}"
+        ok, reason = _all_threads_active(sctx)
+        if not ok:
+            return False, reason
+        plan, msg = spec.parse(op_call)
+        if msg is not None or plan is None:
+            return False, msg
+        for br in buffer_regions(plan):
+            if br.buffer.scope() != "local":
+                return False, f"operand scope {br.buffer.scope()} != local"
+            if br.buffer.layout is None:
+                return False, f"operand {br} has no layout"
+        if spec.check_extras is not None:
+            ok2, reason2 = spec.check_extras(plan.extras, compute_dtype_of(plan))
+            if not ok2:
+                return False, reason2
+        anchor = pick_anchor(plan)
+        ok3, reason3 = _validate_anchor_layout(anchor)
+        if not ok3:
+            return False, reason3
+        ok_scope, reason_scope = _validate_scope_level_anchor(anchor, sctx)
+        if not ok_scope:
+            return False, reason_scope
+        # Shape compat (NumPy-style broadcast): anchor's tensor shape is the
+        # result shape; every operand must broadcast TO anchor.
+        anchor_tshape = _tensor_shape_of(anchor.region)
+        for br in buffer_regions(plan):
+            if br is anchor:
+                continue
+            op_tshape = _tensor_shape_of(br.region)
+            ok_b, reason_b = shape_broadcast_compat(op_tshape, anchor_tshape)
+            if not ok_b:
+                return False, f"shape incompat: {reason_b}"
+        ok4, reason4 = _check_layout_operands_agree(plan, sctx)
+        if not ok4:
+            return False, reason4
+        return True, None
+
+    return check
+
+
+# -----------------------------------------------------------------------------
+# Shared helpers
+# -----------------------------------------------------------------------------
+def _prod(it) -> int:
+    return functools.reduce(operator.mul, it, 1)
+
+
+# -----------------------------------------------------------------------------
+# Main entry
+# -----------------------------------------------------------------------------
+def emit_reg(op_call: TilePrimitiveCall, spec, sctx: DispatchContext) -> PrimFunc:
+    plan, msg = spec.parse(op_call)
+    if msg is not None or plan is None:
+        fail(msg or "parse failed")
+    return _emit_induced(plan, spec, sctx, op_call, pick_anchor(plan))
+
+
+# -----------------------------------------------------------------------------
+# Induced path — anchor with non-trivial layout drives partition
+# -----------------------------------------------------------------------------
+def _strip_thread(layout):
+    """Return a new TileLayout with thread iters removed from shard."""
+    mem_iters = [it for it in layout.shard if not it.axis.is_thread()]
+    return TileLayout.from_iters(mem_iters, list(layout.replica), dict(layout.offset))
+
+
+def _pick_vec_and_carve(spec, op_call, sctx, plan, per_op_mem_layouts):
+    """Pick ``(vec_len, vec_impl, carved_layouts)``.
+
+    Enumerate ``spec.vec_impls`` widest-first. For each candidate ``vec_len``:
+      1. Try ``_carve_tail`` on each operand's mem-only layout (may split a
+         boundary iter so the tail product equals ``vec_len``).
+      2. Verify the carved tail is physically contiguous (stride-1 chain whose
+         product equals ``vec_len``).
+      3. Call ``impl.applies(op_call, sctx, plan)``.
+    First candidate that passes all three on EVERY operand wins. Otherwise
+    fall back to scalar: ``vec_len=1``, ``vec_impl=None``, original (uncarved)
+    layouts.
+    """
+    impls = sorted(getattr(spec, "vec_impls", []), key=lambda i: -i.vec_len)
+    for impl in impls:
+        cand = impl.vec_len
+        carved_try = {}
+        all_ok = True
+        for op_br, layout in per_op_mem_layouts.items():
+            new_iters = _carve_tail(list(layout.shard), cand)
+            if new_iters is None:
+                all_ok = False
+                break
+            new_layout = TileLayout.from_iters(new_iters, list(layout.replica), dict(layout.offset))
+            if not _verify_s_tail_contig(new_layout, cand):
+                all_ok = False
+                break
+            carved_try[op_br] = new_layout
+        if not all_ok:
+            continue
+        ok, _ = impl.applies(op_call, sctx, plan)
+        if not ok:
+            continue
+        return cand, impl, carved_try
+    # Scalar fallback — use uncarved mem layouts as-is.
+    return 1, None, dict(per_op_mem_layouts)
+
+
+def _emit_induced(plan, spec, sctx, op_call, anchor_br) -> PrimFunc:
+    # Every buffer-region operand has a layout (enforced by predicate);
+    # trivial / identity layouts are fine — the algorithm is robust to
+    # layouts with no thread axes (strip is no-op, placeholders empty).
+    layout_others = [br for br in buffer_regions(plan) if br is not anchor_br]
+
+    # Step 1: slice + permute (region offset absorbed into op_p.offset; per-iter
+    # strides reflect post-slice physical addressing). No (st, ext) leaks out.
+    with sctx.target:
+        anchor_p, per_op_aligned = align_operands_to_anchor(anchor_br, layout_others)
+
+    # Step 2: post-align thread-equality check. ``align`` is supposed to
+    # normalize the thread part; we verify it did. (Replica was pre-checked
+    # in the predicate.)
+    def _thread_iters(layout):
+        c = layout.canonicalize()
+        return [(it.axis, int(it.extent), int(it.stride)) for it in c.shard if it.axis.is_thread()]
+
+    anchor_thread = _thread_iters(anchor_p)
+    for op_br, (op_p, _) in per_op_aligned.items():
+        if _thread_iters(op_p) != anchor_thread:
+            fail("thread part mismatch between anchor and operand after alignment")
+
+    # Step 3: drop thread iters; from here on operands have mem-only layouts.
+    per_op_mem = {anchor_br: _strip_thread(anchor_p)}
+    for op_br, (op_p, _) in per_op_aligned.items():
+        per_op_mem[op_br] = _strip_thread(op_p)
+
+    # Step 4: enumerate spec.vec_impls widest-first; try carve tail for each
+    # operand. First candidate that all operands can carve + impl.applies wins.
+    # Otherwise scalar fallback (vec=1, no inner loop).
+    vec_len, vec_impl, per_op_carved = _pick_vec_and_carve(spec, op_call, sctx, plan, per_op_mem)
+
+    # Step 5: totals + emit. per_thread_total = ∏ extents of (carved) mem
+    # layout. All operands have the same per_thread_total (alignment invariant).
+    per_thread_total = _prod(int(it.extent) for it in per_op_carved[anchor_br].shard)
+    outer_total = per_thread_total // vec_len if vec_len > 0 else per_thread_total
+
+    if vec_impl is not None:
+        result = _emit_induced_packed(
+            plan,
+            vec_impl,
+            vec_len,
+            outer_total,
+            per_thread_total,
+            per_op_carved,
+            anchor_br,
+        )
+    else:
+        result = _emit_induced_scalar(
+            plan,
+            spec,
+            outer_total,
+            per_thread_total,
+            per_op_carved,
+            anchor_br,
+        )
+
+    return result
+
+
+def _make_views_meta(per_op_carved, per_thread_total):
+    """Build the per-operand 1D buffer view dict.
+
+    Each view aliases the operand's physical storage as a 1D shape of
+    ``per_thread_total`` elements, with layout = the operand's carved mem-only
+    TileLayout. Scalar indexing into the view goes through this layout's
+    iter strides at codegen time.
+    """
+    return {
+        op_br: T.decl_buffer(
+            (per_thread_total,),
+            op_br.buffer.dtype,
+            op_br.buffer.data,
+            scope="local",
+            layout=per_op_carved[op_br],
+        )
+        for op_br in per_op_carved
+    }
+
+
+# -----------------------------------------------------------------------------
+# Emit — packed (one PTX/CUDA call per outer chunk; no T.vectorized inside)
+# -----------------------------------------------------------------------------
+def _emit_induced_packed(
+    plan, vec_impl, vec_len, outer_total, per_thread_total, per_op_carved, anchor_br
+) -> PrimFunc:
+    extras = plan.extras
+    srcs = plan.srcs
+    dst_br = plan.dst
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        views = T.meta_var(_make_views_meta(per_op_carved, per_thread_total))
+        # Serial loop (not T.unroll): T.unroll materializes each per-iter
+        # ``dst_lane_indices`` / ``src_args`` buffer as a fresh int[1]
+        # declaration, multiplying by outer_total. ptxas unrolls the
+        # static-bound loop without that scratch explosion.
+        for f in range(outer_total):
+            # Pass logical 1D coord; each buffer's own layout maps it to
+            # physical at access time (handles wgmma, broadcast, etc.).
+            dst_lane_indices = [[f * vec_len + k] for k in range(vec_len)]
+            src_args = T.meta_var(
+                [
+                    src.scalar
+                    if src.is_scalar
+                    else (
+                        views[src.buf_region],
+                        [[f * vec_len + k] for k in range(vec_len)],
+                    )
+                    for src in srcs
+                ]
+            )
+            T.evaluate(vec_impl.emit(views[dst_br], dst_lane_indices, src_args, extras))
+
+    return impl
+
+
+# -----------------------------------------------------------------------------
+# Emit — scalar fallback (vec_len = 1; one element per outer iter; no
+# T.vectorized inside, so no codegen vec-packing of multi-dim indices).
+# -----------------------------------------------------------------------------
+def _emit_induced_scalar(
+    plan, spec, outer_total, per_thread_total, per_op_carved, anchor_br
+) -> PrimFunc:
+    extras = plan.extras
+    srcs = plan.srcs
+    dst_br = plan.dst
+    dst_dtype = dst_br.buffer.dtype
+    compute = spec.compute_scalar
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        views = T.meta_var(_make_views_meta(per_op_carved, per_thread_total))
+        # Serial loop (not T.unroll) — see _emit_induced_packed for why.
+        for f in range(outer_total):
+            # Logical 1D coord = f (vec_len = 1 in scalar path); each
+            # buffer's layout maps to physical at access time.
+            src_vals = T.meta_var(
+                [src.scalar if src.is_scalar else views[src.buf_region][f] for src in srcs]
+            )
+            views[dst_br][f] = T.cast(compute(src_vals, extras, dst_dtype), dst_dtype)
+
+    return impl

python/tvm/backend/cuda/operator/tile_primitive/elementwise/register.py

@@ -0,0 +1,61 @@
+# 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.
+
+"""Register each op in ``ALL_OPS`` for both dispatch variants (``reg``, ``smem``).
+
+Mirrors copy PR-640's two-variant model: scope-pair drives the dispatch
+selection, the underlying algorithm (induced vs synthesized) follows.
+"""
+
+from tvm.tirx import PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch
+
+from .ops import ALL_OPS
+from .reg import emit_reg, is_reg_ewise
+from .smem import emit_smem, is_smem_ewise
+
+
+def _register_reg(spec) -> None:
+    @register_dispatch(
+        spec.name,
+        "cuda",
+        variant="reg",
+        priority=10,
+        when=[predicate(f"{spec.name}_reg", is_reg_ewise(spec))],
+    )
+    def _dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _spec=spec) -> PrimFunc:
+        return emit_reg(op, _spec, sctx)
+
+
+def _register_smem(spec) -> None:
+    @register_dispatch(
+        spec.name,
+        "cuda",
+        variant="smem",
+        priority=10,
+        when=[predicate(f"{spec.name}_smem", is_smem_ewise(spec))],
+    )
+    def _dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _spec=spec) -> PrimFunc:
+        return emit_smem(op, _spec, sctx)
+
+
+for _spec in ALL_OPS.values():
+    _register_reg(_spec)
+    _register_smem(_spec)
+
+
+__all__: list[str] = []

python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py

@@ -0,0 +1,264 @@
+# 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.
+
+"""Elementwise dispatch when all operands live in ``shared*``.
+
+Mirrors ``cuda/copy/gmem_smem.py``: no operand carries a per-thread partition,
+so partition is *synthesized* from ``sctx.intra`` (``thread_cnt = ∏ intra``).
+Each thread takes ``ceildiv(total, vec_chunk * thread_cnt)`` strided chunks.
+
+The shared buffers are indexed multi-dim via ``get_indices(fused, dst_st,
+dst_ext)`` and the buffer's own layout resolves to physical addresses at
+codegen time. Packed-vec emit requires the innermost dim to have stride 1
+(non-swizzle slice) so lanes are physically contiguous; checked in
+``_max_layout_vec``.
+"""
+
+from __future__ import annotations
+
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc, TilePrimitiveCall
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.operator.tile_primitive.dispatcher import fail
+
+from ..common import get_indices, get_st_extent, get_thread_cnt
+from ._common import (
+    _TID_AXIS_FOR_SCOPE,
+    _all_threads_active,
+    _axis_decl,
+    _broadcast_indices,
+    _tensor_shape_of,
+    buffer_regions,
+    compute_dtype_of,
+    dtype_bits,
+    emit_scope_sync,
+    fetch_src_value,
+    n_elements,
+    pick_vec_chunk,
+    shape_broadcast_compat,
+)
+
+
+# -----------------------------------------------------------------------------
+# Predicate
+# -----------------------------------------------------------------------------
+def is_smem_ewise(spec):
+    """Predicate factory: dispatch accepted iff all operands in ``shared*``."""
+
+    def check(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]:
+        if not sctx.is_target("cuda"):
+            return False, "non-cuda target"
+        if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"):
+            return False, f"unsupported scope {sctx.scope_kind}"
+        ok, reason = _all_threads_active(sctx)
+        if not ok:
+            return False, reason
+        plan, msg = spec.parse(op_call)
+        if msg is not None or plan is None:
+            return False, msg
+        for br in buffer_regions(plan):
+            if not br.buffer.scope().startswith("shared"):
+                return False, f"operand scope {br.buffer.scope()} != shared*"
+            if br.buffer.layout is None:
+                return False, "shared operand has no layout"
+        if spec.check_extras is not None:
+            ok2, reason2 = spec.check_extras(plan.extras, compute_dtype_of(plan))
+            if not ok2:
+                return False, reason2
+        # NumPy-style right-aligned broadcast: anchor = plan.dst; every src
+        # must be shape-compatible with anchor (extent matches or is 1).
+        anchor_tshape = _tensor_shape_of(plan.dst.region)
+        for s in plan.srcs:
+            if s.buf_region is None:
+                continue
+            src_tshape = _tensor_shape_of(s.buf_region.region)
+            ok_b, reason_b = shape_broadcast_compat(src_tshape, anchor_tshape)
+            if not ok_b:
+                return False, f"shape incompat: {reason_b}"
+        return True, None
+
+    return check
+
+
+# -----------------------------------------------------------------------------
+# vec_chunk selection
+# -----------------------------------------------------------------------------
+def _max_layout_vec(plan, total: int, thread_cnt: int) -> int:
+    """Widest vec_chunk dividing all operands' innermost extents AND
+    ``total / thread_cnt``, within dtype-bit candidates ``{128,64,32,16,8}``."""
+    max_bits = dtype_bits(plan.dst.buffer.dtype)
+    for s in plan.srcs:
+        if s.buf_region is not None:
+            max_bits = max(max_bits, dtype_bits(s.buf_region.buffer.dtype))
+    per_thread = total // thread_cnt if thread_cnt > 0 else total
+    if total % thread_cnt != 0:
+        return 1
+
+    inners = [int(plan.dst.region[-1].extent)]
+    for s in plan.srcs:
+        if s.buf_region is None or s.index_fn is not None:
+            continue
+        inners.append(int(s.buf_region.region[-1].extent))
+
+    for cand_bits in (128, 64, 32, 16, 8):
+        n = cand_bits // max_bits
+        if n <= 0:
+            continue
+        if per_thread % n != 0:
+            continue
+        if all(i % n == 0 for i in inners):
+            return n
+    return 1
+
+
+# -----------------------------------------------------------------------------
+# Main entry
+# -----------------------------------------------------------------------------
+def emit_smem(op_call: TilePrimitiveCall, spec, sctx: DispatchContext) -> PrimFunc:
+    plan, msg = spec.parse(op_call)
+    if msg is not None or plan is None:
+        fail(msg or "parse failed")
+
+    # Use cuda/common.py:get_thread_cnt rather than copy/_common.py:_thread_cnt
+    # — the latter computes ``∏ sctx.intra`` which silently returns 0 for
+    # sub-warp counts at cta scope (warpid extent rounds down to 0). The
+    # former reads launch_params["threadIdx.x"].dom.extent and is correct
+    # for all scopes.
+    thread_cnt = get_thread_cnt(sctx)
+    if thread_cnt is None:
+        fail(f"unsupported scope {sctx.scope_kind} for smem emit")
+    thread_cnt = int(thread_cnt)
+    if thread_cnt <= 0:
+        fail(f"non-positive thread_cnt {thread_cnt}")
+    assert "threadIdx.y" not in sctx.launch_params and "threadIdx.z" not in sctx.launch_params, (
+        "smem emit currently assumes 1D threadIdx"
+    )
+
+    total = n_elements(plan.dst)
+    vec_max = _max_layout_vec(plan, total, thread_cnt)
+    vec_chunk, vec_impl = pick_vec_chunk(spec, op_call, sctx, plan, vec_max)
+
+    if vec_impl is not None:
+        return _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx)
+    return _emit_scalar(plan, spec, vec_chunk, total, thread_cnt, sctx)
+
+
+def _tid_expr(sctx: DispatchContext):
+    """Per-scope tid expr. ``thread`` scope returns 0; collective scopes use
+    ``_axis_decl`` (T.lane_id / T.thread_id_in_wg / threadIdx.x)."""
+    if sctx.scope_kind == "thread":
+        return 0
+    axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind]
+    return _axis_decl(axis_name, sctx)
+
+
+# -----------------------------------------------------------------------------
+# Per-lane src index helper (handles broadcast via right-aligned compat)
+# -----------------------------------------------------------------------------
+def _src_lane_indices(src_br, dst_lane_indices, dst_st, dst_ext, vec_chunk, fused0):
+    """Return the per-lane multi-dim index list for ``src_br``.
+
+    If src's region shape matches dst's, fall through to ``get_indices``
+    (same as the legacy non-broadcast path). Otherwise derive each lane's
+    src index from the corresponding dst lane index via right-aligned
+    broadcast compat.
+    """
+    src_st, src_ext = get_st_extent(src_br)
+    if tuple(int(e) for e in src_ext) == tuple(int(e) for e in dst_ext):
+        return [get_indices(fused0 + k, src_st, src_ext) for k in range(vec_chunk)]
+    return [
+        _broadcast_indices(dst_lane_indices[k], dst_st, dst_ext, src_st, src_ext)
+        for k in range(vec_chunk)
+    ]
+
+
+# -----------------------------------------------------------------------------
+# Emit — packed
+# -----------------------------------------------------------------------------
+def _emit_packed(plan, vec_impl, vec_chunk, total, thread_cnt, sctx) -> PrimFunc:
+    extras = plan.extras
+    srcs = plan.srcs
+    dst_buf = plan.dst.buffer
+    dst_st, dst_ext = get_st_extent(plan.dst)
+    sync = emit_scope_sync(sctx.scope_kind)
+    n_outer = (total + vec_chunk * thread_cnt - 1) // (vec_chunk * thread_cnt)
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        tid = _tid_expr(sctx)
+        for s in T.serial(0, n_outer):
+            # First lane's fused index for this thread, this chunk.
+            fused0 = T.meta_var(s * vec_chunk * thread_cnt + tid * vec_chunk)
+            # Predicate the call (skip the trailing partial chunk).
+            if fused0 + vec_chunk <= total:
+                dst_lane_indices = T.meta_var(
+                    [get_indices(fused0 + k, dst_st, dst_ext) for k in range(vec_chunk)]
+                )
+                src_args = T.meta_var(
+                    [
+                        srcs[i].scalar
+                        if srcs[i].is_scalar
+                        else (
+                            srcs[i].buf_region.buffer,
+                            _src_lane_indices(
+                                srcs[i].buf_region,
+                                dst_lane_indices,
+                                dst_st,
+                                dst_ext,
+                                vec_chunk,
+                                fused0,
+                            ),
+                        )
+                        for i in range(len(srcs))
+                    ]
+                )
+                T.evaluate(vec_impl.emit(dst_buf, dst_lane_indices, src_args, extras))
+        sync()
+
+    return impl
+
+
+# -----------------------------------------------------------------------------
+# Emit — scalar fallback
+# -----------------------------------------------------------------------------
+def _emit_scalar(plan, spec, vec_chunk, total, thread_cnt, sctx) -> PrimFunc:
+    extras = plan.extras
+    srcs = plan.srcs
+    dst_buf = plan.dst.buffer
+    dst_st, dst_ext = get_st_extent(plan.dst)
+    dst_dtype = dst_buf.dtype
+    compute = spec.compute_scalar
+    sync = emit_scope_sync(sctx.scope_kind)
+    n_outer = (total + vec_chunk * thread_cnt - 1) // (vec_chunk * thread_cnt)
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        tid = _tid_expr(sctx)
+        for s in T.serial(0, n_outer):
+            for vec in T.vectorized(vec_chunk):
+                fused = T.meta_var(s * vec_chunk * thread_cnt + tid * vec_chunk + vec)
+                if fused < total:
+                    dst_idx = T.meta_var(get_indices(fused, dst_st, dst_ext))
+                    src_vals = T.meta_var(
+                        [fetch_src_value(src, fused, dst_idx, dst_st, dst_ext) for src in srcs]
+                    )
+                    dst_buf[tuple(dst_idx)] = T.cast(
+                        compute(src_vals, extras, dst_dtype), dst_dtype
+                    )
+        sync()
+
+    return impl

python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/__init__.py

@@ -0,0 +1,40 @@
+# 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.
+
+"""Packed-vector emit functions for elementwise ops.
+
+Each module here exposes one or more ``VecImpl`` instances that an ``OpSpec``
+(in ``ops/``) attaches to its ``vec_impls`` list. ``reg.py``/``smem.py`` then
+pick the widest matching one at dispatch time, mirroring how copy picks
+``copy_{Nb}`` from a menu.
+
+VecImpl emit contract:
+  emit(dst_buf, dst_lane_indices, src_args, extras) -> Expr
+
+  * dst_buf: Buffer
+  * dst_lane_indices: list[list[Expr]] of length ``vec_len``; each entry is the
+                     multi-dim indices for one lane (precomputed by schedule).
+  * src_args[i]: one of
+      - Expr (scalar src — broadcast across all lanes)
+      - tuple (Buffer, list[list[Expr]] of length ``vec_len``) — buffer src
+                                                                with per-lane indices
+  * extras: dict (rounding_mode, etc.)
+
+  Returns the PTX/CUDA call result; the schedule wraps in ``T.evaluate`` at
+  the call site. All Python-side shape branching (scalar vs buffer src) happens
+  in this emit function -- collapses the old 4x2 schema.py factory explosion.
+"""

python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py

@@ -0,0 +1,97 @@
+# 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.
+
+"""Packed f32x2 VecImpls for binary add/sub/mul on sm_100+.
+
+PTX op family: ``{add,sub,mul}.<rm>.ftz.f32x2``. Each call processes 2 f32s
+per operand. The old ``_make_binary_packed_f32x2_factory`` (240+ lines, 8
+``@T.prim_func`` shape combos per op) collapses to one ``emit`` per op
+because operand-shape branching is now Python-level (outside any
+``@T.prim_func``).
+"""
+
+from __future__ import annotations
+
+from tvm.ir.expr import Expr
+from tvm.script import tirx as T
+
+from .._common import dtype_name, scalar_dtype
+from ..ops import VecImpl
+
+
+def _lane(arg, k):
+    """Read lane ``k`` of one operand argument.
+
+    arg is either a scalar Expr (broadcast) or ``(Buffer, lane_indices)``.
+    """
+    if isinstance(arg, tuple):
+        buf, lane_indices = arg
+        return buf[tuple(lane_indices[k])]
+    return arg
+
+
+def _f32x2_applies(op_name):
+    """Predicate: f32 dst+srcs, sm_100+, no broadcasting srcs, two srcs."""
+
+    def applies(op_call, sctx, plan):
+        from ...common import sm_version_ok
+
+        if dtype_name(plan.dst.buffer.dtype) != "float32":
+            return False, "dst dtype not f32"
+        if not sm_version_ok(op_call, sctx, min_version=100)[0]:
+            return False, "sm version < 100"
+        if len(plan.srcs) != 2:
+            return False, "binary requires 2 srcs"
+        for s in plan.srcs:
+            if s.is_scalar:
+                if scalar_dtype(s.scalar) != "float32":
+                    return False, "scalar src dtype not f32"
+            else:
+                if dtype_name(s.buf_region.buffer.dtype) != "float32":
+                    return False, "buffer src dtype not f32"
+                if s.index_fn is not None:
+                    return False, "broadcasting src not supported by f32x2 packed"
+        return True, None
+
+    return applies
+
+
+def _emit_binary_f32x2_for(op_name):
+    op_func = getattr(T.ptx, f"{op_name}_f32x2")
+
+    def emit(dst_buf, dst_lane_indices, src_args, extras) -> Expr:
+        a_arg, b_arg = src_args
+        rm = extras.get("rounding_mode", "rz")
+        return op_func(
+            T.address_of(dst_buf[tuple(dst_lane_indices[0])]),
+            T.cuda.make_float2(_lane(a_arg, 0), _lane(a_arg, 1)),
+            T.cuda.make_float2(_lane(b_arg, 0), _lane(b_arg, 1)),
+            rounding=rm,
+            ftz=True,
+        )
+
+    return emit
+
+
+BINARY_F32X2_IMPLS: dict[str, VecImpl] = {
+    name: VecImpl(
+        vec_len=2,
+        applies=_f32x2_applies(name),
+        emit=_emit_binary_f32x2_for(name),
+    )
+    for name in ("add", "sub", "mul")
+}

python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py

@@ -0,0 +1,92 @@
+# 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.
+
+"""Packed vec_len=2 cast via CUDA pair intrinsics.
+
+Each supported (src_dtype, dst_dtype) pair has a CUDA builtin that converts a
+packed-2 source to a packed-2 destination in one instruction
+(e.g. ``__float22half2_rn``). The intrinsic takes pointers to the first
+element of each packed pair on either side.
+"""
+
+from __future__ import annotations
+
+from tvm.ir.expr import Expr
+from tvm.script import tirx as T
+
+from .._common import dtype_name
+from ..ops import VecImpl
+
+_VEC2_CAST_INTRINSICS = {
+    ("float32", "float16"): "__float22half2_rn",
+    ("float16", "float32"): "__half22float2",
+    ("bfloat16", "float32"): "__bfloat1622float2",
+    ("float32", "bfloat16"): "__float22bfloat162_rn",
+}
+_DTYPE_X2_NAME = {"float32": "float2", "float16": "half2", "bfloat16": "nv_bfloat162"}
+
+
+def _intrinsic_name(src_dtype, dst_dtype):
+    return f"tvm_builtin_cast_{src_dtype}x2_{dst_dtype}x2"
+
+
+def _intrinsic_source(src_dtype, dst_dtype):
+    intrinsic = _VEC2_CAST_INTRINSICS[(src_dtype, dst_dtype)]
+    return (
+        f"\n__forceinline__ __device__ void {_intrinsic_name(src_dtype, dst_dtype)}"
+        f"(void* dst, void* src) {{\n"
+        f"    (({_DTYPE_X2_NAME[dst_dtype]}*)dst)[0] = "
+        f"{intrinsic}((({_DTYPE_X2_NAME[src_dtype]}*)src)[0]);\n"
+        "}\n"
+    )
+
+
+def _cast_vec2_applies(op_call, sctx, plan):
+    if len(plan.srcs) != 1 or plan.srcs[0].is_scalar:
+        return False, "cast requires 1 buffer src"
+    src = plan.srcs[0]
+    if src.index_fn is not None:
+        return False, "broadcasting src not supported by cast vec2"
+    src_dtype = dtype_name(src.buf_region.buffer.dtype)
+    dst_dtype = dtype_name(plan.dst.buffer.dtype)
+    if (src_dtype, dst_dtype) not in _VEC2_CAST_INTRINSICS:
+        return False, f"no vec2 intrinsic for {src_dtype}->{dst_dtype}"
+    return True, None
+
+
+def _emit_cast_vec2(dst_buf, dst_lane_indices, src_args, extras) -> Expr:
+    src_arg = src_args[0]
+    # cast_vec2 requires buffer src (guarded by applies()).
+    assert isinstance(src_arg, tuple), "cast vec2 src must be a buffer"
+    src_buf, src_lane_indices = src_arg
+    src_dtype = dtype_name(src_buf.dtype)
+    dst_dtype = dtype_name(dst_buf.dtype)
+    func_name = _intrinsic_name(src_dtype, dst_dtype)
+    source_code = _intrinsic_source(src_dtype, dst_dtype)
+    return T.cuda.func_call(
+        func_name,
+        T.address_of(dst_buf[tuple(dst_lane_indices[0])]),
+        T.address_of(src_buf[tuple(src_lane_indices[0])]),
+        source_code=source_code,
+    )
+
+
+CAST_VEC2_IMPL = VecImpl(
+    vec_len=2,
+    applies=_cast_vec2_applies,
+    emit=_emit_cast_vec2,
+)

python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.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.
+
+"""Packed f32x2 VecImpl for FMA on sm_100+.
+
+PTX: ``fma.<rm>.ftz.f32x2 d, a, b, c`` — 2 f32 FMAs per call. Same Python-side
+shape collapse as binary_f32x2.
+"""
+
+from __future__ import annotations
+
+from tvm.ir.expr import Expr
+from tvm.script import tirx as T
+
+from .._common import dtype_name, scalar_dtype
+from ..ops import VecImpl
+from .binary_f32x2 import _lane
+
+
+def _fma_f32x2_applies(op_call, sctx, plan):
+    from ...common import sm_version_ok
+
+    if dtype_name(plan.dst.buffer.dtype) != "float32":
+        return False, "dst dtype not f32"
+    if not sm_version_ok(op_call, sctx, min_version=100)[0]:
+        return False, "sm version < 100"
+    if len(plan.srcs) != 3:
+        return False, "fma requires 3 srcs"
+    a, b, c = plan.srcs
+    if a.is_scalar:
+        return False, "fma 'a' must be a buffer (no scalar-a packed FMA)"
+    if dtype_name(a.buf_region.buffer.dtype) != "float32":
+        return False, "src a dtype not f32"
+    if a.index_fn is not None:
+        return False, "broadcasting src a not supported"
+    for s in (b, c):
+        if s.is_scalar:
+            if scalar_dtype(s.scalar) != "float32":
+                return False, "scalar b/c dtype not f32"
+        else:
+            if dtype_name(s.buf_region.buffer.dtype) != "float32":
+                return False, "buffer b/c dtype not f32"
+            if s.index_fn is not None:
+                return False, "broadcasting src b/c not supported"
+    return True, None
+
+
+def _emit_fma_f32x2(dst_buf, dst_lane_indices, src_args, extras) -> Expr:
+    a_arg, b_arg, c_arg = src_args
+    rm = extras.get("rounding_mode", "rz")
+    return T.ptx.fma_f32x2(
+        T.address_of(dst_buf[tuple(dst_lane_indices[0])]),
+        T.cuda.make_float2(_lane(a_arg, 0), _lane(a_arg, 1)),
+        T.cuda.make_float2(_lane(b_arg, 0), _lane(b_arg, 1)),
+        T.cuda.make_float2(_lane(c_arg, 0), _lane(c_arg, 1)),
+        rounding=rm,
+        ftz=True,
+    )
+
+
+FMA_F32X2_IMPL = VecImpl(
+    vec_len=2,
+    applies=_fma_f32x2_applies,
+    emit=_emit_fma_f32x2,
+)

python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.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.
+"""Execution scope utilities for CUDA op dispatches."""
+
+from collections.abc import Callable
+
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def macro_or_prim_func(macro: Callable, need_macro: bool = False) -> Callable:
+    """Wrap a macro in a ``prim_func`` unless the caller explicitly wants the macro."""
+    if need_macro:
+        return macro
+
+    @T.prim_func(check_well_formed=False)
+    def func():
+        macro()
+
+    return func
+
+
+def thread_selector(sctx: DispatchContext, inner_impl, macro: bool = False) -> Callable:
+    """Narrow execution to a single, deterministic thread within ``sctx.exec_scope``.
+
+    The elected thread is stable across invocations so that synchronization
+    primitives (for example PTX ``elect_sync``) behave correctly.
+
+    Parameters
+    ----------
+    sctx : DispatchContext
+        The dispatch context. Only ``sctx.scope_kind`` is consulted; the
+        caller is responsible for having narrowed into the desired scope via an
+        ``if`` guard with a canonical thread-filter predicate before reaching here.
+    inner_impl : T.inline
+        The body to execute inside the selected thread.
+    macro : bool
+        If True, return the macro directly; otherwise wrap it in a ``prim_func``.
+    """
+    assert not isinstance(inner_impl, PrimFunc), "inner_impl must be a macro, not a PrimFunc"
+    name = sctx.scope_kind
+    if name == "thread":
+        return macro_or_prim_func(inner_impl, need_macro=macro)
+    if name == "cta":
+
+        @T.inline()
+        def impl():
+            T.lane_id([32])
+            if T.ptx.elect_sync():
+                inner_impl()
+
+        return macro_or_prim_func(impl, need_macro=macro)
+    if name == "warp":
+
+        @T.inline()
+        def impl():
+            T.lane_id([32])
+            if T.ptx.elect_sync():
+                inner_impl()
+
+        return macro_or_prim_func(impl, need_macro=macro)
+    if name == "warpgroup":
+
+        @T.inline()
+        def impl():
+            warp_id = T.warp_id_in_wg([4])
+            T.lane_id([32])
+            if warp_id == 0:
+                if T.ptx.elect_sync():
+                    inner_impl()
+
+        return macro_or_prim_func(impl, need_macro=macro)
+    raise ValueError(f"thread_selector: unsupported exec_scope {name!r}")
+
+
+def single_thread(op_call: TilePrimitiveCall, sctx: DispatchContext) -> bool:
+    """Predicate for dispatchers that require a single-thread execution scope."""
+    del op_call
+    return sctx.is_thread
+
+
+def exec_scope_ok(
+    op_call: TilePrimitiveCall, sctx: DispatchContext, expected_scopes: list[str]
+) -> tuple[bool, str | None]:
+    """Predicate helper: check that ``sctx.scope_kind`` is in *expected_scopes*."""
+    del op_call
+    ok = sctx.scope_kind in expected_scopes
+    return ok, None if ok else f"unsupported exec_scope {sctx.scope_kind}"

python/tvm/backend/cuda/operator/tile_primitive/gemm/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+
+"""CUDA synchronous ``gemm`` lowerings (warp-level ``mma.sync`` tensor core).
+
+Importing this package registers every synchronous CUDA ``gemm`` dispatch
+candidate as a side effect (each submodule calls ``register_dispatch`` at
+import time).  It is the synchronous counterpart to ``gemm_async``.
+"""
+
+from .mma_m16n8k_ import *

python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py

@@ -0,0 +1,595 @@
+# 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.
+
+"""Warp-level ``mma.sync`` GEMM lowering for the synchronous ``gemm`` op on CUDA."""
+
+from dataclasses import dataclass
+
+from tvm.arith.analyzer import Analyzer
+from tvm.script import tirx as T
+from tvm.tirx import PrimFunc
+from tvm.tirx.layout import TileLayout
+from tvm.tirx.operator.tile_primitive import (
+    DispatchContext,
+    fail,
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+@dataclass(frozen=True)
+class MmaInst:
+    """One concrete mma.sync instruction we can emit. A given shape appears once
+    per dtype signature. Adding an instruction is just adding an entry to
+    MMA_INSTRUCTIONS; the feasibility checks below stay generic. The lane->coord
+    mapping is the fixed m16n8 family structure (hardcoded in the match); the
+    only per-instruction fragment detail is ``k_pack`` (it is not always
+    32/dtype_bits, so it is stored explicitly rather than derived)."""
+
+    name: str  # distinguishing tag, e.g. "m16n8k16.bf16"
+    m: int
+    n: int
+    k: int
+    dtype: tuple[str, str, str, str]  # the single (A, B, C, D) dtype signature
+    k_pack: int  # contiguous-along-K elements packed into one register for A and B
+
+
+# One entry per concrete instruction. The m16n8 family fixes M/N at 16/8; K is
+# 16 or 8. f16/bf16 inputs with f32 accumulation, packing 2 bf16/f16 per b32
+# along K (f16 accumulation, int8, fp8, etc. are added as more entries).
+_BF16_F32 = ("bfloat16", "bfloat16", "float32", "float32")
+_F16_F32 = ("float16", "float16", "float32", "float32")
+MMA_INSTRUCTIONS = (
+    MmaInst("m16n8k16.bf16", 16, 8, 16, _BF16_F32, k_pack=2),
+    MmaInst("m16n8k16.f16", 16, 8, 16, _F16_F32, k_pack=2),
+    MmaInst("m16n8k8.bf16", 16, 8, 8, _BF16_F32, k_pack=2),
+    MmaInst("m16n8k8.f16", 16, 8, 8, _F16_F32, k_pack=2),
+)
+
+
+def _split(grouped, seps):
+    """Split a grouped layout into one shard-only sub-layout per group, plus the
+    layout's offset (returned separately rather than distributed into the subs)."""
+    subs = [
+        TileLayout.from_iters(list(grouped.shard[seps[g] : seps[g + 1]]), [], {})
+        for g in range(len(seps) - 1)
+    ]
+    return subs, dict(grouped.offset)
+
+
+def _combine(layouts, offset):
+    """Concatenate sub-layouts' shards and apply the separately tracked offset."""
+    shard = [it for lay in layouts for it in lay.shard]
+    return TileLayout.from_iters(shard, [], offset)
+
+
+def _canon_perm(iters):
+    """Permutation putting thread iters left, memory iters right; stride-desc within each."""
+    thr = sorted(
+        (i for i, it in enumerate(iters) if it.axis.is_thread()),
+        key=lambda i: -int(iters[i].stride),
+    )
+    mem = sorted(
+        (i for i, it in enumerate(iters) if not it.axis.is_thread()),
+        key=lambda i: -int(iters[i].stride),
+    )
+    return thr + mem
+
+
+def _canon(layout):
+    """Reorder an anchor sub-layout into canonical (thread-left/memory-right) order."""
+    return layout.permute_dims(_canon_perm(list(layout.shard)))
+
+
+def _align(anchor, follower, name):
+    """Regroup follower by anchor's iter extents, then permute its groups to follow
+    the anchor's canonical order (follower keeps its own strides)."""
+    extents = [int(it.extent) for it in anchor.shard]
+    perm = _canon_perm(list(anchor.shard))
+    try:
+        grp, seps = follower.group(extents)
+    except Exception as e:  # follower dim layout incompatible with anchor decomposition
+        fail(f"gemm mma: {name} not alignable to anchor extents {extents}: {e}")
+    return grp.permute_by_groups(seps, perm)
+
+
+def _region_totals(layout):
+    """(product of thread-axis iter extents, product of memory-axis iter extents).
+
+    Computed once on each anchor to fix that dim's thread/memory region lengths;
+    followers reuse the anchor's split rather than re-deriving it from their own
+    iters (which can misclassify -- e.g. B's N, whose register slot is actually a
+    lane iter, would otherwise report memory length 1)."""
+    thr, mem = 1, 1
+    for it in layout.shard:
+        if it.axis.is_thread():
+            thr *= int(it.extent)
+        else:
+            mem *= int(it.extent)
+    return thr, mem
+
+
+def _frag_group(layout, lane, mem, thread_total, mem_total):
+    """Group one logical-dim sub-layout into the fragment shape, then optionally
+    verify it.
+
+    ``lane`` and ``mem`` are lists of ``(extent, stride, want_thread)``;
+    ``thread_total`` and ``mem_total`` are this dim's thread/memory region lengths
+    (from its anchor via _region_totals). The group shape is
+    ``[thread_total // prod(lane), *lane, mem_total // prod(mem), *mem]``: the lane
+    extents are carved off the thread region and the mem extents off the memory
+    region (innermost last, e.g. ``[(reg, ...)]`` for an accumulator dim or
+    ``[(kHi, ...), (k_pack, ...)]`` for A/B's K). The input must already be in
+    canonical (thread-left/memory-right) order.
+
+    Every carved group is verified: it must be a single iter, its axis must be a
+    thread axis when ``want_thread`` else a memory axis (only is_thread is checked,
+    since scope varies the exact thread axis; e.g. B's N register slot is actually
+    a lane, so want_thread=True there), and a non-None ``stride`` pins that iter's
+    stride. Raises on a tiling or verification failure, so a non-matching caller
+    layout is declined via the caller's try/except.
+    """
+    lane_ext = [e for e, _, _ in lane]
+    mem_ext = [e for e, _, _ in mem]
+    lane_prod, mem_prod = 1, 1
+    for e in lane_ext:
+        lane_prod *= e
+    for e in mem_ext:
+        mem_prod *= e
+    grouped, seps = layout.group(
+        [thread_total // lane_prod, *lane_ext, mem_total // mem_prod, *mem_ext]
+    )
+    # group order: [thread_rest, *lane (from idx 1), mem_rest, *mem (after)].
+    specs = [(1 + j, s, t) for j, (_, s, t) in enumerate(lane)]
+    specs += [(2 + len(lane) + j, s, t) for j, (_, s, t) in enumerate(mem)]
+    for idx, stride, want_thread in specs:
+        grp = grouped.shard[seps[idx] : seps[idx + 1]]
+        if len(grp) != 1:
+            raise ValueError(f"frag group {idx} is not a single iter")
+        if int(grp[0].extent) == 1:
+            # An extent-1 group iterates nothing, so its axis/stride is
+            # meaningless and gets dropped downstream (cf. _same_iters /
+            # _reg_layout). This is the kHi == 1 case of m16n8k8: there is a
+            # single high-K register group, which .group() may materialize as a
+            # degenerate split of the (thread) lane axis.
+            continue
+        if grp[0].axis.is_thread() != want_thread:
+            raise ValueError(f"frag group {idx} thread/memory axis mismatch")
+        if stride is not None and int(grp[0].stride) != stride:
+            raise ValueError(f"frag group {idx} stride {int(grp[0].stride)} != {stride}")
+    return grouped, seps
+
+
+def _grp(grouped, seps, i):
+    """The iters of group ``i`` of a grouped layout: ``shard[seps[i]:seps[i+1]]``."""
+    return grouped.shard[seps[i] : seps[i + 1]]
+
+
+def _ext(iters):
+    """Product of the extents of an iter list."""
+    p = 1
+    for it in iters:
+        p *= int(it.extent)
+    return p
+
+
+def _reg_layout(groups, offset):
+    """Per-thread register layout from per-logical-dim iter groups (dropping
+    thread-axis offset terms), plus the matching local-view shape (each dim is
+    the product of that group's extents).
+
+    Extent-1 iters are dropped from the layout: they iterate nothing (offset
+    always 0, so harmless to the mapping) but would otherwise pin a degenerate
+    axis -- e.g. B's N has no real register, so its "register" slot is a single
+    extent-1 lane iter that must not make the register buffer thread-axis."""
+    iters = [it for g in groups for it in g if int(it.extent) != 1]
+    layout = TileLayout.from_iters(
+        iters, [], {ax: v for ax, v in offset.items() if not ax.is_thread()}
+    )
+    return layout, [_ext(g) for g in groups]
+
+
+def _same_iters(a, b):
+    """True iff iter lists ``a`` and ``b`` match elementwise on (extent, stride,
+    axis), ignoring extent-1 iters (they iterate nothing, so their stride/axis is
+    meaningless -- e.g. the degenerate thread-rest left by a group shape's '1')."""
+    a = [it for it in a if int(it.extent) != 1]
+    b = [it for it in b if int(it.extent) != 1]
+    if len(a) != len(b):
+        return False
+    return all(
+        int(x.extent) == int(y.extent)
+        and int(x.stride) == int(y.stride)
+        and x.axis.name == y.axis.name
+        for x, y in zip(a, b, strict=True)
+    )
+
+
+def _full_active_lanes(op: TilePrimitiveCall, sctx: DispatchContext):
+    """The active thread set (sctx.intra) must be complete and un-narrowed.
+
+    mma.sync.aligned is collective over every active thread; an enclosing if
+    that narrows any intra axis makes the .aligned instruction undefined. So
+    each intra axis must be at offset 0 with its full extent: laneid=32,
+    wid_in_wg=4 (warpgroup), and warpid=warps-per-CTA (cta) from the launch
+    config. Any other axis (e.g. cta_id at cluster scope) is not supported.
+    """
+    full = {"laneid": 32, "wid_in_wg": 4}
+    if "warpid" in sctx.intra:
+        tx = sctx.launch_params.get("threadIdx.x")
+        if tx is None:
+            return False, "cta scope needs threadIdx.x in launch_params"
+        try:
+            full["warpid"] = int(tx.dom.extent) // 32
+        except (TypeError, ValueError):
+            return False, f"non-static threadIdx.x extent {tx.dom.extent}"
+    for axis, rng in sctx.intra.items():
+        if axis not in full:
+            return False, f"unsupported active-set axis {axis!r}"
+        extent, offset = int(rng[0]), int(rng[1])
+        if extent != full[axis] or offset != 0:
+            return False, (
+                f"active {axis} is [{offset}, {offset + extent}), need full [0, {full[axis]})"
+            )
+    return True
+
+
+def _no_replica(op: TilePrimitiveCall, sctx: DispatchContext):
+    """All operand layouts must have no replica (no broadcast/duplicated axes)."""
+    for region, name in zip(op.args[:4], ("D", "A", "B", "C")):
+        if region.buffer.layout.replica:
+            return False, f"{name} layout has replica {region.buffer.layout.replica}"
+    return True
+
+
+@register_dispatch(
+    "gemm",
+    "cuda",
+    variant="mma.m16n8k*",
+    priority=10,
+    when=[
+        predicate("full_active_lanes", _full_active_lanes),
+        predicate("no_replica", _no_replica),
+    ],
+)
+def gemm_cuda_mma_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    """``gemm`` -> warp-level ``mma.sync`` of the m16n8k* family.
+
+    This is the ``"mma.m16n8k*"`` variant. It targets the m16n8k* tensor-core
+    instructions -- currently m16n8k16 / m16n8k8 with bf16/f16 inputs and f32
+    accumulation (see MMA_INSTRUCTIONS); other K (and other shapes/dtypes) are
+    added as more entries. Pure-register path: A/B fragments and C/D
+    accumulators all live in registers.
+    """
+    # gemm op args: D = alpha * A @ B + beta * C
+    # D (args[0]) is the output; C (args[3]) is the beta-accumulator input.
+    D_region, A_region, B_region, C_region, transpose_A, transpose_B, alpha, beta = op.args
+    D, A, B, C = D_region.buffer, A_region.buffer, B_region.buffer, C_region.buffer
+
+    # Pure-register mma path: A/B fragments and C/D accumulators all live in
+    # registers ("local"). The caller is responsible for staging A/B into
+    # registers (e.g. via ldmatrix) beforehand.
+    for buf, name in ((D, "D"), (A, "A"), (B, "B"), (C, "C")):
+        if buf.scope() != "local":
+            fail(f"gemm mma requires {name} in register (local) scope, got {buf.scope()}")
+
+    # transpose_A/transpose_B only describe the input's logical orientation; we
+    # normalize to the standard form A=[M,K], B=[K,N] (D/C are always [M,N]).
+    #   transpose_A: False -> buffer is [M,K], True -> [K,M]
+    #   transpose_B: False -> buffer is [K,N], True -> [N,K]
+    # The .row.col K-major requirement is not enforced here -- it is checked
+    # later by the per-instruction fragment match against the real layout.
+    analyzer = Analyzer()
+
+    # mma.sync computes D = A·B + C natively (no scalar scaling), so we support
+    # only alpha=1 and beta in {0, 1}; beta selects whether C is the accumulator
+    # (1 -> c_ptr=C, 0 -> c_ptr=0). General alpha/beta is declined.
+    def _const_scalar(expr):
+        s = analyzer.simplify(expr)
+        try:
+            return float(s.value)
+        except (AttributeError, TypeError, ValueError):
+            return None
+
+    if _const_scalar(alpha) != 1.0:
+        fail(f"gemm mma supports only alpha=1, got alpha={alpha}")
+    if _const_scalar(beta) not in (0.0, 1.0):
+        fail(f"gemm mma supports only beta in {{0, 1}}, got beta={beta}")
+
+    def _mat_extents(region, name):
+        ext = [r.extent for r in region.region if not analyzer.can_prove_equal(r.extent, 1)]
+        if len(ext) != 2:
+            fail(f"gemm mma expects 2D {name}, got non-unit extents {ext}")
+        return ext
+
+    A_ext = _mat_extents(A_region, "A")
+    B_ext = _mat_extents(B_region, "B")
+    D_M, D_N = _mat_extents(D_region, "D")
+    C_M, C_N = _mat_extents(C_region, "C")
+    M, K = (A_ext[1], A_ext[0]) if transpose_A else (A_ext[0], A_ext[1])
+    B_K, N = (B_ext[1], B_ext[0]) if transpose_B else (B_ext[0], B_ext[1])
+    assert analyzer.can_prove_equal(B_K, K), f"gemm mma: A K={K} != B K={B_K}"
+    assert analyzer.can_prove_equal(D_M, M) and analyzer.can_prove_equal(D_N, N), (
+        f"gemm mma: D dims ({D_M}, {D_N}) != (M={M}, N={N})"
+    )
+    assert analyzer.can_prove_equal(C_M, M) and analyzer.can_prove_equal(C_N, N), (
+        f"gemm mma: C dims ({C_M}, {C_N}) != (M={M}, N={N})"
+    )
+
+    # Tiling into instructions needs static extents.
+    def _const(expr, name):
+        try:
+            return int(analyzer.simplify(expr))
+        except (TypeError, ValueError):
+            fail(f"gemm mma needs static {name} extent, got {expr}")
+
+    M, N, K = _const(M, "M"), _const(N, "N"), _const(K, "K")
+
+    # Slice each operand's layout to its region, then group it into its 2D
+    # buffer-order shape; the split below maps those groups to standard
+    # (M,K)/(K,N). group() raises if the layout can't be tiled that way, so a
+    # caller layout that doesn't match the operand shape is declined cleanly.
+    def _slice_group(buf, region, shape2d, name):
+        # slice() itself groups internally, so both slice and group can raise the
+        # ICHECK when the layout can't be tiled as shape2d -- guard both.
+        canon = None
+        try:
+            sliced = buf.layout.slice(buf.shape, region.region)
+            if sliced is not None:
+                canon = sliced.canonicalize()
+        except Exception as e:  # ICHECK failure -> layout not tileable as shape2d
+            fail(f"gemm mma: {name} layout not tileable as {tuple(shape2d)}: {e}")
+        if canon is None:
+            fail(f"gemm mma: cannot slice {name} layout to its region")
+        # All thread iters must share one thread axis (e.g. all laneid). _frag_group
+        # only checks is_thread (not the exact axis, since scope varies), so a layout
+        # mixing two thread axes (e.g. laneid + wid_in_wg) would carve ambiguously --
+        # decline it here, like ldmatrix validating its thread structure.
+        thread_axes = {it.axis.name for it in canon.shard if it.axis.is_thread()}
+        if len(thread_axes) > 1:
+            fail(f"gemm mma: {name} has >1 thread axis {sorted(thread_axes)}, only one supported")
+        try:
+            return canon.group(list(shape2d))
+        except Exception as e:  # ICHECK failure -> layout not tileable as shape2d
+            fail(f"gemm mma: {name} layout not tileable as {tuple(shape2d)}: {e}")
+
+    A_grouped, A_seps = _slice_group(A, A_region, (K, M) if transpose_A else (M, K), "A")
+    B_grouped, B_seps = _slice_group(B, B_region, (N, K) if transpose_B else (K, N), "B")
+    C_grouped, C_seps = _slice_group(C, C_region, (M, N), "C")
+    D_grouped, D_seps = _slice_group(D, D_region, (M, N), "D")
+
+    # Split each operand into per-logical-dim sub-layouts (+ its offset), mapping
+    # the buffer-order subs to standard (M,K)/(K,N) per the transpose flags.
+    (DM, DN), D_off = _split(D_grouped, D_seps)
+    (CM, CN), C_off = _split(C_grouped, C_seps)
+    A_subs, A_off = _split(A_grouped, A_seps)
+    B_subs, B_off = _split(B_grouped, B_seps)
+    AM, AK = (A_subs[1], A_subs[0]) if transpose_A else (A_subs[0], A_subs[1])
+    BK, BN = (B_subs[1], B_subs[0]) if transpose_B else (B_subs[0], B_subs[1])
+
+    # Anchor-align so every operand decomposes each shared logical dim the same
+    # way: M anchor = DM -> AM, CM ; N anchor = DN -> BN, CN ; K anchor = AK -> BK.
+    # _align uses each anchor's raw (pre-canon) per-iter extents to group the
+    # follower and reorders the follower's groups into the anchor's canonical
+    # order, so the followers come out canonical. Canon the anchors themselves
+    # afterwards. Every sub-layout is then in canonical (thread-left/memory-right)
+    # order before the loop, so _frag_group groups directly without re-canon.
+    AM = _align(DM, AM, "A.M")
+    CM = _align(DM, CM, "C.M")
+    BN = _align(DN, BN, "B.N")
+    CN = _align(DN, CN, "C.N")
+    BK = _align(AK, BK, "B.K")
+    DM, DN, AK = _canon(DM), _canon(DN), _canon(AK)
+
+    # Each dim's thread/memory region lengths, fixed once from its anchor (3
+    # anchors x 2 parts = 6 lengths). Every operand of that dim reuses them in
+    # _frag_group, so a follower whose register slot is actually a lane (B's N)
+    # still gets the anchor's memory length instead of its own (mis)classified one.
+    m_thr, m_mem = _region_totals(DM)
+    n_thr, n_mem = _region_totals(DN)
+    k_thr, k_mem = _region_totals(AK)
+
+    # Per-instruction selection: try each candidate in order and use the first
+    # whose shape / dtype / (later) fragment layout all fit. A failing check just
+    # moves on to the next instruction; if none fit, decline.
+    sig = (str(A.dtype), str(B.dtype), str(C.dtype), str(D.dtype))
+    for inst in MMA_INSTRUCTIONS:
+        assert inst.m % 8 == 0 and inst.n % 8 == 0 and inst.k % 8 == 0, (
+            f"mma instruction {inst.name} m/n/k must be multiples of 8"
+        )
+        if M % inst.m or N % inst.n or K % inst.k:
+            continue
+        if sig != inst.dtype:
+            continue
+        # Group every operand into this instruction's fragment shape. The m16n8
+        # lane split is g (8 lanes) along M and t (4 lanes) along N/K:
+        #   C/D accumulator: M = g + 8*rM (inst.m//8 regs),  N = 2*t + rN (inst.n//4 regs)
+        #   A multiplicand:  M as C/D's M,  K = 2*t + p + 8*kHi
+        #   B multiplicand:  K as A's K,    N = g (lane 8, no reg: B has no M so N
+        #                                        reuses the 8-lane g group)
+        # so K's memory tail is [kHi, k_pack] with k_pack the innermost (stride-1)
+        # contiguous pack and kHi = inst.k // (4 * k_pack) high-K register groups.
+        # _frag_group raises if the caller layout can't be tiled or (when any
+        # stride is given) fails the fragment checks -> move on to the next
+        # instruction. lane/mem are [(extent, stride), ...]: the lane stride pins
+        # the laneid stride (g=4, t=1), a mem stride pins a register iter (rN /
+        # k_pack = 1; rM / kHi free = None). C shares D's accumulator fragment and
+        # B's K shares A's K. B's N is the pure 8-lane g group, but aligned to the
+        # accumulator's N (lane 4 + reg 2) it splits into lane g_hi (4, laneid
+        # stride 8) and a "register" g_lo (2, laneid stride 4) -- a lane iter in the
+        # register slot (B's N has no real register). Region lengths come from the
+        # per-dim anchor (m/n/k _thr,_mem), so this split still tiles correctly.
+        # _frag_group(layout, lane, mem, thread_total, mem_total); each carve is
+        # (extent, stride, want_thread): lanes are thread, registers memory, except
+        # B.N's register slot (g_lo) which is itself a lane (want_thread=True).
+        kHi = inst.k // (4 * inst.k_pack)
+        try:
+            DM_g, DM_seps = _frag_group(
+                DM, [(8, 4, True)], [(inst.m // 8, None, False)], m_thr, m_mem
+            )
+            DN_g, DN_seps = _frag_group(DN, [(4, 1, True)], [(inst.n // 4, 1, False)], n_thr, n_mem)
+            CM_g, CM_seps = _frag_group(
+                CM, [(8, 4, True)], [(inst.m // 8, None, False)], m_thr, m_mem
+            )
+            CN_g, CN_seps = _frag_group(CN, [(4, 1, True)], [(inst.n // 4, 1, False)], n_thr, n_mem)
+            AM_g, AM_seps = _frag_group(
+                AM, [(8, 4, True)], [(inst.m // 8, None, False)], m_thr, m_mem
+            )
+            AK_g, AK_seps = _frag_group(
+                AK, [(4, 1, True)], [(kHi, None, False), (inst.k_pack, 1, False)], k_thr, k_mem
+            )
+            BK_g, BK_seps = _frag_group(
+                BK, [(4, 1, True)], [(kHi, None, False), (inst.k_pack, 1, False)], k_thr, k_mem
+            )
+            BN_g, BN_seps = _frag_group(BN, [(4, 8, True)], [(2, 4, True)], n_thr, n_mem)
+        except Exception:
+            continue
+        # M.to (M's warp tiling, group 0) must match across D, A, C so the same
+        # logical M-block lands on the same warp in all three operands.
+        m_to = _grp(DM_g, DM_seps, 0)
+        if not (
+            _same_iters(m_to, _grp(AM_g, AM_seps, 0)) and _same_iters(m_to, _grp(CM_g, CM_seps, 0))
+        ):
+            continue
+        # N.to (N's warp tiling, group 0) must match across D, B, C so the same
+        # logical N-block lands on the same warp in all three operands.
+        n_to = _grp(DN_g, DN_seps, 0)
+        if not (
+            _same_iters(n_to, _grp(BN_g, BN_seps, 0)) and _same_iters(n_to, _grp(CN_g, CN_seps, 0))
+        ):
+            continue
+        # K.to (K's warp tiling, group 0) must match across A, B so the same
+        # logical K-block lands on the same warp in both operands.
+        if not _same_iters(_grp(AK_g, AK_seps, 0), _grp(BK_g, BK_seps, 0)):
+            continue
+        break
+    else:
+        fail(f"no mma instruction fits M={M}, N={N}, K={K}, dtypes={sig}")
+
+    # Per-operand register layout + matching local-view shape, grouped per logical
+    # dim (offset drops thread-axis terms). Iter order = shape dim order:
+    #   D/C -> [M.mo, N.mo, rM, rN]   A -> [M.mo, K.mo, rM, kHi, k_pack]
+    #   B   -> [K.mo, N.mo, kHi, k_pack]   (k_pack innermost / contiguous)
+    D_reg, d_shape = _reg_layout(
+        [
+            _grp(DM_g, DM_seps, 2),
+            _grp(DN_g, DN_seps, 2),
+            _grp(DM_g, DM_seps, 3),
+            _grp(DN_g, DN_seps, 3),
+        ],
+        D_off,
+    )
+    C_reg, c_shape = _reg_layout(
+        [
+            _grp(CM_g, CM_seps, 2),
+            _grp(CN_g, CN_seps, 2),
+            _grp(CM_g, CM_seps, 3),
+            _grp(CN_g, CN_seps, 3),
+        ],
+        C_off,
+    )
+    A_reg, a_shape = _reg_layout(
+        [
+            _grp(AM_g, AM_seps, 2),
+            _grp(AK_g, AK_seps, 2),
+            _grp(AM_g, AM_seps, 3),
+            _grp(AK_g, AK_seps, 3),
+            _grp(AK_g, AK_seps, 4),
+        ],
+        A_off,
+    )
+    B_reg, b_shape = _reg_layout(
+        [
+            _grp(BK_g, BK_seps, 2),
+            _grp(BN_g, BN_seps, 2),
+            _grp(BK_g, BK_seps, 3),
+            _grp(BK_g, BK_seps, 4),
+        ],
+        B_off,
+    )
+
+    # Emit one mma per (m, n) output tile, accumulating over K. The tile / init /
+    # K loops use T.unroll: the UnrollLoop pass fully expands them in TIR (their
+    # bounds are compile-time constants), so the local-buffer indices resolve to
+    # static register slots -- mma register operands must be constant.
+    #
+    # mma is d = a·b + c. D's accumulator is initialized once per output tile --
+    # copying C when beta==1, clearing to 0 when beta==0 -- then every K step
+    # accumulates in place with c = d, giving a single uniform mma form.
+    M_tiles, N_tiles, K_tiles = d_shape[0], d_shape[1], a_shape[1]
+    shape_str = f"m{inst.m}n{inst.n}k{inst.k}"
+    a_type, b_type, c_type, d_type = inst.dtype
+    use_c = _const_scalar(beta) == 1.0
+
+    # Per-register counts in the fixed PTX enumeration order (derived from the
+    # instruction, NOT hardcoded, so m16n8k8 with kHi==1 also works):
+    #   D/C accumulator: rM = inst.m // 8 regs along M, rN = inst.n // 4 along N
+    #                    c_id = 2 * rM + rN  (4 f32 for m16n8k16, also 4 for k8)
+    #   A multiplicand:  rM = inst.m // 8, kHi = inst.k // (4 * inst.k_pack)
+    #                    b32 = rM + 2 * kHi  (4 b32 for k16, 2 b32 for k8)
+    #   B multiplicand:  kHi = inst.k // (4 * inst.k_pack)
+    #                    b32 = kHi          (2 b32 for k16, 1 b32 for k8)
+    n_rM = inst.m // 8
+    n_rN = inst.n // 4
+    n_kHi = inst.k // (4 * inst.k_pack)
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        d_local = D.local(*d_shape, layout=D_reg)
+        c_local = C.local(*c_shape, layout=C_reg)
+        a_local = A.local(*a_shape, layout=A_reg)
+        b_local = B.local(*b_shape, layout=B_reg)
+        for m in T.unroll(M_tiles):
+            for n in T.unroll(N_tiles):
+                # Initialize D[m, n]: copy C (beta==1) or clear to 0 (beta==0).
+                for rM in T.unroll(n_rM):
+                    for rN in T.unroll(n_rN):
+                        if use_c:
+                            d_local[m, n, rM, rN] = c_local[m, n, rM, rN]
+                        else:
+                            d_local[m, n, rM, rN] = T.float32(0)
+                # Accumulate over K in place: d = a·b + d.
+                for k in T.unroll(K_tiles):
+                    # D: 4 f32 in PTX order c_id = 2*rM + rN.
+                    d_ptrs = [
+                        d_local.ptr_to([m, n, rM, rN]) for rM in range(n_rM) for rN in range(n_rN)
+                    ]
+                    # A: b32 regs in PTX order b32 = rM + 2*kHi (kHi outer, rM inner).
+                    a_ptrs = [
+                        a_local.ptr_to([m, k, rM, kHi, 0])
+                        for kHi in range(n_kHi)
+                        for rM in range(n_rM)
+                    ]
+                    # B: b32 regs in PTX order b32 = kHi.
+                    b_ptrs = [b_local.ptr_to([k, n, kHi, 0]) for kHi in range(n_kHi)]
+                    # Accumulate in place into D's own regs: c = d.
+                    T.ptx.mma(
+                        shape_str,
+                        "row",
+                        "col",
+                        d_type,
+                        a_type,
+                        b_type,
+                        c_type,
+                        d_ptrs,
+                        a_ptrs,
+                        b_ptrs,
+                        d_ptrs,
+                    )
+
+    return impl

python/tvm/backend/cuda/operator/tile_primitive/gemm_async/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .tcgen05 import *

python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py

@@ -0,0 +1,62 @@
+# 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.
+
+"""GEMM-related utilities for CUDA op dispatches."""
+
+from tvm.arith.analyzer import Analyzer
+from tvm.tirx import Buffer
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def validate_gemm_op(op_call: TilePrimitiveCall, sctx: DispatchContext) -> bool:
+    """Sanity check for gemm op"""
+    C_buffer_region, A_buffer_region, B_buffer_region = op_call.args[:3]
+    C: Buffer = C_buffer_region.buffer
+    A: Buffer = A_buffer_region.buffer
+    B: Buffer = B_buffer_region.buffer
+    if not (C.layout and A.layout and B.layout and A.dtype == B.dtype):
+        return False
+    # Extract regions and validate dimensions
+    analyzer = Analyzer()
+    C_region, A_region, B_region = (
+        C_buffer_region.region,
+        A_buffer_region.region,
+        B_buffer_region.region,
+    )
+    # Extract extents and validate non-unit dimensions match
+    transA, transB = op_call.args[3:5]
+    C_extent_ = [r.extent for r in C_region if r.extent != 1]
+    A_extent_ = [r.extent for r in A_region if r.extent != 1]
+    B_extent_ = [r.extent for r in B_region if r.extent != 1]
+    assert len(C_extent_) == len(A_extent_) == len(B_extent_) == 2, (
+        "Only 2D C, A, B are supported for gemm"
+    )
+    if transA:
+        A_extent_ = [A_extent_[1], A_extent_[0]]
+    if transB:
+        B_extent_ = [B_extent_[1], B_extent_[0]]
+    # C: MxN, A: MxK, B: NxK
+    if not all(
+        [
+            analyzer.can_prove_equal(C_extent_[0], A_extent_[0]),
+            analyzer.can_prove_equal(C_extent_[1], B_extent_[0]),
+            analyzer.can_prove_equal(A_extent_[1], B_extent_[1]),
+        ]
+    ):
+        return False
+    return True

python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py

@@ -0,0 +1,326 @@
+# 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.
+
+"""Layout analysis utilities for local-memory op dispatches.
+
+Provides functions for analyzing TileLayout thread/local partitions,
+computing local region info, layout signature comparison, and thread
+variable resolution. Used by cast.py, unary.py, and binary.py.
+"""
+
+import functools
+import operator
+from collections import defaultdict
+
+from tvm.arith import Analyzer
+from tvm.tirx.layout import TileLayout
+
+
+def get_sublayout_from_region(layout, buffer_shape, region_st, region_extent):
+    """Get sublayout by slicing the layout with the buffer region.
+
+    Args:
+        layout: The buffer's TileLayout.
+        buffer_shape: The buffer's shape.
+        region_st: Region start indices.
+        region_extent: Region extents.
+
+    Returns:
+        Sublayout if slicing succeeds, otherwise the original layout.
+    """
+    if not layout:
+        return layout
+    region = [(region_st[i], region_st[i] + region_extent[i]) for i in range(len(region_st))]
+    sliced = layout.slice(list(buffer_shape), region)
+    return sliced if sliced is not None else layout
+
+
+def get_layout_thread_local_partition(layout):
+    """Extract thread and local dimension info from layout.
+
+    Returns:
+        tuple | None: On success, (thread_groups, local_dim_indices, local_extents).
+            - thread_groups: dict {axis: (dim_indices, extents)} for each thread axis
+            - local_dim_indices: list of dimension indices for local (memory) axes
+            - local_extents: list of extents for local dimensions
+            Returns None if layout is not supported.
+
+    Validates:
+        - No stride==0 on thread dims (broadcast/overlap = cross-thread semantics)
+        - Local dims may have arbitrary strides (alignment uses actual layout strides)
+        - No thread axes in replica
+
+    Example:
+        Layout (2, 8, 4, 2):(2@warpid, 4@laneid, 1@laneid, 1@m) returns:
+        - thread_groups = {warpid: ([0], [2]), laneid: ([1, 2], [8, 4])}
+        - local_dim_indices = [3], local_extents = [2]
+    """
+    if not isinstance(layout, TileLayout):
+        return None
+
+    shard = getattr(layout, "shard", None)
+    if not shard:
+        return None
+
+    # Partition dimensions into thread and local (memory) axes
+    thread_dim_indices = [i for i, it in enumerate(shard) if it.axis.is_thread()]
+    local_dim_indices = [i for i, it in enumerate(shard) if not it.axis.is_thread()]
+
+    if not thread_dim_indices or not local_dim_indices:
+        return None
+
+    analyzer = Analyzer()
+    for idx in thread_dim_indices:
+        if analyzer.can_prove_equal(shard[idx].stride, 0):
+            return None
+
+    # Replica must not contain thread axes
+    replica = getattr(layout, "replica", None)
+    if replica and any(it.axis.is_thread() for it in replica):
+        return None
+
+    # Group thread dimensions by axis
+    thread_groups_dict = defaultdict(list)
+    for idx in thread_dim_indices:
+        thread_groups_dict[shard[idx].axis].append(idx)
+
+    thread_groups = {}
+
+    for axis, dim_indices in thread_groups_dict.items():
+        dim_indices = sorted(dim_indices)
+        extents = [shard[i].extent for i in dim_indices]
+        thread_groups[axis] = (dim_indices, extents)
+
+    local_extents = [shard[i].extent for i in local_dim_indices]
+    return (thread_groups, local_dim_indices, local_extents)
+
+
+def cast_layout_supported_for_local(layout) -> bool:
+    """Check that layout is valid for local cast (warp/warpgroup/cta/cluster):
+    filter out cross-thread semantics."""
+    return get_layout_thread_local_partition(layout) is not None
+
+
+def get_local_region(orig_layout: TileLayout, buffer_shape, region_st, region_extent):
+    """Compute local storage shape, iteration starts, and extents with validation of region.
+
+    Args:
+        orig_layout: The original (unsliced) TileLayout.
+        buffer_shape: The buffer shape.
+        region_st: Region start in shape space.
+        region_extent: Region extent in shape space.
+
+    Returns:
+        (local_shape, local_st, local_ext), or ([1], [0], [1]) if no local dims.
+        Returns None if the region is invalid (non-contiguous slicing).
+        - local_shape: full storage extents per local dim.
+        - local_st: region start per local dim.
+        - local_ext: region extent per local dim.
+
+    Example:
+        Layout (2, 8, 4, 2):(8@m, 2@laneid, 2@m, 1@m), Shape [16, 8], Region [8:16, :] returns:
+        - local_shape = [2, 8], local_st = [1, 0], local_ext = [1, 8]
+    """
+    grouped, seps = orig_layout.group(list(buffer_shape))
+
+    local_shape = []
+    local_st = []
+    local_ext = []
+    analyzer = Analyzer()
+
+    for d in range(len(buffer_shape)):
+        shard_range = list(range(seps[d], seps[d + 1]))
+        has_local = any(not grouped.shard[s].axis.is_thread() for s in shard_range)
+        if not has_local:
+            continue
+
+        has_thread = any(grouped.shard[s].axis.is_thread() for s in shard_range)
+
+        if not has_thread:
+            # Pure local shape dim: use shape-level values directly.
+            local_shape.append(buffer_shape[d])
+            local_st.append(region_st[d])
+            local_ext.append(region_extent[d])
+        else:
+            # Decompose start element
+            remaining_st = region_st[d]
+            st_coords = []
+            for i, s_idx in enumerate(shard_range):
+                sub_prod = 1
+                for j in range(i + 1, len(shard_range)):
+                    sub_prod = sub_prod * grouped.shard[shard_range[j]].extent
+                st_coords.append(remaining_st // sub_prod)
+                remaining_st = remaining_st % sub_prod
+
+            # Decompose end element
+            remaining_end = region_st[d] + region_extent[d] - 1
+            end_coords = []
+            for i, s_idx in enumerate(shard_range):
+                sub_prod = 1
+                for j in range(i + 1, len(shard_range)):
+                    sub_prod = sub_prod * grouped.shard[shard_range[j]].extent
+                end_coords.append(remaining_end // sub_prod)
+                remaining_end = remaining_end % sub_prod
+
+            # check the rectangularity and contiguity of the sliced region
+            cur_local_shape, cur_local_st, cur_local_end = 1, 0, 0
+            for k in reversed(range(len(st_coords))):
+                if grouped.shard[seps[d] + k].axis.is_thread():
+                    # for thread dims, region must be contiguous and span full extent
+                    if not (
+                        analyzer.can_prove_equal(st_coords[k], 0)
+                        and analyzer.can_prove_equal(
+                            end_coords[k], grouped.shard[seps[d] + k].extent - 1
+                        )
+                    ):
+                        return None
+                else:
+                    if not analyzer.can_prove_equal(end_coords[k] - st_coords[k], 1) and not (
+                        analyzer.can_prove_equal(st_coords[k], 0)
+                        and analyzer.can_prove_equal(
+                            end_coords[k], grouped.shard[seps[d] + k].extent - 1
+                        )
+                    ):
+                        # to ensure contiguity, if the region spans multiple values
+                        # in this dim, it must span the full extent
+                        return None
+                    cur_local_shape *= grouped.shard[seps[d] + k].extent
+                    cur_local_st = cur_local_st * grouped.shard[seps[d] + k].extent + st_coords[k]
+                    cur_local_end = (
+                        cur_local_end * grouped.shard[seps[d] + k].extent + end_coords[k]
+                    )
+
+            # double check the validity of the sliced region
+            assert region_extent[d] == functools.reduce(
+                operator.mul, [end - st + 1 for st, end in zip(st_coords, end_coords)], 1
+            )
+
+            # append the local info without thread dims
+            local_shape.append(cur_local_shape)
+            local_st.append(cur_local_st)
+            local_ext.append(cur_local_end - cur_local_st + 1)
+
+    if not local_shape:
+        return [1], [0], [1]  # treat no local dim case as 1D local shape with 1 element
+    return local_shape, local_st, local_ext
+
+
+def compute_linear_offset(region_st, local_dims, layout):
+    """Compute linear offset using layout's actual strides.
+
+    Physical offset = sum(region_st[dim] * layout.shard[dim].stride) for all local dims.
+    """
+    offset = 0
+    for dim_idx in local_dims:
+        offset = offset + region_st[dim_idx] * layout.shard[dim_idx].stride
+    return offset
+
+
+def _axis_key(axis):
+    if hasattr(axis, "name") and axis.name:
+        return str(axis.name)
+    return str(axis)
+
+
+def layout_signature(layout):
+    """Return semantic signature from canonicalized TileLayout.
+
+    Returns (thread_sig, local_sig, replica_sig).
+    Each sig is a list of (axis_key, extent, stride) in shard/replica order.
+    """
+    if not isinstance(layout, TileLayout):
+        return None
+    shard = getattr(layout, "shard", None)
+    if not shard:
+        return None
+
+    thread_sig = []
+    local_sig = []
+    for it in shard:
+        item = (_axis_key(it.axis), it.extent, it.stride)
+        if it.axis.is_thread():
+            thread_sig.append(item)
+        else:
+            local_sig.append(item)
+
+    replica_sig = []
+    replica = getattr(layout, "replica", None) or []
+    for it in replica:
+        replica_sig.append((_axis_key(it.axis), it.extent, it.stride))
+    return (thread_sig, local_sig, replica_sig)
+
+
+def sig_equal(analyzer: Analyzer, src_sig, dst_sig) -> bool:
+    """Compare two layout signatures with semantic equality (Analyzer).
+
+    Signatures come from layout_signature(layout) and are:
+      (thread_sig, local_sig, replica_sig)
+    Each sig element is (axis_key, extent, stride).
+    """
+    if src_sig is None or dst_sig is None:
+        return False
+
+    src_thread_sig, src_local_sig, src_replica_sig = src_sig
+    dst_thread_sig, dst_local_sig, dst_replica_sig = dst_sig
+
+    if len(src_thread_sig) != len(dst_thread_sig):
+        return False
+    if len(src_local_sig) != len(dst_local_sig):
+        return False
+    if len(src_replica_sig) != len(dst_replica_sig):
+        return False
+
+    def _list_equal(a_list, b_list) -> bool:
+        for (a_key, a_ext, a_str), (b_key, b_ext, b_str) in zip(a_list, b_list):
+            if a_key != b_key:
+                return False
+            if not analyzer.can_prove_equal(a_ext, b_ext):
+                return False
+            if not analyzer.can_prove_equal(a_str, b_str):
+                return False
+        return True
+
+    return (
+        _list_equal(src_thread_sig, dst_thread_sig)
+        and _list_equal(src_local_sig, dst_local_sig)
+        and _list_equal(src_replica_sig, dst_replica_sig)
+    )
+
+
+def resolve_thread_var(axis, sctx):
+    """Map the axis to the corresponding thread variable."""
+    axis_name = getattr(axis, "name", None)
+    if not axis_name:
+        try:
+            axis_name = str(axis)
+        except Exception:
+            axis_name = ""
+
+    for key, itervar in sctx.launch_params.items():
+        if getattr(itervar.var, "name", "") == axis_name:
+            return itervar.var
+
+    if axis_name:
+        axis_name_lower = axis_name.lower()
+        for key in sctx.launch_params:
+            if axis_name_lower in key.lower() or (axis_name == "tx" and "threadIdx.x" in key):
+                return sctx.launch_params[key].var
+
+    if "threadIdx.x" in sctx.launch_params:
+        return sctx.launch_params["threadIdx.x"].var
+
+    return None

python/tvm/backend/cuda/operator/tile_primitive/permute_layout/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+from .warp_xor_swizzle import *

python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py

@@ -0,0 +1,430 @@
+# 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.
+
+"""CUDA permute_layout dispatch: warp register-staged in-place transpose with
+optional per-lane XOR-swizzle to avoid SMEM bank conflicts on the write phase.
+
+The dispatcher reasons about the **layout's shard**, not the buffer's
+declared shape (the two can differ — a buffer with ``shape=(PIPE, M, K)``
+may carry a layout whose shard has more dims internally, with grouping
+mapping shard segments onto buffer dims).  Concretely:
+
+    src_sliced = src.layout.slice(src.shape, region).canonicalize()
+    dst_sliced = dst.layout.slice(dst.shape, region).canonicalize()
+    # If the two sliced shards have different structures (which is common —
+    # a linear layout collapses to 1D under canon while a transposed one
+    # keeps its multi-dim structure), regroup src to dst's shape.
+    if src_sliced.shard != dst_sliced.shard:
+        src_sliced, _ = src_sliced.group(dst.shard.extents)
+    extent  = [int(it.extent) for it in dst_sliced.shard]   # iteration shape
+    src_str = [int(it.stride) for it in src_sliced.shard]
+    dst_str = [int(it.stride) for it in dst_sliced.shard]
+
+The algorithm:
+
+    regs[P]
+    for r in 0..P:
+        j  = r XOR ((lane >> SHIFT) & MASK)
+        i  = lane + j * 32                             # flat logical index
+        idx = decompose(i, extent)                     # iter multi-dim index
+        regs[r] = src[project(idx, src.shape, slice_starts)]
+    warp_sync()
+    for r in 0..P:
+        j  = r XOR ((lane >> SHIFT) & MASK)
+        i  = lane + j * 32
+        idx = decompose(i, extent)
+        dst[project(idx, dst.shape, slice_starts)] = regs[r]
+    warp_sync()
+
+where ``project`` mixed-radix-folds the iter shard dims back onto the
+buffer's iterated slice dims (so the emit's index matches buf.shape rank,
+which TIR's BufferLoad/Store requires).
+
+SHIFT and MASK are chosen by simulating the bank pattern at the **shard
+granularity** (where strides are affine), trying k = 0, 1, …, log2(P)
+and picking the smallest k that makes both phases bank-conflict-free.
+
+Correctness rests on:
+
+* For each lane, ``r ↦ r XOR const`` is a bijection on ``[0, P)``.
+* Therefore (lane, r) ↔ flat over [0, V).
+* Both layouts are verified bijections on the slice (every logical
+  position has a unique byte offset under that layout).
+* The mixed-radix projection from iter shard idx to buf coord is exactly
+  what TIR's BufferLoad does internally when buf.shape rank < shard rank
+  — so iter shard's strides and the buffer-indexed byte offset agree.
+"""
+
+from __future__ import annotations
+
+import math
+
+from tvm.runtime import DataType
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, BufferRegion, IntImm, PrimFunc
+from tvm.tirx.layout import TileLayout, _flatten_coord
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail, register_dispatch
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import get_indices, get_st_extent
+
+# ---------- helpers ----------------------------------------------------------
+
+
+def _as_buffer_and_region(arg):
+    """Normalize a Buffer or BufferRegion to (buffer, start_list, extent_list)."""
+    if isinstance(arg, Buffer):
+        buf = arg
+        extent = list(buf.shape)
+        st = [0] * len(extent)
+    elif isinstance(arg, BufferRegion):
+        buf = arg.buffer
+        st, extent = get_st_extent(arg)
+    else:
+        raise TypeError(f"unexpected permute_layout arg type: {type(arg)}")
+    return buf, list(st), list(extent)
+
+
+def _as_int(x):
+    """Return int(x) if x is int-like, else None."""
+    if isinstance(x, int):
+        return x
+    if isinstance(x, IntImm):
+        return int(x.value)
+    if hasattr(x, "value") and isinstance(x.value, int):
+        return int(x.value)
+    try:
+        return int(x)
+    except (TypeError, ValueError):
+        return None
+
+
+def _layout_shard_int(layout):
+    """Return (extents, strides) as int lists from a TileLayout's shard, or (None, None)."""
+    if not isinstance(layout, TileLayout):
+        return None, None
+    extents, strides = [], []
+    for it in layout.shard:
+        e = _as_int(it.extent)
+        s = _as_int(it.stride)
+        if e is None or s is None:
+            return None, None
+        extents.append(e)
+        strides.append(s)
+    return extents, strides
+
+
+def _decompose_row_major(i, extent):
+    out, rem = [], i
+    for e in reversed(extent):
+        out.append(rem % e)
+        rem //= e
+    return list(reversed(out))
+
+
+def _eval_offset(idx, strides):
+    return sum(i * s for i, s in zip(idx, strides))
+
+
+def _check_bijection(extent, strides):
+    """Iteration extents + strides define a bijection on [0, V)?"""
+    V = math.prod(extent)
+    seen = set()
+    for i in range(V):
+        off = _eval_offset(_decompose_row_major(i, extent), strides)
+        if off in seen:
+            return False
+        seen.add(off)
+    return len(seen) == V
+
+
+def _bank_free(extent, strides, dtype_bytes, P, k):
+    """For every register slot r ∈ [0, P), do the 32 lanes hit 32 distinct banks?"""
+    T, BANKS, BANK_W = 32, 32, 4
+    shift = 5 - k
+    mask = (1 << k) - 1
+    for r in range(P):
+        seen = set()
+        for lane in range(T):
+            j = r ^ ((lane >> shift) & mask)
+            flat = lane + j * T
+            idx = _decompose_row_major(flat, extent)
+            off_bytes = _eval_offset(idx, strides) * dtype_bytes
+            bank = (off_bytes // BANK_W) % BANKS
+            if bank in seen:
+                return False
+            seen.add(bank)
+    return True
+
+
+def _choose_xor_k(extent, src_strides, dst_strides, dtype_bytes, P):
+    max_k = int(math.log2(P)) if P > 0 else 0
+    for k in range(max_k + 1):
+        if _bank_free(extent, src_strides, dtype_bytes, P, k) and _bank_free(
+            extent, dst_strides, dtype_bytes, P, k
+        ):
+            return k
+    return None
+
+
+# ---------- validator + dispatch impl ---------------------------------------
+
+
+def _gather(op_call):
+    op_call = TilePrimitiveCall.downcast(op_call)
+    dst_arg, src_arg = op_call.args[0], op_call.args[1]
+    src_buf, src_st, src_ext = _as_buffer_and_region(src_arg)
+    dst_buf, dst_st, dst_ext = _as_buffer_and_region(dst_arg)
+    return src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext
+
+
+def _why_reject(op_call, sctx):
+    if not sctx.is_warp:
+        return f"scope {sctx.scope_kind!r} is not 'warp'"
+    if "threadIdx.y" in sctx.launch_params or "threadIdx.z" in sctx.launch_params:
+        return "multi-dim threadIdx is not supported"
+
+    src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext = _gather(op_call)
+
+    if src_buf.dtype != dst_buf.dtype:
+        return f"dtype mismatch: dst={dst_buf.dtype} vs src={src_buf.dtype}"
+
+    src_ext_i = [_as_int(e) for e in src_ext]
+    dst_ext_i = [_as_int(e) for e in dst_ext]
+    if None in src_ext_i or None in dst_ext_i:
+        return "extents must be compile-time integers"
+    if src_ext_i != dst_ext_i:
+        return f"slice shape mismatch: src={src_ext_i} vs dst={dst_ext_i}"
+
+    dtype_bytes = DataType(src_buf.dtype).bits // 8
+    if dtype_bytes not in (1, 2, 4, 8, 16):
+        return f"unsupported dtype byte width: {dtype_bytes}"
+
+    if not isinstance(src_buf.layout, TileLayout):
+        return "src buffer's layout is not a plain TileLayout"
+    if not isinstance(dst_buf.layout, TileLayout):
+        return "dst buffer's layout is not a plain TileLayout"
+
+    # Slice + canonicalize both layouts.  The result's shard describes the
+    # iteration domain; runtime starts (like ``ks``) are folded into the
+    # layout's offset, separate from the shard's affine part.
+    src_region = [(s, s + e) for s, e in zip(src_st, src_ext)]
+    dst_region = [(s, s + e) for s, e in zip(dst_st, dst_ext)]
+    src_sliced = src_buf.layout.slice(list(src_buf.shape), src_region)
+    dst_sliced = dst_buf.layout.slice(list(dst_buf.shape), dst_region)
+    if src_sliced is None or dst_sliced is None:
+        return "layout.slice failed"
+    src_sliced = src_sliced.canonicalize()
+    dst_sliced = dst_sliced.canonicalize()
+
+    # Iteration shape: regroup dst onto the iterated buf dims; the result's
+    # shard may stay finer than iter_buf_extents (one buf dim ↔ several shard
+    # dims via seps), which is fine.  Then regroup src to match dst's shard
+    # extents exactly so both phases share the same iteration index space.
+    iter_buf_extents = [e for e in src_ext_i if e != 1]
+    try:
+        dst_grouped, dst_seps = dst_sliced.group(iter_buf_extents)
+        src_grouped, _ = src_sliced.group([int(it.extent) for it in dst_grouped.shard])
+    except Exception as e:
+        return f"layout.group failed: {e}"
+
+    dst_ext_, dst_str_ = _layout_shard_int(dst_grouped)
+    src_ext_, src_str_ = _layout_shard_int(src_grouped)
+    if dst_ext_ is None or src_ext_ is None:
+        return "regrouped layout shard contains non-integer extent/stride"
+    if src_ext_ != dst_ext_:
+        return f"src shard {src_ext_} doesn't match dst shard {dst_ext_} after regrouping"
+
+    extent = dst_ext_
+    V = math.prod(extent)
+    T = 32
+    if V == 0 or V % T != 0:
+        return f"volume {V} not divisible by warp size {T}"
+    P = V // T
+    if P == 0 or (P & (P - 1)) != 0 or P > T:
+        return f"per-thread count {P} must be power of 2 in [1, {T}]"
+
+    if not _check_bijection(extent, src_str_):
+        return "src layout (regrouped) is not a bijection on the slice"
+    if not _check_bijection(extent, dst_str_):
+        return "dst layout is not a bijection on the slice"
+    return None
+
+
+def _impl(op_call, sctx):
+    src_buf, src_st, src_ext, dst_buf, dst_st, dst_ext = _gather(op_call)
+    src_ext_i = [_as_int(e) for e in src_ext]
+
+    src_region = [(s, s + e) for s, e in zip(src_st, src_ext)]
+    dst_region = [(s, s + e) for s, e in zip(dst_st, dst_ext)]
+    src_sliced = src_buf.layout.slice(list(src_buf.shape), src_region).canonicalize()
+    dst_sliced = dst_buf.layout.slice(list(dst_buf.shape), dst_region).canonicalize()
+
+    iter_buf_extents = [e for e in src_ext_i if e != 1]
+    dst_grouped, dst_seps = dst_sliced.group(iter_buf_extents)
+    src_grouped, _ = src_sliced.group([int(it.extent) for it in dst_grouped.shard])
+
+    extent, dst_str_ = _layout_shard_int(dst_grouped)
+    _, src_str_ = _layout_shard_int(src_grouped)
+    V = math.prod(extent)
+    P = V // 32
+    dtype_bytes = DataType(src_buf.dtype).bits // 8
+
+    k_opt = _choose_xor_k(extent, src_str_, dst_str_, dtype_bytes, P)
+    if k_opt is None:
+        fail(f"no XOR-bits k ∈ [0, log2(P)={int(math.log2(P))}] makes both phases bank-free")
+
+    shift = 5 - k_opt
+    mask = (1 << k_opt) - 1
+
+    iter_buf_dims = [i for i, e in enumerate(src_ext_i) if e != 1]
+    seps = list(dst_seps)
+
+    def _project(iter_idx, st_list):
+        buf_idx = list(st_list)
+        for bi in range(len(seps) - 1):
+            lo, hi = seps[bi], seps[bi + 1]
+            flat = _flatten_coord(iter_idx[lo:hi], extent[lo:hi])
+            buf_idx[iter_buf_dims[bi]] = st_list[iter_buf_dims[bi]] + flat
+        return tuple(buf_idx)
+
+    tid_x = sctx.launch_params["threadIdx.x"]
+    dtype = src_buf.dtype
+
+    # Shared 32/64b: base ptr + stride offset avoids buf[] flatten IMAD path.
+    direct = (
+        dtype_bytes in (4, 8)
+        and "shared" in str(src_buf.scope())
+        and "shared" in str(dst_buf.scope())
+    )
+    ptx_t = f"b{dtype_bytes * 8}"
+    # ld/st only move bits, so use an unsigned container of the matching width:
+    # ``ptx.ld(..., bN)`` rejects a float return dtype, and the permute is a pure
+    # byte shuffle, so a float32/float64 tile loads/stores correctly as uint.
+    bits_dtype = f"uint{dtype_bytes * 8}"
+
+    def _iter_off(iter_idx, strides):
+        return sum(iter_idx[d] * strides[d] for d in range(len(strides)))
+
+    # fmt: off
+    if direct:
+        @T.prim_func
+        def impl():
+            warp_size = T.meta_var(32)
+            lane_id = T.meta_var(tid_x % warp_size)
+            regs = T.alloc_buffer((P,), bits_dtype, scope="local")
+            base_src = T.meta_var(src_buf.ptr_to(list(src_st)))
+            base_dst = T.meta_var(dst_buf.ptr_to(list(dst_st)))
+            # Phase 1: read via L_src
+            for r in T.unroll(0, P):
+                j = T.meta_var(r ^ ((lane_id >> shift) & mask))
+                flat = T.meta_var(lane_id + j * warp_size)
+                iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
+                off = T.meta_var(_iter_off(iter_idx, src_str_))
+                ptr = T.meta_var(T.ptr_byte_offset(base_src, off * dtype_bytes, dtype))
+                regs[r] = T.ptx.ld(ptr, bits_dtype, ptx_t, space="shared")
+            T.cuda.warp_sync()
+            # Phase 2: write via L_dst
+            for r in T.unroll(0, P):
+                j = T.meta_var(r ^ ((lane_id >> shift) & mask))
+                flat = T.meta_var(lane_id + j * warp_size)
+                iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
+                off = T.meta_var(_iter_off(iter_idx, dst_str_))
+                ptr = T.meta_var(T.ptr_byte_offset(base_dst, off * dtype_bytes, dtype))
+                T.evaluate(T.ptx.st(ptr, regs[r], space="shared", ptx_type=ptx_t))
+            T.cuda.warp_sync()
+    else:
+        @T.prim_func
+        def impl():
+            warp_size = T.meta_var(32)
+            lane_id = T.meta_var(tid_x % warp_size)
+            regs = T.alloc_buffer((P,), dtype, scope="local")
+            # Phase 1: read via L_src
+            for r in T.unroll(0, P):
+                j = T.meta_var(r ^ ((lane_id >> shift) & mask))
+                flat = T.meta_var(lane_id + j * warp_size)
+                iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
+                src_idx = T.meta_var(_project(iter_idx, src_st))
+                regs[r] = src_buf[tuple(src_idx)]
+            T.cuda.warp_sync()
+            # Phase 2: write via L_dst
+            for r in T.unroll(0, P):
+                j = T.meta_var(r ^ ((lane_id >> shift) & mask))
+                flat = T.meta_var(lane_id + j * warp_size)
+                iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent))
+                dst_idx = T.meta_var(_project(iter_idx, dst_st))
+                dst_buf[tuple(dst_idx)] = regs[r]
+            T.cuda.warp_sync()
+    # fmt: on
+    return impl
+
+
+# === Variant: permute_layout/warp_xor_swizzle (priority=20) ============
+#
+# When: warp scope; matching dst/src dtype + slice shape; both buffers carry
+# a plain TileLayout; after slice + canonicalize (and regrouping src to dst's
+# structure if needed), the iteration extents form a power-of-2 ≤32 elements
+# per lane; both layouts are bijections on the slice; and there exists an
+# XOR-bits ``k`` that makes both phases bank-conflict-free.
+#
+# Buffer ``shape`` rank does NOT need to equal layout ``shard`` rank — the
+# dispatcher uses the layout shard for iteration (after slice+canon) and
+# projects back onto ``buf.shape`` via mixed-radix grouping for the emit.
+#
+# Before (TilePrimitiveCall):
+#     with T.warp():
+#         # SFA_smem: u32 (PIPE, BLK_SFA//32, 32), layout shard 4D
+#         #   (PIPE, BLK_SFA//128, 4, 32) strides (BLK_SFA, 128, 32, 1)
+#         # SFA_post: same shape; layout shard 4D, strides (BLK_SFA, 128, 1, 4)
+#         Tx.permute_layout(SFA_post[ks, :, :], SFA_smem[ks, :, :])
+#
+# After (BLK_SFA=128, P=4, k=2, shift=3):
+#     lane_id = threadIdx.x % 32
+#     regs = T.alloc_buffer((4,), "uint32", scope="local")
+#     for r in T.unroll(4):
+#         j = r ^ ((lane_id >> 3) & 0x3)
+#         flat = lane_id + j * 32
+#         (g, l) = decompose(flat, extent=[4, 32])
+#         regs[r] = src[ks, g, l]
+#     T.cuda.warp_sync()
+#     for r in T.unroll(4):
+#         j = r ^ ((lane_id >> 3) & 0x3)
+#         flat = lane_id + j * 32
+#         (g, l) = decompose(flat, extent=[4, 32])
+#         dst[ks, g, l] = regs[r]
+#     T.cuda.warp_sync()
+@register_dispatch(
+    "permute_layout",
+    "cuda",
+    variant="warp_xor_swizzle",
+    priority=20,
+)
+def permute_layout_dispatch(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    reason = _why_reject(op, sctx)
+    if reason is not None:
+        fail(reason)
+    return _impl(op, sctx)
+
+
+__all__ = [
+    "_bank_free",
+    "_check_bijection",
+    "_choose_xor_k",
+    "_decompose_row_major",
+    "_eval_offset",
+    "permute_layout_dispatch",
+]

python/tvm/backend/cuda/operator/tile_primitive/reduction/__init__.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
+# under the License.
+
+from .local import *
+from .shared import *
+from .sm100_packed import *

python/tvm/backend/cuda/operator/tile_primitive/reduction/local.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.
+
+"""CUDA reduction operator dispatch: local-memory variant.
+
+Registered ops: sum, max, min.
+
+When: dst and src are both local-scope buffers with matching dtype, on CUDA.
+
+(A) Thread scope -- sequential per-element reduction
+    (_emit_reduction_local_thread_wise):
+
+Before:
+    Tx.sum(B_local[0:2, 0:3], A_local[0:2, 0:3, 0:4], [-1], False)
+
+After (scheduled PrimFunc, spatial_len=6, reduction_len=4):
+    for spa in range(6):
+        B_local[spa] = T.float32(0.0)                      # init (skipped if accum)
+        for red in range(4):
+            B_local[spa] = B_local[spa] + A_local[spa * 4 + red]
+
+(B) Warp/Warpgroup scope -- layout-driven reduction
+    (_emit_reduction_local_view):
+    Requires TileLayout with valid thread-partition. Decomposes layout to
+    identify thread-local elements, then optionally shuffles partial sums.
+
+    thread_reduce=False: local-only, no shuffle (warp and warpgroup).
+    thread_reduce=True: local reduction + cross-thread shfl_xor steps (warp only).
+    accum=True + shuffle: saves old dst before reduce+shuffle, combines after (warp only).
+
+Before:
+    Tx.warp.sum(red_view[0:16, 0:4], acc_view[0:16, 0:128], [-1], False,
+                   thread_reduce=True)
+
+After (scheduled PrimFunc, local_total=2, local_red=32, 2 shuffle steps):
+    src_local = acc_view.view(64)
+    dst_local = red_view.view(2)
+    for spa in range(2):
+        dst_local[spa] = T.float32(0.0)
+        for red in range(32):
+            dst_local[spa] = dst_local[spa] + src_local[...]
+        dst_local[spa] = dst_local[spa] + shfl_xor(..., 1, 32, 32)
+        dst_local[spa] = dst_local[spa] + shfl_xor(..., 2, 32, 32)
+"""
+
+import functools
+import operator
+from typing import Any
+
+from tvm.arith.analyzer import Analyzer
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx.layout import TileLayout, laneid
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import get_indices, get_st_extent
+from ..layout_utils import get_local_region, get_sublayout_from_region
+from .utils import (
+    _REDUCE_OP_TO_STR,
+    _analyze_axes,
+    _analyze_layout_dims,
+    _build_local_dim_map,
+    _compute_shuffle_masks,
+    _match_reduction_storage_scope,
+    _reduction_args,
+    _validate_reduction_layout,
+    reduce_default_value_table,
+    reduce_op_table,
+)
+
+
+def _analyze_shuffle_reduce(src_layout, dst_layout):
+    """Analyze src/dst layouts for laneid shard->replica reduce pattern.
+
+    Returns (reduce_width, local_elems) if the pattern matches, or None.
+    - reduce_width: number of lanes participating in each group's reduction
+    - local_elems: per-thread element count (product of non-laneid shard extents)
+    """
+    if src_layout.is_swizzle() or dst_layout.is_swizzle():
+        return None
+
+    src_canon = src_layout.canonicalize()
+    dst_canon = dst_layout.canonicalize()
+
+    # Extract laneid iters from shard and replica
+    src_laneid_shard = [it for it in src_canon.shard if it.axis == laneid]
+    dst_laneid_replica = [it for it in dst_canon.replica if it.axis == laneid]
+
+    # src shard must contain laneid (data distributed across lanes)
+    if not src_laneid_shard:
+        return None
+    # dst replica must contain laneid (result broadcast to lanes)
+    if not dst_laneid_replica:
+        return None
+
+    # laneid span must be 32 (full warp)
+    src_laneid_span = 1 + sum(abs(int(it.stride)) * (int(it.extent) - 1) for it in src_laneid_shard)
+    if src_laneid_span != 32:
+        return None
+
+    reduce_width = functools.reduce(operator.mul, [int(it.extent) for it in dst_laneid_replica], 1)
+    if reduce_width <= 0 or reduce_width > 32 or (reduce_width & (reduce_width - 1)) != 0:
+        return None  # must be power of 2
+
+    # local_elems = product of non-laneid shard extents in src
+    src_non_laneid = [it for it in src_canon.shard if it.axis != laneid]
+    local_elems = functools.reduce(operator.mul, [int(it.extent) for it in src_non_laneid], 1)
+
+    return reduce_width, local_elems
+
+
+def _gen_warp_shuffle_reduce(src, dst, reduce_width, local_elems, accum, op_type, init_value):
+    """Generate warp shuffle reduce codegen for laneid shard->replica pattern.
+
+    Unified for both full warp (reduce_width=32) and partial warp (e.g. reduce_width=8).
+    """
+    is_same_buffer = src.same_as(dst)
+    op_str = _REDUCE_OP_TO_STR[op_type]
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        src_local = src.local(local_elems)
+        dst_local = dst.local(local_elems)
+        for k in T.serial(local_elems):
+            if not is_same_buffer:
+                dst_local[k] = src_local[k]
+            dst_local[k] = T.cuda.warp_reduce(dst_local[k], op_str, reduce_width)
+    # fmt: on
+
+    return impl
+
+
+def validate_reduction_local(
+    op: TilePrimitiveCall, sctx: DispatchContext
+) -> tuple[bool, str | None]:
+    """Validate reduction in local memory."""
+    op = TilePrimitiveCall.downcast(op)
+    dst_br, src_br = op.output, op.input
+    dst, src = dst_br.buffer, src_br.buffer
+
+    if not (src.scope() == "local" and dst.scope() == "local" and sctx.is_target("cuda")):
+        return False, "expected local scope and CUDA target"
+    if src.dtype != dst.dtype:
+        return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
+
+    if sctx.is_thread:
+        return True, None  # thread-wise reduction
+    elif sctx.scope_kind in ["warp", "warpgroup"]:
+        if not sctx.is_warp and op.config.get("thread_reduce", False):
+            return (
+                False,
+                "thread_reduce=True is only supported in warp scope; "
+                "warpgroup local reduction is thread-local only",
+            )
+        # VIEW: need layouts and layout analysis
+        if not (src.layout and dst.layout):
+            return False, "layouts required for view-based local reduction"
+        if not (isinstance(src.layout, TileLayout) and isinstance(dst.layout, TileLayout)):
+            return False, "TileLayout required for view-based local reduction"
+        if src.layout.is_swizzle() or dst.layout.is_swizzle():
+            return False, "swizzle layout unsupported for local reduction"
+
+        analyzer = Analyzer()
+
+        # Validate get_local_region succeeds for both
+        src_st, src_extent = get_st_extent(src_br)
+        dst_st, dst_extent = get_st_extent(dst_br)
+
+        if sctx.is_warp:
+            # Check for laneid shard->replica shuffle reduce pattern first.
+            # This pattern has laneid in dst replica (broadcast), which the
+            # general validation below would reject.
+            shuffle_info = _analyze_shuffle_reduce(src.layout, dst.layout)
+            if shuffle_info is not None:
+                return True, None
+
+        for layout, buf, st, ext, name in [
+            (src.layout, src, src_st, src_extent, "src"),
+            (dst.layout, dst, dst_st, dst_extent, "dst"),
+        ]:
+            for it in layout.shard:
+                if it.axis.is_thread() and analyzer.can_prove_equal(it.stride, 0):
+                    return False, f"thread dim with zero stride in {name}"
+            replica = getattr(layout, "replica", None) or []
+            if any(it.axis.is_thread() for it in replica):
+                return False, f"thread axis in replica for {name}"
+            if get_local_region(layout, list(buf.shape), st, ext) is None:
+                return False, f"get_local_region failed for {name}"
+
+        # Validate layout compatibility
+        # Spatial dims match, reduce dims in dst have local_extent==1
+        reduce_axes = tuple(int(a) for a in op.reduce_axes)
+        src_ndim = len(src_br.region)
+        try:
+            reduce_dims, _ = _analyze_axes(src_ndim, reduce_axes)
+        except AssertionError as e:
+            return False, str(e)
+        src_sliced = get_sublayout_from_region(src.layout, src.shape, src_st, src_extent)
+        dst_sliced = get_sublayout_from_region(dst.layout, dst.shape, dst_st, dst_extent)
+        ok, msg = _validate_reduction_layout(
+            src_sliced, dst_sliced, list(src_extent), list(dst_extent), reduce_dims
+        )
+        return ok, msg
+    else:
+        return False, f"unsupported exec_scope {sctx.scope_kind} for local reduction"
+
+
+def _emit_reduction_local_thread_wise(
+    dst_br: BufferRegion,
+    src_br: BufferRegion,
+    accum: bool,
+    reduce_op: ReduceOpType,
+    reduce_dims: list[int],
+    spatial_dims: list[int],
+) -> PrimFunc:
+    dst, src = dst_br.buffer, src_br.buffer
+    dtype = src.dtype
+    src_st, src_extent = get_st_extent(src_br)
+    dst_st, dst_extent = get_st_extent(dst_br)
+    src_ndim = len(src_extent)
+    spa_extents = [src_extent[d] for d in spatial_dims]
+    red_extents = [src_extent[d] for d in reduce_dims]
+    spatial_len = functools.reduce(operator.mul, spa_extents, 1)
+    reduction_len = functools.reduce(operator.mul, red_extents, 1)
+
+    op_func = reduce_op_table.get(reduce_op)
+    assert op_func is not None
+    init_value = reduce_default_value_table(dtype).get(reduce_op)
+
+    def get_src_indices(spa_fused, red_fused):
+        spa_indices = []
+        rem = spa_fused
+        for e in reversed(spa_extents):
+            spa_indices.append(rem % e)
+            rem //= e
+        spa_indices.reverse()
+
+        red_indices = []
+        rem = red_fused
+        for e in reversed(red_extents):
+            red_indices.append(rem % e)
+            rem //= e
+        red_indices.reverse()
+
+        full = [None] * src_ndim
+        for i, d in enumerate(spatial_dims):
+            full[d] = spa_indices[i] + src_st[d]
+        for i, d in enumerate(reduce_dims):
+            full[d] = red_indices[i] + src_st[d]
+        return full
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        for spa in T.serial(spatial_len):
+            dst_idx = T.meta_var(get_indices(spa, dst_st, dst_extent))
+            if not accum:
+                dst[tuple(dst_idx)] = init_value
+            for red in T.serial(reduction_len):
+                src_idx = T.meta_var(get_src_indices(spa, red))
+                dst[tuple(dst_idx)] = op_func(dst[tuple(dst_idx)], src[tuple(src_idx)])
+    # fmt: on
+
+    return impl
+
+
+def _emit_reduction_local_view(
+    dst_br: BufferRegion,
+    src_br: BufferRegion,
+    accum: bool,
+    reduce_op: ReduceOpType,
+    config: dict[str, Any],
+    reduce_dims: set[int],
+    spatial_dims: list[int],
+    src_local_info,
+    dst_local_info,
+    shuffle_masks: list[int],
+) -> PrimFunc:
+    dst, src = dst_br.buffer, src_br.buffer
+    dtype = src.dtype
+
+    op_func = reduce_op_table.get(reduce_op)
+    assert op_func is not None
+    init_value = reduce_default_value_table(dtype).get(reduce_op)
+
+    src_local_shape, src_local_st, src_local_ext = src_local_info
+    dst_local_shape, dst_local_st, dst_local_ext = dst_local_info
+
+    # Build maps from original dim index to position in get_local_region output
+    src_dim_map = _build_local_dim_map(src.layout, list(src.shape))
+    dst_dim_map = _build_local_dim_map(dst.layout, list(dst.shape))
+
+    # Only include reduction dims that have local parts in src
+    src_ndim = len(src_br.region)
+    reduce_local_dims = [d for d in reduce_dims if src_dim_map[d] is not None]
+    reduction_local_ext = [src_local_ext[src_dim_map[d]] for d in reduce_local_dims]
+    reduction_local_st = [src_local_st[src_dim_map[d]] for d in reduce_local_dims]
+
+    reduction_local_total = functools.reduce(operator.mul, reduction_local_ext, 1)
+    dst_local_total = functools.reduce(operator.mul, dst_local_ext, 1)
+
+    def _get_src_local_index(dst_fused, red_fused):
+        """Compute src local multi-dim index from dst fused index and reduction fused index."""
+        dst_indices = get_indices(dst_fused, dst_local_st, dst_local_ext)
+        red_indices = get_indices(red_fused, reduction_local_st, reduction_local_ext)
+
+        # Interleave into src local indices (skipping pure-thread dims)
+        src_local = []
+        ri = 0
+        for d in range(src_ndim):
+            if src_dim_map[d] is None:
+                continue  # pure-thread in src, not in src.local()
+            if d in reduce_dims:
+                src_local.append(red_indices[ri])
+                ri += 1
+            else:
+                # Spatial dim: use corresponding dst local position
+                src_local.append(dst_indices[dst_dim_map[d]])
+
+        return src_local
+
+    # is_same_buffer = src.same_as(dst)
+    shuffle = bool(config.get("thread_reduce", False))
+    in_place = dst.same_as(src)
+
+    def shuffle_data(mask, dst_local, dst_idx):
+        @T.inline
+        def inner_shuffle(v, shuffle_mask):
+            dst_local[tuple(dst_idx)] = op_func(
+                v, T.tvm_warp_shuffle_xor(mask, v, shuffle_mask, 32, 32)
+            )
+
+        for i in range(len(shuffle_masks)):
+            inner_shuffle(dst_local[tuple(dst_idx)], shuffle_masks[i])
+
+    need_save_accum = accum and shuffle
+
+    # fmt: off
+    if need_save_accum:
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            src_local = src.local(*src_local_shape)
+            dst_local = dst.local(*dst_local_shape)
+            old_val = T.alloc_buffer([1], dtype, scope="local")
+
+            for spa in T.serial(dst_local_total):
+                dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))
+                old_val[0] = dst_local[tuple(dst_idx)]
+                if not in_place:
+                    dst_local[tuple(dst_idx)] = init_value
+                    for red in T.serial(reduction_local_total):
+                        src_idx = T.meta_var(_get_src_local_index(spa, red))
+                        dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)])  # noqa: E501
+                if shuffle:
+                    mask = T.tvm_warp_activemask()
+                    shuffle_data(mask, dst_local, dst_idx)
+                dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], old_val[0])
+    else:
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            src_local = src.local(*src_local_shape)
+            dst_local = dst.local(*dst_local_shape)
+
+            for spa in T.serial(dst_local_total):
+                dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))
+                if not in_place:
+                    if not accum:
+                        dst_local[tuple(dst_idx)] = init_value
+                    for red in T.serial(reduction_local_total):
+                        src_idx = T.meta_var(_get_src_local_index(spa, red))
+                        dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)])  # noqa: E501
+                if shuffle:
+                    mask = T.tvm_warp_activemask()
+                    shuffle_data(mask, dst_local, dst_idx)
+    # fmt: on
+
+    return impl
+
+
+def reduction_local_impl(
+    op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext
+) -> PrimFunc | None:
+    dst_br, src_br, reduce_axes, accum, config = _reduction_args(op)
+    src_ndim = len(src_br.region)
+    reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
+
+    if sctx.is_thread:
+        return _emit_reduction_local_thread_wise(
+            dst_br, src_br, accum, op_type, reduce_dims, spatial_dims
+        )
+    elif sctx.scope_kind in ["warp", "warpgroup"]:
+        src = src_br.buffer
+        dst = dst_br.buffer
+
+        if sctx.is_warp:
+            # --- Try laneid shard->replica shuffle reduce ---
+            shuffle_info = _analyze_shuffle_reduce(src.layout, dst.layout)
+            if shuffle_info is not None:
+                reduce_width, local_elems = shuffle_info
+                if op_type not in _REDUCE_OP_TO_STR:
+                    fail(f"unsupported reduce op: {op_type}")
+                dtype = src.dtype
+                init_value = reduce_default_value_table(dtype).get(op_type)
+                return _gen_warp_shuffle_reduce(
+                    src, dst, reduce_width, local_elems, accum, op_type, init_value
+                )
+        elif config.get("thread_reduce", False):
+            fail(
+                "thread_reduce=True is only supported in warp scope; "
+                "warpgroup local reduction is thread-local only"
+            )
+
+        # --- Existing WGMMA layout path below ---
+        src_st, src_extent = get_st_extent(src_br)
+        dst_st, dst_extent = get_st_extent(dst_br)
+
+        src_local_info = get_local_region(src.layout, list(src.shape), src_st, src_extent)
+        dst_local_info = get_local_region(dst.layout, list(dst.shape), dst_st, dst_extent)
+        assert src_local_info is not None and dst_local_info is not None
+
+        src_dim_info = _analyze_layout_dims(src.layout, list(src.shape))
+        shuffle_masks = (
+            _compute_shuffle_masks(src_dim_info, reduce_dims)
+            if config.get("thread_reduce", False)
+            else []
+        )
+
+        return _emit_reduction_local_view(
+            dst_br,
+            src_br,
+            accum,
+            op_type,
+            config,
+            reduce_dims,
+            spatial_dims,
+            src_local_info,
+            dst_local_info,
+            shuffle_masks,
+        )
+    else:
+        fail(f"unsupported exec_scope {sctx.scope_kind} for reduction_local_impl")
+
+
+# ---------------------------------------------------------------------------
+# Registration: local memory reduction (priority=10)
+# ---------------------------------------------------------------------------
+
+for op_name, op_type in [
+    ("sum", ReduceOpType.SUM),
+    ("max", ReduceOpType.MAX),
+    ("min", ReduceOpType.MIN),
+]:
+
+    @register_dispatch(
+        op_name,
+        "cuda",
+        variant="local",
+        priority=10,
+        when=[
+            predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["local"]),
+            predicate("local_valid", validate_reduction_local),
+        ],
+    )
+    def _local_dispatch(op: TilePrimitiveCall, sctx: DispatchContext, _op_type=op_type) -> PrimFunc:
+        op = TilePrimitiveCall.downcast(op)
+        return reduction_local_impl(op, _op_type, sctx)

python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py

@@ -0,0 +1,298 @@
+# 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.
+
+"""CUDA reduction operator dispatch: shared-memory variant.
+
+Registered ops: sum, max, min.
+
+When: dst and src are both shared-memory buffers, exec scope is one of
+{cta, warpgroup, warp, thread}, threadIdx.x bound, reduce axes valid.
+
+(A) CTA/warpgroup/warp scope -- adaptive-group shuffle tree
+    (_emit_reduction_shared_cta):
+    group_size = min(next_power_of_2(reduction_len), 32).
+    Each group of threads reduces one spatial position via shfl_xor.
+
+Before:
+    Tx.cta.sum(B_smem[0:4], A_smem[0:4, 0:8], [-1], False)
+
+After (scheduled PrimFunc, group_size=8, spatial_par=4):
+    thread_data[0] = T.float32(0.0)
+    thread_data[0] = thread_data[0] + A_smem[tid_in_scope]  # gather
+    # log2(8) = 3 shuffle-xor steps with width=8
+    thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 1, 8, 32)
+    thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 2, 8, 32)
+    thread_data[0] = thread_data[0] + shfl_xor(thread_data[0], 4, 8, 32)
+    if tid_in_scope % 8 == 0:
+        B_smem[tid_in_scope // 8] = thread_data[0]
+
+(B) Thread scope -- sequential loop (_emit_reduction_shared_thread):
+
+Before:
+    if tid == 65:
+        Tx.sum(B_smem[0:4], A_smem[0:4, 0:8], [-1], False)
+
+After (scheduled PrimFunc):
+    for spa in range(4):
+        B_smem[spa] = T.float32(0.0)                       # init (skipped if accum)
+        for red in range(8):
+            B_smem[spa] = B_smem[spa] + A_smem[spa * 8 + red]
+"""
+
+import functools
+import math
+import operator
+
+from tvm.arith.analyzer import Analyzer
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext, fail
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import get_indices, get_st_extent, next_power_of_2
+from .utils import (
+    _analyze_axes,
+    _match_reduction_storage_scope,
+    _reduction_args,
+    build_src_indices,
+    reduce_default_value_table,
+    reduce_op_table,
+)
+
+
+def validate_reduction_shared(
+    op: TilePrimitiveCall, sctx: DispatchContext
+) -> tuple[bool, str | None]:
+    """Validate reduction in shared memory."""
+    if sctx.scope_kind not in ["cta", "warpgroup", "warp", "thread"]:
+        return False, f"unsupported exec_scope {sctx.scope_kind} for shared reduction"
+
+    op = TilePrimitiveCall.downcast(op)
+    dst, src = op.output.buffer, op.input.buffer
+    if not (src.scope().startswith("shared") and dst.scope().startswith("shared")):
+        return False, "expected shared scope for both src and dst"
+    if src.dtype != dst.dtype:
+        return False, f"dtype mismatch: src={src.dtype} dst={dst.dtype}"
+
+    if "threadIdx.x" not in sctx.launch_params:
+        return False, "threadIdx.x not in launch_params"
+    if "threadIdx.y" in sctx.launch_params or "threadIdx.z" in sctx.launch_params:
+        return False, "multi-dimensional thread binding not supported for shared reduction"
+
+    reduce_axes = tuple(int(a) for a in op.reduce_axes)
+    src_region = op.input.region
+    dst_region = op.output.region
+    src_ndim = len(src_region)
+    try:
+        reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
+    except AssertionError as e:
+        return False, str(e)
+
+    # Validate dst shape matches spatial dims of src
+    src_extent = [r.extent for r in src_region]
+    dst_extent = [r.extent for r in dst_region]
+    expected_dst_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
+    actual_dst_len = functools.reduce(operator.mul, dst_extent, 1)
+    analyzer = Analyzer()
+    if not analyzer.can_prove_equal(expected_dst_len, actual_dst_len):
+        return (False, f"dst size {actual_dst_len} != expected spatial size {expected_dst_len}")
+
+    return True, None
+
+
+def _emit_reduction_shared_cta(
+    dst_br: BufferRegion,
+    src_br: BufferRegion,
+    accum: bool,
+    reduce_op: ReduceOpType,
+    sctx: DispatchContext,
+    reduce_dims: list[int],
+    spatial_dims: list[int],
+) -> PrimFunc:
+    exec_scope_name = sctx.scope_kind
+
+    def get_thread_cnt():
+        if exec_scope_name == "cta":
+            return sctx.launch_params["threadIdx.x"].dom.extent
+        elif exec_scope_name == "warpgroup":
+            return 128
+        elif exec_scope_name == "warp":
+            return 32
+        elif exec_scope_name == "thread":
+            return 1
+
+    thread_cnt = get_thread_cnt()
+    dst, src = dst_br.buffer, src_br.buffer
+    src_st, src_extent = get_st_extent(src_br)
+    dst_st, dst_extent = get_st_extent(dst_br)
+    dtype = src.dtype
+
+    # Compute spatial/reduction from the explicit axes
+    spatial_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
+    reduction_len = functools.reduce(operator.mul, [src_extent[d] for d in reduce_dims], 1)
+
+    op_func = reduce_op_table.get(reduce_op)
+    assert op_func is not None
+    init_value = reduce_default_value_table(dtype).get(reduce_op)
+
+    # Adaptive group size: nearest power-of-2 for reduction length, capped at warp size and thread count. # noqa: E501
+    group_size = min(next_power_of_2(int(reduction_len)), 32, int(thread_cnt))
+    group_size = max(group_size, 1)  # ensure at least 1
+    n_shuffles = int(math.log2(group_size)) if group_size > 1 else 0
+    spatial_par = int(thread_cnt) // group_size
+
+    def get_tid_in_scope():
+        tx_var = sctx.launch_params["threadIdx.x"].var
+        if exec_scope_name == "cta":
+            return tx_var
+        elif exec_scope_name in ("warp", "warpgroup"):
+            return tx_var % thread_cnt
+        elif exec_scope_name == "thread":
+            return 0
+
+    def shuffle_data(thread_data):
+        @T.inline
+        def inner_shuffle(mask, v, shuffle_mask):
+            v[0] = op_func(v[0], T.tvm_warp_shuffle_xor(mask, v[0], shuffle_mask, group_size, 32))
+
+        if n_shuffles > 0:
+            mask = T.tvm_warp_activemask()
+            for i in range(n_shuffles):
+                inner_shuffle(mask, thread_data, 1 << i)
+
+    @T.inline
+    def sync():
+        if exec_scope_name == "cta":
+            T.cuda.cta_sync()
+        elif exec_scope_name == "warpgroup":
+            T.cuda.warpgroup_sync(8)  # TODO: fix this hardcoded value
+        elif exec_scope_name == "warp":
+            T.cuda.warp_sync()
+        elif exec_scope_name == "thread":
+            pass
+
+    # fmt: off
+    @T.prim_func
+    def impl():
+        tid_in_scope = get_tid_in_scope()
+        thread_data = T.alloc_buffer([1], dtype=dtype, scope="local")
+        group_id = T.meta_var(T.floordiv(tid_in_scope, group_size))
+        lane_in_grp = T.meta_var(tid_in_scope % group_size)
+        for step in T.serial(T.ceildiv(spatial_len, spatial_par)):
+            spa_fused = T.meta_var(step * spatial_par + group_id)
+            if spa_fused < spatial_len:
+                thread_data[0] = init_value
+                for t in T.serial(T.ceildiv(reduction_len, group_size)):
+                    red_fused = T.meta_var(t * group_size + lane_in_grp)
+                    if red_fused < reduction_len:
+                        src_indices = T.meta_var(build_src_indices(spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st))  # noqa: E501
+                        thread_data[0] = op_func(thread_data[0], src[tuple(src_indices)])
+                shuffle_data(thread_data)
+                if lane_in_grp == 0:
+                    dst_indices = T.meta_var(get_indices(spa_fused, dst_st, dst_extent))
+                    dst[tuple(dst_indices)] = T.if_then_else(T.bool(accum), op_func(dst[tuple(dst_indices)], thread_data[0]), thread_data[0])  # noqa: E501
+
+        sync()
+    # fmt: on
+
+    return impl
+
+
+def _emit_reduction_shared_thread(
+    dst_br: BufferRegion,
+    src_br: BufferRegion,
+    accum: bool,
+    reduce_op: ReduceOpType,
+    sctx: DispatchContext,
+    reduce_dims: list[int],
+    spatial_dims: list[int],
+) -> PrimFunc:
+    dst, src = dst_br.buffer, src_br.buffer
+    src_st, src_extent = get_st_extent(src_br)
+    dst_st, dst_extent = get_st_extent(dst_br)
+    dtype = src.dtype
+
+    # Compute spatial/reduction from the explicit axes
+    spatial_len = functools.reduce(operator.mul, [src_extent[d] for d in spatial_dims], 1)
+    reduction_len = functools.reduce(operator.mul, [src_extent[d] for d in reduce_dims], 1)
+
+    op_func = reduce_op_table.get(reduce_op)
+    assert op_func is not None
+    init_value = reduce_default_value_table(dtype).get(reduce_op)
+
+    @T.prim_func
+    def impl():
+        for spa_fused in T.serial(spatial_len):
+            dst_indices = T.meta_var(get_indices(spa_fused, dst_st, dst_extent))
+            if not accum:
+                dst[tuple(dst_indices)] = init_value
+            for red_fused in T.serial(reduction_len):
+                src_indices = T.meta_var(
+                    build_src_indices(
+                        spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st
+                    )
+                )
+                dst[tuple(dst_indices)] = op_func(dst[tuple(dst_indices)], src[tuple(src_indices)])
+
+    return impl
+
+
+def reduction_shared_impl(
+    op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext
+) -> PrimFunc | None:
+    dst_br, src_br, reduce_axes, accum, config = _reduction_args(op)
+    src_ndim = len(src_br.region)
+    reduce_dims, spatial_dims = _analyze_axes(src_ndim, reduce_axes)
+    if sctx.scope_kind in ["cta", "warpgroup", "warp"]:
+        return _emit_reduction_shared_cta(
+            dst_br, src_br, accum, op_type, sctx, reduce_dims, spatial_dims
+        )
+    elif sctx.is_thread:
+        return _emit_reduction_shared_thread(
+            dst_br, src_br, accum, op_type, sctx, reduce_dims, spatial_dims
+        )
+    else:
+        fail(f"unsupported exec_scope {sctx.scope_kind} for reduction_shared_impl")
+
+
+# ---------------------------------------------------------------------------
+# Registration: shared memory reduction (priority=10)
+# ---------------------------------------------------------------------------
+
+for op_name, op_type in [
+    ("sum", ReduceOpType.SUM),
+    ("max", ReduceOpType.MAX),
+    ("min", ReduceOpType.MIN),
+]:
+
+    @register_dispatch(
+        op_name,
+        "cuda",
+        variant="shared",
+        priority=10,
+        when=[
+            predicate("storage_scope", _match_reduction_storage_scope, expected_scope=["shared*"]),
+            predicate("shared_valid", validate_reduction_shared),
+        ],
+    )
+    def _shared_dispatch(
+        op: TilePrimitiveCall, sctx: DispatchContext, _op_type=op_type
+    ) -> PrimFunc:
+        op = TilePrimitiveCall.downcast(op)
+        return reduction_shared_impl(op, _op_type, sctx)

python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py

@@ -0,0 +1,249 @@
+# 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.
+
+"""CUDA reduction operator dispatch: SM100+ packed optimized variant.
+
+Registered ops: sum, max, min.
+
+When: thread scope, all local buffers, float32, 1D src with len >= 8,
+SM100+ (uses packed PTX instructions not available on older GPUs).
+
+Before (TilePrimitiveCall -- sum example):
+    Tx.sum(dst_local[0:1], src_local[0:32])   # float32, reduce 32 -> 1 (thread scope)
+
+After -- packed_add_sum (uses add.f32x2 to reduce pairs):
+    # Iteratively reduce: 32 -> 16 -> 8 -> 4 -> 2 -> 1
+    # Each step: add.f32x2 combines adjacent pairs
+    for i in T.serial(16):
+        T.cuda.func_call("add_f32x2", &buf[i*2], &buf[i*2], &buf[i*2+2])
+    # ... repeat halving until scalar result
+    dst_local[0] = buf[0]
+
+After -- 3input_maxmin (uses 3-input PTX max/min):
+    # Tree reduction with 3-input instructions:
+    # max(a, b, c) in one PTX instruction
+    for i in T.serial(n // 3):
+        T.cuda.func_call("max3_f32", &buf[i*3], &buf[i*3+1], &buf[i*3+2])
+
+With accum=True: accumulator folded into first element/pair of the reduction.
+"""
+
+import functools
+import operator
+
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion, PrimFunc
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import sm_version_ok
+from ..exec_scope_utils import exec_scope_ok
+from .utils import (
+    _dst_len_ok,
+    _dtype_ok,
+    _local_scope_match,
+    _reduction_len_ok,
+    _src_ndim_ok,
+    reduce_op_table,
+)
+
+
+def _emit_reduction_local_thread_packed_add_sum(
+    dst_buffer_region: BufferRegion,
+    src_buffer_region: BufferRegion,
+    accum: bool,
+    reduce_op: ReduceOpType,
+    sctx: DispatchContext,
+) -> PrimFunc:
+    dst, src = dst_buffer_region.buffer, src_buffer_region.buffer
+    src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
+    dtype = src.dtype
+
+    src_extent = [r.extent for r in src_region]
+    [r.extent for r in dst_region]
+    src_st = [r.min for r in src_region]
+    dst_st = [r.min for r in dst_region]
+
+    reduction_len = functools.reduce(operator.mul, src_extent, 1)
+
+    src_base = src_st[0]
+    num_full_chunks = reduction_len // 8
+    remainder = reduction_len % 8
+    remainder_base = num_full_chunks * 8
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        local_sum = T.alloc_buffer([8], dtype, scope="local")
+        # First pass: copy first 8 elements (with optional accumulator)
+        for i in T.unroll(8):
+            if accum and i == 0:
+                local_sum[i] = src[src_base + i] + dst[tuple(dst_st)]
+            else:
+                local_sum[i] = src[src_base + i]
+
+        # Process remaining full chunks of 8
+        for outer in T.serial(num_full_chunks - 1):
+            for j in T.unroll(4):
+                T.ptx.add_f32x2(
+                    T.address_of(local_sum[2 * j]),
+                    T.cuda.make_float2(local_sum[2 * j], local_sum[2 * j + 1]),
+                    T.cuda.make_float2(
+                        src[src_base + 8 * (outer + 1) + 2 * j],
+                        src[src_base + 8 * (outer + 1) + 2 * j + 1],
+                    ),
+                    ftz=True,
+                )
+
+        # Handle remainder elements (0 to 7)
+        for i in T.serial(remainder):
+            local_sum[0] = local_sum[0] + src[src_base + remainder_base + i]
+
+        # Final packed add sum: 8 -> 4 -> 2 -> 1
+        T.ptx.add_f32x2(
+            T.address_of(local_sum[0]),
+            T.cuda.make_float2(local_sum[0], local_sum[1]),
+            T.cuda.make_float2(local_sum[2], local_sum[3]),
+            ftz=True,
+        )
+        T.ptx.add_f32x2(
+            T.address_of(local_sum[4]),
+            T.cuda.make_float2(local_sum[4], local_sum[5]),
+            T.cuda.make_float2(local_sum[6], local_sum[7]),
+            ftz=True,
+        )
+        T.ptx.add_f32x2(
+            T.address_of(local_sum[0]),
+            T.cuda.make_float2(local_sum[0], local_sum[1]),
+            T.cuda.make_float2(local_sum[4], local_sum[5]),
+            ftz=True,
+        )
+        dst[tuple(dst_st)] = local_sum[0] + local_sum[1]
+    # fmt: on
+
+    return impl
+
+
+def _emit_reduction_local_thread_3input_maxmin(
+    dst_buffer_region: BufferRegion,
+    src_buffer_region: BufferRegion,
+    accum: bool,
+    reduce_op: ReduceOpType,
+    sctx: DispatchContext,
+) -> PrimFunc:
+    dst, src = dst_buffer_region.buffer, src_buffer_region.buffer
+    src_region, dst_region = src_buffer_region.region, dst_buffer_region.region
+    dtype = src.dtype
+
+    src_extent = [r.extent for r in src_region]
+    src_st = [r.min for r in src_region]
+    dst_st = [r.min for r in dst_region]
+
+    reduction_len = functools.reduce(operator.mul, src_extent, 1)
+
+    op_func = reduce_op_table[reduce_op]
+    reduce3_func = T.ptx.reduce3_max_f32 if reduce_op == ReduceOpType.MAX else T.ptx.reduce3_min_f32
+
+    src_base = src_st[0]
+    num_full_chunks = reduction_len // 8
+    remainder = reduction_len % 8
+    remainder_base = num_full_chunks * 8
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        temp = T.alloc_buffer([4], dtype, scope="local")
+        # First pass: process first 8 elements into 4 temps
+        for i in T.unroll(4):
+            if accum and i == 0:
+                temp[i] = reduce3_func(src[src_base + 2 * i], src[src_base + 2 * i + 1], dst[tuple(dst_st)])  # noqa: E501
+            else:
+                temp[i] = op_func(src[src_base + 2 * i], src[src_base + 2 * i + 1])
+
+        # Process remaining full chunks of 8
+        for outer in T.serial(num_full_chunks - 1):
+            for i in T.unroll(4):
+                temp[i] = reduce3_func(
+                    temp[i],
+                    src[src_base + 8 * (outer + 1) + 2 * i],
+                    src[src_base + 8 * (outer + 1) + 2 * i + 1],
+                )
+
+        # Process remainder elements (0 to 7 elements)
+        for i in T.serial(remainder):
+            temp[0] = op_func(temp[0], src[src_base + remainder_base + i])
+
+        # Final merge: combine 4 temps into result
+        dst[tuple(dst_st)] = op_func(temp[0], temp[1])
+        dst[tuple(dst_st)] = reduce3_func(dst[tuple(dst_st)], temp[2], temp[3])
+    # fmt: on
+
+    return impl
+
+
+def _sm100_packed_add_sum_impl(op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext):
+    op = TilePrimitiveCall.downcast(op)
+    return _emit_reduction_local_thread_packed_add_sum(op.output, op.input, op.accum, op_type, sctx)
+
+
+def _sm100_3input_maxmin_impl(op: TilePrimitiveCall, op_type: ReduceOpType, sctx: DispatchContext):
+    op = TilePrimitiveCall.downcast(op)
+    return _emit_reduction_local_thread_3input_maxmin(op.output, op.input, op.accum, op_type, sctx)
+
+
+_optimized_local_reduction_predicates = [
+    predicate("exec_scope", exec_scope_ok, expected_scopes=["thread"]),
+    predicate("local_scope", _local_scope_match),
+    predicate("dst_len", _dst_len_ok, expected_len=1),
+    predicate("src_ndim", _src_ndim_ok, expected_ndim=1),
+    predicate("dtype", _dtype_ok, expected_dtype="float32"),
+    predicate("sm_version", sm_version_ok, min_version=100),
+    predicate("reduction_len", _reduction_len_ok, min_len=8),
+]
+
+_optimized_impl_table = {
+    ReduceOpType.SUM: ("packed_add_sum", _sm100_packed_add_sum_impl),
+    ReduceOpType.MAX: ("3input_maxmin", _sm100_3input_maxmin_impl),
+    ReduceOpType.MIN: ("3input_maxmin", _sm100_3input_maxmin_impl),
+}
+
+
+# ---------------------------------------------------------------------------
+# Registration: SM100+ optimized local reduction (priority=20)
+# ---------------------------------------------------------------------------
+
+for op_name, op_type in [
+    ("sum", ReduceOpType.SUM),
+    ("max", ReduceOpType.MAX),
+    ("min", ReduceOpType.MIN),
+]:
+    variant_name, optimized_impl = _optimized_impl_table[op_type]
+
+    @register_dispatch(
+        op_name,
+        "cuda",
+        variant=variant_name,
+        priority=20,
+        when=_optimized_local_reduction_predicates,
+    )
+    def _optimized_dispatch(
+        op: TilePrimitiveCall, sctx: DispatchContext, _impl=optimized_impl, _op_type=op_type
+    ) -> PrimFunc:
+        op = TilePrimitiveCall.downcast(op)
+        return _impl(op, _op_type, sctx)

python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py

@@ -0,0 +1,262 @@
+# 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.
+
+"""Shared helpers for reduction operator dispatches on CUDA targets."""
+
+import functools
+import math
+import operator
+
+from tvm.arith.analyzer import Analyzer
+from tvm.script import tirx as T
+from tvm.tirx import BufferRegion
+from tvm.tirx.operator.tile_primitive import DispatchContext
+from tvm.tirx.operator.tile_primitive.common import ReduceOpType
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import match_scope
+
+reduce_op_table = {
+    ReduceOpType.SUM: lambda a, b: a + b,
+    ReduceOpType.MAX: T.max,
+    ReduceOpType.MIN: T.min,
+}
+
+
+def reduce_default_value_table(dtype):
+    return {
+        ReduceOpType.SUM: 0.0,
+        ReduceOpType.MAX: T.min_value(dtype),
+        ReduceOpType.MIN: T.max_value(dtype),
+    }
+
+
+def _reduction_args(
+    op: TilePrimitiveCall,
+) -> tuple[BufferRegion, BufferRegion, tuple[int, ...], bool, dict]:
+    """Parse ReduceOp -> (dst, src, reduce_axes, accum, config)."""
+    op = TilePrimitiveCall.downcast(op)
+    dst = op.output
+    src = op.input
+    reduce_axes = tuple(int(a) for a in op.reduce_axes)
+    accum = op.accum
+    config = op.config
+    return dst, src, reduce_axes, accum, config
+
+
+def _match_reduction_storage_scope(
+    op: TilePrimitiveCall, sctx: DispatchContext, expected_scope: list[str]
+) -> tuple[bool, str | None]:
+    """Check that dst and src scopes match one of the expected patterns."""
+    op = TilePrimitiveCall.downcast(op)
+    dst_scope = op.output.buffer.scope()
+    src_scope = op.input.buffer.scope()
+
+    ok = any(match_scope(dst_scope, p) and match_scope(src_scope, p) for p in expected_scope)
+    msg = f"storage scope mismatch: dst {dst_scope}, src {src_scope}; expected {expected_scope}"
+    return (ok, None if ok else msg)
+
+
+def _analyze_axes(src_ndim: int, reduce_axes: tuple[int, ...]) -> tuple[list[int], list[int]]:
+    """Normalize negative axes -> (reduce_dim_set, spatial_dim_list)."""
+    reduce_dims = set()
+    for ax in reduce_axes:
+        a = ax if ax >= 0 else ax + src_ndim
+        assert 0 <= a < src_ndim, f"reduce axis {ax} out of range for ndim={src_ndim}"
+        reduce_dims.add(a)
+    spatial_dims = [d for d in range(src_ndim) if d not in reduce_dims]
+    return sorted(reduce_dims), spatial_dims
+
+
+def _analyze_layout_dims(layout, shape):
+    """layout.group(shape) -> decompose each dim into thread/local iters.
+
+    Returns list of per-dim (thread_extent, local_extent, thread_strides):
+        thread_extent = product of thread iter extents in this dim
+        local_extent  = product of local iter extents in this dim
+        thread_strides = list of (stride, extent) for thread iters in this dim
+    """
+    grouped, seps = layout.group(list(shape))
+    result = []
+    for d in range(len(shape)):
+        shard_range = list(range(seps[d], seps[d + 1]))
+        thread_extent = 1
+        local_extent = 1
+        thread_strides = []
+        for s_idx in shard_range:
+            it = grouped.shard[s_idx]
+            if it.axis.is_thread():
+                thread_extent *= it.extent
+                thread_strides.append((it.stride, it.extent))
+            else:
+                local_extent *= it.extent
+        result.append((thread_extent, local_extent, thread_strides))
+    return result
+
+
+def _compute_shuffle_masks(dim_info, reduce_dims: set[int]) -> list[int]:
+    """From reduction dims' thread iter (stride, extent) pairs, compute XOR masks.
+
+    For each thread iter in a reduction dim:
+        masks += [stride * 2^i for i in range(log2(extent))]
+    Sorted ascending.
+    """
+    masks = []
+    for d in reduce_dims:
+        _, _, thread_strides = dim_info[d]
+        for stride, extent in thread_strides:
+            ext_int = int(extent) if hasattr(extent, "__int__") else extent
+            n_bits = int(math.log2(ext_int))
+            for i in range(n_bits):
+                stride_int = int(stride) if hasattr(stride, "__int__") else stride
+                masks.append(stride_int * (1 << i))
+    masks.sort()
+    return masks
+
+
+def _build_local_dim_map(layout, buffer_shape):
+    """Map original dim index to position in get_local_region output (None if pure-thread)."""
+    grouped, seps = layout.group(list(buffer_shape))
+    dim_map = {}
+    local_pos = 0
+    for d in range(len(buffer_shape)):
+        shard_range = list(range(seps[d], seps[d + 1]))
+        has_local = any(not grouped.shard[s].axis.is_thread() for s in shard_range)
+        if has_local:
+            dim_map[d] = local_pos
+            local_pos += 1
+        else:
+            dim_map[d] = None
+    return dim_map
+
+
+def _validate_reduction_layout(
+    src_layout, dst_layout, src_shape, dst_shape, reduce_dims: list[int]
+) -> tuple[bool, str | None]:
+    """Validate that spatial dims of src/dst have matching thread+local structure,
+    and that reduction dims in dst have local_extent == 1.
+    """
+    src_dim_info = _analyze_layout_dims(src_layout, src_shape)
+    dst_dim_info = _analyze_layout_dims(dst_layout, dst_shape)
+    analyzer = Analyzer()
+
+    # Spatial dims: src/dst must match in both thread and local extents.
+    # Reduce dims: src/dst thread extent must match, and dst local extent must be 1.
+
+    # get expected simplified dst layout
+    expected_dst_dim = []
+    for src_idx in range(len(src_shape)):
+        if analyzer.can_prove_equal(src_dim_info[src_idx][0], 1) and analyzer.can_prove_equal(
+            src_dim_info[src_idx][1], 1
+        ):
+            continue  # skip if extent=1
+        if src_idx in reduce_dims:  # reduce dims
+            if not analyzer.can_prove_equal(src_dim_info[src_idx][0], 1):
+                expected_dst_dim.append((src_dim_info[src_idx][0], 1))
+        else:  # spatial dims
+            expected_dst_dim.append((src_dim_info[src_idx][0], src_dim_info[src_idx][1]))
+
+    # check dst layout
+    check_idx = 0
+    for dst_idx in range(len(dst_shape)):
+        if analyzer.can_prove_equal(dst_dim_info[dst_idx][0], 1) and analyzer.can_prove_equal(
+            dst_dim_info[dst_idx][1], 1
+        ):
+            continue
+        if not (
+            analyzer.can_prove_equal(dst_dim_info[dst_idx][0], expected_dst_dim[check_idx][0])
+            and analyzer.can_prove_equal(dst_dim_info[dst_idx][1], expected_dst_dim[check_idx][1])
+        ):
+            return False, "mismatch dst/src layout for reduction"
+        check_idx += 1
+    if check_idx != len(expected_dst_dim):
+        return False, "mismatch dst/src layout for reduction"
+    return True, None
+
+
+def build_src_indices(spa_fused, red_fused, spatial_dims, reduce_dims, src_extent, src_st):
+    """Combine spatial and reduction indices into full src index tuple."""
+
+    # Build index helpers that work with the explicit axis split
+    def get_spatial_or_reduction_src_indices(spa_or_red_fused, is_spatial):
+        dims = spatial_dims if is_spatial else reduce_dims
+        spa_extents = [src_extent[d] for d in dims]
+        indices = []
+        rem = spa_or_red_fused
+        for e in reversed(spa_extents):
+            indices.append(rem % e)
+            rem //= e
+        indices.reverse()
+        return [idx + src_st[d] for idx, d in zip(indices, dims)]
+
+    spa_vals = get_spatial_or_reduction_src_indices(spa_fused, is_spatial=True)
+    red_vals = get_spatial_or_reduction_src_indices(red_fused, is_spatial=False)
+    full = [None] * len(src_extent)
+    for i, d in enumerate(spatial_dims):
+        full[d] = spa_vals[i]
+    for i, d in enumerate(reduce_dims):
+        full[d] = red_vals[i]
+    return full
+
+
+_REDUCE_OP_TO_STR = {ReduceOpType.SUM: "sum", ReduceOpType.MAX: "max", ReduceOpType.MIN: "min"}
+
+
+def _dtype_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_dtype: str):
+    op = TilePrimitiveCall.downcast(op)
+    dtype = op.input.buffer.dtype
+    ok = dtype == expected_dtype
+    return (ok, None if ok else f"dtype {dtype} != {expected_dtype}")
+
+
+def _reduction_len_ok(op: TilePrimitiveCall, sctx: DispatchContext, min_len: int):
+    op = TilePrimitiveCall.downcast(op)
+    src_extent = [r.extent for r in op.input.region]
+    reduction_len = functools.reduce(operator.mul, src_extent, 1)
+    ok = reduction_len >= min_len
+    return (ok, None if ok else f"reduction_len {reduction_len} < {min_len}")
+
+
+def _dst_len_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_len: int):
+    op = TilePrimitiveCall.downcast(op)
+    dst_extent = [r.extent for r in op.output.region]
+    dst_len = functools.reduce(operator.mul, dst_extent, 1)
+    ok = dst_len == expected_len
+    return (ok, None if ok else f"dst_len {dst_len} != {expected_len}")
+
+
+def _src_ndim_ok(op: TilePrimitiveCall, sctx: DispatchContext, expected_ndim: int):
+    op = TilePrimitiveCall.downcast(op)
+    src_extent = [r.extent for r in op.input.region]
+    ok = len(src_extent) == expected_ndim
+    return (ok, None if ok else f"src ndim {len(src_extent)} != {expected_ndim}")
+
+
+def _local_scope_match(op: TilePrimitiveCall, sctx: DispatchContext):
+    op = TilePrimitiveCall.downcast(op)
+    src, dst = op.input.buffer, op.output.buffer
+    ok = all(
+        [
+            src.scope() == "local",
+            dst.scope() == "local",
+            src.dtype == dst.dtype,
+            sctx.is_target("cuda"),
+        ]
+    )
+    if not ok:
+        return (False, "src/dst must be local scope with matching dtype on CUDA")
+    return (True, None)

python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py

@@ -0,0 +1,117 @@
+# 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.
+
+"""TMA (Tensor Memory Accelerator) utilities for CUDA op dispatches."""
+
+import copy
+from enum import Enum
+
+import tvm
+from tvm.arith.analyzer import Analyzer
+from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout
+
+
+class SwizzleMode(Enum):
+    """The swizzle mode of the TMA."""
+
+    SWIZZLE_NONE = 0
+    SWIZZLE_32B_ATOM = 1
+    SWIZZLE_64B_ATOM = 2
+    SWIZZLE_128B_ATOM = 3
+
+
+def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> SwizzleLayout:
+    """Generate the MMA-compatible shared-memory atom layout."""
+    bits = tvm.DataType(dtype).bits
+    if isinstance(swizzle_mode, int):
+        swizzle_mode = SwizzleMode(swizzle_mode)
+    return SwizzleLayout(
+        per_element=(128 // bits).bit_length() - 1, swizzle_len=swizzle_mode.value, atom_len=3
+    )
+
+
+def mma_atom_shape(dtype: str, swizzle_mode: SwizzleMode | int, shape: list[int] | None = None):
+    """Generate the MMA-compatible shared-memory atom shape."""
+    bits = tvm.DataType(dtype).bits
+    if isinstance(swizzle_mode, int):
+        swizzle_mode = SwizzleMode(swizzle_mode)
+    atom_shape = {
+        SwizzleMode.SWIZZLE_32B_ATOM: [8, 256],
+        SwizzleMode.SWIZZLE_64B_ATOM: [8, 512],
+        SwizzleMode.SWIZZLE_128B_ATOM: [8, 1024],
+    }[swizzle_mode]
+    atom_shape[-1] //= bits
+    if shape is None:
+        return atom_shape
+    atom_shape = [1] * (len(shape) - len(atom_shape)) + atom_shape
+    return atom_shape
+
+
+def mma_shared_layout(dtype: str, swizzle_mode: SwizzleMode | int, shape) -> Layout:
+    """Generate the MMA-compatible shared-memory layout for shape and dtype.
+
+    It uses a default tiling strategy to tile the TMA atom layout into the shared memory.
+    """
+    if isinstance(swizzle_mode, int):
+        swizzle_mode = SwizzleMode(swizzle_mode)
+    if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
+        return TileLayout(S[tuple(shape)]).canonicalize()
+    atom_shape = mma_atom_shape(dtype, swizzle_mode, shape)
+    layout = mma_atom_layout(dtype, swizzle_mode)
+    tile_to_shape = copy.copy(atom_shape)
+    tile_to_shape[-2] = shape[-2]
+    return layout.tile_to(tile_to_shape, atom_shape).tile_to(shape, tile_to_shape).canonicalize()
+
+
+# Backward-compatible aliases kept during the alloc_mma migration.
+tma_atom_layout = mma_atom_layout
+tma_atom_shape = mma_atom_shape
+tma_shared_layout = mma_shared_layout
+
+
+def tma_atom_compatible(dst_shape, dst_st, dst_extent, atom_shape):
+    """Check if the copy region in dst is compatible with the TMA atom shape."""
+    analyzer = Analyzer()
+    for i, _ in enumerate(dst_st):
+        if any(
+            not analyzer.can_prove_equal(x % atom_shape[i], 0)
+            for x in [dst_shape[i], dst_st[i], dst_extent[i]]
+        ):
+            return False
+    return True
+
+
+def get_swizzle_mode_from_layout(layout: Layout) -> SwizzleMode | None:
+    """Extract swizzle mode from a shared memory layout."""
+    if isinstance(layout, ComposeLayout):
+        swizzle = layout.swizzle  # SwizzleLayout is named 'swizzle' in ComposeLayout
+        swizzle_len = swizzle.swizzle_len
+    elif isinstance(layout, SwizzleLayout):
+        swizzle_len = layout.swizzle_len
+    elif isinstance(layout, TileLayout):
+        # TileLayout without SwizzleLayout means no swizzle (mode 0)
+        return SwizzleMode.SWIZZLE_NONE
+    else:
+        return None
+
+    # Map swizzle_len to SwizzleMode
+    return {
+        0: SwizzleMode.SWIZZLE_NONE,
+        1: SwizzleMode.SWIZZLE_32B_ATOM,
+        2: SwizzleMode.SWIZZLE_64B_ATOM,
+        3: SwizzleMode.SWIZZLE_128B_ATOM,
+    }.get(swizzle_len)

python/tvm/backend/cuda/script.py

@@ -0,0 +1,573 @@
+# 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.
+"""CUDA TVMScript namespaces."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from tvm.backend.cuda import op as _cuda_op
+from tvm.tirx import Buffer
+from tvm.tirx import op as _tir_op
+from tvm.tirx.script.builder.ir import _dtype_forward, _op_wrapper
+
+# pylint: disable=protected-access
+
+
+def _ptx_ldg32(reg, guard, addr, local_addr):
+    if isinstance(addr, Buffer):
+        addr = addr[0]
+    return _tir_op.call_intrin(reg.ty, "tirx.ptx.ldg32", reg, guard, addr, local_addr)
+
+
+_ptx_ldg32.__tir_op_name__ = "ptx.ldg32"
+
+
+class PTXNamespace:
+    """The PTX instruction submodule."""
+
+    def __init__(self):
+        self.ldg32 = _ptx_ldg32
+        self.ldmatrix = _dtype_forward(_cuda_op.ptx_ldmatrix)
+        # Apache-compatible variant. Same lowered intrinsic as
+        # ``ldmatrix`` but accepts the historical ``(trans, num, dtype,
+        # local_ptr, local_offset, smem_ptr, smem_offset)`` form. Coexists
+        # with the fork-native version so upstream-derived tests keep
+        # working without rewriting their tirx code.
+        self.ldmatrix_legacy = _dtype_forward(_cuda_op.ptx_ldmatrix_legacy)
+        self.stmatrix = _op_wrapper(_cuda_op.ptx_stmatrix)
+        self.setmaxnreg: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_setmaxnreg)
+        self.elect_sync: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_elect_sync)
+        self.clc_try_cancel = _op_wrapper(_cuda_op.ptx_clc_try_cancel)
+        self.clc_query_cancel = _op_wrapper(_cuda_op.ptx_clc_query_cancel)
+        self.fetch_register: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_fetch_register)
+        self.ld = _op_wrapper(_cuda_op.ptx_ld)
+        self.ld_acquire = _op_wrapper(_cuda_op.ptx_ld_acquire)
+        self.ld_relaxed = _op_wrapper(_cuda_op.ptx_ld_relaxed)
+        self.ld_volatile = _op_wrapper(_cuda_op.ptx_ld_volatile)
+        self.ld_mmio = _op_wrapper(_cuda_op.ptx_ld_mmio)
+        self.ld_global_acquire = _op_wrapper(_cuda_op.ptx_ld_global_acquire)
+        self.red_scalar = _op_wrapper(_cuda_op.ptx_red_scalar)
+        self.atom_scalar = _op_wrapper(_cuda_op.ptx_atom_scalar)
+        self.prefetch_tensormap = _op_wrapper(_cuda_op.ptx_prefetch_tensormap)
+        self.mbarrier_test_wait_parity = _op_wrapper(_cuda_op.ptx_mbarrier_test_wait_parity)
+        self.cp_async_bulk_g2s_cta = _op_wrapper(_cuda_op.ptx_cp_async_bulk_g2s_cta)
+        self.cp_async_bulk_g2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_g2s_cluster)
+        self.cp_async_bulk_s2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2s_cluster)
+        self.cp_async_bulk_s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2g)
+        self.st = _op_wrapper(_cuda_op.ptx_st)
+        self.st_relaxed = _op_wrapper(_cuda_op.ptx_st_relaxed)
+        self.st_release = _op_wrapper(_cuda_op.ptx_st_release)
+        self.st_volatile = _op_wrapper(_cuda_op.ptx_st_volatile)
+        self.st_mmio = _op_wrapper(_cuda_op.ptx_st_mmio)
+        self.st_bulk = _op_wrapper(_cuda_op.ptx_st_bulk)
+        self.fns_b32 = _op_wrapper(_cuda_op.ptx_fns_b32)
+        self.add_rn_f32_bf16 = _op_wrapper(_cuda_op.ptx_add_rn_f32_bf16)
+        self.mapa = _op_wrapper(_cuda_op.ptx_mapa)
+        self.map_shared_rank = _op_wrapper(_cuda_op.ptx_map_shared_rank)
+        self.any_sync = _op_wrapper(_cuda_op.ptx_any_sync)
+        # Math operations
+        self.exp2 = _op_wrapper(_cuda_op.ptx_exp2)
+        self.rcp = _op_wrapper(_cuda_op.ptx_rcp)
+        self.reduce3_min_f32 = _op_wrapper(_cuda_op.ptx_reduce3_min_f32)
+        self.reduce3_max_f32 = _op_wrapper(_cuda_op.ptx_reduce3_max_f32)
+        # add/sub/mul/fma DPS form: (d_addr, a, b[, c], *, rounding, ftz[, sat])
+        self.add_f32 = _op_wrapper(_cuda_op.ptx_add_f32)
+        self.add_f32x2 = _op_wrapper(_cuda_op.ptx_add_f32x2)
+        self.add_f64 = _op_wrapper(_cuda_op.ptx_add_f64)
+        self.sub_f32 = _op_wrapper(_cuda_op.ptx_sub_f32)
+        self.sub_f32x2 = _op_wrapper(_cuda_op.ptx_sub_f32x2)
+        self.sub_f64 = _op_wrapper(_cuda_op.ptx_sub_f64)
+        self.mul_f32 = _op_wrapper(_cuda_op.ptx_mul_f32)
+        self.mul_f32x2 = _op_wrapper(_cuda_op.ptx_mul_f32x2)
+        self.mul_f64 = _op_wrapper(_cuda_op.ptx_mul_f64)
+        self.fma_f32 = _op_wrapper(_cuda_op.ptx_fma_f32)
+        self.fma_f32x2 = _op_wrapper(_cuda_op.ptx_fma_f32x2)
+        self.fma_f64 = _op_wrapper(_cuda_op.ptx_fma_f64)
+        self.max_f32 = _op_wrapper(_cuda_op.ptx_max_f32)
+        self.mma = MmaNamespace()
+        self.cp_async = CpAsyncNamespace()
+        self.wgmma = WgmmaNamespace()
+        self.mbarrier = MbarrierNamespace()
+        self.tcgen05 = Tcgen05Namespace()
+        self.bar = BarNamespace()
+        self.barrier = BarrierNamespace()
+        self.fence = FenceNamespace()
+        self.griddepcontrol = GriddepcontrolNamespace()
+
+
+class MmaNamespace:
+    """The MMA instruction submodule."""
+
+    def __init__(self):
+        self.sp = _dtype_forward(_cuda_op.ptx_mma_sp)
+        # Apache-compatible variant of ptx_mma. Coexists with the
+        # fork-native ``__call__`` form (``T.ptx.mma(...)``).
+        self.legacy = _dtype_forward(_cuda_op.ptx_mma_legacy)
+        # __call__ corresponds to ptx_mma
+        self.__tir_call_op_name__ = "ptx_mma"
+
+    def __call__(self, *args, **kwds):
+        return _dtype_forward(_cuda_op.ptx_mma)(*args, **kwds)
+
+
+class CpAsyncNamespace:
+    """The CpAsync instruction submodule."""
+
+    def __init__(self):
+        self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_commit_group)
+        self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_wait_group)
+        # Legacy variant: takes (dst_ptr, dst_offset, src_ptr, src_offset,
+        # cp_size). Offsets are folded into the pointers; coexists with
+        # the fork-native ``__call__`` form.
+        self.legacy = _dtype_forward(_cuda_op.ptx_cp_async_legacy)
+        self.bulk = CpAsyncBulkNamespace()
+        self.mbarrier = CpAsyncMbarrierNamespace()
+
+    def __call__(self, *args, **kwds):
+        # Accept the legacy 6-arg form ``(elem_dtype, dst, dst_off, src,
+        # src_off, cp_size)`` that the printer round-trips for the raw
+        # ``tirx.ptx.cp_async`` Call emitted by
+        # ``tvm.backend.cuda.transform.InjectPTXAsyncCopy``. The pass-emitted
+        # Call has 5 args (no ``tvm_access_ptr`` fold) and a
+        # per-element-dtype Call.dtype, so build it directly.
+        if len(args) == 6 and isinstance(args[0], str) and "dtype" not in kwds:
+            import tvm
+
+            elem_dtype, dst, dst_off, src, src_off, cp_size = args
+            return tvm.ir.Call(
+                tvm.ir.Op.get("tirx.ptx.cp_async_raw"),
+                [dst, dst_off, src, src_off, cp_size],
+                ret_ty=tvm.ir.PrimType(elem_dtype),
+            )
+        return _dtype_forward(_cuda_op.ptx_cp_async)(*args, **kwds)
+
+    # __call__ corresponds to ptx_cp_async
+    __tir_call_op_name__ = "ptx_cp_async"
+
+
+class CpAsyncBulkNamespace:
+    """The CpAsyncBulk instruction submodule."""
+
+    def __init__(self):
+        self.commit_group = _op_wrapper(_cuda_op.ptx_cp_async_bulk_commit_group)
+        self.wait_group = _op_wrapper(_cuda_op.ptx_cp_async_bulk_wait_group)
+        self.tensor = CpAsyncBulkTensorNamespace()
+        self.s2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_shared_to_cluster)
+
+    def __call__(self, *args, **kwds):
+        return _dtype_forward(_cuda_op.ptx_cp_async_bulk)(*args, **kwds)
+
+    # __call__ corresponds to ptx_cp_async_bulk
+    __tir_call_op_name__ = "ptx_cp_async_bulk"
+
+
+class CpAsyncBulkTensorNamespace:
+    """The CpAsyncBulkTensor instruction submodule."""
+
+    def __init__(self):
+        self.g2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster)
+        self.g2c_tile_gather4 = _op_wrapper(
+            _cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster
+        )
+        self.s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global)
+        self.s2g_reduce = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global_reduce)
+        self.g2c_prefetch = _op_wrapper(
+            _cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch
+        )
+
+    @staticmethod
+    def g2c_bar_addr(
+        dim,
+        dst_ptr,
+        bar_addr,
+        tensormap_addr,
+        cta_mask,
+        cta_group,
+        cache_hint,
+        *coords,
+        cache_policy=None,
+    ):
+        _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP)
+        cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy)
+        return _tir_op.call_intrin(
+            "",
+            "tirx.ptx.cp_async_bulk_tensor_global_to_cluster",
+            dim,
+            dst_ptr,
+            bar_addr,
+            tensormap_addr,
+            cta_mask,
+            cta_group,
+            cache_policy,
+            int(has_cache_policy),
+            1,
+            *coords,
+        )
+
+    @staticmethod
+    def g2c_tile_gather4_bar_addr(
+        dim,
+        dst_ptr,
+        bar_addr,
+        tensormap_addr,
+        cta_mask,
+        cta_group,
+        cache_hint,
+        *coords,
+        cache_policy=None,
+    ):
+        _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP)
+        cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy)
+        return _tir_op.call_intrin(
+            "",
+            "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster",
+            dim,
+            dst_ptr,
+            bar_addr,
+            tensormap_addr,
+            cta_mask,
+            cta_group,
+            cache_policy,
+            int(has_cache_policy),
+            1,
+            *coords,
+        )
+
+
+class CpAsyncMbarrierNamespace:
+    """The CpAsyncMbarrier instruction submodule."""
+
+    def __init__(self):
+        self.arrive = _op_wrapper(_cuda_op.ptx_cp_async_mbarrier_arrive)
+
+
+class WgmmaNamespace:
+    """The WGMMA instruction submodule."""
+
+    def __init__(self):
+        self.fence: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_wgmma_fence)
+        self.commit_group = _op_wrapper(_cuda_op.ptx_wgmma_commit_group)
+        self.wait_group = _op_wrapper(_cuda_op.ptx_wgmma_wait_group)
+        self.noop_barrier = _op_wrapper(_cuda_op.ptx_wgmma_noop_barrier)
+        self.mma_async = WgmmaMmaAsyncNamespace()
+        self.encode_matrix_descriptor = _op_wrapper(_cuda_op.ptx_wgmma_encode_matrix_descriptor)
+
+
+class WgmmaMmaAsyncNamespace:
+    """The WGMMA MMAAsync instruction submodule."""
+
+    def __init__(self):
+        self.ss = _op_wrapper(_cuda_op.ptx_wgmma_mma_async_ss)
+        self.rs = _op_wrapper(_cuda_op.ptx_wgmma_mma_async_rs)
+
+
+class MbarrierNamespace:
+    """The Mbarrier instruction submodule."""
+
+    def __init__(self):
+        self.init = _op_wrapper(_cuda_op.ptx_mbarrier_init)
+        self.try_wait = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait)
+        self.try_wait_once = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_once)
+        self.try_wait_acquire_cluster = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_acquire_cluster)
+        self.arrive = MbarrierArriveNamespace()
+
+
+class MbarrierArriveNamespace:
+    """The Mbarrier Arrive instruction submodule."""
+
+    def __init__(self):
+        self.expect_tx = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_expect_tx)
+        self.cluster_count = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_cluster_count)
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.ptx_mbarrier_arrive)(*args, **kwds)
+
+    # __call__ corresponds to ptx_mbarrier_arrive
+    __tir_call_op_name__ = "ptx_mbarrier_arrive"
+
+
+class Tcgen05Namespace:
+    """The Tcgen05 instruction submodule."""
+
+    def __init__(self):
+        self.alloc = _op_wrapper(_cuda_op.ptx_tcgen05_alloc)
+        self.dealloc = _op_wrapper(_cuda_op.ptx_tcgen05_dealloc)
+        self.relinquish_alloc_permit = _op_wrapper(_cuda_op.ptx_tcgen05_relinquish_alloc_permit)
+        self.encode_matrix_descriptor = _op_wrapper(_cuda_op.ptx_tcgen05_encode_matrix_descriptor)
+        self.encode_instr_descriptor = _op_wrapper(_cuda_op.ptx_tcgen05_encode_instr_descriptor)
+        self.encode_instr_descriptor_block_scaled = _op_wrapper(
+            _cuda_op.ptx_tcgen05_encode_instr_descriptor_block_scaled
+        )
+        self.ld = _op_wrapper(_cuda_op.ptx_tcgen05_ld)
+        self.st = _op_wrapper(_cuda_op.ptx_tcgen05_st)
+        self.cp = _op_wrapper(_cuda_op.ptx_tcgen05_cp)
+        self.shift = _op_wrapper(_cuda_op.ptx_tcgen05_shift)
+        self.commit = _op_wrapper(_cuda_op.ptx_tcgen05_commit)
+        self.wait = Tcgen05WaitNamespace()
+        self.mma = Tcgen05MmaNamespace()
+        self.fence = Tcgen05FenceNamespace()
+
+
+class Tcgen05FenceNamespace:
+    """The Tcgen05 Fence instruction submodule."""
+
+    def __init__(self):
+        self.before_thread_sync = _op_wrapper(_cuda_op.ptx_tcgen05_fence_before_thread_sync)
+        self.after_thread_sync = _op_wrapper(_cuda_op.ptx_tcgen05_fence_after_thread_sync)
+
+
+class Tcgen05MmaNamespace:
+    """The Tcgen05 MMA instruction submodule."""
+
+    def __init__(self):
+        self.block_scale = _op_wrapper(_cuda_op.ptx_tcgen05_mma_block_scale)
+        self.sp = Tcgen05MmaSpNamespace()
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.ptx_tcgen05_mma)(*args, **kwds)
+
+    # __call__ corresponds to ptx_tcgen05_mma
+    __tir_call_op_name__ = "ptx_tcgen05_mma"
+
+
+class Tcgen05MmaSpNamespace:
+    """Tcgen05 Sparse MMA instruction submodule."""
+
+    def __init__(self):
+        self.block_scale = _op_wrapper(_cuda_op.ptx_tcgen05_mma_sp_block_scale)
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.ptx_tcgen05_mma_sp)(*args, **kwds)
+
+    # __call__ corresponds to ptx_tcgen05_mma_sp
+    __tir_call_op_name__ = "ptx_tcgen05_mma_sp"
+
+
+class Tcgen05WaitNamespace:
+    """The Tcgen05 Wait instruction submodule."""
+
+    def __init__(self):
+        self.ld = _op_wrapper(_cuda_op.ptx_tcgen05_wait_ld)
+        self.st = _op_wrapper(_cuda_op.ptx_tcgen05_wait_st)
+
+
+class BarNamespace:
+    """The Bar instruction submodule."""
+
+    def __init__(self):
+        self.arrive = _op_wrapper(_cuda_op.ptx_bar_arrive)
+        self.sync = _op_wrapper(_cuda_op.ptx_bar_sync)
+
+
+class BarrierNamespace:
+    """The Barrier instruction submodule."""
+
+    def __init__(self):
+        self.cluster = BarrierClusterNamespace()
+
+
+class BarrierClusterNamespace:
+    """The BarrierCluster instruction submodule."""
+
+    def __init__(self):
+        self.arrive = _op_wrapper(_cuda_op.ptx_barrier_cluster_arrive)
+        self.wait = _op_wrapper(_cuda_op.ptx_barrier_cluster_wait)
+
+
+class FenceNamespace:
+    """PTX fence instruction submodule."""
+
+    def __init__(self):
+        self.proxy_async = _op_wrapper(_cuda_op.ptx_fence_proxy_async)
+        self.mbarrier_init = _op_wrapper(_cuda_op.ptx_fence_mbarrier_init)
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.ptx_fence)(*args, **kwds)
+
+    __tir_call_op_name__ = "ptx_fence"
+
+
+class GriddepcontrolNamespace:
+    """PTX griddepcontrol instruction submodule (sm_90+)."""
+
+    def __init__(self):
+        self.wait = _op_wrapper(_cuda_op.ptx_griddepcontrol_wait)
+        self.launch_dependents = _op_wrapper(_cuda_op.ptx_griddepcontrol_launch_dependents)
+
+
+class CUDANamespace:
+    """The CUDA intrinsics submodule."""
+
+    def __init__(self):
+        self.atomic_add = _op_wrapper(_cuda_op.cuda_atomic_add)
+        self.thread_fence = _op_wrapper(_cuda_op.cuda_thread_fence)
+        self.warpgroup_sync = _op_wrapper(_cuda_op.cuda_warpgroup_sync)
+        self.warp_sync = _op_wrapper(_cuda_op.cuda_warp_sync)
+        self.warp_reduce = _op_wrapper(_cuda_op.cuda_warp_reduce)
+        self.warp_sum = _op_wrapper(_cuda_op.cuda_warp_sum)
+        self.warp_max = _op_wrapper(_cuda_op.cuda_warp_max)
+        self.warp_min = _op_wrapper(_cuda_op.cuda_warp_min)
+        self.cta_reduce = _op_wrapper(_cuda_op.cuda_cta_reduce)
+        self.cta_sum = _op_wrapper(_cuda_op.cuda_cta_sum)
+        self.cta_max = _op_wrapper(_cuda_op.cuda_cta_max)
+        self.cta_min = _op_wrapper(_cuda_op.cuda_cta_min)
+        self.cta_sync = _op_wrapper(_cuda_op.cuda_cta_sync)
+        self.grid_sync = _op_wrapper(_cuda_op.cuda_grid_sync)
+        self.cluster_sync = _op_wrapper(_cuda_op.cuda_cluster_sync)
+        self.thread_rank = _op_wrapper(_cuda_op.cuda_thread_rank)
+        self.trap_when_assert_failed = _op_wrapper(_cuda_op.cuda_trap_when_assert_failed)
+        self.runtime_instr_desc = _op_wrapper(_cuda_op.cuda_runtime_instr_desc)
+        self.half2float = _op_wrapper(_cuda_op.cuda_half2float)
+        self.bfloat162float = _op_wrapper(_cuda_op.cuda_bfloat162float)
+        self.float22half2 = _op_wrapper(_cuda_op.cuda_float22half2)
+        self.half8tofloat8 = _op_wrapper(_cuda_op.cuda_half8tofloat8)
+        self.float8tohalf8 = _op_wrapper(_cuda_op.cuda_float8tohalf8)
+        self.syncthreads_and = _op_wrapper(_cuda_op.cuda_syncthreads_and)
+        self.syncthreads_or = _op_wrapper(_cuda_op.cuda_syncthreads_or)
+        self.nano_sleep = _op_wrapper(_cuda_op.cuda_nano_sleep)
+        self.atomic_cas = _op_wrapper(_cuda_op.cuda_atomic_cas)
+        self.func_call = _op_wrapper(_cuda_op.cuda_func_call)
+        self.printf = _op_wrapper(_cuda_op.cuda_printf)
+        self.ldg = _op_wrapper(_cuda_op.cuda_ldg)
+        self.get_tmem_addr = _op_wrapper(_cuda_op.cuda_get_tmem_addr)
+        self.cvta_generic_to_shared = _op_wrapper(_cuda_op.cuda_cvta_generic_to_shared)
+        self.smem_addr_from_uint64 = _op_wrapper(_cuda_op.cuda_smem_addr_from_uint64)
+        self.sm100_tma_2sm_mbarrier_addr = _op_wrapper(_cuda_op.cuda_sm100_tma_2sm_mbarrier_addr)
+        self.uint_as_float = _op_wrapper(_cuda_op.cuda_uint_as_float)
+        self.float_as_uint = _op_wrapper(_cuda_op.cuda_float_as_uint)
+        self.ballot_sync = _op_wrapper(_cuda_op.cuda_ballot_sync)
+        self.ffs_u32 = _op_wrapper(_cuda_op.cuda_ffs_u32)
+        self.reduce_add_sync_u32 = _op_wrapper(_cuda_op.cuda_reduce_add_sync_u32)
+        self.reduce_min_sync_u32 = _op_wrapper(_cuda_op.cuda_reduce_min_sync_u32)
+        self.clock64 = _op_wrapper(_cuda_op.cuda_clock64)
+        self.make_float2 = _op_wrapper(_cuda_op.cuda_make_float2)
+        self.float2_x = _op_wrapper(_cuda_op.cuda_float2_x)
+        self.float2_y = _op_wrapper(_cuda_op.cuda_float2_y)
+        self.fmul2_rn = _op_wrapper(_cuda_op.cuda_fmul2_rn)
+        self.fadd2_rn = _op_wrapper(_cuda_op.cuda_fadd2_rn)
+        self.float22bfloat162_rn = _op_wrapper(_cuda_op.cuda_float22bfloat162_rn)
+        self.float22bfloat162_rn_from_float2 = _op_wrapper(
+            _cuda_op.cuda_float22bfloat162_rn_from_float2
+        )
+        self.bfloat1622float2 = _op_wrapper(_cuda_op.cuda_bfloat1622float2)
+        self.hmin2 = _op_wrapper(_cuda_op.cuda_hmin2)
+        self.hmax2 = _op_wrapper(_cuda_op.cuda_hmax2)
+        self.fp8x4_e4m3_from_float4 = _op_wrapper(_cuda_op.cuda_fp8x4_e4m3_from_float4)
+        self.timer_init = _op_wrapper(_cuda_op.timer_init_cuda)
+        self.timer_start = _op_wrapper(_cuda_op.timer_start_cuda)
+        self.timer_end = _op_wrapper(_cuda_op.timer_end_cuda)
+        self.timer_finalize = _op_wrapper(_cuda_op.timer_finalize_cuda)
+        self.mma_store = _dtype_forward(_cuda_op.mma_store)
+        self.mma_fill = _dtype_forward(_cuda_op.mma_fill)
+        self.mma_store_legacy = _dtype_forward(_cuda_op.mma_store_legacy)
+        self.mma_fill_legacy = _dtype_forward(_cuda_op.mma_fill_legacy)
+        setattr(self, "__shfl_sync", self._shfl_sync)
+        setattr(self, "__shfl_up_sync", self._shfl_up_sync)
+        setattr(self, "__shfl_down_sync", self._shfl_down_sync)
+        setattr(self, "__shfl_xor_sync", self._shfl_xor_sync)
+        setattr(self, "__activemask", self._activemask)
+
+    @staticmethod
+    def _shfl_sync(mask, var, lane, width):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_sync", mask, var, lane, width)
+
+    @staticmethod
+    def _shfl_up_sync(mask, var, delta, width):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_up_sync", mask, var, delta, width)
+
+    @staticmethod
+    def _shfl_down_sync(mask, var, delta, width):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_down_sync", mask, var, delta, width)
+
+    @staticmethod
+    def _shfl_xor_sync(mask, var, lane_mask, width):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.cuda.__shfl_xor_sync", mask, var, lane_mask, width)
+
+    @staticmethod
+    def _activemask():
+        return _tir_op.call_intrin("uint32", "tirx.cuda.__activemask")
+
+
+class NVSHMEMNamespace:
+    """The NVSHMEM intrinsics submodule."""
+
+    def __init__(self):
+        self.my_pe = _op_wrapper(_cuda_op.nvshmem_my_pe)
+        self.n_pes = _op_wrapper(_cuda_op.nvshmem_n_pes)
+        self.signal_op = _op_wrapper(_cuda_op.nvshmem_signal_op)
+        self.wait_until = _op_wrapper(_cuda_op.nvshmem_wait_until)
+        self.quiet = _op_wrapper(_cuda_op.nvshmem_quiet)
+        self.fence = _op_wrapper(_cuda_op.nvshmem_fence)
+        self.barrier_all = _op_wrapper(_cuda_op.nvshmem_barrier_all)
+        self.getmem_nbi = NVSHMEMGetMemNBINamespace()
+        self.putmem_nbi = NVSHMEMPutMemNBINamespace()
+        self.putmem_signal_nbi = NVSHMEMPutMemSignalNBINamespace()
+
+
+class NVSHMEMGetMemNBINamespace:
+    """The NVSHMEM GetMemNBI intrinsics submodule."""
+
+    def __init__(self):
+        self.warp = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_warp)
+        self.block = _op_wrapper(_cuda_op.nvshmem_getmem_nbi_block)
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.nvshmem_getmem_nbi)(*args, **kwds)
+
+    # __call__ corresponds to nvshmem_getmem_nbi
+    __tir_call_op_name__ = "nvshmem_getmem_nbi"
+
+
+class NVSHMEMPutMemNBINamespace:
+    """The NVSHMEM PutMemNBI intrinsics submodule."""
+
+    def __init__(self):
+        self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_warp)
+        self.block = _op_wrapper(_cuda_op.nvshmem_putmem_nbi_block)
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.nvshmem_putmem_nbi)(*args, **kwds)
+
+    # __call__ corresponds to nvshmem_putmem_nbi
+    __tir_call_op_name__ = "nvshmem_putmem_nbi"
+
+
+class NVSHMEMPutMemSignalNBINamespace:
+    """The NVSHMEM PutMemSignalNBI intrinsics submodule."""
+
+    def __init__(self):
+        self.warp = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_warp)
+        self.block = _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi_block)
+
+    def __call__(self, *args, **kwds):
+        return _op_wrapper(_cuda_op.nvshmem_putmem_signal_nbi)(*args, **kwds)
+
+    # __call__ corresponds to nvshmem_putmem_signal_nbi
+    __tir_call_op_name__ = "nvshmem_putmem_signal_nbi"
+
+
+__all__ = ["CUDANamespace", "NVSHMEMNamespace", "PTXNamespace"]

python/tvm/backend/cuda/target_tags.py

@@ -0,0 +1,374 @@
+# 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.
+"""NVIDIA CUDA target tags."""
+
+from tvm.target import register_tag
+
+
+def _register_cuda_tag(name, arch, shared_mem=49152, regs=65536, **extra):
+    config = {
+        "kind": "cuda",
+        "keys": ["cuda", "gpu"],
+        "arch": arch,
+        "max_shared_memory_per_block": shared_mem,
+        "max_threads_per_block": 1024,
+        "thread_warp_size": 32,
+        "registers_per_block": regs,
+    }
+    config.update(extra)
+    register_tag(name, config)
+
+
+def _register_jetson_tag(name, arch, mcpu, num_cores, regs=65536):
+    register_tag(
+        name,
+        {
+            "kind": "cuda",
+            "arch": arch,
+            "max_shared_memory_per_block": 49152,
+            "max_threads_per_block": 1024,
+            "thread_warp_size": 32,
+            "registers_per_block": regs,
+            "host": {
+                "kind": "llvm",
+                "mtriple": "aarch64-linux-gnu",
+                "mcpu": mcpu,
+                "num-cores": num_cores,
+            },
+        },
+    )
+
+
+# =====================================================================
+# Data center / Tesla GPUs
+# =====================================================================
+_register_cuda_tag("nvidia/nvidia-a100", "sm_80", l2_cache_size_bytes=41943040)
+_register_cuda_tag("nvidia/nvidia-h100", "sm_90a", l2_cache_size_bytes=52428800)
+_register_cuda_tag("nvidia/nvidia-b100", "sm_100a", l2_cache_size_bytes=52428800)
+_register_cuda_tag("nvidia/nvidia-a40", "sm_86")
+_register_cuda_tag("nvidia/nvidia-a30", "sm_80")
+_register_cuda_tag("nvidia/nvidia-a10", "sm_86")
+_register_cuda_tag("nvidia/nvidia-a10g", "sm_86")
+_register_cuda_tag("nvidia/nvidia-a16", "sm_86")
+_register_cuda_tag("nvidia/nvidia-a2", "sm_86")
+_register_cuda_tag("nvidia/nvidia-t4", "sm_75")
+_register_cuda_tag("nvidia/nvidia-v100", "sm_70")
+_register_cuda_tag("nvidia/tesla-p100", "sm_60")
+_register_cuda_tag("nvidia/tesla-p40", "sm_61")
+_register_cuda_tag("nvidia/tesla-p4", "sm_61")
+_register_cuda_tag("nvidia/tesla-m60", "sm_52")
+_register_cuda_tag("nvidia/tesla-m40", "sm_52")
+_register_cuda_tag("nvidia/tesla-k80", "sm_37")
+_register_cuda_tag("nvidia/tesla-k40", "sm_35")
+_register_cuda_tag("nvidia/tesla-k20", "sm_35")
+_register_cuda_tag("nvidia/tesla-k10", "sm_30")
+_register_cuda_tag("nvidia/tesla-c2075", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/tesla-c2050", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/tesla-c2070", "sm_20", regs=32768)
+
+# =====================================================================
+# Quadro / RTX professional desktop GPUs
+# =====================================================================
+_register_cuda_tag("nvidia/rtx-a6000", "sm_86")
+_register_cuda_tag("nvidia/quadro-rtx-8000", "sm_75")
+_register_cuda_tag("nvidia/quadro-rtx-6000", "sm_75")
+_register_cuda_tag("nvidia/quadro-rtx-5000", "sm_75")
+_register_cuda_tag("nvidia/quadro-rtx-4000", "sm_75")
+_register_cuda_tag("nvidia/quadro-gv100", "sm_70")
+_register_cuda_tag("nvidia/quadro-gp100", "sm_60")
+_register_cuda_tag("nvidia/quadro-p6000", "sm_61")
+_register_cuda_tag("nvidia/quadro-p5000", "sm_61")
+_register_cuda_tag("nvidia/quadro-p4000", "sm_61")
+_register_cuda_tag("nvidia/quadro-p2200", "sm_61")
+_register_cuda_tag("nvidia/quadro-p2000", "sm_61")
+_register_cuda_tag("nvidia/quadro-p1000", "sm_61")
+_register_cuda_tag("nvidia/quadro-p620", "sm_61")
+_register_cuda_tag("nvidia/quadro-p600", "sm_61")
+_register_cuda_tag("nvidia/quadro-p400", "sm_61")
+_register_cuda_tag("nvidia/quadro-m6000-24gb", "sm_52")
+_register_cuda_tag("nvidia/quadro-m6000", "sm_52")
+_register_cuda_tag("nvidia/quadro-k6000", "sm_35")
+_register_cuda_tag("nvidia/quadro-m5000", "sm_52")
+_register_cuda_tag("nvidia/quadro-k5200", "sm_35")
+_register_cuda_tag("nvidia/quadro-k5000", "sm_30")
+_register_cuda_tag("nvidia/quadro-m4000", "sm_52")
+_register_cuda_tag("nvidia/quadro-k4200", "sm_30")
+_register_cuda_tag("nvidia/quadro-k4000", "sm_30")
+_register_cuda_tag("nvidia/quadro-m2000", "sm_52")
+_register_cuda_tag("nvidia/quadro-k2200", "sm_50")
+_register_cuda_tag("nvidia/quadro-k2000", "sm_30")
+_register_cuda_tag("nvidia/quadro-k2000d", "sm_30")
+_register_cuda_tag("nvidia/quadro-k1200", "sm_50")
+_register_cuda_tag("nvidia/quadro-k620", "sm_50")
+_register_cuda_tag("nvidia/quadro-k600", "sm_30")
+_register_cuda_tag("nvidia/quadro-k420", "sm_30")
+_register_cuda_tag("nvidia/quadro-410", "sm_30")
+_register_cuda_tag("nvidia/quadro-plex-7000", "sm_20", regs=32768)
+
+# =====================================================================
+# Quadro / RTX professional mobile GPUs (Turing)
+# =====================================================================
+_register_cuda_tag("nvidia/rtx-5000", "sm_75")
+_register_cuda_tag("nvidia/rtx-4000", "sm_75")
+_register_cuda_tag("nvidia/rtx-3000", "sm_75")
+_register_cuda_tag("nvidia/t2000", "sm_75")
+_register_cuda_tag("nvidia/t1000", "sm_75")
+_register_cuda_tag("nvidia/p620", "sm_61")
+_register_cuda_tag("nvidia/p520", "sm_61")
+
+# =====================================================================
+# Quadro professional mobile GPUs (Pascal / Maxwell)
+# =====================================================================
+_register_cuda_tag("nvidia/quadro-p5200", "sm_61")
+_register_cuda_tag("nvidia/quadro-p4200", "sm_61")
+_register_cuda_tag("nvidia/quadro-p3200", "sm_61")
+_register_cuda_tag("nvidia/quadro-p3000", "sm_61")
+_register_cuda_tag("nvidia/quadro-p500", "sm_61")
+_register_cuda_tag("nvidia/quadro-m5500m", "sm_52")
+_register_cuda_tag("nvidia/quadro-m2200", "sm_52")
+_register_cuda_tag("nvidia/quadro-m1200", "sm_50")
+_register_cuda_tag("nvidia/quadro-m620", "sm_52")
+_register_cuda_tag("nvidia/quadro-m520", "sm_50")
+
+# =====================================================================
+# Quadro professional mobile GPUs (Kepler / Maxwell)
+# =====================================================================
+_register_cuda_tag("nvidia/quadro-k6000m", "sm_30")
+_register_cuda_tag("nvidia/quadro-k5200m", "sm_30")
+_register_cuda_tag("nvidia/quadro-k5100m", "sm_30")
+_register_cuda_tag("nvidia/quadro-m5000m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k500m", "sm_30")
+_register_cuda_tag("nvidia/quadro-k4200m", "sm_30")
+_register_cuda_tag("nvidia/quadro-k4100m", "sm_30")
+_register_cuda_tag("nvidia/quadro-m4000m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k3100m", "sm_30")
+_register_cuda_tag("nvidia/quadro-m3000m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k2200m", "sm_30")
+_register_cuda_tag("nvidia/quadro-k2100m", "sm_30")
+_register_cuda_tag("nvidia/quadro-m2000m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k1100m", "sm_30")
+_register_cuda_tag("nvidia/quadro-m1000m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k620m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k610m", "sm_35")
+_register_cuda_tag("nvidia/quadro-m600m", "sm_50")
+_register_cuda_tag("nvidia/quadro-k510m", "sm_35")
+_register_cuda_tag("nvidia/quadro-m500m", "sm_50")
+
+# =====================================================================
+# NVS cards
+# =====================================================================
+_register_cuda_tag("nvidia/nvidia-nvs-810", "sm_50")
+_register_cuda_tag("nvidia/nvidia-nvs-510", "sm_30")
+_register_cuda_tag("nvidia/nvidia-nvs-315", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/nvidia-nvs-310", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/nvs-5400m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/nvs-5200m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/nvs-4200m", "sm_21", regs=32768)
+
+# =====================================================================
+# GeForce RTX 50-series desktop
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-rtx-5060-ti", "sm_120", l2_cache_size_bytes=33554432)
+
+# =====================================================================
+# GeForce RTX 40-series desktop
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-rtx-4090", "sm_89", l2_cache_size_bytes=75497472)
+
+# =====================================================================
+# GeForce RTX 30-series desktop
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-rtx-3090-ti", "sm_86")
+_register_cuda_tag("nvidia/geforce-rtx-3090", "sm_86")
+_register_cuda_tag("nvidia/geforce-rtx-3080-ti", "sm_86")
+_register_cuda_tag("nvidia/geforce-rtx-3080", "sm_86")
+_register_cuda_tag("nvidia/geforce-rtx-3070-ti", "sm_86")
+_register_cuda_tag("nvidia/geforce-rtx-3070", "sm_86")
+_register_cuda_tag("nvidia/geforce-rtx-3060", "sm_86")
+
+# =====================================================================
+# GeForce RTX 20-series / TITAN (Turing)
+# =====================================================================
+_register_cuda_tag("nvidia/nvidia-titan-rtx", "sm_75")
+_register_cuda_tag("nvidia/geforce-rtx-2080-ti", "sm_75")
+_register_cuda_tag("nvidia/geforce-rtx-2080", "sm_75")
+_register_cuda_tag("nvidia/geforce-rtx-2070", "sm_75")
+_register_cuda_tag("nvidia/geforce-rtx-2060", "sm_75")
+
+# =====================================================================
+# GeForce TITAN / GTX 10-series (Pascal)
+# =====================================================================
+_register_cuda_tag("nvidia/nvidia-titan-v", "sm_70")
+_register_cuda_tag("nvidia/nvidia-titan-xp", "sm_61")
+_register_cuda_tag("nvidia/nvidia-titan-x", "sm_61")
+_register_cuda_tag("nvidia/geforce-gtx-1080-ti", "sm_61")
+_register_cuda_tag("nvidia/geforce-gtx-1080", "sm_61")
+_register_cuda_tag("nvidia/geforce-gtx-1070-ti", "sm_61")
+_register_cuda_tag("nvidia/geforce-gtx-1070", "sm_61")
+_register_cuda_tag("nvidia/geforce-gtx-1060", "sm_61")
+_register_cuda_tag("nvidia/geforce-gtx-1050", "sm_61")
+
+# =====================================================================
+# GeForce GTX 900/700 series desktop (Maxwell / Kepler)
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-gtx-titan-x", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-titan-z", "sm_35")
+_register_cuda_tag("nvidia/geforce-gtx-titan-black", "sm_35")
+_register_cuda_tag("nvidia/geforce-gtx-titan", "sm_35")
+_register_cuda_tag("nvidia/geforce-gtx-980-ti", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-980", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-970", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-960", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-950", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-780-ti", "sm_35")
+_register_cuda_tag("nvidia/geforce-gtx-780", "sm_35")
+_register_cuda_tag("nvidia/geforce-gtx-770", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-760", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-750-ti", "sm_50")
+_register_cuda_tag("nvidia/geforce-gtx-750", "sm_50")
+_register_cuda_tag("nvidia/geforce-gtx-690", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-680", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-670", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-660-ti", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-660", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-650-ti-boost", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-650-ti", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-650", "sm_30")
+
+# =====================================================================
+# GeForce GTX 500/400 series desktop (Fermi)
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-gtx-560-ti", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-550-ti", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-460", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gts-450", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-590", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-580", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-570", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-480", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-470", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-465", "sm_20", regs=32768)
+
+# =====================================================================
+# GeForce GT desktop (Kepler / Fermi)
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-gt-740", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-730", "sm_35")
+_register_cuda_tag("nvidia/geforce-gt-730-ddr3,128bit", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-720", "sm_35")
+_register_cuda_tag("nvidia/geforce-gt-705", "sm_35")
+_register_cuda_tag("nvidia/geforce-gt-640-gddr5", "sm_35")
+_register_cuda_tag("nvidia/geforce-gt-640-gddr3", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-630", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-620", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-610", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-520", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-440", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-430", "sm_21", regs=32768)
+
+# =====================================================================
+# GeForce notebook GPUs (Maxwell / Kepler)
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-gtx-980m", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-970m", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-965m", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-960m", "sm_50")
+_register_cuda_tag("nvidia/geforce-gtx-950m", "sm_50")
+_register_cuda_tag("nvidia/geforce-940m", "sm_50")
+_register_cuda_tag("nvidia/geforce-930m", "sm_50")
+_register_cuda_tag("nvidia/geforce-920m", "sm_35")
+_register_cuda_tag("nvidia/geforce-910m", "sm_52")
+_register_cuda_tag("nvidia/geforce-gtx-880m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-870m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-860m-sm-30", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-860m-sm-50", "sm_50")
+_register_cuda_tag("nvidia/geforce-gtx-850m", "sm_50")
+_register_cuda_tag("nvidia/geforce-840m", "sm_50")
+_register_cuda_tag("nvidia/geforce-830m", "sm_50")
+_register_cuda_tag("nvidia/geforce-820m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-800m", "sm_21", regs=32768)
+
+# =====================================================================
+# GeForce notebook GPUs (Kepler / Fermi, older)
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-gtx-780m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-770m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-765m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-760m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-680mx", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-680m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-675mx", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-675m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-670mx", "sm_30")
+_register_cuda_tag("nvidia/geforce-gtx-670m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-660m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-755m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-750m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-650m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-745m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-645m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-740m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-730m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-640m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-640m-le", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-735m", "sm_30")
+_register_cuda_tag("nvidia/geforce-gt-635m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-630m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-625m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-720m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-620m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-710m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-705m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-610m", "sm_21", regs=32768)
+
+# =====================================================================
+# GeForce notebook GPUs (Fermi, GTX 5xx/4xxM)
+# =====================================================================
+_register_cuda_tag("nvidia/geforce-gtx-580m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-570m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-560m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-555m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-550m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-540m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-525m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-520mx", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-520m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-485m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-470m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-460m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-445m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-435m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-420m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gt-415m", "sm_21", regs=32768)
+_register_cuda_tag("nvidia/geforce-gtx-480m", "sm_20", regs=32768)
+_register_cuda_tag("nvidia/geforce-410m", "sm_21", regs=32768)
+
+# =====================================================================
+# Jetson boards (simple, no host)
+# =====================================================================
+_register_cuda_tag("nvidia/jetson-nano", "sm_53", regs=32768)
+_register_cuda_tag("nvidia/jetson-tx2", "sm_62", regs=32768)
+_register_cuda_tag("nvidia/jetson-tx1", "sm_53", regs=32768)
+_register_cuda_tag("nvidia/tegra-x1", "sm_53", regs=32768)
+
+# =====================================================================
+# Jetson boards (with LLVM host)
+# =====================================================================
+_register_jetson_tag("nvidia/jetson-agx-xavier", "sm_72", "carmel", 8)
+_register_jetson_tag("nvidia/jetson-orin-nano", "sm_87", "carmel", 6)
+_register_jetson_tag("nvidia/jetson-agx-orin-32gb", "sm_87", "cortex-a78", 8)
+_register_jetson_tag("nvidia/jetson-agx-orin-64gb", "sm_87", "cortex-a78", 12)

python/tvm/backend/hexagon/__init__.py

@@ -0,0 +1,51 @@
+# 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.
+"""Hexagon-owned backend hooks."""
+
+from importlib import import_module
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.base import _LOADED_LIBS
+
+_LAZY_SUBMODULES = {"target_tags"}
+
+
+def register_backend():
+    """Register Hexagon-owned Python semantics."""
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Runtime sidecars only need registration side effects; libtvm_runtime is global.
+        _LOADED_LIBS["tvm_runtime_hexagon"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_hexagon",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    import_module(f"{__name__}.target_tags")
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["register_backend", "target_tags"]

python/tvm/backend/hexagon/target_tags.py

@@ -0,0 +1,56 @@
+# 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.
+"""Qualcomm Hexagon target tags."""
+
+from tvm.target import register_tag
+
+_ONE_MB = 2**20
+
+_HEXAGON_VERSIONS = {
+    "v65": {"vtcm": _ONE_MB // 4, "mattr": ["+hvxv65", "+hvx-length128b"]},
+    "v66": {"vtcm": _ONE_MB // 4, "mattr": ["+hvxv66", "+hvx-length128b"]},
+    "v68": {
+        "vtcm": 4 * _ONE_MB,
+        "mattr": ["+hvxv68", "+hvx-length128b", "+hvx-qfloat", "-hvx-ieee-fp"],
+        "llvm-options": ["-force-hvx-float"],
+    },
+    "v69": {
+        "vtcm": 8 * _ONE_MB,
+        "mattr": ["+hvxv69", "+hvx-length128b", "+hvx-qfloat", "-hvx-ieee-fp"],
+    },
+    "v73": {
+        "vtcm": 8 * _ONE_MB,
+        "mattr": ["+hvxv73", "+hvx-length128b", "+hvx-qfloat", "-hvx-ieee-fp"],
+    },
+    "v75": {
+        "vtcm": 8 * _ONE_MB,
+        "mattr": ["+hvxv75", "+hvx-length128b", "+hvx-qfloat", "-hvx-ieee-fp"],
+    },
+}
+
+for _ver, _info in _HEXAGON_VERSIONS.items():
+    _config = {
+        "kind": "hexagon",
+        "mtriple": "hexagon",
+        "mcpu": "hexagon" + _ver,
+        "mattr": _info["mattr"],
+        "num-cores": 4,
+        "vtcm-capacity": _info["vtcm"],
+    }
+    if "llvm-options" in _info:
+        _config["llvm-options"] = _info["llvm-options"]
+    register_tag("qcom/hexagon-" + _ver, _config)

python/tvm/backend/loader.py

@@ -0,0 +1,187 @@
+# 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.
+"""Backend loading and public alias support."""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import sys
+import types
+from typing import Any
+
+_LOADED_BACKENDS: dict[str, Any] = {}
+
+
+class _AliasModule(types.ModuleType):
+    """Module object that exposes a backend module under a public alias."""
+
+    def __init__(self, fullname: str, module):
+        super().__init__(fullname, getattr(module, "__doc__", None))
+        self.__dict__["__tvm_backend_module__"] = module
+        self.__dict__["__package__"] = fullname.rpartition(".")[0]
+        if hasattr(module, "__all__"):
+            self.__dict__["__all__"] = module.__all__
+        if hasattr(module, "__path__"):
+            self.__dict__["__path__"] = []
+
+    def __getattr__(self, name: str):
+        return getattr(self.__dict__["__tvm_backend_module__"], name)
+
+    def __setattr__(self, name: str, value):
+        setattr(self.__dict__["__tvm_backend_module__"], name, value)
+
+    def __delattr__(self, name: str):
+        delattr(self.__dict__["__tvm_backend_module__"], name)
+
+    def __dir__(self):
+        return sorted(set(super().__dir__()) | set(dir(self.__dict__["__tvm_backend_module__"])))
+
+
+class _AliasLoader:
+    """Loader that returns an already-resolved module for an alias spec."""
+
+    def __init__(self, fullname: str, module):
+        self._fullname = fullname
+        self._module = module
+
+    def create_module(self, spec):
+        return _get_alias_module(self._fullname, self._module)
+
+    def exec_module(self, module):
+        _set_module_alias(self._fullname, self._module)
+        return None
+
+    def is_package(self, fullname):
+        return hasattr(self._module, "__path__")
+
+
+def _redirect_tirx_backend_alias(fullname: str) -> str | None:
+    prefix = "tvm.tirx."
+    if not fullname.startswith(prefix):
+        return None
+    rest = fullname[len(prefix) :]
+    backend_name, sep, tail = rest.partition(".")
+    if not sep or backend_name not in _LOADED_BACKENDS:
+        return None
+    return f"tvm.backend.{backend_name}.{tail}"
+
+
+class _BackendAliasFinder:
+    """Redirect ``tvm.tirx.<backend>.*`` imports to ``tvm.backend.<backend>.*``."""
+
+    @classmethod
+    def find_spec(cls, fullname, path, target=None):
+        redirected = _redirect_tirx_backend_alias(fullname)
+        if redirected is None:
+            return None
+        module = importlib.import_module(redirected)
+        _set_module_alias(fullname, module)
+        loader = _AliasLoader(fullname, module)
+        spec = importlib.util.spec_from_loader(
+            fullname, loader, is_package=hasattr(module, "__path__")
+        )
+        if spec is not None and hasattr(module, "__path__"):
+            spec.submodule_search_locations = []
+        return spec
+
+
+if not any(isinstance(finder, _BackendAliasFinder) for finder in sys.meta_path):
+    sys.meta_path.insert(0, _BackendAliasFinder())
+
+
+def _get_alias_module(alias: str, module):
+    existing = sys.modules.get(alias)
+    if (
+        isinstance(existing, _AliasModule)
+        and existing.__dict__.get("__tvm_backend_module__") is module
+    ):
+        return existing
+    return _AliasModule(alias, module)
+
+
+def _set_module_alias(alias: str, module, *, direct: bool = False) -> None:
+    alias_module = module if direct else _get_alias_module(alias, module)
+    sys.modules[alias] = alias_module
+    parent_name, _, child_name = alias.rpartition(".")
+    parent = sys.modules.get(parent_name)
+    if parent is not None:
+        setattr(parent, child_name, alias_module)
+
+
+def _alias_loaded_backend_modules(name: str) -> None:
+    backend_prefix = f"tvm.backend.{name}"
+    public_prefix = f"tvm.tirx.{name}"
+    for module_name, module in sorted(list(sys.modules.items())):
+        if module_name == backend_prefix or module_name.startswith(f"{backend_prefix}."):
+            public_name = f"{public_prefix}{module_name[len(backend_prefix) :]}"
+            _set_module_alias(public_name, module, direct=module_name == backend_prefix)
+
+
+def _import_backend(name: str):
+    module_name = f"tvm.backend.{name}"
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError as err:
+        if err.name == module_name:
+            raise ImportError(
+                f"Cannot load TVM backend {name!r}: expected Python package {module_name!r}. "
+                "Install the backend package or check the backend name."
+            ) from err
+        raise
+
+
+def load(name: str) -> None:
+    """Load a backend's Python registration hooks.
+
+    Loading is idempotent.  A backend package must live at ``tvm.backend.<name>``
+    and expose ``register_backend()``.
+    """
+
+    if name in _LOADED_BACKENDS:
+        return None
+
+    module = _import_backend(name)
+    register_backend = getattr(module, "register_backend", None)
+    if register_backend is None:
+        raise AttributeError(f"Backend package 'tvm.backend.{name}' has no register_backend()")
+
+    import tvm.tirx as tirx  # pylint: disable=import-outside-toplevel
+
+    setattr(tirx, name, module)
+    sys.modules[f"tvm.tirx.{name}"] = module
+    _LOADED_BACKENDS[name] = module
+    try:
+        register_backend()
+        _alias_loaded_backend_modules(name)
+    except Exception:
+        _LOADED_BACKENDS.pop(name, None)
+        if getattr(tirx, name, None) is module:
+            delattr(tirx, name)
+        if sys.modules.get(f"tvm.tirx.{name}") is module:
+            sys.modules.pop(f"tvm.tirx.{name}", None)
+        raise
+    return None
+
+
+def is_loaded(name: str) -> bool:
+    """Return whether a backend has been loaded."""
+
+    return name in _LOADED_BACKENDS
+
+
+__all__ = ["is_loaded", "load"]

python/tvm/backend/metal/__init__.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.
+"""Metal-owned TIRx modules."""
+
+from importlib import import_module
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.base import _LOADED_LIBS
+
+_LAZY_SUBMODULES = {"op", "script", "target_tags"}
+
+
+def _detect_target_from_device(dev):
+    from tvm.target import Target  # pylint: disable=import-outside-toplevel
+
+    return Target(
+        {
+            "kind": "metal",
+            "max_shared_memory_per_block": 32768,
+            "max_threads_per_block": dev.max_threads_per_block,
+            "thread_warp_size": dev.warp_size,
+        }
+    )
+
+
+def register_backend():
+    """Register Metal-owned Python semantics."""
+    from tvm.target.detect_target import register_device_target_detector
+    from tvm.tirx.script.builder import ir as builder_ir  # pylint: disable=import-outside-toplevel
+
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Runtime sidecars only need registration side effects; libtvm_runtime is global.
+        _LOADED_LIBS["tvm_runtime_metal"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_metal",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    register_device_target_detector("metal", _detect_target_from_device)
+    for name, namespace in script_namespaces().items():
+        builder_ir.register_script_namespace(name, namespace)
+    import_module(f"{__name__}.target_tags")
+
+
+def script_namespaces(**_):
+    """Return Metal-owned TVMScript namespaces."""
+    from .script import MetalNamespace  # pylint: disable=import-outside-toplevel
+
+    return {"metal": MetalNamespace()}
+
+
+def script_namespace(**kwargs):
+    """Return the Metal TVMScript namespace object."""
+    return script_namespaces(**kwargs)["metal"]
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "op",
+    "register_backend",
+    "script",
+    "script_namespace",
+    "script_namespaces",
+    "target_tags",
+]

python/tvm/backend/metal/op.py

@@ -0,0 +1,84 @@
+# 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.
+"""Metal-owned TIR intrinsic builders."""
+
+from __future__ import annotations
+
+from tvm.tirx.op import call_intrin
+
+
+def make_filled_simdgroup_matrix(d, index, value, col=8, row=8):
+    """Create a filled SIMDGroup matrix."""
+
+    return call_intrin("void", "tirx.make_filled_simdgroup_matrix", d, index, value, col, row)
+
+
+def simdgroup_load(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+    """Load data from device or threadgroup memory to simdgroup."""
+
+    return call_intrin(
+        "void",
+        "tirx.simdgroup_load",
+        d,
+        index,
+        ptr,
+        stride,
+        col,
+        row,
+        transpose_matrix,
+    )
+
+
+def simdgroup_store(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+    """Store data from simdgroup to device or threadgroup memory."""
+
+    return call_intrin(
+        "void",
+        "tirx.simdgroup_store",
+        d,
+        index,
+        ptr,
+        stride,
+        col,
+        row,
+        transpose_matrix,
+    )
+
+
+def simdgroup_multiply_accumulate(d, index_d, a, index_a, b, index_b, c, index_c):
+    """Multiply and accumulate two matrices in simdgroup."""
+
+    return call_intrin(
+        "void",
+        "tirx.simdgroup_multiply_accumulate",
+        d,
+        index_d,
+        a,
+        index_a,
+        b,
+        index_b,
+        c,
+        index_c,
+    )
+
+
+__all__ = [
+    "make_filled_simdgroup_matrix",
+    "simdgroup_load",
+    "simdgroup_multiply_accumulate",
+    "simdgroup_store",
+]

python/tvm/backend/metal/script.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal TVMScript namespace."""
+
+from __future__ import annotations
+
+from tvm.backend.metal import op as _metal_op
+from tvm.tirx import Buffer
+from tvm.tirx import op as _tir_op
+from tvm.tirx.script.builder.ir import _op_wrapper
+
+
+class MetalNamespace:
+    """The Metal intrinsics submodule."""
+
+    def __init__(self):
+        self.make_filled_simdgroup_matrix = _op_wrapper(_metal_op.make_filled_simdgroup_matrix)
+        self.simdgroup_load = _op_wrapper(_metal_op.simdgroup_load)
+        self.simdgroup_store = _op_wrapper(_metal_op.simdgroup_store)
+        self.simdgroup_multiply_accumulate = _op_wrapper(_metal_op.simdgroup_multiply_accumulate)
+
+    @staticmethod
+    def simd_shuffle(var, lane):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle", var, lane)
+
+    @staticmethod
+    def simd_shuffle_up(var, delta):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_up", var, delta)
+
+    @staticmethod
+    def simd_shuffle_down(var, delta):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_down", var, delta)
+
+
+__all__ = ["MetalNamespace"]

python/tvm/backend/metal/target_tags.py

@@ -0,0 +1,50 @@
+# 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.
+"""Apple Metal GPU target tags."""
+
+from tvm.target import register_tag
+
+_METAL_HOST_TRIPLE = "arm64-apple-macos"
+
+
+def _register_metal_tag(name, max_threads, shared_mem, warp_size, mcpu):
+    try:
+        from tvm.target.codegen import llvm_is_valid_cpu
+
+        if not llvm_is_valid_cpu(mcpu, _METAL_HOST_TRIPLE):
+            return
+    except Exception:  # pylint: disable=broad-except
+        pass  # LLVM not available; register unconditionally
+    register_tag(
+        name,
+        {
+            "kind": "metal",
+            "max_threads_per_block": max_threads,
+            "max_shared_memory_per_block": shared_mem,
+            "thread_warp_size": warp_size,
+            "host": {
+                "kind": "llvm",
+                "mtriple": _METAL_HOST_TRIPLE,
+                "mcpu": mcpu,
+            },
+        },
+    )
+
+
+_register_metal_tag("apple/m1-gpu", 1024, 32768, 32, "apple-m1")
+_register_metal_tag("apple/m1-gpu-restricted", 256, 32768, 32, "apple-m1")
+_register_metal_tag("apple/m2-gpu", 1024, 32768, 32, "apple-m2")

python/tvm/backend/opencl/__init__.py

@@ -0,0 +1,58 @@
+# 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.
+"""OpenCL-owned backend hooks."""
+
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.base import _LOADED_LIBS
+
+
+def _detect_target_from_device(dev):
+    from tvm.target import Target  # pylint: disable=import-outside-toplevel
+
+    return Target(
+        {
+            "kind": "opencl",
+            "max_shared_memory_per_block": dev.max_shared_memory_per_block,
+            "max_threads_per_block": dev.max_threads_per_block,
+            "thread_warp_size": dev.warp_size,
+        }
+    )
+
+
+def register_backend():
+    """Register OpenCL-owned Python semantics."""
+    from tvm.target.detect_target import register_device_target_detector
+
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Runtime sidecars only need registration side effects; libtvm_runtime is global.
+        _LOADED_LIBS["tvm_runtime_opencl"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_opencl",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    register_device_target_detector("opencl", _detect_target_from_device)
+    return None
+
+
+__all__ = ["register_backend"]

python/tvm/backend/rocm/__init__.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.
+"""ROCm-owned TIRx modules."""
+
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.base import _LOADED_LIBS
+
+
+def _detect_target_from_device(dev):
+    from tvm.target import Target  # pylint: disable=import-outside-toplevel
+
+    return Target(
+        {
+            "kind": "rocm",
+            "mtriple": "amdgcn-amd-amdhsa-hcc",
+            "max_shared_memory_per_block": dev.max_shared_memory_per_block,
+            "max_threads_per_block": dev.max_threads_per_block,
+            "thread_warp_size": dev.warp_size,
+        }
+    )
+
+
+def register_backend():
+    """Register ROCm-owned Python semantics."""
+    from tvm.target.detect_target import register_device_target_detector
+
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Runtime sidecars only need registration side effects; libtvm_runtime is global.
+        _LOADED_LIBS["tvm_runtime_rocm"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_rocm",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    register_device_target_detector("rocm", _detect_target_from_device)
+    return None
+
+
+__all__ = ["register_backend"]

python/tvm/backend/trn/__init__.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.
+"""Trainium-owned TIRx modules."""
+
+from importlib import import_module
+
+_LAZY_SUBMODULES = {"layout", "op", "operator", "pipeline", "script", "target_tags", "transform"}
+
+
+def register_backend():
+    """Register Trainium-owned Python semantics."""
+    from tvm.tirx import compilation_pipeline  # pylint: disable=import-outside-toplevel
+    from tvm.tirx.script.builder import ir as builder_ir  # pylint: disable=import-outside-toplevel
+
+    for name, namespace in script_namespaces().items():
+        builder_ir.register_script_namespace(name, namespace)
+
+    import_module(f"{__name__}.operator.tile_primitive")
+    trn_pipeline = import_module(f"{__name__}.pipeline")
+    import_module(f"{__name__}.target_tags")
+    import_module(f"{__name__}.transform")
+    compilation_pipeline.register_tir_pipeline("trn", trn_pipeline.trn_pipeline)
+
+
+def script_namespace(op_wrapper=None):
+    """Return the Trainium TVMScript namespace object."""
+    from .script import NKINamespace  # pylint: disable=import-outside-toplevel
+
+    return NKINamespace(op_wrapper)
+
+
+def script_namespaces(op_wrapper=None, **_):
+    """Return Trainium-owned TVMScript namespaces."""
+    return {"nki": script_namespace(op_wrapper)}
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "layout",
+    "op",
+    "operator",
+    "pipeline",
+    "register_backend",
+    "script",
+    "script_namespace",
+    "script_namespaces",
+    "target_tags",
+    "transform",
+]

python/tvm/backend/trn/layout.py

@@ -0,0 +1,123 @@
+# 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.
+"""Trainium-specific TIRx layout helpers."""
+
+from __future__ import annotations
+
+import functools
+import operator
+import re
+
+import tvm
+from tvm.tirx.expr import Expr
+from tvm.tirx.layout import Axis, Iter, Layout, S, TileLayout
+
+_TRN_MEMORY_AXES = {"F", "P", "Bank"}
+_PSUM_MAX_ELEM_PER_BANK = 512
+
+
+def is_trainium_layout(layout: Layout | None) -> bool:
+    """Return whether a layout uses only Trainium memory axes."""
+    if not isinstance(layout, TileLayout):
+        return False
+    return not any(
+        iter.axis.is_memory() and iter.axis.name not in _TRN_MEMORY_AXES for iter in layout.shard
+    )
+
+
+def trainium_layout(annotation: str, shape: tuple[Expr], is_psum: bool = False) -> TileLayout:
+    """Create a Trainium tile layout from a PF annotation string and logical shape."""
+    analyzer = tvm.arith.Analyzer()
+    assert re.fullmatch(r"[PF]*", annotation), (
+        f"annotation {annotation} must be a string of 'P' and 'F'"
+    )
+    assert len(annotation) == len(shape), (
+        f"annotation {annotation} and shape {shape} must have the same length"
+    )
+    num_p_dim = annotation.count("P")
+    if num_p_dim == 1:
+        p_idx = annotation.index("P")
+        p_dim = shape[p_idx]
+        assert analyzer.can_prove(p_dim <= 128 or p_dim % 128 == 0), (
+            f"There is only 1 P in the annotation. Partition size {p_dim} must be less than "
+            "or equal to 128 or a multiple of 128"
+        )
+        if analyzer.can_prove(p_dim > 128):
+            annotation = "F" + annotation
+            shape = (p_dim // 128, *shape[:p_idx], 128, *shape[p_idx + 1 :])
+    elif num_p_dim > 1:
+        p_dim_prod = functools.reduce(
+            operator.mul, [s for s, c in zip(shape, annotation) if c == "P"]
+        )
+        assert analyzer.can_prove(p_dim_prod <= 128), (
+            f"There are {num_p_dim} Ps in the annotation. Partition size {p_dim_prod} must be "
+            "less than or equal to 128"
+        )
+
+    f_shape = [s for s, c in zip(shape, annotation) if c == "F"]
+    p_shape = [s for s, c in zip(shape, annotation) if c == "P"]
+    f_strides = Layout._get_default_strides(f_shape, 1)  # pylint: disable=protected-access
+    p_strides = Layout._get_default_strides(p_shape, 1)  # pylint: disable=protected-access
+    f_tile_layout = TileLayout(S[tuple(f_shape) : tuple(s @ Axis.F for s in f_strides)])
+    p_tile_layout = TileLayout(S[tuple(p_shape) : tuple(s @ Axis.P for s in p_strides)])
+    result = []
+    f_index = p_index = 0
+
+    for char in annotation:
+        if char == "F":
+            result.append(f_tile_layout.shard[f_index])
+            f_index += 1
+        else:
+            result.append(p_tile_layout.shard[p_index])
+            p_index += 1
+    if num_p_dim == 1 and analyzer.can_prove(p_dim > 128):
+        higher_p = result[0]
+        result = result[1:]
+        result = [*result[:p_idx], higher_p, *result[p_idx:]]
+
+    res = TileLayout.from_iters(result, [], {})
+    if is_psum:
+        res = to_psum_layout(res)
+    return res
+
+
+def to_psum_layout(layout: TileLayout) -> TileLayout:
+    """Convert a Trainium sbuf layout to its psum physical-bank layout."""
+    analyzer = tvm.arith.Analyzer()
+    shard = []
+    for iter in layout.shard:
+        if iter.axis.name == "F":
+            if analyzer.can_prove(iter.stride % _PSUM_MAX_ELEM_PER_BANK == 0):
+                stride = analyzer.simplify(iter.stride // _PSUM_MAX_ELEM_PER_BANK)
+                shard.append(Iter(iter.extent, stride, Axis.get("Bank")))
+            elif analyzer.can_prove(_PSUM_MAX_ELEM_PER_BANK % iter.stride == 0):
+                c = analyzer.simplify(_PSUM_MAX_ELEM_PER_BANK // iter.stride)
+                if analyzer.can_prove(iter.extent < c):
+                    shard.append(iter)
+                elif analyzer.can_prove(iter.extent % c == 0):
+                    shard.append(Iter(analyzer.simplify(iter.extent // c), 1, Axis.get("Bank")))
+                    shard.append(Iter(c, iter.stride, Axis.get("F")))
+                else:
+                    raise ValueError(f"layout {layout} can not be converted to psum layout")
+            else:
+                raise ValueError(f"layout {layout} can not be converted to psum layout")
+        else:
+            shard.append(iter)
+    return TileLayout.from_iters(shard, [], {})
+
+
+__all__ = ["is_trainium_layout", "to_psum_layout", "trainium_layout"]

python/tvm/backend/trn/op.py

@@ -0,0 +1,153 @@
+# 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.
+"""Trainium-owned NKI intrinsic Python wrappers."""
+
+from __future__ import annotations
+
+from tvm.tirx.op import call_intrin
+
+
+def nki_load(res, data):
+    return call_intrin("", "tirx.nki.load", res, data)
+
+
+def nki_store(res, data):
+    return call_intrin("", "tirx.nki.store", res, data)
+
+
+def nki_tensor_copy(res, data):
+    return call_intrin("", "tirx.nki.tensor_copy", res, data)
+
+
+def nki_matmul(res, lhs, rhs, accum=True):
+    return call_intrin("", "tirx.nki.matmul", res, lhs, rhs, accum)
+
+
+def nki_activation(result, data, opcode, bias=0.0, scale=1.0):
+    return call_intrin("", "tirx.nki.activation", result, data, opcode, bias, scale)
+
+
+def nki_reciprocal(result, data):
+    return call_intrin("", "tirx.nki.reciprocal", result, data)
+
+
+def nki_tensorreduce(result, data, opcode, negate, *axes):
+    return call_intrin("", "tirx.nki.tensorreduce", result, data, opcode, negate, *axes)
+
+
+def nki_tensortensor(result, operand0, operand1, opcode):
+    return call_intrin("", "tirx.nki.tensortensor", result, operand0, operand1, opcode)
+
+
+def nki_tensorscalar(result, operand0, operand1, opcode, reverse=False):
+    return call_intrin("", "tirx.nki.tensorscalar", result, operand0, operand1, opcode, reverse)
+
+
+def nki_memset(result, value):
+    return call_intrin("", "tirx.nki.memset", result, value)
+
+
+def nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias=0.0, scale=1.0):
+    return call_intrin(
+        "",
+        "tirx.nki.activation_reduce",
+        reduce_res,
+        act_res,
+        data,
+        opcode,
+        reduce_opcode,
+        bias,
+        scale,
+    )
+
+
+def nki_tensorscalar_reduce(
+    reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse=False
+):
+    return call_intrin(
+        "",
+        "tirx.nki.tensorscalar_reduce",
+        reduce_res,
+        tensorscalar_res,
+        operand0,
+        operand1,
+        opcode,
+        reduce_opcode,
+        reverse,
+    )
+
+
+def nki_identity(result, size):
+    return call_intrin("", "tirx.nki.identity", result, size)
+
+
+def nki_scalar_tensor_tensor(
+    result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+):
+    return call_intrin(
+        "",
+        "tirx.nki.scalar_tensor_tensor",
+        result,
+        data,
+        operand0,
+        operand1,
+        opcode0,
+        opcode1,
+        reverse0,
+        reverse1,
+    )
+
+
+def nki_scalar_tensor_scalar(
+    result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
+):
+    return call_intrin(
+        "",
+        "tirx.nki.scalar_tensor_scalar",
+        result,
+        data,
+        operand0,
+        operand1,
+        opcode0,
+        opcode1,
+        reverse0,
+        reverse1,
+    )
+
+
+def nki_affine_select(result, pred, true_value, false_value):
+    return call_intrin("", "tirx.nki.affine_select", result, pred, true_value, false_value)
+
+
+__all__ = [
+    "nki_activation",
+    "nki_activation_reduce",
+    "nki_affine_select",
+    "nki_identity",
+    "nki_load",
+    "nki_matmul",
+    "nki_memset",
+    "nki_reciprocal",
+    "nki_scalar_tensor_scalar",
+    "nki_scalar_tensor_tensor",
+    "nki_store",
+    "nki_tensor_copy",
+    "nki_tensorreduce",
+    "nki_tensorscalar",
+    "nki_tensorscalar_reduce",
+    "nki_tensortensor",
+]

python/tvm/backend/trn/operator/__init__.py

@@ -0,0 +1,22 @@
+# 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.
+"""Trainium backend operator package.
+
+Loaded by the Trainium backend registration hook.
+"""
+
+__all__ = ["tile_primitive"]

python/tvm/backend/trn/operator/tile_primitive/__init__.py

@@ -0,0 +1,25 @@
+# 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.
+
+from .binary import *
+from .compose_op import *
+from .copy import *
+from .gemm import *
+from .private_alloc import *
+from .reduction import *
+from .select import *
+from .unary import *