chore: import upstream snapshot with attribution

tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py

@@ -0,0 +1,249 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring
+"""Tests for the priority=0 ``copy/fallback`` dispatch — scalar single-thread
+emit picked when every higher-priority variant rejects.
+
+The cases here are *intentionally* shaped so ``gmem_smem`` rejects (region
+element count doesn't divide ``thread_cnt``) and ``reg`` / ``ld_stmatrix`` /
+... don't apply (scope pair mismatch). The dispatcher should land on
+fallback, the emit should pick one active thread, and the round-trip
+``A_gmem → A_smem → B_gmem`` should match.
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+
+# Force the fallback dispatch to register before any test compiles a kernel.
+# Without this import, in fresh pytest workers the `copy/fallback` variant
+# isn't yet registered when the dispatcher snapshots its registry.
+from tvm.tirx.cuda.operator.tile_primitive.copy import fallback as _fallback_module  # noqa: F401
+from tvm.tirx.layout import S, TileLayout
+
+
+def _round_trip_shapes_and_threads():
+    """Cases where ``gmem_smem`` rejects on ``n_elements % thread_cnt``.
+
+    Per task: ``(scope, n_threads, shape, why_fallback)``. ``shape`` is small
+    enough that scalar emit is fine, and chosen so the higher-priority
+    variants can't accept it (size doesn't divide thread_cnt).
+    """
+    return [
+        # warp scope, 32 threads, 24 elements (4x6) → 24 % 32 != 0.
+        ("warp", 32, (4, 6), "4*6=24 ∤ 32"),
+        # warp scope, 32 threads, 8 elements (1x8) → 8 % 32 != 0.
+        ("warp", 32, (1, 8), "1*8=8 ∤ 32"),
+        # warpgroup scope, 128 threads, 24 elements (4x6) → 24 % 128 != 0.
+        ("warpgroup", 128, (4, 6), "4*6=24 ∤ 128"),
+        # warpgroup scope, 128 threads, 32 elements (4x8) → 32 % 128 != 0.
+        ("warpgroup", 128, (4, 8), "4*8=32 ∤ 128"),
+        # cta scope, 256 threads, 32 elements (4x8) → 32 % 256 != 0.
+        ("cta", 256, (4, 8), "4*8=32 ∤ 256"),
+        # cta scope, 1024 threads, 64 elements (8x8) → 64 % 1024 != 0.
+        # Mimics the test_partial_reduction sparse-write-back pattern.
+        ("cta", 1024, (8, 8), "8*8=64 ∤ 1024"),
+    ]
+
+
+def _build_round_trip_kernel(scope, n_threads, shape, dtype):
+    """``Tx.copy(A_smem, A); Tx.copy(B, A_smem)`` at the given scope. Both
+    copies hit the same predicates; both should fall to fallback."""
+    s_layout = TileLayout(S[shape])
+    full = tuple(slice(0, d) for d in shape)
+
+    # Each scope variant inserts an explicit ``cta_sync`` between the two
+    # copies — fallback's emit no longer sneaks one in, so the writer/reader
+    # pair on ``A_smem`` would otherwise race.
+    if scope == "warp":
+
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, shape, dtype)
+            B = T.match_buffer(B_ptr, shape, dtype)
+            T.device_entry()
+            T.cta_id([1])
+            T.lane_id([32])
+            T.thread_id([n_threads])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+            Tx.warp.copy(A_smem[full], A[full])
+            T.cuda.cta_sync()
+            Tx.warp.copy(B[full], A_smem[full])
+
+    elif scope == "warpgroup":
+
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, shape, dtype)
+            B = T.match_buffer(B_ptr, shape, dtype)
+            T.device_entry()
+            T.cta_id([1])
+            T.warpgroup_id([n_threads // 128])
+            T.warp_id_in_wg([4])
+            T.lane_id([32])
+            T.thread_id_in_wg([128])
+            T.thread_id([n_threads])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+            Tx.wg.copy(A_smem[full], A[full])
+            T.cuda.cta_sync()
+            Tx.wg.copy(B[full], A_smem[full])
+
+    elif scope == "cta":
+
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, shape, dtype)
+            B = T.match_buffer(B_ptr, shape, dtype)
+            T.device_entry()
+            T.cta_id([1])
+            T.warp_id([n_threads // 32])
+            T.lane_id([32])
+            T.thread_id([n_threads])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+            Tx.cta.copy(A_smem[full], A[full])
+            T.cuda.cta_sync()
+            Tx.cta.copy(B[full], A_smem[full])
+
+    else:
+        raise ValueError(f"unsupported scope {scope!r}")
+
+    return kernel
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+@pytest.mark.parametrize(
+    "scope,n_threads,shape,why",
+    [
+        pytest.param(s, n, sh, w, id=f"{s}-{n}-{'x'.join(map(str, sh))}")
+        for s, n, sh, w in _round_trip_shapes_and_threads()
+    ],
+)
+def test_fallback_round_trip(scope, n_threads, shape, why):
+    """End-to-end: compile + run + compare. Failure means either the
+    dispatcher didn't pick fallback (silent crash earlier) or fallback's
+    emit is wrong (mismatch on B vs A)."""
+    del why
+    dtype = "float32"
+    kernel = _build_round_trip_kernel(scope, n_threads, shape, dtype)
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.warns(UserWarning, match="copy/fallback"):
+        mod = tvm.IRModule({"main": kernel})
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    A_np = tvm.testing.generate_random_array(dtype, shape)
+    B_np = np.zeros(shape, dtype=np_dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, dev)
+        B = tvm.runtime.tensor(B_np, dev)
+        compiled(A, B)
+        np.testing.assert_array_equal(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_fallback_thread_scope():
+    """``T.thread()`` — single thread, no gate. Either ``gmem_smem`` picks
+    it up (n_elements % 1 == 0) or ``fallback`` does — both end up emitting
+    a sensible single-thread copy. We only check the round trip is correct,
+    not which variant fired."""
+    shape = (4, 6)
+    dtype = "float32"
+    s_layout = TileLayout(S[shape])
+    full = tuple(slice(0, d) for d in shape)
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, dtype)
+        B = T.match_buffer(B_ptr, shape, dtype)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([1])
+        A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+        Tx.copy(A_smem[full], A[full])
+        T.cuda.cta_sync()
+        Tx.copy(B[full], A_smem[full])
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": kernel})
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    A_np = tvm.testing.generate_random_array(dtype, shape)
+    B_np = np.zeros(shape, dtype=np_dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, dev)
+        B = tvm.runtime.tensor(B_np, dev)
+        compiled(A, B)
+        np.testing.assert_array_equal(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_fallback_emits_gate():
+    """Compiled CUDA source must contain a single-thread gate so only one
+    active thread executes the scalar copy (not all of them, which would
+    work but be racy + wasteful and indicate gate elision)."""
+    shape = (4, 6)
+    dtype = "float32"
+    s_layout = TileLayout(S[shape])
+    full = tuple(slice(0, d) for d in shape)
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, dtype)
+        B = T.match_buffer(B_ptr, shape, dtype)
+        T.device_entry()
+        T.cta_id([1])
+        T.warp_id([8])  # 256 threads => 8 warps
+        T.lane_id([32])
+        T.thread_id([256])
+        A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+        Tx.cta.copy(A_smem[full], A[full])
+        Tx.cta.copy(B[full], A_smem[full])
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.warns(UserWarning, match="copy/fallback"):
+        mod = tvm.IRModule({"main": kernel})
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+    src = "".join(im.inspect_source() for im in compiled.mod.imports)
+    # The gate compiles to something like ``if (((int)threadIdx.x) == 0)``.
+    # We don't pin the exact spelling; just require an equality predicate
+    # against threadIdx.x somewhere in the source.
+    assert "threadIdx.x" in src
+    # At least one ``== 0`` (or ``== <literal>``) comparison must exist for
+    # the single-thread gate.
+    assert "== 0" in src, "fallback emit didn't produce a tid==0 gate; src:\n" + src[:2000]
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py

@@ -0,0 +1,578 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring
+"""Round-trip tests for the ``gmem_smem`` copy dispatch (synthesized partition).
+
+Pipeline: A_gmem --G2S--> A_smem --S2G--> B_gmem. If either direction is
+wrong the round trip leaves B mismatched against A.
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
+
+
+def _build_kernel(scope, n_threads, shape, dtype):
+    s_layout = TileLayout(S[shape])
+    full_slices = tuple(slice(0, d) for d in shape)
+
+    if scope == "warpgroup":
+
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, shape, dtype)
+            B = T.match_buffer(B_ptr, shape, dtype)
+            T.device_entry()
+            T.cta_id([1])
+            T.warpgroup_id([n_threads // 128])
+            T.warp_id_in_wg([4])
+            T.lane_id([32])
+            T.thread_id_in_wg([128])
+            T.thread_id([n_threads])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+            Tx.wg.copy(A_smem[full_slices], A[full_slices])
+            T.cuda.cta_sync()
+            Tx.wg.copy(B[full_slices], A_smem[full_slices])
+
+    elif scope == "warp":
+
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, shape, dtype)
+            B = T.match_buffer(B_ptr, shape, dtype)
+            T.device_entry()
+            T.cta_id([1])
+            T.lane_id([32])
+            T.thread_id([n_threads])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+            Tx.warp.copy(A_smem[full_slices], A[full_slices])
+            T.cuda.cta_sync()
+            Tx.warp.copy(B[full_slices], A_smem[full_slices])
+
+    elif scope == "cta":
+
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, shape, dtype)
+            B = T.match_buffer(B_ptr, shape, dtype)
+            T.device_entry()
+            T.cta_id([1])
+            T.warp_id([n_threads // 32])
+            T.lane_id([32])
+            T.thread_id([n_threads])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+            Tx.cta.copy(A_smem[full_slices], A[full_slices])
+            T.cuda.cta_sync()
+            Tx.cta.copy(B[full_slices], A_smem[full_slices])
+    else:
+        raise ValueError(f"unsupported scope {scope!r}")
+
+    return kernel
+
+
+# (scope, n_threads, shape) — shape chosen so total / T / vec_len > 1 with at
+# least one outer round, and total is divisible by T*vec_len.
+TASKS = [
+    ("warp", 32, (32, 32)),  # 1024 total, T=32, vec 8 → outer 4
+    ("warp", 32, (32, 64)),  # 2048 total, outer 8
+    ("warpgroup", 128, (128, 32)),  # 4096, outer 4
+    ("warpgroup", 128, (128, 64)),  # 8192, outer 8
+    ("warpgroup", 128, (256, 16)),  # 4096, outer 4
+    ("cta", 256, (256, 32)),  # 8192, T=256, vec 8 → outer 4
+    ("cta", 256, (512, 16)),  # 8192, outer 4
+]
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+@pytest.mark.parametrize(
+    "scope,n_threads,shape",
+    [pytest.param(*t, id=f"{t[0]}-{t[1]}-{'x'.join(map(str, t[2]))}") for t in TASKS],
+)
+@pytest.mark.parametrize("dtype", ["float16", "float32", "uint8"])
+def test_gmem_smem_roundtrip(scope, n_threads, shape, dtype):
+    kernel = _build_kernel(scope, n_threads, shape, dtype)
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": kernel})
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    A_np = tvm.testing.generate_random_array(dtype, shape)
+    B_np = np.zeros(shape, dtype=np_dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, dev)
+        B = tvm.runtime.tensor(B_np, dev)
+        compiled(A, B)
+        np.testing.assert_array_equal(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ----------------------------------------------------------------------------
+# Migrated from test_copy_sync.py: sync G↔S copy via the user-facing
+# Tx.copy() (which dispatches to gmem_smem).
+# ----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "task",
+    [
+        # A[0:128, 0:32] -> A_smem[0:128, 0:32] -> B[0:128, 0:32]
+        (
+            (128, 32),
+            (128, 32),
+            ((0, 128), (0, 32)),
+            32,
+            TileLayout(S[128, 32]),
+            TileLayout(S[128, 32]),
+            TileLayout(S[128, 32]),
+        ),
+        # A[32:64, 32:64] -> A_smem[0:32, 0:32] -> B[32:64, 32:64]
+        (
+            (64, 64),
+            (32, 32),
+            ((32, 64), (32, 64)),
+            32,
+            TileLayout(S[64, 64]),
+            TileLayout(S[64, 64]),
+            TileLayout(S[32, 32]),
+        ),
+        # A[0:1, 0:32, 0:32] -> A_smem[0:32, 0:32] -> B[0:1, 0:32, 0:32]
+        (
+            (4, 32, 32),
+            (32, 32),
+            ((0, 1), (0, 32), (0, 32)),
+            32,
+            TileLayout(S[4, 32, 32]),
+            TileLayout(S[4, 32, 32]),
+            TileLayout(S[32, 32]),
+        ),
+        # A[0:8, 0:8] -> A_smem[0:8, 0:8] -> B[0:8, 0:8]
+        (
+            (16, 16),
+            (8, 8),
+            ((0, 8), (0, 8)),
+            32,
+            TileLayout(S[16, 16]),
+            TileLayout(S[16, 16]),
+            TileLayout(S[8, 8]),
+        ),
+        # A[32:96, 256:512] -> A_smem[0:32, 0:256] -> B[32:96, 256:512] (swizzled)
+        (
+            (96, 512),
+            (32, 256),
+            ((16, 48), (256, 512)),
+            32,
+            TileLayout(S[96, 512]),
+            TileLayout(S[96, 512]),
+            ComposeLayout(SwizzleLayout(3, 3, 3), TileLayout(S[8, 64]))
+            .tile_to((16, 128), (8, 64))
+            .tile_to((32, 256), (16, 128)),
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+@pytest.mark.parametrize(
+    "dtype", ["int8", "float8_e4m3fn", "float8_e5m2", "float16", "bfloat16", "float32"]
+)
+@pytest.mark.parametrize("scope", ["cta", "thread"])
+def test_copy_g2s_s2g(task, dtype, scope):
+    g_shape, s_shape, g_region, thread_cnt, layoutA, layoutB, layoutS = task
+
+    r_smem = tuple(slice(None) for _ in range(len(s_shape)))
+    r_gmem = tuple(slice(g_region[i][0], g_region[i][1]) for i in range(len(g_shape)))
+
+    if scope == "thread":
+        thread_cnt = 1
+
+    @T.prim_func
+    def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA)
+        B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB)
+
+        T.device_entry()
+        T.cta_id([2])
+        T.thread_id([thread_cnt])
+
+        A_smem = T.alloc_buffer(s_shape, dtype, scope="shared", layout=layoutS)
+        # `scope` is parametrized at runtime; select the scope namespace
+        # dynamically (T.cta / T.thread) instead of a literal prefix.
+        getattr(Tx, scope).copy(A_smem[r_smem], A[r_gmem])
+        T.cuda.cta_sync()
+        getattr(Tx, scope).copy(B[r_gmem], A_smem[r_smem])
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": copy_sync})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        np.random.seed(0)
+        A_np = tvm.testing.generate_random_array(dtype, g_shape)
+        B_np = np.zeros(g_shape, dtype=np_dtype)
+
+        B_ref = B_np.copy()
+        B_ref[r_gmem] = A_np[r_gmem]
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            np.testing.assert_allclose(B_ref, B.numpy())
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ----------------------------------------------------------------------------
+# Regression tests for known correctness gaps in ``align_layouts_gs``.
+#
+# These are intentionally algorithm-level (no GPU runtime) and currently
+# XFAIL; flipping them to passing is the contract for the upcoming swizzle /
+# alignment fix.
+# ----------------------------------------------------------------------------
+
+
+def _align(
+    g_layout, g_shape, s_layout, s_shape, elem_bits, thread_cnt, g_region=None, s_region=None
+):
+    from tvm.tirx.cuda.operator.tile_primitive.copy._common import align_layouts_gs
+
+    target = tvm.target.Target("cuda")
+    if g_region is None:
+        g_region = [(0, d) for d in g_shape]
+    if s_region is None:
+        s_region = [(0, d) for d in s_shape]
+    with target:
+        return align_layouts_gs(
+            g_layout,
+            g_shape,
+            g_region,
+            s_layout,
+            s_shape,
+            s_region,
+            elem_bits,
+            thread_cnt,
+        )
+
+
+@pytest.mark.xfail(
+    reason="align_layouts_gs ignores swizzle chunk size; "
+    "_extract_tile strips the swizzle wrap before vec_len pick."
+)
+@pytest.mark.parametrize("per_element,expected_max_vec", [(2, 4), (1, 2), (0, 1)])
+def test_swizzled_smem_vec_len_must_fit_chunk(per_element, expected_max_vec):
+    """``SwizzleLayout(per_element, ...)`` keeps the bottom ``per_element``
+    bits unswizzled. vec must stay within that chunk or it crosses an XOR
+    boundary and reads/writes the wrong physical bytes."""
+    shape = (32, 32)  # 1024 fp16 elements total
+    g_layout = TileLayout(S[shape])
+    s_layout = ComposeLayout(SwizzleLayout(per_element, 3, 3), TileLayout(S[shape]))
+    _g, _s, vec_len = _align(g_layout, shape, s_layout, shape, elem_bits=16, thread_cnt=32)
+    chunk_elems = 1 << per_element
+    assert vec_len <= chunk_elems, (
+        f"vec_len={vec_len} crosses swizzle chunk size={chunk_elems} "
+        f"(SwizzleLayout(per_element={per_element}, ...))"
+    )
+
+
+def test_unaligned_strides_must_clamp_vec_len():
+    """G layout with row stride 20 (non-multiple of vec_len=8) → tid=2's
+    base offset = 20 elements * 2 bytes = 40 bytes, which is not 16-byte
+    aligned for a 128-bit vec ld/st (uint4 reinterpret crashes)."""
+    shape = (2, 16)
+    # row stride 20 (instead of 16) — leaves 4-elem gap between rows.
+    g_layout = TileLayout(S[(2, 16) : (20, 1)])
+    s_layout = TileLayout(S[(2, 16) : (20, 1)])
+    _g, s_p, vec_len = _align(g_layout, shape, s_layout, shape, elem_bits=16, thread_cnt=4)
+    # All non-vec strides must be multiples of vec_len so per-thread / per-round
+    # starting offset stays vec-aligned. vec iter is s_p.shard[-1] (always
+    # stride=1 by construction).
+    for it in s_p.shard[:-1]:
+        stride = int(it.stride)
+        assert stride % vec_len == 0, (
+            f"stride={stride} not a multiple of vec_len={vec_len}; "
+            f"per-thread / per-round offset will be misaligned for the vec ld/st"
+        )
+
+
+def test_unaligned_region_offset_must_clamp_vec_len():
+    """Slicing the gmem region at a non-vec-aligned column (e.g. col 3 in
+    fp16) means the per-thread base offset starts at 3 elements = 6 bytes,
+    which is not 16/8/4-byte aligned — vec_len must drop to 1."""
+    shape = (4, 16)
+    g_layout = TileLayout(S[(4, 32)])  # full buffer is 4x32 fp16
+    s_layout = TileLayout(S[(4, 16)])
+    # Take cols [3, 19) — start offset 3 (odd for any vec_len > 1 in fp16).
+    g_region = [(0, 4), (3, 19)]
+    s_region = [(0, 4), (0, 16)]
+    _g, _s, vec_len = _align(
+        g_layout,
+        (4, 32),
+        s_layout,
+        (4, 16),
+        elem_bits=16,
+        thread_cnt=4,
+        g_region=g_region,
+        s_region=s_region,
+    )
+    assert 3 % vec_len == 0, (
+        f"vec_len={vec_len} doesn't divide the region's starting column 3; "
+        f"per-thread base offset will be misaligned for the vec ld/st"
+    )
+
+
+def test_swizzled_smem_emit_must_be_swizzle_aware():
+    """Codegen-level: emitted S address should go through the SwizzleLayout's
+    Apply so the XOR scrambling is honored. Currently emit uses
+    ``s_buf.ptr_to([0,..,0]) + linear_offset`` which only matches a
+    non-swizzled storage layout."""
+    import tvm
+    from tvm.script import tirx as T
+    from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
+
+    shape = (128, 32)
+    s_layout = ComposeLayout(SwizzleLayout(3, 3, 3), TileLayout(S[shape]))
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float16")
+        T.device_entry()
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        T.warp_id_in_wg([4])
+        T.lane_id([32])
+        T.thread_id_in_wg([128])
+        T.thread_id([128])
+        A_smem = T.alloc_buffer(shape, "float16", scope="shared", layout=s_layout)
+        Tx.wg.copy(A_smem[0:128, 0:32], A[0:128, 0:32])
+
+    # NB: pin sm_90 explicitly — the default cuda target falls back to sm_50
+    # when no GPU is detected, which nvcc 13+ rejects. Codegen happens before
+    # nvcc; if the whole tvm.compile pipeline fails, we never see the source.
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_90"})
+    with target:
+        mod = tvm.IRModule({"main": kernel})
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    src = "".join(im.inspect_source() for im in compiled.mod.imports)
+
+    # If emit is swizzle-aware, two ways it shows up in the generated
+    # CUDA:
+    #   1. fallback path emits ``swizzle.apply(linear)``, which lowers
+    #      to a ``^`` (XOR) somewhere in the S-offset computation
+    #      (typically on a separate ``s_off_ptr[0] = ...`` line, not on
+    #      the ``tvm_builtin_pointer_offset`` line itself).
+    #   2. fast path precomputes a ``signed_strides[N]`` register array
+    #      (one per binary outer iter), so each per-iter offset is a
+    #      sum of those strides — fingerprintable by the ``1 - 2 *``
+    #      sign-computation idiom emit_init writes.
+    # XOR-less code paired with no signed_strides init means swizzle
+    # was silently dropped.
+    has_xor = "^" in src
+    has_signed_strides_init = "1 - 2 *" in src or "(1 - 2 *" in src
+    assert has_xor or has_signed_strides_init, (
+        "emitted s_ptr address shows no swizzle handling — no XOR (fallback "
+        "path) and no signed_strides init (fast path)"
+    )
+
+
+def test_layout_permute_copy_preserves_smem_strides():
+    """Regression for the MMA-style K-tiled SMEM layout (``tcgen05_mma_ss_no_tma``):
+
+    ``Tx.copy(A_smem, A)`` where A is plain row-major and A_smem uses the
+    K-tiled MMA layout. The two layouts cover the same byte range but map
+    ``(i, j)`` to *different* physical offsets:
+
+      A      : ``A[i, j]      → i*K + j``         (row-major)
+      A_smem : ``A_smem[i, j] → i*8 + (j//8)*1024 + (j%8)``  (K-tiled)
+
+    Earlier ``align_layouts_gs`` sorted+canonicalized BOTH sides
+    independently. A_smem's three iters all chain by stride
+    (8*128 == 1024, 1*8 == 8), so ``FuseContiguousShardIters`` collapsed
+    A_smem to ``[(8192, 1)]`` — same as A's canonical form. The partition
+    then synthesized identical ``s_p`` and ``g_p`` strides, ``apply``
+    emitted ``tid*8`` on *both* sides, and the copy treated A_smem as
+    row-major. MMA descriptors that re-read A_smem with the K-tiled
+    formula then saw 99% wrong elements.
+
+    Fix: ``align_layouts_gs`` only sorts+canonicalizes G. S is grouped by
+    G's pre-sort iter extents and permuted by G's stride-desc permutation;
+    S's per-iter strides are preserved end-to-end.
+
+    This test asserts the structural property: ``s_p.apply`` for any
+    ``tid > 0`` must produce an offset that differs from G's row-major
+    ``tid * vec_len`` — proving S kept its K-tiled stride 1024.
+    """
+    from tvm.tirx import Var as _TirVar
+    from tvm.tirx.expr import IntImm as _IntImm
+
+    M, K = 128, 64
+    # Plain row-major GMEM.
+    g_layout = TileLayout(S[M, K])
+    # K-tiled SMEM: 3D shape (M, K//8, 8) with strides (8, M*8, 1) — the
+    # MMA descriptor's canonical SWIZZLE=0 layout from
+    # tests/.../codegen/test_codegen_blackwell.py::test_tcgen05_mma_ss_no_tma.
+    s_layout = TileLayout(S[(M, K // 8, 8) : (8, M * 8, 1)])
+
+    g_p, s_p, vec_len = _align(
+        g_layout,
+        (M, K),
+        s_layout,
+        (M, K),
+        elem_bits=16,
+        thread_cnt=128,
+    )
+
+    # vec_len must reach 8 (fp16 → 128-bit vec ld/st).
+    assert vec_len == 8, f"expected vec_len=8 for K-tiled fp16 MMA layout, got {vec_len}"
+
+    # S must keep at least one iter with stride 1024 (the K-tile jump
+    # between 8-elem columns). After the fix, s_p.shard has 4 iters with
+    # strides [128, 8, 1024, 1]; the old (broken) code collapsed to 3
+    # iters all matching g_p's row-major strides [1024, 8, 1].
+    s_strides = [int(it.stride) for it in s_p.shard]
+    assert 1024 in s_strides, (
+        f"s_p strides {s_strides} lost the K-tile stride 1024 — "
+        f"align_layouts_gs collapsed A_smem to row-major and the copy "
+        f"will write A_smem in the wrong layout"
+    )
+
+    # Codegen-level check: s_p.apply on (f=0, tid, v=0) must depend on
+    # ``tid % 8`` (the K-tile jump), not just ``tid * 8`` (row-major).
+    # We pin this by evaluating apply for a couple of concrete tids.
+    target = tvm.target.Target("cuda")
+    with target:
+        apply_shape = [_IntImm("int32", 8), _IntImm("int32", 128), _IntImm("int32", 8)]
+        tid_var = _TirVar("tid", "int32")
+        s_off_expr = s_p.apply(
+            _IntImm("int32", 0),
+            tid_var,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )["m"]
+        g_off_expr = g_p.apply(
+            _IntImm("int32", 0),
+            tid_var,
+            _IntImm("int32", 0),
+            shape=apply_shape,
+        )["m"]
+
+    # G is row-major: g_off(tid) = tid * 8.
+    # S is K-tiled : s_off(tid) = (tid // 8) * 8 + (tid % 8) * 1024.
+    # For tid=1 the two MUST differ — they're identical iff S was
+    # collapsed to row-major (the regression).
+    from tvm.tirx import stmt_functor
+
+    analyzer = tvm.arith.Analyzer()
+    s_off_at_1 = analyzer.simplify(
+        stmt_functor.substitute(s_off_expr, {tid_var: _IntImm("int32", 1)})
+    )
+    g_off_at_1 = analyzer.simplify(
+        stmt_functor.substitute(g_off_expr, {tid_var: _IntImm("int32", 1)})
+    )
+    assert int(s_off_at_1) == 1024, (
+        f"s_p.apply at tid=1 produced offset {s_off_at_1}, expected 1024 "
+        f"(K-tile jump). S was collapsed to row-major somewhere."
+    )
+    assert int(g_off_at_1) == 8, (
+        f"g_p.apply at tid=1 produced offset {g_off_at_1}, expected 8 (row-major)"
+    )
+
+
+# ----------------------------------------------------------------------------
+# Fast-path firing test (positive). Pairs with the var_bounds wiring inside
+# ``gmem_smem._emit_gmem_smem``.
+#
+# Setup: warp-scope 32x64 fp16 G2S/S2G with 128b swizzled SMEM. The outer
+# iter stride is ``thread_cnt * vec_len = 32 * 8 = 256``, which puts the
+# binary-split bj's at {5, 6, 7} — well above the swizzle XOR region (so
+# Case 1.D, signed_stride = +T). The (C1) analyzer check
+# ``bit_bj(s_off // C) == 0`` needs the placeholder var bounded to
+# laneid ∈ [0, 32); the dispatch passes ``var_bounds`` so it can discharge,
+# recognizer accepts, and emit lowers to the
+# ``base_off + sum_j bit_j(f) · signed_strides[j]`` precomputed form.
+# ----------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_gmem_smem_swizzle_fast_path_fires_with_var_bounds():
+    """Warp-scope 32x64 fp16 G2S/S2G with 128b swizzled SMEM. Fast path
+    must fire: a 3-slot ``v_<n>[]`` signed_strides buffer + bit-select adds
+    per outer iter, no per-iter ``swizzle.apply`` XOR splice in the hot path."""
+    import re
+
+    swizzle = SwizzleLayout(3, 3, 3)
+    shape = (32, 64)
+    g_layout = TileLayout(S[shape])
+    s_layout = ComposeLayout(swizzle, TileLayout(S[shape]))
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float16", layout=g_layout)
+        B = T.match_buffer(B_ptr, shape, "float16", layout=g_layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.lane_id([32])
+        T.thread_id([32])
+        smem = T.alloc_buffer(shape, "float16", scope="shared", layout=s_layout)
+        Tx.warp.copy(smem, A[:, :])
+        T.cuda.cta_sync()
+        Tx.warp.copy(B[:, :], smem)
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": kernel})
+        ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = ex.mod.imports[0].inspect_source()
+
+    bitsel = re.findall(r"& 1\) \* v_\d+\[", src)
+    v_decls = re.findall(r"alignas\(\d+\) int v_\d+\[(\d+)\]", src)
+    assert bitsel, (
+        "expected fast-path ``(bit & 1) * v_<n>[i]`` adds; if missing, "
+        "var_bounds wiring may have regressed"
+    )
+    assert "3" in v_decls, (
+        f"expected at least one 3-slot signed_strides buffer for bjs "
+        f"[7, 6, 5]; got decl sizes {v_decls}"
+    )
+
+    # Round-trip correctness.
+    A_np = np.arange(32 * 64, dtype="float16").reshape(shape)
+    B_np = np.zeros(shape, dtype="float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        ex(A, B)
+        np.testing.assert_allclose(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py

@@ -0,0 +1,544 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring
+"""Round-trip tests for the ``ldstmatrix`` copy dispatch.
+
+Pipeline:
+  ld direction: A_gmem → A_smem (per-thread init) → R_local (T.copy dispatch
+                under test) → B_gmem (per-thread write).
+  st direction: A_gmem → R_local (per-thread init) → A_smem (T.copy dispatch
+                under test) → B_gmem (per-thread write).
+
+Both directions must round-trip ``A == B``. Layout strides are constructed
+so that:
+  - trans=False S layout matches step-9's row-major spec (8→p, 4→2, num→q, 2→1).
+  - trans=True  S layout matches step-9's col-major spec (8→1, 4→2p, num→q, 2→p).
+
+Uniform shape ``(*scope_outer, 8, 4, num, 2)`` is used for every num (including
+num=1, which gets an extent-1 placeholder for the num atom).
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout, laneid, tid_in_wg, tx
+
+
+def _compile_src(kernel):
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": kernel})
+    with target:
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    return compiled, compiled.mod.imports[0].inspect_source()
+
+
+# ---------------------------------------------------------------------------
+# Layout builders.
+# ---------------------------------------------------------------------------
+def _r_layout_warp(num):
+    return TileLayout(S[(8, 4, num, 2) : (4 @ laneid, 1 @ laneid, 2, 1)])
+
+
+def _r_layout_warpgroup(num):
+    return TileLayout(S[(4, 8, 4, num, 2) : (32 @ tid_in_wg, 4 @ tid_in_wg, 1 @ tid_in_wg, 2, 1)])
+
+
+def _r_layout_cta(num):
+    return TileLayout(S[(4, 8, 4, num, 2) : (32 @ tx, 4 @ tx, 1 @ tx, 2, 1)])
+
+
+def _s_layout_warp(num, trans):
+    if not trans:
+        return TileLayout(S[(8, 4, num, 2) : (num * 8, 2, 8, 1)])
+    return TileLayout(S[(8, 4, num, 2) : (1, 2 * num * 8, 8, num * 8)])
+
+
+def _s_layout_warpgroup_or_cta(num, trans):
+    if not trans:
+        return TileLayout(S[(4, 8, 4, num, 2) : (64 * num, num * 8, 2, 8, 1)])
+    return TileLayout(S[(4, 8, 4, num, 2) : (64 * num, 1, 16 * num, 8, 8 * num)])
+
+
+# 128b swizzle for fp16 (p=3 ⇒ 8 fp16 chunk; sw=at=3 ⇒ 8-row swizzle period).
+_SWIZZLE_128B = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
+
+
+def _maybe_wrap_swizzle(tile_layout, enable: bool):
+    if not enable:
+        return tile_layout
+    return ComposeLayout(_SWIZZLE_128B, tile_layout)
+
+
+# ---------------------------------------------------------------------------
+# Warp scope kernel builder.
+# ---------------------------------------------------------------------------
+def _build_warp_kernel(num, direction, trans, swizzle=False):
+    r_layout = _r_layout_warp(num)
+    s_layout = _maybe_wrap_swizzle(_s_layout_warp(num, trans), swizzle)
+    s_shape = (8, 4, num, 2)
+    full = (slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
+    M, N = 8, num * 8
+
+    def _coord(row, cp, t, w):
+        # Map per-thread layout coord (row, cp, t, w) to gmem (row, col).
+        if not trans:
+            return row, t * 8 + cp * 2 + w
+        return cp * 2 + w, row + t * 8
+
+    # fmt: off
+    if direction == "ld":
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), "float16")
+            B = T.match_buffer(B_ptr, (M, N), "float16")
+            T.device_entry()
+            T.cta_id([1])
+            T.lane_id([32])
+            tid = T.thread_id([32])
+            A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
+            row = tid // 4
+            cp = tid % 4
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(row, cp, t, w)
+                    A_smem[row, cp, t, w] = A[gr, gc]
+            T.cuda.cta_sync()
+            R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
+            Tx.warp.copy(R_local[full], A_smem[full])
+            r_view = R_local.local()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(row, cp, t, w)
+                    B[gr, gc] = r_view[t * 2 + w]
+    else:  # direction == "st"
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), "float16")
+            B = T.match_buffer(B_ptr, (M, N), "float16")
+            T.device_entry()
+            T.cta_id([1])
+            T.lane_id([32])
+            tid = T.thread_id([32])
+            A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
+            row = tid // 4
+            cp = tid % 4
+            R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
+            r_view = R_local.local()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(row, cp, t, w)
+                    r_view[t * 2 + w] = A[gr, gc]
+            Tx.warp.copy(A_smem[full], R_local[full])
+            T.cuda.cta_sync()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(row, cp, t, w)
+                    B[gr, gc] = A_smem[row, cp, t, w]
+    # fmt: on
+    return kernel, (M, N)
+
+
+# ---------------------------------------------------------------------------
+# Warpgroup scope kernel builder. 4 warps stacked vertically.
+# ---------------------------------------------------------------------------
+def _build_warpgroup_kernel(num, direction, trans, swizzle=False):
+    r_layout = _r_layout_warpgroup(num)
+    s_layout = _maybe_wrap_swizzle(_s_layout_warpgroup_or_cta(num, trans), swizzle)
+    s_shape = (4, 8, 4, num, 2)
+    full = (slice(0, 4), slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
+    M, N = 32, num * 8
+
+    def _coord(wid, row, cp, t, w):
+        if not trans:
+            return wid * 8 + row, t * 8 + cp * 2 + w
+        return wid * 8 + cp * 2 + w, row + t * 8
+
+    # fmt: off
+    if direction == "ld":
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), "float16")
+            B = T.match_buffer(B_ptr, (M, N), "float16")
+            T.device_entry()
+            T.cta_id([1])
+            T.warpgroup_id([1])
+            T.warp_id_in_wg([4])
+            T.lane_id([32])
+            T.thread_id_in_wg([128])
+            tid = T.thread_id([128])
+            A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
+            wid = tid // 32
+            lid = tid % 32
+            row = lid // 4
+            cp = lid % 4
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    A_smem[wid, row, cp, t, w] = A[gr, gc]
+            T.cuda.cta_sync()
+            R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
+            Tx.wg.copy(R_local[full], A_smem[full])
+            r_view = R_local.local()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    B[gr, gc] = r_view[t * 2 + w]
+    else:
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), "float16")
+            B = T.match_buffer(B_ptr, (M, N), "float16")
+            T.device_entry()
+            T.cta_id([1])
+            T.warpgroup_id([1])
+            T.warp_id_in_wg([4])
+            T.lane_id([32])
+            T.thread_id_in_wg([128])
+            tid = T.thread_id([128])
+            A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
+            wid = tid // 32
+            lid = tid % 32
+            row = lid // 4
+            cp = lid % 4
+            R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
+            r_view = R_local.local()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    r_view[t * 2 + w] = A[gr, gc]
+            Tx.wg.copy(A_smem[full], R_local[full])
+            T.cuda.cta_sync()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    B[gr, gc] = A_smem[wid, row, cp, t, w]
+    # fmt: on
+    return kernel, (M, N)
+
+
+# ---------------------------------------------------------------------------
+# CTA scope kernel builder. Same geometry as warpgroup, but R uses ``tx``.
+# ---------------------------------------------------------------------------
+def _build_cta_kernel(num, direction, trans, swizzle=False):
+    r_layout = _r_layout_cta(num)
+    s_layout = _maybe_wrap_swizzle(_s_layout_warpgroup_or_cta(num, trans), swizzle)
+    s_shape = (4, 8, 4, num, 2)
+    full = (slice(0, 4), slice(0, 8), slice(0, 4), slice(0, num), slice(0, 2))
+    M, N = 32, num * 8
+
+    def _coord(wid, row, cp, t, w):
+        if not trans:
+            return wid * 8 + row, t * 8 + cp * 2 + w
+        return wid * 8 + cp * 2 + w, row + t * 8
+
+    # fmt: off
+    if direction == "ld":
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), "float16")
+            B = T.match_buffer(B_ptr, (M, N), "float16")
+            T.device_entry()
+            T.cta_id([1])
+            T.warp_id([4])
+            T.lane_id([32])
+            tid = T.thread_id([128])
+            A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
+            wid = tid // 32
+            lid = tid % 32
+            row = lid // 4
+            cp = lid % 4
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    A_smem[wid, row, cp, t, w] = A[gr, gc]
+            T.cuda.cta_sync()
+            R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
+            Tx.cta.copy(R_local[full], A_smem[full])
+            r_view = R_local.local()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    B[gr, gc] = r_view[t * 2 + w]
+    else:
+        @T.prim_func
+        def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), "float16")
+            B = T.match_buffer(B_ptr, (M, N), "float16")
+            T.device_entry()
+            T.cta_id([1])
+            T.warp_id([4])
+            T.lane_id([32])
+            tid = T.thread_id([128])
+            A_smem = T.alloc_buffer(s_shape, "float16", scope="shared", layout=s_layout)
+            wid = tid // 32
+            lid = tid % 32
+            row = lid // 4
+            cp = lid % 4
+            R_local = T.alloc_buffer(s_shape, "float16", scope="local", layout=r_layout)
+            r_view = R_local.local()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    r_view[t * 2 + w] = A[gr, gc]
+            Tx.cta.copy(A_smem[full], R_local[full])
+            T.cuda.cta_sync()
+            for t in range(num):
+                for w in range(2):
+                    gr, gc = _coord(wid, row, cp, t, w)
+                    B[gr, gc] = A_smem[wid, row, cp, t, w]
+    # fmt: on
+    return kernel, (M, N)
+
+
+_BUILDERS = {
+    "warp": _build_warp_kernel,
+    "warpgroup": _build_warpgroup_kernel,
+    "cta": _build_cta_kernel,
+}
+
+
+@pytest.mark.parametrize("scope", ["warp", "warpgroup", "cta"])
+@pytest.mark.parametrize("trans", [False, True])
+@pytest.mark.parametrize("direction", ["ld", "st"])
+@pytest.mark.parametrize("num", [1, 2, 4])
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_ldstmatrix(scope, trans, direction, num):
+    kernel, (M, N) = _BUILDERS[scope](num, direction, trans)
+    compiled, src = _compile_src(kernel)
+
+    inst = "ldmatrix" if direction == "ld" else "stmatrix"
+    trans_inst = ".trans" if trans else ""
+    expected = f"{inst}.sync.aligned.m8n8.x{num}{trans_inst}.shared.b16"
+    assert expected in src, f"{expected} not emitted; src=\n{src}"
+
+    A_np = np.arange(M * N, dtype="float16").reshape(M, N)
+    B_np = np.zeros((M, N), dtype="float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        compiled(A, B)
+        np.testing.assert_allclose(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Swizzled-S round-trip. Verifies the ldstmatrix dispatch's swizzle fast
+# path (when recognized) and slow path (fallback) both produce correct
+# A == B. The 128b swizzle (p=sw=at=3) is the most common fp16 SMEM
+# swizzle; with it the dispatch's per-tile S offset goes through
+# ``swizzle.apply`` (or its precomputed signed-stride lowering) per mm.
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("scope", ["warp", "warpgroup", "cta"])
+@pytest.mark.parametrize("trans", [False, True])
+@pytest.mark.parametrize("direction", ["ld", "st"])
+@pytest.mark.parametrize("num", [1, 2, 4])
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_ldstmatrix_swizzle(scope, trans, direction, num):
+    kernel, (M, N) = _BUILDERS[scope](num, direction, trans, swizzle=True)
+    compiled, src = _compile_src(kernel)
+
+    inst = "ldmatrix" if direction == "ld" else "stmatrix"
+    trans_inst = ".trans" if trans else ""
+    expected = f"{inst}.sync.aligned.m8n8.x{num}{trans_inst}.shared.b16"
+    assert expected in src, f"{expected} not emitted; src=\n{src}"
+
+    A_np = np.arange(M * N, dtype="float16").reshape(M, N)
+    B_np = np.zeros((M, N), dtype="float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        compiled(A, B)
+        np.testing.assert_allclose(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Multi-iter outer (m_outer > 1) fast-path round-trip. The existing 128b
+# swizzle tests above all have m_outer = 1 (only base_off, no signed_strides
+# bits). These two cover the non-trivial outer iter case:
+#
+#   pow2 case (32x64):  R outermost mem ext = 4 (pow2).
+#                       m_outer iters on S 6D seg 3 = [(4, 512), (2, 32)].
+#                       Binary split → bjs [7, 6, 2]. All BitIter.
+#                       signed_strides buffer has 3 slots.
+#
+#   linear case (40x64): R outermost mem ext = 5 (non-pow2). Stride lands
+#                       the outermost S 6D seg 3 iter at (5, 512). 512 is
+#                       exactly 2^(p+at+sw) = swizzle period → Case 1.D pure,
+#                       so the LinearIter relaxation accepts it.
+#                       outer_iters = [LinearIter(5, 512), BitIter(2, ...)].
+#                       Inner BitIter (bj=2 Case 1.A) is the only slot in
+#                       signed_strides; the outer LinearIter contributes
+#                       ``c * 512`` per mm as a compile-time constant.
+# ---------------------------------------------------------------------------
+def _build_multi_iter_kernel(outer_ext: int):
+    """Warp + R=(outer_ext, 8, 2, 4, 4, 2):(16, 4@laneid, 8, 2, 1@laneid, 1)
+    + 333 swizzle on S. Mem strides 16/8/2/1 on extents outer_ext/2/4/2 are
+    bijective for outer_ext ∈ {4, 5}: max = (outer_ext-1)*16 + 8 + 6 + 1
+    = outer_ext*16 - 1, matching extent product = 16*outer_ext."""
+    shape = (outer_ext, 8, 2, 4, 4, 2)
+    r_layout = TileLayout(S[shape : (16, 4 @ laneid, 8, 2, 1 @ laneid, 1)])
+    s_layout = SwizzleLayout(3, 3, 3)
+    full = tuple(slice(0, e) for e in shape)
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float16")
+        B = T.match_buffer(B_ptr, shape, "float16")
+        T.device_entry()
+        T.cta_id([1])
+        T.lane_id([32])
+        tid = T.thread_id([32])
+        A_smem = T.alloc_buffer(shape, "float16", scope="shared", layout=s_layout)
+        for a in range(outer_ext):
+            for c in range(2):
+                for d in range(4):
+                    for e in range(2):
+                        A_smem[a, tid // 4, c, d, tid % 4, e] = A[a, tid // 4, c, d, tid % 4, e]
+        T.cuda.cta_sync()
+        R_local = T.alloc_buffer(shape, "float16", scope="local", layout=r_layout)
+        Tx.warp.copy(R_local[full], A_smem[full])
+        r_view = R_local.local()
+        for a in range(outer_ext):
+            for c in range(2):
+                for d in range(4):
+                    for e in range(2):
+                        B[a, tid // 4, c, d, tid % 4, e] = r_view[a * 16 + c * 8 + d * 2 + e]
+
+    return kernel, shape
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_ldstmatrix_swizzle_multi_iter_pow2():
+    """32x64 fp16 warp; outer m_outer split into multiple BitIters (no
+    LinearIter). Fast path must fire with a 3-slot signed_strides buffer."""
+    import re
+
+    kernel, shape = _build_multi_iter_kernel(outer_ext=4)
+    compiled, src = _compile_src(kernel)
+    assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src
+
+    # Fast-path fingerprint: 3-slot signed_strides + bit-select uses.
+    assert re.search(r"alignas\(\d+\) int v_\d+\[3\]", src), (
+        "expected 3-slot signed_strides buffer for bjs [7, 6, 2]"
+    )
+    bitsel = re.findall(r"& 1\) \* v_\d+\[", src)
+    assert bitsel, "fast-path bit-select pattern '& 1) * v_<n>[' missing"
+
+    n_elem = 1
+    for e in shape:
+        n_elem *= e
+    A_np = np.arange(n_elem, dtype="float16").reshape(shape)
+    B_np = np.zeros(shape, dtype="float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        compiled(A, B)
+        np.testing.assert_allclose(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_ldstmatrix_tcgen05_warpgroup_atom_emits_ldmatrix():
+    """Regression: a warpgroup ``.16x256b`` tcgen05 register atom loaded from a
+    128B-swizzled SMEM tile must dispatch to ``ldmatrix.x4``.
+
+    The atom's laneid is split across two shard iters (stride 4 and stride 1)
+    plus ``wid_in_wg`` — it only fuses to a single ``tid_in_wg`` thread axis
+    after the slice. With the redundant pre-slice ``canonicalize()`` removed,
+    Step 2 relies on ``TileLayout.Slice``'s built-in global pre-canonicalize, so
+    the slice no longer leaves an ill-formed split-laneid sub-layout that
+    ``GetScope`` would reject (which silently fell back to a scalar reg path).
+    Compile-only (no GPU): asserts the instruction appears in generated source.
+    """
+    from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
+    from tvm.tirx.layout import tcgen05_atom_layout
+
+    m, k = 64, 64
+    # bf16 payload carrying the *fp32* atom layout (mirrors the tf32-prenorm
+    # cast warp's ``a_bf16``: one element per 32-bit slot).
+    reg_layout = tcgen05_atom_layout("16x256b", (m, k), "float32")
+    smem_layout = mma_shared_layout("bfloat16", 3, (m, k))
+
+    @T.prim_func
+    def kernel(s_ptr: T.handle) -> None:
+        smem = T.match_buffer(s_ptr, (m, k), "bfloat16", scope="shared", layout=smem_layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        T.warp_id_in_wg([4])
+        T.lane_id([32])
+        T.thread_id_in_wg([128])
+        a_reg = T.alloc_buffer((m, k), "bfloat16", scope="local", layout=reg_layout)
+        Tx.wg.copy(a_reg, smem)
+
+    _, src = _compile_src(kernel)
+    assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src, (
+        f"warpgroup tcgen05 atom did not emit ldmatrix.x4; src=\n{src}"
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_ldstmatrix_swizzle_multi_iter_linear():
+    """40x64 fp16 warp; outer ext=5 is non-pow2 but stride lands on swizzle
+    period (Case 1.D pure) so the LinearIter relaxation fires. Pattern has
+    a 1-slot signed_strides (inner BitIter bj=2 Case 1.A); outer iter
+    contributes ``c * 512`` per mm as a compile-time constant."""
+    import re
+
+    kernel, shape = _build_multi_iter_kernel(outer_ext=5)
+    compiled, src = _compile_src(kernel)
+    assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src
+
+    # Fast-path fingerprint: 1-slot signed_strides (just the inner BitIter).
+    assert re.search(r"alignas\(\d+\) int v_\d+\[1\]", src), (
+        "expected 1-slot signed_strides buffer (only the inner Case-1.A bj=2)"
+    )
+    bitsel = re.findall(r"& 1\) \* v_\d+\[", src)
+    assert bitsel, "fast-path bit-select pattern missing"
+
+    n_elem = 1
+    for e in shape:
+        n_elem *= e
+    A_np = np.arange(n_elem, dtype="float16").reshape(shape)
+    B_np = np.zeros(shape, dtype="float16")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        compiled(A, B)
+        np.testing.assert_allclose(B.numpy(), A_np)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)

tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py

@@ -0,0 +1,629 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring
+"""Round-trip tests for the ``reg`` copy dispatch.
+
+R = per-thread local (register). The dispatch handles round-trips between R
+and any non-R buffer (``shared*`` or ``global``); ``non_r_scope`` parametrize
+toggles which side is exercised.
+
+Self-contained: each thread direct-stores its row into the non-R buffer (no
+G2S / G2L dispatch needed because each thread writes its own address), the
+dispatch does the inbound copy into R and the outbound copy back, then each
+thread reads its row into ``B``. Round-trip mismatch ⇒ at least one direction
+is wrong.
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+from tvm.tirx.layout import S, TileLayout, laneid, tid_in_wg, tx
+
+
+def _r_layout(scope, shape):
+    if scope == "warpgroup":
+        return TileLayout(S[shape : (1 @ tid_in_wg, 1)])
+    if scope == "warp":
+        return TileLayout(S[shape : (1 @ laneid, 1)])
+    if scope == "cta":
+        return TileLayout(S[shape : (1 @ tx, 1)])
+    raise ValueError(f"unsupported scope {scope!r}")
+
+
+def _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope):
+    """Build a kernel that round-trips data through R via ``non_r_scope``.
+
+    ``non_r_scope == "shared"``: ``A_smem`` is allocated inside the kernel.
+    Kernel signature: ``kernel(B_ptr)``.
+
+    ``non_r_scope == "global"``: a separate gmem ``A`` is the staging area.
+    Kernel signature: ``kernel(A_ptr, B_ptr)``.
+    """
+    shape = (n_threads, k)
+    full_slices = (slice(0, n_threads), slice(0, k))
+    r_layout = _r_layout(scope, shape)
+
+    if non_r_scope == "shared":
+        s_layout = TileLayout(S[shape])
+
+        if scope == "warpgroup":
+
+            @T.prim_func
+            def kernel(B_ptr: T.handle) -> None:
+                B = T.match_buffer(B_ptr, shape, dtype)
+                T.device_entry()
+                T.cta_id([1])
+                T.warpgroup_id([n_threads // 128])
+                T.warp_id_in_wg([4])
+                T.lane_id([32])
+                T.thread_id_in_wg([128])
+                tid = T.thread_id([n_threads])
+                A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+                for kk in range(k):
+                    A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
+                T.cuda.cta_sync()
+                R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
+                Tx.wg.copy(R_local[full_slices], A_smem[full_slices])
+                for kk in range(k):
+                    A_smem[tid, kk] = T.cast(0, dtype)
+                T.cuda.cta_sync()
+                Tx.wg.copy(A_smem[full_slices], R_local[full_slices])
+                T.cuda.cta_sync()
+                for kk in range(k):
+                    B[tid, kk] = A_smem[tid, kk]
+
+        elif scope == "warp":
+
+            @T.prim_func
+            def kernel(B_ptr: T.handle) -> None:
+                B = T.match_buffer(B_ptr, shape, dtype)
+                T.device_entry()
+                T.cta_id([1])
+                T.lane_id([32])
+                tid = T.thread_id([n_threads])
+                A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+                for kk in range(k):
+                    A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
+                T.cuda.cta_sync()
+                R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
+                Tx.warp.copy(R_local[full_slices], A_smem[full_slices])
+                for kk in range(k):
+                    A_smem[tid, kk] = T.cast(0, dtype)
+                T.cuda.cta_sync()
+                Tx.warp.copy(A_smem[full_slices], R_local[full_slices])
+                T.cuda.cta_sync()
+                for kk in range(k):
+                    B[tid, kk] = A_smem[tid, kk]
+
+        elif scope == "cta":
+
+            @T.prim_func
+            def kernel(B_ptr: T.handle) -> None:
+                B = T.match_buffer(B_ptr, shape, dtype)
+                T.device_entry()
+                T.cta_id([1])
+                T.warp_id([n_threads // 32])
+                T.lane_id([32])
+                tid = T.thread_id([n_threads])
+                A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=s_layout)
+                for kk in range(k):
+                    A_smem[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
+                T.cuda.cta_sync()
+                R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
+                Tx.cta.copy(R_local[full_slices], A_smem[full_slices])
+                for kk in range(k):
+                    A_smem[tid, kk] = T.cast(0, dtype)
+                T.cuda.cta_sync()
+                Tx.cta.copy(A_smem[full_slices], R_local[full_slices])
+                T.cuda.cta_sync()
+                for kk in range(k):
+                    B[tid, kk] = A_smem[tid, kk]
+
+        return kernel
+
+    if non_r_scope == "global":
+        if scope == "warpgroup":
+
+            @T.prim_func
+            def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+                A = T.match_buffer(A_ptr, shape, dtype)
+                B = T.match_buffer(B_ptr, shape, dtype)
+                T.device_entry()
+                T.cta_id([1])
+                T.warpgroup_id([n_threads // 128])
+                T.warp_id_in_wg([4])
+                T.lane_id([32])
+                T.thread_id_in_wg([128])
+                tid = T.thread_id([n_threads])
+                for kk in range(k):
+                    A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
+                T.cuda.cta_sync()
+                R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
+                Tx.wg.copy(R_local[full_slices], A[full_slices])
+                for kk in range(k):
+                    A[tid, kk] = T.cast(0, dtype)
+                T.cuda.cta_sync()
+                Tx.wg.copy(A[full_slices], R_local[full_slices])
+                T.cuda.cta_sync()
+                for kk in range(k):
+                    B[tid, kk] = A[tid, kk]
+
+        elif scope == "warp":
+
+            @T.prim_func
+            def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+                A = T.match_buffer(A_ptr, shape, dtype)
+                B = T.match_buffer(B_ptr, shape, dtype)
+                T.device_entry()
+                T.cta_id([1])
+                T.lane_id([32])
+                tid = T.thread_id([n_threads])
+                for kk in range(k):
+                    A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
+                T.cuda.cta_sync()
+                R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
+                Tx.warp.copy(R_local[full_slices], A[full_slices])
+                for kk in range(k):
+                    A[tid, kk] = T.cast(0, dtype)
+                T.cuda.cta_sync()
+                Tx.warp.copy(A[full_slices], R_local[full_slices])
+                T.cuda.cta_sync()
+                for kk in range(k):
+                    B[tid, kk] = A[tid, kk]
+
+        elif scope == "cta":
+
+            @T.prim_func
+            def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+                A = T.match_buffer(A_ptr, shape, dtype)
+                B = T.match_buffer(B_ptr, shape, dtype)
+                T.device_entry()
+                T.cta_id([1])
+                T.warp_id([n_threads // 32])
+                T.lane_id([32])
+                tid = T.thread_id([n_threads])
+                for kk in range(k):
+                    A[tid, kk] = T.cast(tid * 100 + kk + 1, dtype)
+                T.cuda.cta_sync()
+                R_local = T.alloc_buffer(shape, dtype, scope="local", layout=r_layout)
+                Tx.cta.copy(R_local[full_slices], A[full_slices])
+                for kk in range(k):
+                    A[tid, kk] = T.cast(0, dtype)
+                T.cuda.cta_sync()
+                Tx.cta.copy(A[full_slices], R_local[full_slices])
+                T.cuda.cta_sync()
+                for kk in range(k):
+                    B[tid, kk] = A[tid, kk]
+
+        return kernel
+
+    raise ValueError(f"unsupported non_r_scope {non_r_scope!r}")
+
+
+def _expected(shape, dtype):
+    n, k = shape
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    out = np.empty(shape, dtype=np_dtype)
+    for t in range(n):
+        for kk in range(k):
+            out[t, kk] = (t * 100 + kk + 1) % 256 if dtype == "uint8" else t * 100 + kk + 1
+    return out
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+@pytest.mark.parametrize("non_r_scope", ["shared", "global"])
+@pytest.mark.parametrize(
+    "scope,n_threads,k",
+    [
+        ("warpgroup", 128, 16),
+        ("warpgroup", 128, 32),
+        ("warpgroup", 128, 8),
+        ("warp", 32, 8),
+        ("warp", 32, 16),
+        ("cta", 256, 8),
+        ("cta", 256, 16),
+    ],
+)
+@pytest.mark.parametrize("dtype", ["float16", "float32", "uint8"])
+def test_reg_roundtrip(scope, n_threads, k, dtype, non_r_scope):
+    shape = (n_threads, k)
+    kernel = _build_roundtrip_kernel(scope, n_threads, k, dtype, non_r_scope)
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": kernel})
+        compiled = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    B_np = np.zeros(shape, dtype=np_dtype)
+    expected = _expected(shape, dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        B = tvm.runtime.tensor(B_np, dev)
+        if non_r_scope == "shared":
+            compiled(B)
+        else:
+            A_np = np.zeros(shape, dtype=np_dtype)
+            A = tvm.runtime.tensor(A_np, dev)
+            compiled(A, B)
+        np.testing.assert_array_equal(B.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ----------------------------------------------------------------------------
+# Migrated from test_copy_sync.py: sync G↔L copy via Tx.copy() (L = local =
+# per-thread register, so it dispatches to the reg variant).
+# ----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "task",
+    [
+        # A[3:4, 8:16, 8:16] -> A_local[0:8, 0:8] -> B[3:4, 8:16, 8:16]
+        (
+            (4, 16, 16),  # g_shape
+            (8, 8),  # l_shape
+            ((3, 4), (8, 16), (8, 16)),  # g_region
+            1,  # thread_cnt
+            TileLayout(S[4, 16, 16]),  # layoutA
+            TileLayout(S[4, 16, 16]),  # layoutB
+            TileLayout(S[8, 8]),  # layoutLocal
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+@pytest.mark.parametrize(
+    "dtype", ["int8", "float8_e4m3fn", "float8_e5m2", "float16", "bfloat16", "float32"]
+)
+def test_copy_g2l_l2g_vec_load(task, dtype):
+    g_shape, l_shape, g_region, thread_cnt, layoutA, layoutB, layoutLocal = task
+
+    r_lmem = tuple(slice(None) for _ in range(len(l_shape)))
+    r_gmem = tuple(slice(g_region[i][0], g_region[i][1]) for i in range(len(g_shape)))
+
+    @T.prim_func
+    def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA)
+        B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB)
+
+        T.device_entry()
+        T.cta_id([2])
+        T.thread_id([thread_cnt])
+        A_local = T.alloc_buffer(l_shape, dtype, scope="local", layout=layoutLocal)
+        Tx.copy(A_local[r_lmem], A[r_gmem])
+        Tx.copy(B[r_gmem], A_local[r_lmem])
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": copy_sync})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        np.random.seed(0)
+        A_np = tvm.testing.generate_random_array(dtype, g_shape)
+        B_np = np.zeros(g_shape, dtype=np_dtype)
+
+        B_ref = B_np.copy()
+        B_ref[r_gmem] = A_np[r_gmem]
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            np.testing.assert_allclose(B_ref, B.numpy())
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_reg_copy_wg_local_to_swizzled_shared_uses_swizzle_fastpath():
+    """Regression: R→S copy where R has a ``wg_local_layout`` (thread iter
+    ``1 @ tid_in_wg``) must pick the widest vec PTX ``st.shared.v4`` AND use the
+    swizzle fast path (precomputed ``signed_strides`` + per-iter
+    bit-select), not the per-iter ``swizzle.apply()`` fallback.
+
+    Two distinct bugs this test guards against:
+
+    (1) ``_choose_vec_len`` used to include R-side thread-iter strides in
+    its alignment check. ``wg_local_layout``'s thread iter has stride 1;
+    a vec=8 (16-byte) alignment check on ``1 % 8 != 0`` would reject
+    every wider variant and fall to scalar ``copy_16b``. Thread-axis
+    strides are partition-coord (virtual), not storage-physical, so they
+    must be excluded.
+
+    (2) Even at the widest vec, if the outer loop is a runtime serial
+    (Python ``range`` doesn't actually unroll in TVMScript) the swizzle
+    fast path's per-iter constant-fold can't kick in and the
+    ``tvm_builtin_pointer_offset`` swizzle XOR ends up recomputed every
+    iteration. Loop must be ``T.unroll``.
+    """
+    from tvm.tirx.layout import SwizzleLayout, wg_local_layout
+
+    N_THREADS, EPI_N = 128, 64
+    g_shape = (N_THREADS, EPI_N)
+    g_layout = TileLayout(S[g_shape])
+    # 128b swizzle on the SMEM side (per_element=3 ⇒ 8 fp16 atom width).
+    smem_layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3)
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, "float16", layout=g_layout)
+        B = T.match_buffer(B_ptr, g_shape, "float16", layout=g_layout)
+
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([N_THREADS])
+        tid = T.thread_id_in_wg([N_THREADS])
+        reg = T.alloc_buffer(g_shape, "float16", scope="local", layout=wg_local_layout(EPI_N))
+        smem = T.alloc_buffer(g_shape, "float16", scope="shared", layout=smem_layout)
+
+        # Populate the per-thread slice via .local() (decomposes the wg
+        # thread-axis layout into a per-thread 1D view).
+        reg_local = reg.local(EPI_N)
+        for i in T.serial(EPI_N):
+            reg_local[i] = A[tid, i]
+        Tx.wg.copy(smem, reg)
+        T.cuda.cta_sync()
+        for i in T.serial(EPI_N):
+            B[tid, i] = smem[tid, i]
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": kernel})
+        ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = ex.mod.imports[0].inspect_source()
+
+    # (1) Widest variant: 8 fp16 elements per call (16 bytes → v4.u32 st).
+    assert "tvm_builtin_ptx_st" in src, (
+        "expected PTX st in generated CUDA, alignment check fell back to a narrower variant"
+    )
+    assert "st.shared.v4" in src, "expected 128b vector store (st.shared.v4.u32)"
+    assert "tvm_builtin_copy_" not in src, (
+        "copy_xxb helpers appeared — reg dispatch should use PTX ld/st only"
+    )
+    # (2) Swizzle fast path fingerprint:
+    #   * emit_init allocates a size-N int buffer of "signed strides".
+    #   * emit_iter_offset uses bit-select * signed-stride: ``(bit) * v[i]``
+    #     where ``bit = (f >> M) & 1``.
+    # The fallback (per-iter ``swizzle.apply(s_off + ds_per_iter)``) has no
+    # such bit-select * signed-stride pattern.
+    import re
+
+    bitsel_pattern = re.findall(r"& 1\) \* v_\d+\[", src)
+    assert bitsel_pattern, (
+        "fast-path bit-select pattern '& 1) * v_<n>[' not found; "
+        "looks like emit_iter_offset's fast path didn't fire."
+    )
+
+
+# --- tcgen05 D epilogue deposit (tf32_hc_prenorm_gemm) -----------------------
+# Production op: ``Tx.warpgroup.copy(smem_cd_mma, d_reg)`` after ``tcgen05.ld``
+# pulls the M=64 accumulator fragment from TMEM into ``d_reg``, then deposits it
+# into 128B-swizzled MMA SMEM for the subsequent TMA store to gmem D.
+_TCGEN05_D_ATOM = "16x256b"
+_TCGEN05_D_SHAPE = (64, 64)
+_TCGEN05_D_DTYPE = "float32"
+_TCGEN05_D_SWIZZLE = 3  # SwizzleMode 128B → mma_shared_layout(..., 3, shape)
+_TCGEN05_D_SLICE = (slice(0, 64), slice(0, 64))
+
+
+def _tcgen05_d_epilogue_layouts():
+    from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
+    from tvm.tirx.layout import tcgen05_atom_layout
+
+    m, n = _TCGEN05_D_SHAPE
+    reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE)
+    smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n))
+    return m, n, reg_layout, smem_layout
+
+
+def _build_tcgen05_d_epilogue_deposit():
+    """``Tx.wg.copy(smem[slice], d_reg[slice])``: R (tcgen05 atom) → S (128B swizzle)."""
+    from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
+    from tvm.tirx.layout import tcgen05_atom_layout
+
+    m, n = _TCGEN05_D_SHAPE
+    smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n))
+    reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE)
+    sl_m, sl_n = _TCGEN05_D_SLICE
+
+    @T.prim_func
+    def deposit(
+        d_reg: T.Buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout),
+    ) -> None:
+        smem_cd_mma = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="shared", layout=smem_layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        T.warp_id_in_wg([4])
+        T.lane_id([32])
+        T.thread_id_in_wg([128])
+        Tx.wg.copy(smem_cd_mma[sl_m, sl_n], d_reg[sl_m, sl_n])
+
+    return deposit
+
+
+def _compile_tcgen05_d_epilogue_deposit():
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": _build_tcgen05_d_epilogue_deposit()})
+        return tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+
+def test_reg_copy_tcgen05_d_epilogue_deposit_layout_pairing():
+    """Pre-fix bug: canonical ``r_p`` collapses atom ``m`` groups and drops S pairings.
+
+    Copy: ``Tx.wg.copy(smem_cd_mma[0:64,0:64], d_reg[0:64,0:64])`` (R→S).
+    ``d_reg``: ``(64,64)`` fp32 ``tcgen05_atom_layout("16x256b", ...)``.
+    ``smem_cd_mma``: ``(64,64)`` fp32 ``mma_shared_layout(..., swizzle=128B)``.
+    """
+    from tvm.backend.cuda.operator.tile_primitive.copy.reg import (
+        _split_thread_loop,
+        align_layouts_raw,
+    )
+    from tvm.tirx.exec_scope import ExecScope
+    from tvm.tirx.operator.tile_primitive import DispatchContext
+
+    m, n, reg_layout, smem_layout = _tcgen05_d_epilogue_layouts()
+    region = [(0, m), (0, n)]
+    sctx = DispatchContext(
+        tvm.target.Target("cuda"), ExecScope("warpgroup"), {}, {}, scope_kind="warpgroup"
+    )
+    with sctx.target:
+        r_sliced = reg_layout.slice([m, n], region)
+        s_sliced = smem_layout.slice([m, n], region)
+        r_p, s_p, s_seps, r_perm = align_layouts_raw(r_sliced, s_sliced, region)
+
+    r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps)
+    r_iters_bug, s_groups_bug = _split_thread_loop(r_p, s_p, s_seps)
+    mem_extents = [int(it.extent) for it in r_iters]
+    bug_extents = [int(it.extent) for it in r_iters_bug]
+
+    # Fixed path: 3 register m-groups stay 1:1 with 3 S-side groups.
+    assert mem_extents == [8, 2, 2]
+    assert len(r_iters) == len(s_groups) == 3
+    # Pre-fix path (what _split_thread_loop used to take): 1 fused m-iter, only
+    # the first S group is paired — the other two are silently dropped.
+    assert bug_extents == [32]
+    assert len(r_iters_bug) == len(s_groups_bug) == 1
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_reg_copy_tcgen05_d_epilogue_deposit_codegen():
+    """``reg`` dispatch lowers the D epilogue deposit; pre-fix loop only covered half.
+
+    Without the ``r_perm`` fix the emitted outer loop ran ``f < 8`` (one fused
+    register group) instead of ``f < 16`` (three atom m-groups x vec tail), and
+    the swizzled SMEM stores landed in the wrong (row, col) slots.
+    """
+    import re
+
+    ex = _compile_tcgen05_d_epilogue_deposit()
+    src = ex.mod.imports[0].inspect_source()
+
+    assert "copy/fallback" not in src, "reg dispatch must not fall back to scalar copy"
+    assert "tvm_builtin_copy_" not in src, "reg dispatch should emit PTX ld/st only"
+    assert "tvm_builtin_ptx_st" in src
+    assert "st.shared.v2.u32" in src, "fp32 vec=2 → 8B shared store per outer iter"
+
+    loop = re.search(r"for \(int f = 0; f < (\d+)", src)
+    assert loop is not None, "expected reg copy outer loop in generated CUDA"
+    assert loop.group(1) == "16", (
+        f"fixed pairing emits 16 outer stores (pre-fix bug collapsed to 8); got f < {loop.group(1)}"
+    )
+
+
+def _tcgen05_16x256b_row_col(tid_wg: T.int32, lane: T.int32, reg_idx: T.int32):
+    """Map ``(tid_in_wg, reg)`` → logical ``(row, col)`` for ``.16x256b`` fp32 atom."""
+    t0 = lane & T.int32(3)
+    t1 = lane >> 2
+    v0p = reg_idx & T.int32(1)
+    va = (reg_idx >> 1) & T.int32(1)
+    vb = reg_idx >> 2
+    wid = tid_wg >> 5
+    row = t1 + T.int32(8) * va + T.int32(16) * wid
+    col = v0p + T.int32(2) * t0 + T.int32(8) * vb
+    return row, col
+
+
+def _build_tcgen05_d_epilogue_deposit_roundtrip():
+    """Fill ``d_reg``, R→S deposit, S→R reload, dump via ``.local()`` to gmem."""
+    from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout
+    from tvm.tirx.layout import tcgen05_atom_layout
+
+    m, n = _TCGEN05_D_SHAPE
+    regs_per_thread = 32  # ``.16x256b.x8`` fp32: 4 regs/slot x rep=8
+    smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n))
+    reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE)
+    sl_m, sl_n = _TCGEN05_D_SLICE
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (m, n), _TCGEN05_D_DTYPE)
+        B = T.match_buffer(B_ptr, (m, n), _TCGEN05_D_DTYPE)
+        T.device_entry()
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        T.warp_id_in_wg([4])
+        T.lane_id([32])
+        tid_wg = T.thread_id_in_wg([128])
+        lane = T.lane_id([32])
+        d_reg = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout)
+        d_reg_out = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout)
+        smem_cd_mma = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="shared", layout=smem_layout)
+        reg_in = d_reg.local(regs_per_thread)
+        reg_out = d_reg_out.local(regs_per_thread)
+        for r in T.serial(regs_per_thread):
+            row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32"))
+            reg_in[r] = A[row, col]
+        Tx.wg.copy(smem_cd_mma[sl_m, sl_n], d_reg[sl_m, sl_n])
+        T.cuda.cta_sync()
+        Tx.wg.copy(d_reg_out[sl_m, sl_n], smem_cd_mma[sl_m, sl_n])
+        for r in T.serial(regs_per_thread):
+            row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32"))
+            B[row, col] = reg_out[r]
+
+    return kernel
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+def test_reg_copy_tcgen05_d_epilogue_deposit_gpu():
+    """GPU: R→S deposit + S→R reload must preserve the Layout-F register tile.
+
+    Host fills gmem ``A[row,col]=row*100+col``; each thread scatters into ``d_reg``
+    via the ``.16x256b`` (tid,reg)→(row,col) map, runs production
+    ``Tx.wg.copy(smem_cd_mma, d_reg)`` then the inverse
+    ``Tx.wg.copy(d_reg_out, smem_cd_mma)``, and dumps ``d_reg_out`` back to gmem
+    ``B`` through ``.local()``. Pre-fix pairing dropped 2/3 of the S groups —
+    ``max|B-A|`` was hundreds, not 0.
+    """
+    m, n = _TCGEN05_D_SHAPE
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(
+            tvm.IRModule({"main": _build_tcgen05_d_epilogue_deposit_roundtrip()}),
+            target=target,
+            tir_pipeline="tirx",
+        )
+
+    rows = np.arange(m, dtype=np.int32)[:, None]
+    cols = np.arange(n, dtype=np.int32)[None, :]
+    a_np = (rows * 100 + cols).astype(np.float32)
+    b_np = np.zeros((m, n), dtype=np.float32)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        a = tvm.runtime.tensor(a_np, dev)
+        b = tvm.runtime.tensor(b_np, dev)
+        mod(a, b)
+        np.testing.assert_allclose(b.numpy(), a_np, rtol=0, atol=0)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py

@@ -0,0 +1,443 @@
+# 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.
+
+"""Tests for the generic swizzle-aware iter pattern in
+``cuda/copy/_swizzle_iter.py``.
+
+Two layers:
+
+* **Recognizer tests** check that ``try_recognize`` returns the expected
+  ``SwizzlePattern`` (or rejects) for each of conditions (a)+(b)+(c).
+* **Numeric correctness tests** verify the proof empirically: for many
+  ``(M0, k)`` samples, the formula
+  ``apply(M0) + sum_{j : bit_j(k)=1} signed_strides[j]``
+  equals ``apply(M0 + ds_k)`` computed by the layout's own Apply formula.
+  Plus a per-thread-sign-matters test that would fail for a constant-sign
+  implementation, ensuring the test isn't trivially satisfied.
+
+All algorithm-level (no GPU needed). End-to-end emit is tested in
+``test_gmem_smem.py::test_swizzled_smem_emit_must_be_swizzle_aware``.
+"""
+
+import pytest
+
+import tvm
+from tvm.tirx import Var as _TirVar
+from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
+    get_swizzle,
+    try_recognize,
+)
+from tvm.tirx.expr import IntImm as _IntImm
+from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout
+
+# ----------------------------------------------------------------------------
+# Pure-Python reference: SwizzleLayout's Apply, plus the proof's formula.
+# Used as ground truth — both must agree for the proof to hold.
+# ----------------------------------------------------------------------------
+
+
+def py_swizzle_apply(M: int, p: int, sw: int, at: int) -> int:
+    """Pure-Python reimplementation of SwizzleLayoutNode::Apply (swizzle_inner=True):
+    phys = swz_q * C + (M mod C)
+    q = M / C; swz_q = q XOR ((q & outer_mask) >> at)
+    """
+    C = 1 << p
+    q = M // C
+    outer_mask = ((1 << sw) - 1) << at
+    swz_q = q ^ ((q & outer_mask) >> at)
+    return swz_q * C + (M % C)
+
+
+def py_signed_strides(
+    M0: int, p: int, sw: int, at: int, bit_positions: list[int], iter_strides_elems: list[int]
+) -> list[int]:
+    """Pure-Python reimplementation of emit_init's formula. Mirrors:
+    if bj >= sw: sigma_bj = +1                              (mid_bits)
+    else       : sigma_bj = 1 - 2 * bit_(at+bj)(M0/C)        (chunk_bits)
+    signed_strides[j] = sigma_bj * iter_strides_elems[j]
+    """
+    C = 1 << p
+    q = M0 // C
+    out: list[int] = []
+    for bj, stride in zip(bit_positions, iter_strides_elems):
+        if bj >= sw:
+            out.append(stride)
+        else:
+            row_bit = (q >> (at + bj)) & 1
+            sigma = 1 - 2 * row_bit
+            out.append(sigma * stride)
+    return out
+
+
+def py_iter_offset(base_off: int, k: int, signed_strides: list[int]) -> int:
+    """Formula sum: base_off + sum_{j : bit_(n-1-j)(k)=1} signed_strides[j]."""
+    n = len(signed_strides)
+    off = base_off
+    for j in range(n):
+        if (k >> (n - 1 - j)) & 1:
+            off += signed_strides[j]
+    return off
+
+
+def py_outer_ds(k: int, iter_extents: list[int], iter_strides: list[int]) -> int:
+    """Decode a flat outer index k into per-iter coords (matching
+    _flat_outer_coords) and sum coord_i * stride_i for the corresponding ds."""
+    coords: list[int] = []
+    rem = k
+    for ext in reversed(iter_extents):
+        coords.append(rem % ext)
+        rem //= ext
+    coords.reverse()
+    return sum(c * s for c, s in zip(coords, iter_strides))
+
+
+# ----------------------------------------------------------------------------
+# Recognizer tests — verify try_recognize accepts / rejects under (a)+(b)+(c).
+# ----------------------------------------------------------------------------
+
+
+def test_get_swizzle_extracts_from_compose():
+    sw = SwizzleLayout(3, 3, 3)
+    assert get_swizzle(sw) is not None
+    assert get_swizzle(ComposeLayout(sw, TileLayout(S[(64, 64)]))) is not None
+    assert get_swizzle(TileLayout(S[(64, 64)])) is None
+
+
+def test_recognize_nvfp4_case():
+    """nvfp4's epilogue: SwizzleLayout(3,3,3), iter extents [2,2,2] strides
+    [8,16,32], M0 = tid * 64 (each thread starts at col 0 of one row;
+    row_stride 64 = 8 chunks, ensures chunk bits of M0/C are zero for all
+    iter bit positions)."""
+    sw = SwizzleLayout(3, 3, 3)
+    tid = _TirVar("tid", "int32")
+    # M0 = tid * 64 → M0/C = tid * 8 → bits 0,1,2 are 0 (since multiplied by 8).
+    M0 = tid * _IntImm("int32", 64)
+    pat = try_recognize(sw, [2, 2, 2], [8, 16, 32], M0)
+    assert pat is not None
+    assert pat.bit_positions == [0, 1, 2]
+    assert pat.iter_strides_elems == [8, 16, 32]
+    assert pat.n_binary_iters == 3
+
+
+def test_recognize_binary_split():
+    """A single outer iter with extent=4 stride=8 splits into two binary
+    iters with strides 16 and 8 (outermost first, matching _flat_outer_coords)."""
+    sw = SwizzleLayout(3, 3, 3)
+    tid = _TirVar("tid", "int32")
+    M0 = tid * _IntImm("int32", 64)
+    pat = try_recognize(sw, [4], [8], M0)
+    assert pat is not None
+    # Split: stride 8*2 = 16 (outermost), stride 8 (innermost) → bits [1, 0]
+    assert pat.bit_positions == [1, 0]
+    assert pat.iter_strides_elems == [16, 8]
+
+
+def test_recognize_mid_bits():
+    """SwizzleLayout(p=4, sw=2, at=4): chunk bits [0,2), mid bits [2,4),
+    row bits [4,6). An iter at bj=2 lives in mid_bits → sigma is always +1
+    (i.e., the recognizer accepts and the sign formula won't read row bits)."""
+    sw = SwizzleLayout(4, 2, 4)  # C=16, mid_bits cover bits 2..3
+    tid = _TirVar("tid", "int32")
+    # M0/C must have bit 2 == 0. Pick row_stride = 64 (= 4*C) so M0/C = tid*4
+    # which has zeros at bit 0,1, and bit 2 is bit 0 of tid... hmm that varies.
+    # Use row_stride = 128 (= 8*C, contributes 4 to M0/C per tid → bit 2 of M0/C
+    # depends on whether tid is even/odd — not zero. Instead use row_stride such
+    # that M0/C is provably 0 mod 8 = 0 at bits 0..2: row_stride = 256 (= 16*C)
+    # → M0/C = tid*16 → bits 0..3 all 0. iter_mask = bit 2, divisor = C*4 = 64.
+    M0 = tid * _IntImm("int32", 256)
+    pat = try_recognize(sw, [2], [64], M0)  # stride 64 = C * 2^2 → bj=2 (mid)
+    assert pat is not None
+    assert pat.bit_positions == [2]
+    assert pat.iter_strides_elems == [64]
+
+
+def test_reject_not_chunk_aligned():
+    """Condition (a): stride must be a multiple of C."""
+    sw = SwizzleLayout(3, 3, 3)  # C=8
+    tid = _TirVar("tid", "int32")
+    M0 = tid * _IntImm("int32", 64)
+    # stride 4 is not a multiple of C=8 → reject.
+    assert try_recognize(sw, [2], [4], M0) is None
+
+
+def test_reject_carries_into_row_bits():
+    """Condition (b): bj < at. A binary iter with stride C * 2^at lands at
+    bj=at, which would change the row bits → reject."""
+    sw = SwizzleLayout(3, 3, 3)  # at=3, so max bj = 2
+    tid = _TirVar("tid", "int32")
+    M0 = tid * _IntImm("int32", 64)
+    # Strides 8,16,32 OK (bj=0,1,2); 64 → bj=3 → reject.
+    assert try_recognize(sw, [2, 2, 2, 2], [8, 16, 32, 64], M0) is None
+
+
+def test_reject_chunk_overlap():
+    """Condition (c): (M0/C) must have 0 bits at all iter-bit positions per
+    thread. If M0 = tid * 8 (so M0/C = tid), then bit 0 of M0/C is bit 0 of
+    tid — analyzer can't prove this is 0 across all threads, so reject."""
+    sw = SwizzleLayout(3, 3, 3)  # C=8
+    tid = _TirVar("tid", "int32")
+    # M0 = tid * C = tid * 8 → M0/C = tid → bit 0 NOT provably zero.
+    M0 = tid * _IntImm("int32", 8)
+    assert try_recognize(sw, [2], [8], M0) is None
+
+
+def test_recognize_no_outer_iters():
+    """Degenerate case: no outer iter at all. Recognizer returns a trivial
+    pattern (empty bit_positions). Emit will use base_off alone."""
+    sw = SwizzleLayout(3, 3, 3)
+    tid = _TirVar("tid", "int32")
+    M0 = tid * _IntImm("int32", 64)
+    pat = try_recognize(sw, [], [], M0)
+    assert pat is not None
+    assert pat.n_binary_iters == 0
+
+
+# ----------------------------------------------------------------------------
+# Numeric correctness — the PROOF. The formula must equal apply(M0 + ds_k)
+# for all sampled (M0, k) and for non-trivial M0 values per thread.
+# ----------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "p,sw,at,iter_extents,iter_strides,row_stride",
+    [
+        # nvfp4-like (p=sw=at=3, 3 binary iters covering one swizzle row)
+        (3, 3, 3, [2, 2, 2], [8, 16, 32], 64),
+        # single binary iter at chunk_bit position
+        (3, 3, 3, [2], [8], 64),
+        # split-from-extent-4 (one outer becomes two binary)
+        (3, 3, 3, [4], [8], 64),
+        # mid_bits region
+        (4, 2, 4, [2], [64], 256),
+        # mix: one chunk_bit + one mid_bit
+        (
+            3,
+            2,
+            4,
+            [2, 2],
+            [8, 32],
+            256,
+        ),  # C=8, sw=2, at=4 → bj_max=3 for stride 64 → use 32 (bj=2 in mid)
+    ],
+)
+def test_formula_matches_apply_under_conditions(
+    p,
+    sw,
+    at,
+    iter_extents,
+    iter_strides,
+    row_stride,
+):
+    """For every (M0, k) sample, the signed-strides formula must equal
+    py_swizzle_apply(M0 + ds_k). Sweeps multiple per-thread M0 values to
+    catch any per-thread-sign bug (a constant-sign impl would fail here)."""
+    swizzle = SwizzleLayout(p, sw, at)
+    tid = _TirVar("tid", "int32")
+    M0_template = tid * _IntImm("int32", row_stride)
+    pat = try_recognize(swizzle, iter_extents, iter_strides, M0_template)
+    assert pat is not None, (
+        f"recognizer rejected supposedly-valid case p={p},sw={sw},at={at} "
+        f"iter_extents={iter_extents} iter_strides={iter_strides} row_stride={row_stride}"
+    )
+
+    total_iters = 1
+    for ext in iter_extents:
+        total_iters *= ext
+
+    # Per-thread sweep: pick concrete tid values that span a few rows of the
+    # swizzle atom. Tid = 0 alone would hide the sign issue (M0/C row bits all
+    # 0 → all signs +1); larger tids exercise the sign-flip branches.
+    for tid_val in [0, 1, 3, 5, 7, 13, 21]:
+        M0 = tid_val * row_stride
+        base_off = py_swizzle_apply(M0, p, sw, at)
+        ss = py_signed_strides(
+            M0,
+            p,
+            sw,
+            at,
+            pat.bit_positions,
+            pat.iter_strides_elems,
+        )
+        for k in range(total_iters):
+            ds_k = py_outer_ds(k, iter_extents, iter_strides)
+            ground_truth = py_swizzle_apply(M0 + ds_k, p, sw, at)
+            formula = py_iter_offset(base_off, k, ss)
+            assert formula == ground_truth, (
+                f"formula mismatch: p={p},sw={sw},at={at} "
+                f"iter_extents={iter_extents} iter_strides={iter_strides} "
+                f"tid={tid_val} M0={M0} k={k} ds_k={ds_k} "
+                f"apply(M0+ds_k)={ground_truth} formula={formula} "
+                f"signed_strides={ss}"
+            )
+
+
+def test_per_thread_sign_actually_varies():
+    """Guard against a 'constant +1 stride' bug: for the nvfp4-like case,
+    different tids MUST produce different signed_strides[0] (since the
+    formula is sigma_0 = 1 - 2 * bit_(at)(M0/C) and that bit toggles with
+    tid). If a buggy impl always returned +stride, this test would catch it."""
+    p, sw, at = 3, 3, 3
+    row_stride = 64  # M0/C = tid * 8 → bit (at=3) = bit 0 of tid
+    # tid=0 → M0/C bit 3 = 0 → sigma = +1; tid=1 → bit 3 = 1 → sigma = -1
+    ss_even = py_signed_strides(0 * row_stride, p, sw, at, [0], [8])
+    ss_odd = py_signed_strides(1 * row_stride, p, sw, at, [0], [8])
+    assert ss_even != ss_odd, "per-thread sign formula degenerated to constant — proof / impl bug"
+    assert ss_even == [8]
+    assert ss_odd == [-8]
+
+
+# ----------------------------------------------------------------------------
+# Fallback path — when recognizer rejects, per-iter swizzle.apply gives the
+# right answer. This is trivial (we delegate to layout.apply) but documents
+# the contract.
+# ----------------------------------------------------------------------------
+
+
+def test_recognize_linear_iter_pure_case_1d():
+    """Outer iter with non-pow2 ext is accepted IF its stride is a multiple
+    of the swizzle period 2^(p+at+sw) (pure Case 1.D, swizzle has no XOR
+    effect). The iter is stored as a LinearIter (no bit decomposition).
+    """
+    from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
+        _BitIter,
+        _LinearIter,
+    )
+
+    p, sw, at = 3, 3, 3
+    swizzle = SwizzleLayout(p, sw, at)
+    period = 1 << (p + at + sw)  # 512
+    # Outer iter (ext=3, stride=period) — non-pow2 but pure Case 1.D.
+    # Inner iter (ext=2, stride=8) — pow2, Case 1.A (bj=0).
+    pat = try_recognize(swizzle, [3, 2], [period, 8], _IntImm("int32", 0))
+    assert pat is not None
+    assert len(pat.outer_iters) == 2
+    # Outermost (index 0) corresponds to first input iter = the linear one.
+    assert isinstance(pat.outer_iters[0], _LinearIter)
+    assert pat.outer_iters[0].ext == 3
+    assert pat.outer_iters[0].stride == period
+    # Innermost (index 1) is the binary-split iter.
+    assert isinstance(pat.outer_iters[1], _BitIter)
+    assert pat.outer_iters[1].ext == 2
+    assert pat.outer_iters[1].n_bits == 1
+    assert pat.outer_iters[1].slot_start == 0
+    # bit_positions / iter_strides_elems only contain the binary iter's bit.
+    assert pat.bit_positions == [0]  # 8/8 = 2^0
+    assert pat.iter_strides_elems == [8]
+
+
+def test_reject_non_pow2_ext_not_case_1d():
+    """Non-pow2 ext where stride is NOT in pure Case 1.D regime — reject.
+    stride=64 = 2^(p+at) = one atom row, which is Case 1.C (in [at, at+sw))
+    territory and the XOR depends on M0, so the linear path is unsafe."""
+    swizzle = SwizzleLayout(3, 3, 3)
+    pat = try_recognize(swizzle, [3], [64], _IntImm("int32", 0))
+    assert pat is None
+
+
+def test_emit_mixed_linear_bit_correctness():
+    """Brute-force: for a mixed (LinearIter outer, BitIter inner) pattern,
+    emit_iter_offset's prediction must equal the actual swizzle output for
+    every (tid, k) — including the non-pow2 outer extent's coord 2."""
+    from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
+        _LinearIter,
+    )
+
+    p, sw, at = 3, 3, 3
+    swizzle = SwizzleLayout(p, sw, at)
+    period = 1 << (p + at + sw)  # 512
+    iter_extents, iter_strides = [3, 2], [period, 8]
+    tid = _TirVar("tid", "int32")
+    # Inner iter bj=0 in [0, sw); (C1) needs bit_0(M0/C) = 0 ⇒ M0/C even
+    # ⇒ M0 multiple of 16. So row_stride = 16.
+    M0_template = tid * _IntImm("int32", 16)
+    pat = try_recognize(swizzle, iter_extents, iter_strides, M0_template)
+    assert pat is not None
+
+    def py_emit(pattern, signed_strides, base_off, k):
+        off = base_off
+        remaining = k
+        for it in reversed(pattern.outer_iters):
+            c = remaining % it.ext
+            remaining = remaining // it.ext
+            if isinstance(it, _LinearIter):
+                off += c * it.stride
+            else:
+                for b in range(it.n_bits):
+                    bit_pos = it.n_bits - 1 - b
+                    slot = it.slot_start + b
+                    if (c >> bit_pos) & 1:
+                        off += signed_strides[slot]
+        return off
+
+    total_k = iter_extents[0] * iter_extents[1]
+    for tid_val in [0, 1, 5, 7, 13]:
+        M0 = tid_val * 16
+        base_off = py_swizzle_apply(M0, p, sw, at)
+        ss = py_signed_strides(
+            M0,
+            p,
+            sw,
+            at,
+            pat.bit_positions,
+            pat.iter_strides_elems,
+        )
+        for k in range(total_k):
+            ds_k = py_outer_ds(k, iter_extents, iter_strides)
+            ground_truth = py_swizzle_apply(M0 + ds_k, p, sw, at)
+            formula = py_emit(pat, ss, base_off, k)
+            assert formula == ground_truth, (
+                f"mixed mismatch: tid={tid_val} M0={M0} k={k} ds_k={ds_k} "
+                f"truth={ground_truth} formula={formula} ss={ss}"
+            )
+
+
+def test_fallback_path_when_recognizer_rejects():
+    """The recognizer should reject when (c) fails, and the resulting
+    fallback emit (swizzle.apply per iter) is the correct path. This test
+    proves the rejection and demonstrates that the swizzled offset really
+    differs from the linear offset for the rejected case — so a buggy
+    `linear-offset-without-XOR` emit (the pre-fix behavior) would give the
+    wrong answer on at least one (tid, k) sample. The fallback emit, by
+    construction, delegates to swizzle.apply and is thus correct."""
+    p, sw, at = 3, 3, 3
+    swizzle = SwizzleLayout(p, sw, at)
+    tid = _TirVar("tid", "int32")
+    M0_template = tid * _IntImm("int32", 8)  # (c) fails: bit 0 of M0/C = bit 0 of tid
+    pat = try_recognize(swizzle, [2], [8], M0_template)
+    assert pat is None, "recognizer must reject when (c) fails"
+
+    # Demonstrate the swizzled offset differs from linear for at least one
+    # (tid, k) — proves the swizzle is actually non-trivial here and the
+    # broken linear-offset emit would give the wrong physical address.
+    iter_extents, iter_strides = [2], [8]
+    diverging_samples = 0
+    for tid_val in range(16):
+        M0 = tid_val * 8
+        for k in range(2):
+            ds_k = py_outer_ds(k, iter_extents, iter_strides)
+            linear = M0 + ds_k
+            swizzled = py_swizzle_apply(linear, p, sw, at)
+            if swizzled != linear:
+                diverging_samples += 1
+    assert diverging_samples > 0, (
+        "no (tid, k) sample shows swizzled != linear — the swizzle is a "
+        "no-op for this layout, so the test isn't catching anything"
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()