chore: import upstream snapshot with attribution

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,
+)