chore: import upstream snapshot with attribution

python/tvm/tirx/operator/__init__.py

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# `tile_primitive` defines Python Op classes (`Zero(UnaryOp)`, etc.) whose
+# class bodies call `Op.get("tirx.<name>")` at class-definition time, which
+# requires the compiler-side FFI. Load it lazily so that
+# `tvm.tirx.operator.intrinsics._common` (pure data) and other runtime-safe
+# submodules can be imported under `TVM_USE_RUNTIME_LIB=1`, matching apache's
+# discipline for `tvm.tirx`.
+def __getattr__(name):
+    # `from . import tile_primitive` here would recurse: Python's import
+    # machinery does `getattr(self, 'tile_primitive')` to see if the submodule
+    # is already loaded, which goes back through this __getattr__. Use
+    # importlib.import_module to bypass attribute lookup; it sets the attribute
+    # on the parent package as a side effect, so subsequent lookups go through
+    # the normal attribute path, not this __getattr__.
+    import sys  # pylint: disable=import-outside-toplevel
+    from importlib import import_module  # pylint: disable=import-outside-toplevel
+
+    tp_qualname = f"{__name__}.tile_primitive"
+    tile_primitive = sys.modules.get(tp_qualname) or import_module(tp_qualname)
+    if hasattr(tile_primitive, name):
+        return getattr(tile_primitive, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["get_tirx_op"]

python/tvm/tirx/operator/intrinsics/_common.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.
+"""Shared enum / value tables for PTX intrinsic schemas and user wrappers.
+
+Single source of truth. Both ``tvm.tirx.op`` (user wrappers that validate
+arguments via ``_choice``) and ``tvm.tirx.cuda.operator.intrinsics.*``
+(schema declarations using ``Choice(choices=...)`` / ``IntAttr(choices=...)``)
+import from here.
+
+Adding a new modifier value requires changing exactly one place.
+"""
+
+# Memory ordering / scope -----------------------------------------------------
+FENCE_SEM = ("sc", "acq_rel")
+FENCE_SCOPE = ("cta", "cluster", "gpu", "sys")
+FENCE_PROXY_ASYNC_SPACE = ("", "global", "shared::cta", "shared::cluster")
+CLUSTER_BARRIER_SEM = ("", "release", "relaxed")
+
+# CTA group (used by tcgen05 and TMA) -----------------------------------------
+TCGEN05_CTA_GROUP = (1, 2)
+
+# NVSHMEM ---------------------------------------------------------------------
+NVSHMEM_CMP = ("eq", "ne", "gt", "ge", "lt", "le")
+NVSHMEM_SIG_OP = ("set", "add")
+
+# Floating-point rounding -----------------------------------------------------
+F32X2_ROUND = ("rz", "rn", "rm", "rp")
+
+# cp.async (non-bulk) ---------------------------------------------------------
+CP_ASYNC_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal")
+CP_ASYNC_PREFETCH_SIZE = (-1, 64, 128, 256)
+CP_ASYNC_FILL_MODE = ("", "zero")
+
+# cp.async.bulk (TMA) ---------------------------------------------------------
+CP_ASYNC_BULK_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal", "evict_last_use")
+CP_ASYNC_BULK_RED_OP = ("add", "min", "max", "inc", "dec", "and", "or", "xor")
+
+# ldmatrix / stmatrix ---------------------------------------------------------
+LDMATRIX_DTYPE = (".b16", ".b8")
+LDMATRIX_NUM = (1, 2, 4)
+
+# tcgen05.cp ------------------------------------------------------------------
+TCGEN05_CP_SHAPES = ("32x128b", "4x256b", "128x128b", "128x256b", "64x128b")
+TCGEN05_CP_MULTICAST = ("", "warpx4", "warpx2::02_13", "warpx2::01_23")
+TCGEN05_CP_DECOMPRESS = ("", "b8x16.b4x16_p64", "b8x16.b6x16_p32")
+
+# tcgen05.ld / tcgen05.st -----------------------------------------------------
+TCGEN05_LDST_SHAPES = ("16x32bx2", "16x64b", "16x128b", "16x256b", "32x32b")

python/tvm/tirx/operator/tile_primitive/__init__.py

@@ -0,0 +1,30 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT 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: I001
+
+# Op class declarations (Add, Sub, Gemm, ...) — must run first so their
+# `op = Op.get("tirx.<name>")` registrations execute before any dispatch
+# code refers to the same ops.
+from .ops import *
+
+# Dispatch infrastructure. Per-backend schedule registrations are loaded via
+# ``tvm.backend.load(<name>)``.
+from .dispatcher import fail, list_registered_schedules, predicate, register_dispatch
+from .registry import DispatchContext
+
+__all__ = ["DispatchContext", "fail", "list_registered_schedules", "predicate", "register_dispatch"]

python/tvm/tirx/operator/tile_primitive/common.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.
+"""TIRx operator dispatch common utilities."""
+
+from enum import Enum
+
+
+class MapOpType(Enum):
+    """Enumeration of common unary and binary operator types."""
+
+    ADD = 0
+    SUB = 1
+    MUL = 2
+    FDIV = 3
+    ZERO = 4
+    SQRT = 5
+    RECIPROCAL = 6
+    FILL = 7
+    MAX = 8
+    MIN = 9
+    EXP = 10
+    EXP2 = 11
+    SILU = 12
+
+
+class ReduceOpType(Enum):
+    """Enumeration of common reduce operator types."""
+
+    SUM = 0
+    MAX = 1
+    MIN = 2

python/tvm/tirx/operator/tile_primitive/dispatch_context.py

@@ -0,0 +1,209 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""TIRx operator dispatch context."""
+
+from tvm_ffi import register_object
+
+from tvm.ir import Range
+from tvm.runtime import Object, Scriptable
+from tvm.target import Target
+from tvm.tirx import Buffer, IterVar, Stmt, Var, _ffi_api
+from tvm.tirx.exec_scope import ExecScope
+
+
+@register_object("tirx.DispatchContext")
+class DispatchContext(Object, Scriptable):
+    """DispatchContext node.
+
+    Parameters
+    ----------
+    target : Target
+        The target of the dispatch context.
+
+    exec_scope : ExecScope
+        The execution scope of the dispatch context.
+
+    launch_params : Dict[str, Expr]
+        The launch parameters of the dispatch context.
+
+    var_range_map : Dict[Var, Range]
+        A map from loop variables to their ranges.
+
+    callbacks : Dict[str, Object]
+        The callbacks of the dispatch context.
+
+    shared_state : Dict[str, Object]
+        Shared state persisting across dispatch calls within a single lowering pass.
+    """
+
+    target: Target
+    exec_scope: ExecScope
+    launch_params: dict[str, IterVar]
+    var_range_map: dict[Var, Range]
+    alloc_only: bool
+    callbacks: dict[str, Object]
+    shared_state: dict[str, Object]
+    inter: dict[str, list]
+    intra: dict[str, list]
+    scope_kind: str
+
+    kPrivateAlloc = "private_alloc"
+    kDeviceInitStmt = "device_init_stmt"
+    kHostInitStmt = "host_init_stmt"
+    kPostBufferDefStmt = "post_buffer_def_stmt"
+
+    def __init__(
+        self,
+        target: Target,
+        exec_scope: ExecScope,
+        launch_params: dict[str, IterVar],
+        var_range_map: dict[Var, Range],
+        alloc_only: bool = False,
+        callbacks: dict[str, Object] = {},
+        shared_state: dict[str, Object] = {},
+        inter: dict[str, list] | None = None,
+        intra: dict[str, list] | None = None,
+        scope_kind: str = "",
+    ) -> None:
+        self.__init_handle_by_constructor__(
+            _ffi_api.DispatchContext,  # pylint: disable=no-member
+            target,
+            exec_scope,
+            launch_params,
+            var_range_map,
+            alloc_only,
+            callbacks,
+            shared_state,
+            inter or {},
+            intra or {},
+            scope_kind,
+        )
+
+    def add_alloc_buffer(self, buffer: Buffer) -> None:
+        """Add an allocated buffer to the dispatch context.
+           Can be called only if alloc_only is True.
+           The buffer will be added to the workspace of operator (the key in the workspace is the buffer name).
+
+        Parameters
+        ----------
+        buffer : Buffer
+            The buffer to be added.
+        """  # noqa: E501
+        _ffi_api.DispatchContextAddAllocBuffer(self, buffer)  # pylint: disable=no-member
+
+    def add_init_stmt(self, stmt: Stmt, host: bool = False) -> None:
+        """Add an initialization statement to the dispatch context.
+           Device initialization statements is only allowed if alloc_only is True.
+           Host initialization statements will be ignored if alloc_only is True.
+           The statements will be added to the beginning of the kernel.
+
+        Parameters
+        ----------
+        stmt : Stmt
+            The initialization statement to be added.
+        host : bool
+            Whether the statement is a host statement.
+            If True, the statement will be added to the host code (before the kernel).
+            If False, the statement will be added to the kernel body (at the beginning of the kernel).
+        """  # noqa: E501
+        _ffi_api.DispatchContextAddInitStmt(self, stmt, host)  # pylint: disable=no-member
+
+    def add_post_buffer_def_stmt(self, buffer: Buffer, stmt: Stmt) -> None:
+        """Add a statement to be inserted after a buffer's definition (DeclBuffer/AllocBuffer).
+
+        Parameters
+        ----------
+        buffer : Buffer
+            The buffer whose definition scope the statement should appear in.
+        stmt : Stmt
+            The statement to be inserted.
+        """
+        _ffi_api.DispatchContextAddPostBufferDefStmt(self, buffer, stmt)  # pylint: disable=no-member
+
+    def cache_get(self, key: str) -> Object | None:
+        """Look up a cached value by key.
+
+        Parameters
+        ----------
+        key : str
+            Cache key (built by the caller from construction parameters).
+
+        Returns
+        -------
+        Optional[Object]
+            The cached value, or None on miss.
+        """
+        return _ffi_api.DispatchContextSharedStateGet(self, key)
+
+    def cache_set(self, key: str, value: Object) -> None:
+        """Store a value in the cross-dispatch cache.
+
+        Parameters
+        ----------
+        key : str
+            Cache key (built by the caller from construction parameters).
+        value : Object
+            The object to cache (e.g. a Buffer or Var).
+        """
+        _ffi_api.DispatchContextSharedStateSet(self, key, value)
+
+    def is_cuda(self) -> bool:
+        """Check if the target is CUDA."""
+        return self.target.kind.name == "cuda"
+
+    def is_trn(self) -> bool:
+        """Check if the target is Trainium."""
+        return self.target.kind.name == "trn"
+
+    def is_target(self, name: str) -> bool:
+        """Check if the target kind matches ``name``."""
+        return self.target.kind.name == name
+
+    # -- scope predicates ----------------------------------------------------
+    #
+    # Each ``is_<scope>`` returns True iff the op site is at that scope kind.
+    # Backed by ``self.scope_kind``, which 1-1 maps to a canonical intra
+    # TileLayout shape:
+    #   thread     -> {}
+    #   warp       -> {laneid}
+    #   warpgroup  -> {laneid, wid_in_wg}
+    #   cta        -> {laneid, warpid}
+    #   cluster    -> {laneid, warpid, cta_id}
+    #
+    # Prefer these predicates over raw ``self.scope_kind == "..."`` comparisons
+    # so dispatchers that later need stricter intra/inter shape checks can
+    # tighten the predicate body without touching every call site.
+
+    @property
+    def is_thread(self) -> bool:
+        return self.scope_kind == "thread"
+
+    @property
+    def is_warp(self) -> bool:
+        return self.scope_kind == "warp"
+
+    @property
+    def is_warpgroup(self) -> bool:
+        return self.scope_kind == "warpgroup"
+
+    @property
+    def is_cta(self) -> bool:
+        return self.scope_kind == "cta"
+
+    @property
+    def is_cluster(self) -> bool:
+        return self.scope_kind == "cluster"

python/tvm/tirx/operator/tile_primitive/dispatcher.py

@@ -0,0 +1,329 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Rich dispatcher for TIRx operator dispatchs.
+
+This module adds a structured dispatch table with predicates and
+deterministic failure reporting via exceptions.
+"""
+
+from __future__ import annotations
+
+import traceback
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from tvm.ir import Op
+from tvm.tirx import PrimFunc
+from tvm.tirx.operator import get_tirx_op
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from .dispatch_context import DispatchContext
+
+
+class DispatchFail(RuntimeError):
+    """Raised by variants or predicates to provide a reasoned failure."""
+
+
+@dataclass
+class Predicate:
+    """A named predicate. The callable can return:
+
+    - bool
+    - (bool, str) where the second element is an optional reason on failure
+    - raise DispatchFail(reason)
+    """
+
+    name: str
+    fn: Callable[[TilePrimitiveCall, DispatchContext], Any]
+    kwargs: dict[str, Any]
+
+    def evaluate(
+        self, op_call: TilePrimitiveCall, sctx: DispatchContext
+    ) -> tuple[bool, str | None]:
+        try:
+            out = self.fn(op_call, sctx, **self.kwargs)
+            if isinstance(out, tuple):
+                ok, reason = out
+                return bool(ok), (str(reason) if not ok and reason is not None else None)
+            return bool(out), None
+        except DispatchFail as e:  # surface explicit failure reasons
+            return False, str(e)
+        except Exception as e:  # unexpected predicate exception
+            return False, f"predicate exception: {type(e).__name__}: {e}"
+
+
+def predicate(
+    name: str, fn: Callable[[TilePrimitiveCall, DispatchContext], Any], **kwargs
+) -> Predicate:
+    """Wrap a callable into a named predicate."""
+
+    return Predicate(name=name, fn=fn, kwargs=kwargs)
+
+
+def fail(reason: str) -> None:
+    """Helper for schedule variants to explain why they decline to handle the op."""
+
+    raise DispatchFail(reason)
+
+
+@dataclass
+class DispatchCase:
+    variant: str
+    priority: int
+    preds: list[Predicate]
+    # Impl must either return a PrimFunc or raise DispatchFail
+    impl: Callable[[TilePrimitiveCall, DispatchContext], PrimFunc]
+
+
+# Keyed by (Op, target_kind)
+_DISPATCH_TABLE: dict[tuple[Op, str], list[DispatchCase]] = {}
+
+
+def _target_kind_name(sctx: DispatchContext) -> str:
+    """Normalize target kind to a stable dispatch key."""
+
+    kind = getattr(getattr(sctx, "target", None), "kind", None)
+    return getattr(kind, "name", str(kind))
+
+
+def register_dispatch(
+    op_name: str,
+    target_kind: str,
+    *,
+    variant: str,
+    priority: int = 0,
+    when: list[Predicate] | None = None,
+):
+    """Decorator to add a dispatch case for an op/target pair.
+
+    Cases with higher priority run earlier. When list predicates must all pass.
+    The impl must return a PrimFunc on success, and must NOT return None.
+    To decline handling, raise `fail("reason")` (or `DispatchFail`).
+    """
+
+    op = get_tirx_op(op_name)
+
+    def decorator(impl: Callable[[TilePrimitiveCall, DispatchContext], Any]):
+        # Wrap impl to forbid returning None; require raise-or-PrimFunc
+        def wrapped_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+            res = impl(op_call, sctx)
+            if res is None:
+                # Enforce raise-or-PrimFunc contract for schedule implementations
+                raise DispatchFail(
+                    "impl returned None; schedule must return PrimFunc or raise fail()"
+                )
+            return res  # type: ignore[return-value]
+
+        cases = _DISPATCH_TABLE.setdefault((op, target_kind), [])
+        cases.append(
+            DispatchCase(variant=variant, priority=priority, preds=when or [], impl=wrapped_impl)
+        )
+        return impl
+
+    return decorator
+
+
+def list_registered_schedules() -> dict[str, dict[str, list[str]]]:
+    """Return a mapping: op_name -> target_kind -> [variant names]."""
+
+    out: dict[str, dict[str, list[str]]] = {}
+    for (op, tgt), cases in _DISPATCH_TABLE.items():
+        name = op.name
+        out.setdefault(name, {}).setdefault(tgt, [])
+        # keep insertion order by default; sort by priority desc for readability
+        for c in sorted(cases, key=lambda x: (-x.priority, x.variant)):
+            out[name][tgt].append(c.variant)
+    return out
+
+
+def _format_opcall(op_call: TilePrimitiveCall) -> str:
+    """Return a readable representation of the failing opcall."""
+    # Prefer TVMScript or IR text printer if available on this object
+    try:
+        script_method = getattr(op_call, "script", None)
+        if callable(script_method):
+            try:
+                return str(script_method())
+            except TypeError:
+                # Some versions may require keyword args; fall back safely
+                return str(script_method())
+        astext_method = getattr(op_call, "astext", None)
+        if callable(astext_method):
+            return str(astext_method())
+    except Exception:
+        pass
+    try:
+        s = str(op_call)
+        # constrain extremely long single-line prints from repr
+        return s
+    except Exception:
+        pass
+    try:
+        args_len = len(getattr(op_call, "args", []))
+    except Exception:
+        args_len = -1
+    try:
+        op_name = op_call.op.name  # type: ignore[attr-defined]
+    except Exception:
+        op_name = "<unknown-op>"
+    return f"op={op_name}, args={args_len}"
+
+
+def _format_failure_table(header: str, rows: list[tuple[str, list[str]]]) -> str:
+    """Format failures into a readable ASCII table.
+
+    Parameters
+    ----------
+    header : str
+        The header line describing the op/target
+    rows : List[Tuple[str, str, Optional[str]]]
+        Each row is (variant_label, error_summary, traceback_str)
+
+    Returns
+    -------
+    str
+        The formatted report string
+    """
+    # Compute column widths
+    variant_header = "Variant"
+    error_header = "Error"
+    variant_col_w = (
+        max(len(variant_header), *(len(v) for (v, _) in rows)) if rows else len(variant_header)
+    )
+    # Error column width needs to consider multi-line cells
+    if rows:
+        error_col_w = max(
+            len(error_header), *(max(len(line) for line in errs) for (_, errs) in rows)
+        )
+    else:
+        error_col_w = len(error_header)
+
+    def hline(sep: str = "+") -> str:
+        return f"{sep}{'-' * (variant_col_w + 2)}{sep}{'-' * (error_col_w + 2)}{sep}"
+
+    lines: list[str] = [header]
+    if not rows:
+        # No rows; keep the header only
+        return "\n".join(lines)
+
+    # Table header
+    lines.append(hline("+"))
+    lines.append(f"| {variant_header.ljust(variant_col_w)} | {error_header.ljust(error_col_w)} |")
+    lines.append(hline("+"))
+
+    # Rows (support multi-line Error column)
+    for variant, errs in rows:
+        if not errs:
+            errs = [""]
+        for i, err_line in enumerate(errs):
+            v_text = variant if i == 0 else ""
+            lines.append(f"| {v_text.ljust(variant_col_w)} | {err_line.ljust(error_col_w)} |")
+    lines.append(hline("+"))
+
+    return "\n".join(lines)
+
+
+def run_dispatch(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None:
+    """Run structured dispatch.
+
+    Returns a PrimFunc on success. Otherwise, raises RuntimeError with
+    an aggregated reason report.
+    """
+
+    target_kind = _target_kind_name(sctx)
+    key = (op_call.op, target_kind)
+    cases = _DISPATCH_TABLE.get(key)
+    if not cases:
+        header = f"TIRx schedule dispatch failed: op={op_call.op.name} target={target_kind}"
+        report = _format_failure_table(header, [])
+        # Append a simple reason when there are no variants at all
+        report = "\n".join([report, "no registered variants for this op/target"])
+        raise RuntimeError(report)
+
+    # Collect structured failure rows: (variant_label, error_lines)
+    # error_lines: [summary, traceback lines...]
+    failure_rows: list[tuple[str, list[str]]] = []
+    last_exception: BaseException | None = None
+
+    # If explicit dispatch is set, filter to that variant only
+    forced_variant = getattr(op_call, "dispatch", None)
+    if forced_variant is not None:
+        cases = [c for c in cases if c.variant == forced_variant]
+        if not cases:
+            msg_header = f"TIRx schedule dispatch failed: op={op_call.op.name} target={target_kind}"
+            table = _format_failure_table(msg_header, [])
+            msg = "\n".join([table, f"no variant named '{forced_variant}' is registered"])
+            raise RuntimeError(msg)
+
+    for case in sorted(cases, key=lambda c: (-c.priority, c.variant)):
+        # evaluate predicates
+        pred_ok = True
+        pred_msgs: list[str] = []
+        for pred in case.preds:
+            ok, reason = pred.evaluate(op_call, sctx)
+            if not ok:
+                pred_ok = False
+                msg = f"rejected: {pred.name}"
+                if reason:
+                    msg += f" — {reason}"
+                pred_msgs.append(msg)
+        if not pred_ok:
+            # Include the offending TilePrimitiveCall IR in the error cell
+            op_str = _format_opcall(op_call)
+            op_lines = [line.rstrip("\n") for line in str(op_str).splitlines()] if op_str else []
+            failure_rows.append(
+                (
+                    f"{case.variant} (prio={case.priority})",
+                    ["; ".join(pred_msgs), "opcall:", *op_lines],
+                )
+            )
+            continue
+
+        # run impl
+        try:
+            res = case.impl(op_call, sctx)
+            # Defensive check in case a legacy impl bypassed the wrapper
+            if res is None:  # pragma: no cover - legacy guard
+                raise DispatchFail("impl returned None (legacy behavior not allowed)")
+            return res
+        except DispatchFail as e:
+            op_str = _format_opcall(op_call)
+            op_lines = [line.rstrip("\n") for line in str(op_str).splitlines()] if op_str else []
+            failure_rows.append(
+                (
+                    f"{case.variant} (prio={case.priority})",
+                    [f"declined — {e!s}", "opcall:", *op_lines],
+                )
+            )
+        except Exception as e:  # keep searching other variants
+            exc_summary = f"exception — {type(e).__name__}: {e}"
+            tb_str = "".join(traceback.format_exception(type(e), e, e.__traceback__))
+            # Expand traceback into lines
+            tb_lines = [line.rstrip("\n") for line in tb_str.splitlines()]
+            op_str = _format_opcall(op_call)
+            op_lines = [line.rstrip("\n") for line in str(op_str).splitlines()] if op_str else []
+            error_lines = [exc_summary, "opcall:", *op_lines, *tb_lines]
+            failure_rows.append((f"{case.variant} (prio={case.priority})", error_lines))
+            last_exception = e
+
+    # no success
+    header = f"TIRx schedule dispatch failed: op={op_call.op.name} target={target_kind}"
+    report = _format_failure_table(header, failure_rows)
+    if last_exception is not None:
+        raise RuntimeError(report) from last_exception
+    raise RuntimeError(report)

python/tvm/tirx/operator/tile_primitive/ops.py

@@ -0,0 +1,567 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT 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 TIR operator."""
+
+from tvm.ir import Op
+from tvm.tirx import Expr
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def get_tirx_op(op_name: str):
+    assert isinstance(op_name, str)
+    return Op.get("tirx.tile." + op_name)
+
+
+class ArgProperty:
+    def __init__(self, index):
+        self.index = index
+
+    def __get__(self, obj, objtype=None):
+        assert obj is not None, "TilePrimitiveCall cannot be None"
+        return obj.args[self.index]
+
+
+### Base Operator Classes ###
+class UnaryOp(TilePrimitiveCall):
+    """Base class for unary operators: unary(output, input).
+
+    Unary operators take a single input tensor and produce a single output tensor.
+    """
+
+    scalar_input = False
+    output = ArgProperty(0)
+    input = ArgProperty(1)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expression (input) of the operator."""
+        return [self.input]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expression (output) of the operator."""
+        return [self.output]
+
+
+class UnaryOpWithBiasScale(UnaryOp):
+    """Extended unary operator with bias and scale parameters: unary_with_bias_scale(output, input, bias, scale).
+
+    These operators support additional bias and scale parameters for more complex operations (only on trn).
+    output = unary(input * scale + bias)
+    """  # noqa: E501
+
+    bias = ArgProperty(2)
+    scale = ArgProperty(3)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.input, self.bias, self.scale]
+
+
+class BinaryOp(TilePrimitiveCall):
+    """Base class for binary operators: binary(output, input0, input1).
+
+    Binary operators take two input tensors and produce a single output tensor.
+    """
+
+    lhs = ArgProperty(1)
+    rhs = ArgProperty(2)
+    output = ArgProperty(0)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.lhs, self.rhs]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expression (output) of the operator."""
+        return [self.output]
+
+
+class ReduceOp(TilePrimitiveCall):
+    """Base class for reduction operators: reduce(output, input, reduce_axes, accum).
+
+    Reduction operators reduce one or more dimensions of the input tensor.
+    """
+
+    input = ArgProperty(1)
+    output = ArgProperty(0)
+    reduce_axes = ArgProperty(2)
+    accum = ArgProperty(3)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expression (input) of the operator."""
+        return [self.input]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expression (output) of the operator."""
+        return [self.output]
+
+
+### Schedule Operators ###
+class Zero(UnaryOp):
+    """Zero out all elements in src and store to dst."""
+
+    op = get_tirx_op("zero")
+
+
+class Sqrt(UnaryOpWithBiasScale):
+    """Compute square root of all elements in src and store to dst.
+
+    If bias and scale are provided: dst = sqrt(src * scale + bias)
+    """
+
+    op = get_tirx_op("sqrt")
+
+
+class Fill(UnaryOp):
+    """Fill dst with a scalar value."""
+
+    op = get_tirx_op("fill")
+    scalar_input = True
+
+
+class Add(BinaryOp):
+    """Add src1 and src2 element-wise and store to dst."""
+
+    op = get_tirx_op("add")
+
+
+class Sub(BinaryOp):
+    """Subtract src2 from src1 element-wise and store to dst."""
+
+    op = get_tirx_op("sub")
+
+
+class Mul(BinaryOp):
+    """Multiply src1 and src2 element-wise and store to dst."""
+
+    op = get_tirx_op("mul")
+
+
+class FDiv(BinaryOp):
+    """Divide src1 by src2 element-wise using floating point division and store to dst."""
+
+    op = get_tirx_op("fdiv")
+
+
+class FMA(TilePrimitiveCall):
+    """Fused multiply-add: output = input * scale + bias.
+
+    fma(output, input, scale, bias)
+
+    scale and bias can each be either a BufferRegion or a Expr scalar.
+    """
+
+    op = get_tirx_op("fma")
+
+    output = ArgProperty(0)
+    input = ArgProperty(1)
+    scale = ArgProperty(2)
+    bias = ArgProperty(3)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.input, self.scale, self.bias]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expression (output) of the operator."""
+        return [self.output]
+
+
+class Cast(UnaryOp):
+    """Cast src to dst."""
+
+    op = get_tirx_op("cast")
+
+
+class Copy(TilePrimitiveCall):
+    """Copy all elements from src to dst.
+
+    Args:
+        dst: Destination buffer region
+        src: Source buffer region
+    """
+
+    op = get_tirx_op("copy")
+
+    dst = ArgProperty(0)
+    src = ArgProperty(1)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.src]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expressions (outputs) of the operator."""
+        return [self.dst]
+
+
+class CopyAsync(TilePrimitiveCall):
+    """Copy all elements from src to dst asynchronously.
+
+    Args:
+        dst: Destination buffer region
+        src: Source buffer region
+    """
+
+    op = get_tirx_op("copy_async")
+
+    dst = ArgProperty(0)
+    src = ArgProperty(1)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.src]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expressions (outputs) of the operator."""
+        return [self.dst]
+
+
+class Gemm(TilePrimitiveCall):
+    """General matrix multiplication: D = A * B * alpha + C * beta.
+
+    Args:
+        D: Output matrix
+        A: First input matrix
+        B: Second input matrix
+        C: Third input matrix (for bias)
+        transpose_A: Whether to transpose A
+        transpose_B: Whether to transpose B
+        alpha: Scalar multiplier for A*B
+        beta: Scalar multiplier for C
+    """
+
+    op = get_tirx_op("gemm")
+    output = ArgProperty(0)
+    lhs = ArgProperty(1)
+    rhs = ArgProperty(2)
+    bias = ArgProperty(3)
+    transpose_A = ArgProperty(4)
+    transpose_B = ArgProperty(5)
+    alpha = ArgProperty(6)
+    beta = ArgProperty(7)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source matrices."""
+        return [self.lhs, self.rhs, self.bias]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination matrix."""
+        return [self.output]
+
+
+class GemmAsync(TilePrimitiveCall):
+    """General matrix multiplication asynchronously.
+
+    Supports two arg layouts:
+    - Regular (6 args): C, A, B, transA, transB, accum
+    - Block-scaled (8 args): C, A, B, SFA, SFB, transA, transB, accum
+    """
+
+    op = get_tirx_op("gemm_async")
+    output = ArgProperty(0)
+    lhs = ArgProperty(1)
+    rhs = ArgProperty(2)
+
+    @property
+    def is_block_scaled(self) -> bool:
+        """Whether this is a block-scaled MMA operation."""
+        return len(self.args) == 8
+
+    @property
+    def sfa(self):
+        """Get the scale factor buffer for A (None for regular MMA)."""
+        return self.args[3] if self.is_block_scaled else None
+
+    @property
+    def sfb(self):
+        """Get the scale factor buffer for B (None for regular MMA)."""
+        return self.args[4] if self.is_block_scaled else None
+
+    @property
+    def transA(self):
+        return self.args[5] if self.is_block_scaled else self.args[3]
+
+    @property
+    def transB(self):
+        return self.args[6] if self.is_block_scaled else self.args[4]
+
+    @property
+    def accum(self):
+        return self.args[7] if self.is_block_scaled else self.args[5]
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source matrices (including scale factors if block-scaled)."""
+        srcs = [self.lhs, self.rhs]
+        if self.is_block_scaled:
+            srcs.extend([self.sfa, self.sfb])
+        return srcs
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination matrix."""
+        return [self.output]
+
+
+class Sum(ReduceOp):
+    """Sum elements in src along specified axes and store in dst."""
+
+    op = get_tirx_op("sum")
+
+
+class Max(ReduceOp):
+    """Compute maximum value in src along specified axes and store in dst."""
+
+    op = get_tirx_op("max")
+
+
+class Min(ReduceOp):
+    """Compute minimum value in src along specified axes and store in dst."""
+
+    op = get_tirx_op("min")
+
+
+class Reciprocal(UnaryOp):
+    """Compute reciprocal (1/x) for all elements in src and store to dst."""
+
+    op = get_tirx_op("reciprocal")
+
+
+class SiLU(UnaryOp):
+    """Compute SiLU (x * sigmoid(x)) for all elements in src and store to dst."""
+
+    op = get_tirx_op("silu")
+
+
+class Memset(UnaryOp):
+    """Set all elements in dst to a specified value."""
+
+    op = get_tirx_op("memset")
+    scalar_input = True
+
+
+class Maximum(BinaryOp):
+    """Compute element-wise maximum of src1 and src2 and store to dst."""
+
+    op = get_tirx_op("maximum")
+
+
+class Minimum(BinaryOp):
+    """Compute element-wise minimum of src1 and src2 and store to dst."""
+
+    op = get_tirx_op("minimum")
+
+
+class Exp(UnaryOpWithBiasScale):
+    """Compute exponential (e^x) of all elements in src and store to dst.
+
+    If bias and scale are provided: dst = exp(src * scale + bias)
+    """
+
+    op = get_tirx_op("exp")
+
+
+class Exp2(UnaryOpWithBiasScale):
+    """Compute base-2 exponential (2^x) of all elements in src and store to dst.
+
+    If bias and scale are provided: dst = exp2(src * scale + bias)
+    """
+
+    op = get_tirx_op("exp2")
+
+
+class Select(BinaryOp):
+    """Select elements from src1 or src2 based on the predicate.
+
+    select(dst, src1, src2, predicate)
+    """
+
+    op = get_tirx_op("select")
+    predicate = ArgProperty(3)
+
+
+### Compose Ops ###
+class BinaryReduce(TilePrimitiveCall):
+    """Combine a binary operation with a reduction operation.
+
+    binary_reduce(binary_output, reduce_output, binary_input1, binary_input2, binary_op, reduce_op, reduce_axes, )
+    """  # noqa: E501
+
+    op = get_tirx_op("binary_reduce")
+
+    binary_output = ArgProperty(0)
+    reduce_output = ArgProperty(1)
+    binary_input1 = ArgProperty(2)
+    binary_input2 = ArgProperty(3)
+    binary_op = ArgProperty(4)
+    reduce_op = ArgProperty(5)
+    reduce_axes = ArgProperty(6)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.binary_input1, self.binary_input2]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expressions (outputs) of the operator."""
+        return [self.binary_output, self.reduce_output]
+
+
+class UnaryReduce(TilePrimitiveCall):
+    """Combine a unary operation with a reduction operation.
+
+    unary_reduce(unary_output, reduce_output, unary_input, unary_op, reduce_op, bias, scale, reduce_axes)
+    """  # noqa: E501
+
+    op = get_tirx_op("unary_reduce")
+
+    unary_output = ArgProperty(0)
+    reduce_output = ArgProperty(1)
+    unary_input = ArgProperty(2)
+    unary_op = ArgProperty(3)
+    reduce_op = ArgProperty(4)
+    bias = ArgProperty(5)
+    scale = ArgProperty(6)
+    reduce_axes = ArgProperty(7)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.unary_input, self.bias, self.scale]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expressions (outputs) of the operator."""
+        return [self.unary_output, self.reduce_output]
+
+
+class BinaryChain(TilePrimitiveCall):
+    """Chain multiple binary operations together.
+
+    binary_chain(output, data, operand0, operand1, op0, op1, reverse1)
+
+    if not reverse1:
+        output = (operand0 op0 data) op1 operand1
+    else:
+        output = operand1 op1 (operand0 op0 data)
+    """
+
+    op = get_tirx_op("binary_chain")
+
+    output = ArgProperty(0)
+    data = ArgProperty(1)
+    operand0 = ArgProperty(2)
+    operand1 = ArgProperty(3)
+    op0 = ArgProperty(4)
+    op1 = ArgProperty(5)
+    reverse1 = ArgProperty(6)
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        return [self.data, self.operand0, self.operand1]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expressions (outputs) of the operator."""
+        return [self.output]
+
+
+class ReduceNegate(ReduceOp):
+    """
+    Negate the result of a reduction operation.
+
+    reduce_negate(output, input, reduce_axes, accum, reduce_op)
+    """
+
+    op = get_tirx_op("reduce_negate")
+
+    reduce_op = ArgProperty(4)
+
+
+class ComposeOp(TilePrimitiveCall):
+    """Generic operator for composition of multiple operations.
+
+    Must be lowered to specific compose operations before operator-level passes.
+    """
+
+    # TODO: add a pass to lower generic compose_op to specific compose ops
+
+    op = get_tirx_op("compose_op")
+
+    @property
+    def srcs(self) -> list[Expr]:
+        """Get the source expressions (inputs) of the operator."""
+        raise NotImplementedError(
+            "Generic compose_op must be lowered to specific compose ops before operator-level passes"  # noqa: E501
+        )
+
+    @property
+    def dsts(self) -> list[Expr]:
+        """Get the destination expressions (outputs) of the operator."""
+        raise NotImplementedError(
+            "Generic compose_op must be lowered to specific compose ops before operator-level passes"  # noqa: E501
+        )
+
+
+class PermuteLayout(TilePrimitiveCall):
+    """Move data so the buffer's bytes are arranged under a different layout.
+
+    Logical shape is preserved; only the byte placement changes. ``dst`` and
+    ``src`` carry their own TileLayouts; on lowering, the dispatcher reads
+    those layouts and emits a register-staged warp transpose, optionally
+    inserting a bank-conflict-avoiding XOR-swizzle on the per-lane register
+    slots.
+
+    Args: ``permute_layout(dst_region, src_region)``.
+    ``dst`` and ``src`` may alias the same underlying SMEM (in-place).
+    """
+
+    op = get_tirx_op("permute_layout")
+
+    @property
+    def dst(self) -> Expr:
+        return self.args[0]
+
+    @property
+    def src(self) -> Expr:
+        return self.args[1]
+
+    @property
+    def srcs(self) -> list[Expr]:
+        return [self.src]
+
+    @property
+    def dsts(self) -> list[Expr]:
+        return [self.dst]

