chore: import upstream snapshot with attribution

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

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

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

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

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

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

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

@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Scalar single-thread copy fallback (priority=0)."""
+
+import warnings
+
+import tvm
+from tvm.script import tirx as T
+from tvm.tirx import Buffer, PrimFunc
+from tvm.tirx.operator.tile_primitive.dispatcher import (
+    predicate,
+    register_dispatch,
+)
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ._common import _TID_AXIS_FOR_SCOPE
+from .reg import _axis_decl
+from .utils import _is_valid_copy
+
+
+def _region_st_extent(buffer_region):
+    region = buffer_region.region
+    return [r.min for r in region], [r.extent for r in region]
+
+
+def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    op_call = TilePrimitiveCall.downcast(op_call)
+    src: Buffer = op_call.src.buffer
+    dst: Buffer = op_call.dst.buffer
+    src_st, src_extent = _region_st_extent(op_call.src)
+    dst_st, dst_extent = _region_st_extent(op_call.dst)
+
+    warnings.warn(
+        f"copy/fallback (scalar single-thread) picked for {src.scope()} -> "
+        f"{dst.scope()} at scope_kind={sctx.scope_kind}; all faster variants "
+        f"rejected.",
+        stacklevel=2,
+    )
+
+    def _copy_body(dst_buf, src_buf):
+        dst_indices = [i for i in range(len(dst_buf.shape)) if dst_extent[i] != 1]
+        src_indices = [i for i in range(len(src_buf.shape)) if src_extent[i] != 1]
+        assert len(dst_indices) == len(src_indices)
+        copy_extents = [dst_extent[i] for i in dst_indices]
+
+        def _dst_coord(lvs):
+            if isinstance(lvs, tvm.tirx.Var):
+                lvs = [lvs]
+            coord = list(dst_st)
+            for k, lv in enumerate(lvs):
+                coord[dst_indices[k]] += lv
+            return coord
+
+        def _src_coord(lvs):
+            if isinstance(lvs, tvm.tirx.Var):
+                lvs = [lvs]
+            coord = list(src_st)
+            for k, lv in enumerate(lvs):
+                coord[src_indices[k]] += lv
+            return coord
+
+        with T.grid(*copy_extents) as lvs:
+            T.buffer_store(dst_buf, src_buf[tuple(_src_coord(lvs))], _dst_coord(lvs))
+
+    scope_kind = sctx.scope_kind
+
+    if scope_kind == "thread":
+
+        @T.prim_func(check_well_formed=False)
+        def impl():
+            _copy_body(dst, src)
+
+        return impl
+
+    tid_axis_name = _TID_AXIS_FOR_SCOPE[scope_kind]
+    # first-active tid = composition of per-axis offsets (radix-32, since a warp is 32 lanes)
+    first_tid = int(sctx.intra["laneid"][1])
+    if scope_kind == "warpgroup":
+        first_tid += 32 * int(sctx.intra["wid_in_wg"][1])
+    elif scope_kind == "cta":
+        first_tid += 32 * int(sctx.intra["warpid"][1])
+
+    @T.prim_func(check_well_formed=False)
+    def impl():
+        tid = _axis_decl(tid_axis_name, sctx)
+        if tid == first_tid:
+            _copy_body(dst, src)
+
+    return impl
+
+
+@register_dispatch(
+    "copy",
+    "cuda",
+    variant="fallback",
+    priority=0,
+    when=[predicate("validate_copy_op", _is_valid_copy)],
+)
+def copy_schedule_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
+    return _emit_fallback(op_call, sctx)

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

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

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

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

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

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

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

@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Shared helpers for copy operator dispatches on CUDA targets."""
+
+from collections.abc import Iterable
+
+from tvm.tirx import Buffer
+from tvm.tirx.operator.tile_primitive.registry import DispatchContext
+from tvm.tirx.stmt import TilePrimitiveCall
+
+from ..common import match_scope, validate_copy_op
+
+
+def _is_valid_smem_tmem_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    """Validate smem->tmem copy operation.
+
+    The new tcgen05.cp.32x128b.warpx4 dispatch requires the destination tmem
+    buffer to declare warpx4 broadcast as ``R[4 : 32@TLane]``. The legacy
+    128-row dispatch path (no replica) goes through a separate code path and
+    is not handled here.
+    """
+    dst_region, src_region = op_call.args[:2]
+    src: Buffer = src_region.buffer
+    dst: Buffer = dst_region.buffer
+    if not (src.scope().startswith("shared") and dst.scope() == "tmem"):
+        return (False, f"expected shared->tmem, got {src.scope()}->{dst.scope()}")
+    if not (src.layout and dst.layout):
+        return (False, "both buffers must have layouts")
+    if dst.allocated_addr is None:
+        return (False, "tmem buffer must have allocated_addr")
+    # Require warpx4 router on TMEM side so this dispatch only handles the
+    # 32x128b.warpx4 case; other shapes (128x256b/128x128b etc.) fall back
+    # to the legacy dispatch.
+    rep = dst.layout.replica
+    if not (
+        len(rep) == 1
+        and int(rep[0].extent) == 4
+        and int(rep[0].stride) == 32
+        and "TLane" in str(rep[0].axis)
+    ):
+        return (False, f"requires R[4:32@TLane] on tmem, got replica={list(rep)}")
+    return (True, None)
+
+
+def _single_thread_exec(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    """Predicate: exec scope must be single-thread."""
+    exec_scope = sctx.scope_kind
+    ok = exec_scope == "thread"
+    return (ok, None if ok else f"expected thread exec_scope, got {exec_scope}")
+
+
+DEFAULT_ALLOWED_PAIRS: tuple[tuple[str, str], ...] = (
+    ("global", "shared*"),
+    ("shared*", "global"),
+    ("global", "local"),
+    ("local", "global"),
+    ("shared*", "local"),
+    ("local", "shared*"),
+)
+
+
+def _scope_allowed(
+    op_call: TilePrimitiveCall,
+    sctx: DispatchContext,
+    allowed_pairs: Iterable[tuple[str, str]] = DEFAULT_ALLOWED_PAIRS,
+):
+    op_call = TilePrimitiveCall.downcast(op_call)
+    dst_buffer_region, src_buffer_region = (op_call.dst, op_call.src)
+    src_scope = src_buffer_region.buffer.scope()
+    dst_scope = dst_buffer_region.buffer.scope()
+    ok = any(
+        (
+            match_scope(src_scope, src_pat) and match_scope(dst_scope, dst_pat)
+            for src_pat, dst_pat in allowed_pairs
+        )
+    )
+    if not ok:
+        allowed_str = ", ".join((f"{a}->{b}" for a, b in allowed_pairs))
+        return (
+            False,
+            f"unsupported memory scopes src={src_scope} dst={dst_scope}; allowed: {allowed_str}",
+        )
+    return (True, None)
+
+
+def _is_valid_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
+    return (validate_copy_op(op_call, sctx), "validate_copy_op failed")