python/tvm/tirx/operator/tile_primitive/registry.py

@@ -0,0 +1,66 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""TIRx operator dispatch registry.
+
+All operator dispatch is handled by the rich dispatcher. This module exposes
+the global entry `tirx.f_op_dispatcher` used by the C++ lowering pass to query a
+dispatch result.
+"""
+
+from tvm_ffi import register_global_func
+
+from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+# Note: legacy `register_schedule` is intentionally removed.
+
+
+@register_global_func("tirx.f_op_dispatcher")
+def f_op_dispatcher(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    """Find and return a schedule for the operator.
+
+    Parameters
+    ----------
+    op_call : TilePrimitiveCall
+        The operator to be scheduled
+    sctx : DispatchContext
+        The dispatch context
+
+    Returns
+    -------
+    Optional[PrimFunc]
+        The result of the operator implementation
+    """
+    assert sctx.target is not None, "Target not found"
+    (op_call.op, str(sctx.target.kind))
+
+    # Use rich dispatcher for all dispatching
+    try:
+        from .dispatcher import run_dispatch  # local import to avoid cycles
+    except Exception:  # pragma: no cover - fallback if import fails
+        run_dispatch = None  # type: ignore
+
+    if run_dispatch is not None:
+        try:
+            res = run_dispatch(op_call, sctx)
+        except Exception:
+            # propagate exceptions from dispatcher
+            raise
+        if res is not None:
+            return res
+    # Dispatcher reports errors on failure; unreachable on success
+    return None