chore: import upstream snapshot with attribution

tests/python/tirx/__init__.py

@@ -0,0 +1,16 @@
+# 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/python/tirx/codegen/test_codegen_ampere.py

@@ -0,0 +1,220 @@
+# 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
+"""Codegen tests for Ampere (sm_80) warp-level ``mma.sync`` tensor cores.
+
+These exercise the ``T.ptx.mma`` intrinsic directly (not via the gemm
+dispatch). ``ptx.mma`` takes one pointer per 32-bit register for each operand
+(``d_ptrs`` / ``a_ptrs`` / ``b_ptrs`` / ``c_ptrs``), enumerated in the fixed
+PTX register order, so the b32 registers may be scattered in the register file
+while the two packed fp16/bf16 within a b32 stay contiguous. For m16n8k{8,16}
+with f32 accumulation the per-lane register counts are:
+
+    A: 2 inputs per b32 -> k16: 4 b32 (regs 0,2,4,6); k8: 2 b32 (regs 0,2)
+    B: 2 inputs per b32 -> k16: 2 b32 (regs 0,2);     k8: 1 b32 (reg 0)
+    D/C: 4 f32 accumulator registers (0,1,2,3)
+"""
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def _get_source(func: tvm.tirx.PrimFunc):
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source()
+    return src, mod
+
+
+def _np_in(dtype):
+    if dtype == "bfloat16":
+        return __import__("ml_dtypes").bfloat16
+    return np.float16
+
+
+def _run_mma(mod, K, no_c_ptr, np_in):
+    """Run an m16n8kK mma kernel and check D == A @ B (+ C) against numpy."""
+    np.random.seed(0)
+    A_np = np.random.randn(16, K).astype(np_in)
+    B_np = np.random.randn(K, 8).astype(np_in)
+    C_np = np.random.randn(16, 8).astype(np.float32)
+    ref = A_np.astype(np.float32) @ B_np.astype(np.float32)
+    if not no_c_ptr:
+        ref = ref + C_np
+
+    def run_and_check():
+        dev = tvm.device("cuda")
+        D = tvm.runtime.tensor(np.zeros((16, 8), np.float32), device=dev)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        C = tvm.runtime.tensor(C_np, device=dev)
+        mod(D, A, B, C)
+        np.testing.assert_allclose(D.numpy(), ref, atol=1e-2, rtol=1e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("a_type", ["float16", "bfloat16"])
+@pytest.mark.parametrize("no_c_ptr", [False, True])
+def test_ptx_mma_m16n8k16(a_type, no_c_ptr):
+    """m16n8k16 row.col mma, f32 accumulate: A is 16x16 (4 b32/lane), B is 16x8
+    as [K, N] (2 b32/lane), D/C is 16x8 (4 f32/lane)."""
+    if a_type == "bfloat16":
+        pytest.importorskip("ml_dtypes")
+    b_type = a_type
+
+    # fmt: off
+    @T.prim_func
+    def main(
+        D: T.Buffer((16, 8), "float32"),
+        A: T.Buffer((16, 16), a_type),
+        B: T.Buffer((16, 8), b_type),
+        C: T.Buffer((16, 8), "float32"),
+    ):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        D_local = T.alloc_local([4], "float32")
+        A_local = T.alloc_local([8], a_type)
+        B_local = T.alloc_local([4], b_type)
+        C_local = T.alloc_local([4], "float32")
+
+        @T.inline
+        def G2L(buf_local, buf_global, block_8x8, mode="row"):
+            if mode == "row":
+                for i in range(block_8x8):
+                    row = T.meta_var(i % 2 * 8 + tx // 4)
+                    col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
+                    for j in range(2):
+                        buf_local[i * 2 + j] = buf_global[row, col + j]
+            elif mode == "col":
+                for i in range(block_8x8):
+                    row = T.meta_var(i % 2 * 8 + (tx % 4) * 2)
+                    col = T.meta_var(i // 2 * 8 + tx // 4)
+                    for j in range(2):
+                        buf_local[i * 2 + j] = buf_global[row + j, col]
+
+        G2L(D_local, D, 2)
+        G2L(A_local, A, 4)
+        G2L(B_local, B, 2, "col")
+        G2L(C_local, C, 2)
+
+        # One pointer per b32 register, in PTX order: A=4, B=2, D/C=4.
+        d_ptrs = [D_local.ptr_to([i]) for i in range(4)]
+        a_ptrs = [A_local.ptr_to([2 * i]) for i in range(4)]
+        b_ptrs = [B_local.ptr_to([2 * i]) for i in range(2)]
+        if no_c_ptr:
+            T.ptx.mma("m16n8k16", "row", "col", "float32", a_type, b_type, "float32",
+                       d_ptrs, a_ptrs, b_ptrs)
+        else:
+            c_ptrs = [C_local.ptr_to([i]) for i in range(4)]
+            T.ptx.mma("m16n8k16", "row", "col", "float32", a_type, b_type, "float32",
+                       d_ptrs, a_ptrs, b_ptrs, c_ptrs)
+
+        for i in range(2):
+            row = T.meta_var(i % 2 * 8 + tx // 4)
+            col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
+            for j in range(2):
+                D[row, col + j] = D_local[i * 2 + j]
+    # fmt: on
+
+    src, mod = _get_source(main)
+    assert "mma.sync.aligned.m16n8k16.row.col" in src
+    _run_mma(mod, 16, no_c_ptr, _np_in(a_type))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("a_type", ["float16", "bfloat16"])
+@pytest.mark.parametrize("no_c_ptr", [False, True])
+def test_ptx_mma_m16n8k8(a_type, no_c_ptr):
+    """m16n8k8 row.col mma, f32 accumulate: A is 16x8 (2 b32/lane), B is 8x8
+    as [K, N] (1 b32/lane), D/C is 16x8 (4 f32/lane)."""
+    if a_type == "bfloat16":
+        pytest.importorskip("ml_dtypes")
+    b_type = a_type
+
+    # fmt: off
+    @T.prim_func
+    def main(
+        D: T.Buffer((16, 8), "float32"),
+        A: T.Buffer((16, 8), a_type),
+        B: T.Buffer((8, 8), b_type),
+        C: T.Buffer((16, 8), "float32"),
+    ):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        D_local = T.alloc_local([4], "float32")
+        A_local = T.alloc_local([4], a_type)
+        B_local = T.alloc_local([2], b_type)
+        C_local = T.alloc_local([4], "float32")
+
+        @T.inline
+        def G2L(buf_local, buf_global, block_8x8, mode="row"):
+            if mode == "row":
+                for i in range(block_8x8):
+                    row = T.meta_var(i % 2 * 8 + tx // 4)
+                    col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
+                    for j in range(2):
+                        buf_local[i * 2 + j] = buf_global[row, col + j]
+            elif mode == "col":
+                for i in range(block_8x8):
+                    row = T.meta_var(i % 2 * 8 + (tx % 4) * 2)
+                    col = T.meta_var(i // 2 * 8 + tx // 4)
+                    for j in range(2):
+                        buf_local[i * 2 + j] = buf_global[row + j, col]
+
+        G2L(D_local, D, 2)
+        G2L(A_local, A, 2)
+        G2L(B_local, B, 1, "col")
+        G2L(C_local, C, 2)
+
+        # One pointer per b32 register, in PTX order: A=2, B=1, D/C=4.
+        d_ptrs = [D_local.ptr_to([i]) for i in range(4)]
+        a_ptrs = [A_local.ptr_to([2 * i]) for i in range(2)]
+        b_ptrs = [B_local.ptr_to([0])]
+        if no_c_ptr:
+            T.ptx.mma("m16n8k8", "row", "col", "float32", a_type, b_type, "float32",
+                       d_ptrs, a_ptrs, b_ptrs)
+        else:
+            c_ptrs = [C_local.ptr_to([i]) for i in range(4)]
+            T.ptx.mma("m16n8k8", "row", "col", "float32", a_type, b_type, "float32",
+                       d_ptrs, a_ptrs, b_ptrs, c_ptrs)
+
+        for i in range(2):
+            row = T.meta_var(i % 2 * 8 + tx // 4)
+            col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
+            for j in range(2):
+                D[row, col + j] = D_local[i * 2 + j]
+    # fmt: on
+
+    src, mod = _get_source(main)
+    assert "mma.sync.aligned.m16n8k8.row.col" in src
+    _run_mma(mod, 8, no_c_ptr, _np_in(a_type))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/codegen/test_codegen_blackwell.py

@@ -0,0 +1,413 @@
+# 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
+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
+
+
+def _get_source(func: tvm.tirx.PrimFunc) -> str:
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source()
+    return src, mod
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_tmem_alloc_dealloc_relinquish():
+    N_COLS = 512
+    cta_group = 1
+
+    # fmt: off
+    @T.prim_func
+    def test_tmem(A: T.Buffer((16, 16), "float16")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([4])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([128])
+        # tmem_addr = T.alloc_buffer((1,), "uint32", scope="shared", align=8)
+        tmem_addr = T.shared_scalar("uint32")
+
+        # alloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=N_COLS, cta_group=cta_group)
+        T.cuda.cta_sync()
+
+        # dealloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+            T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
+    # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        src, _ = _get_source(test_tmem)
+        assert f"tcgen05.alloc.cta_group::{cta_group}.sync.aligned.shared::cta.b32" in src
+        assert f"tcgen05.dealloc.cta_group::{cta_group}.sync.aligned.b32" in src
+        assert f"tcgen05.relinquish_alloc_permit.cta_group::{cta_group}.sync.aligned" in src
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_mbarrier_try_wait_once_codegen():
+    # fmt: off
+    @T.prim_func
+    def test_try_wait_once(A: T.Buffer((16, 16), "float16")):
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([128])
+        bar = T.shared_scalar("uint64")
+        T.evaluate(T.ptx.mbarrier.try_wait_once(T.address_of(bar), 0, 0))
+    # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        src, _ = _get_source(test_try_wait_once)
+        assert "mbarrier.try_wait.parity.shared::cta.b64" in src
+        assert "selp.u32" in src
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_fence_before_after_thread_sync():
+    # fmt: off
+    @T.prim_func
+    def test_fence(A: T.Buffer((16, 16), "float16")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([4])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([128])
+        T.ptx.tcgen05.fence.before_thread_sync()
+        T.ptx.bar.sync(0, 32)
+        T.ptx.tcgen05.fence.after_thread_sync()
+    # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        src, _ = _get_source(test_fence)
+        assert "tcgen05.fence::after_thread_sync" in src
+        assert "tcgen05.fence::before_thread_sync" in src
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_tcgen05_ld_st_roundtrip():
+    HEIGHT = 128
+    WIDTH = 256
+    N_COLS = 512
+    REPEAT_NUM = 1
+    cta_group = 1
+
+    # fmt: off
+    @T.prim_func
+    def test_ld_st(A: T.Buffer((HEIGHT, WIDTH), "float32"), B: T.Buffer((HEIGHT, WIDTH), "float32")):  # noqa: E501
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([4])
+        lane_id = T.lane_id([32])
+        tx = T.thread_id([128])
+        reg = T.alloc_buffer((WIDTH,), "float32", scope="local")
+        # tmem_addr = T.alloc_buffer((1,), "uint32", scope="shared", align=8)
+        tmem_addr = T.shared_scalar("uint32")
+
+        # alloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=N_COLS, cta_group=cta_group)
+        T.cuda.cta_sync()
+        # GMEM -> RF
+        for i in range(WIDTH):
+            reg[i] = A[tx, i]
+        # RF -> TMEM
+        for i in range(WIDTH):
+            T.ptx.tcgen05.st(tmem_addr, reg[i], shape="32x32b", num=REPEAT_NUM, row=warp_id * 32, col=i)  # noqa: E501
+        T.ptx.tcgen05.wait.st()
+        T.cuda.cta_sync()
+        # reset RF
+        for i in range(WIDTH):
+            reg[i] = 0.0
+        T.cuda.cta_sync()
+        # TMEM -> RF
+        T.ptx.tcgen05.fence.after_thread_sync()
+        for i in range(WIDTH):
+            T.ptx.tcgen05.ld(tmem_addr, reg[i], shape="32x32b", num=REPEAT_NUM, row=warp_id * 32, col=i)  # noqa: E501
+        T.ptx.tcgen05.wait.ld()
+        # RF -> GMEM
+        for i in range(WIDTH):
+            B[tx, i] = reg[i]
+
+        # dealloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+            T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
+    # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        src, mod = _get_source(test_ld_st)
+        assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
+        assert "tcgen05.st.sync.aligned.32x32b.x1.b32" in src
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_np = np.random.randn(HEIGHT, WIDTH).astype("float32")
+        B_np = np.zeros((HEIGHT, WIDTH), dtype="float32")
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        mod(A, B)
+        np.testing.assert_allclose(A.numpy(), B.numpy())
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_tcgen05_cp_ld_roundtrip():
+    dtype = "float32"
+    dtype_bits = tvm.DataType(dtype).bits
+    HEIGHT = 128
+    WIDTH = 64
+    N_COLS = 512
+    REPEAT_NUM = 1
+    SWIZZLE = 0
+    A_layout = T.TileLayout(T.S[(HEIGHT, WIDTH // 4, 4) : (4, HEIGHT * 4, 1)])
+    ldo, sdo = 128, 8
+    cta_group = 1
+
+    # fmt: off
+    @T.prim_func
+    def test_cp_ld(A: T.Buffer((HEIGHT, WIDTH), dtype, layout=T.TileLayout(T.S[(HEIGHT, WIDTH // 4, 4) : (4, HEIGHT * 4, 1)])),  # noqa: E501
+                   B: T.Buffer((HEIGHT, WIDTH), dtype, layout=T.TileLayout(T.S[(HEIGHT, WIDTH // 4, 4) : (4, HEIGHT * 4, 1)]))):  # noqa: E501
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([4])
+        lane_id = T.lane_id([32])
+        tx = T.thread_id([128])
+        A_smem = T.alloc_buffer((HEIGHT, WIDTH), dtype, scope="shared", layout=A_layout)
+        reg = T.alloc_buffer((WIDTH,), dtype, scope="local")
+        # tmem_addr = T.alloc_buffer((1,), "uint32", scope="shared", align=8)
+        tmem_addr = T.shared_scalar("uint32")
+        descA = T.alloc_buffer((1,), "uint64", scope="local")
+        bar = T.alloc_buffer((1,), "uint64", scope="shared", align=8)
+        phase = T.alloc_buffer((1,), "int32", scope="local")
+
+        # alloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=N_COLS, cta_group=cta_group)
+        T.cuda.cta_sync()
+        Tx.cta.copy(A_smem[:, :], A[:, :])
+        T.ptx.fence.proxy_async("shared::cta")
+        T.cuda.cta_sync()
+        # reset RF
+        for i in range(WIDTH):
+            reg[i] = 0.0
+        # SMEM -> TMEM (cp)
+        phase[0] = 0
+        if tx == 0:
+            T.ptx.mbarrier.init(bar.data, 1)
+            for k in range(dtype_bits * WIDTH // 256):
+                T.ptx.tcgen05.encode_matrix_descriptor(descA.data, A_smem.access_ptr("r", offset=A_smem.elem_offset_of([0, k * 8])), ldo=ldo, sdo=sdo, swizzle=SWIZZLE)  # noqa: E501
+                T.ptx.tcgen05.cp(tmem_addr, descA[0], shape="128x256b", cta_group=cta_group, col=k * 256 // 32)  # noqa: E501
+            T.ptx.tcgen05.commit(bar.data, cta_group)
+        T.ptx.mbarrier.try_wait(bar.data, phase[0])
+        phase[0] = phase[0] ^ 1
+        T.cuda.cta_sync()
+        # TMEM -> RF (ld)
+        T.ptx.tcgen05.fence.after_thread_sync()
+        for i in range(WIDTH):
+            T.ptx.tcgen05.ld(tmem_addr, reg[i], shape="32x32b", num=REPEAT_NUM, row=warp_id * 32, col=i)  # noqa: E501
+        T.ptx.tcgen05.wait.ld()
+        # RF -> GMEM
+        for i in range(WIDTH):
+            B[tx, i] = reg[i]
+
+        # dealloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+            T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
+    # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        src, mod = _get_source(test_cp_ld)
+        assert "tcgen05.cp.cta_group::1.128x256b" in src
+        assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_np = np.random.randn(HEIGHT, WIDTH).astype(dtype)
+        B_np = np.zeros((HEIGHT, WIDTH), dtype=dtype)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        mod(A, B)
+        np.testing.assert_allclose(A.numpy(), B.numpy())
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("swizzle", [0, 1, 2, 3])
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_tcgen05_mma_ss_no_tma(swizzle):
+    d_type, a_type, b_type = "float32", "float16", "float16"
+    M, N, K = 128, 128, 64
+    MMA_K = 16
+    N_COLS = 512
+    REPEAT_NUM = 1
+    SWIZZLE = swizzle
+    cta_group = 1
+
+    if SWIZZLE == 0:
+        A_layout = T.TileLayout(T.S[(M, K // 8, 8) : (8, M * 8, 1)])
+        B_layout = T.TileLayout(T.S[(N, K // 8, 8) : (8, N * 8, 1)])
+        ldo, sdo = 128, 8
+    elif SWIZZLE == 1:
+        A_layout = T.ComposeLayout(
+            T.SwizzleLayout(3, 1, 3, swizzle_inner=True),
+            T.TileLayout(T.S[(M, K // 16, 16) : (16, M * 16, 1)]),
+        )
+        B_layout = T.ComposeLayout(
+            T.SwizzleLayout(3, 1, 3, swizzle_inner=True),
+            T.TileLayout(T.S[(N, K // 16, 16) : (16, N * 16, 1)]),
+        )
+        ldo, sdo = 256, 16
+    elif SWIZZLE == 2:
+        A_layout = T.ComposeLayout(
+            T.SwizzleLayout(3, 2, 3, swizzle_inner=True),
+            T.TileLayout(T.S[(M, K // 32, 32) : (32, M * 32, 1)]),
+        )
+        B_layout = T.ComposeLayout(
+            T.SwizzleLayout(3, 2, 3, swizzle_inner=True),
+            T.TileLayout(T.S[(N, K // 32, 32) : (32, N * 32, 1)]),
+        )
+        ldo, sdo = 512, 32
+    elif SWIZZLE == 3:
+        A_layout = T.ComposeLayout(
+            T.SwizzleLayout(3, 3, 3, swizzle_inner=True),
+            T.TileLayout(T.S[(M, 1, 64) : (64, M * 64, 1)]),
+        )
+        B_layout = T.ComposeLayout(
+            T.SwizzleLayout(3, 3, 3, swizzle_inner=True),
+            T.TileLayout(T.S[(N, 1, 64) : (64, N * 64, 1)]),
+        )
+        ldo, sdo = 1, 64
+    else:
+        raise ValueError(f"Invalid swizzle: {SWIZZLE}")
+
+    dyn_smem_bytes = 1024 + (M * K + N * K) * 2
+
+    # fmt: off
+    @T.prim_func
+    def test_mma_ss_no_tma(A: T.Buffer((M, K), a_type, layout=T.TileLayout(T.S[M, K])),
+                           B: T.Buffer((N, K), b_type, layout=T.TileLayout(T.S[N, K])),
+                           C: T.Buffer((M, N), d_type)):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([4])
+        lane_id = T.lane_id([32])
+        tx = T.thread_id([128])
+        dyn = T.alloc_buffer((dyn_smem_bytes,), "uint8", scope="shared")
+        tmem_addr = T.decl_scalar("uint32", dyn.data, scope="shared", elem_offset=0)
+        A_smem = T.decl_buffer((M, K), a_type, dyn.data, elem_offset=256, layout=A_layout)
+        B_smem = T.decl_buffer((N, K), b_type, dyn.data, elem_offset=256 + M*K, layout=B_layout)
+        bar = T.decl_buffer((1,), "uint64", dyn.data, scope="shared", elem_offset=8)
+
+        reg = T.alloc_buffer((N,), d_type, scope="local")
+        descA = T.alloc_buffer((1,), "uint64", scope="local")
+        descB = T.alloc_buffer((1,), "uint64", scope="local")
+        descI = T.alloc_buffer((1,), "uint32", scope="local")
+        phase = T.alloc_buffer((1,), "int32", scope="local")
+
+        # alloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=N_COLS, cta_group=cta_group)
+        T.cuda.cta_sync()
+        for i in range(N):
+            reg[i] = 0.0
+        Tx.cta.copy(A_smem[:, :], A[:, :])
+        Tx.cta.copy(B_smem[:, :], B[:, :])
+        T.ptx.fence.proxy_async("shared::cta")
+        T.cuda.cta_sync()
+        # MMA
+        phase[0] = 0
+        if tx == 0:
+            T.ptx.mbarrier.init(bar.data, 1)
+            T.ptx.tcgen05.encode_instr_descriptor(descI.data, d_dtype=d_type, a_dtype=a_type, b_dtype=b_type, M=M, N=N, K=MMA_K, trans_a=False, trans_b=False, n_cta_groups=cta_group)  # noqa: E501
+            for k in range(K // MMA_K):
+                T.ptx.tcgen05.encode_matrix_descriptor(descA.data, A_smem.access_ptr("r", offset=A_smem.elem_offset_of([0, k * MMA_K])), ldo=ldo, sdo=sdo, swizzle=SWIZZLE)  # noqa: E501
+                T.ptx.tcgen05.encode_matrix_descriptor(descB.data, B_smem.access_ptr("r", offset=B_smem.elem_offset_of([0, k * MMA_K])), ldo=ldo, sdo=sdo, swizzle=SWIZZLE)  # noqa: E501
+                if k == 0:
+                    T.ptx.tcgen05.mma(tmem_addr, descA[0], descB[0], descI[0], d_dtype=d_type, a_dtype=a_type, b_dtype=b_type, use_a_tmem=False, cta_group=cta_group, enable_input_d=0)  # noqa: E501
+                else:
+                    T.ptx.tcgen05.mma(tmem_addr, descA[0], descB[0], descI[0], d_dtype=d_type, a_dtype=a_type, b_dtype=b_type, use_a_tmem=False, cta_group=cta_group, enable_input_d=1)  # noqa: E501
+            T.ptx.tcgen05.commit(bar.data, cta_group)
+        T.ptx.mbarrier.try_wait(bar.data, phase[0])
+        phase[0] = phase[0] ^ 1
+        T.cuda.cta_sync()
+
+        # TMEM -> RF
+        T.ptx.tcgen05.fence.after_thread_sync()
+        for i in range(N):
+            T.ptx.tcgen05.ld(tmem_addr, reg[i], shape="32x32b", num=REPEAT_NUM, row=warp_id * 32, col=i)  # noqa: E501
+        T.ptx.tcgen05.wait.ld()
+        # RF -> GMEM
+        for i in range(N):
+            C[tx, i] = reg[i]
+
+        # dealloc TMEM
+        if warp_id == 0:
+            T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+            T.ptx.tcgen05.dealloc(tmem_addr, n_cols=N_COLS, cta_group=cta_group)
+    # fmt: on
+
+    import torch
+
+    torch.manual_seed(42)
+    target = tvm.target.Target("cuda")
+    with target:
+        src, mod = _get_source(test_mma_ss_no_tma)
+        print(src)
+        assert "tcgen05.mma.cta_group::1.kind::f16" in src
+        assert "tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64" in src
+        assert "tcgen05.ld.sync.aligned.32x32b.x1.b32" in src
+        assert "tcgen05.wait::ld.sync.aligned" in src
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_torch = torch.rand((M, K), dtype=torch.float16)
+        B_torch = torch.rand((N, K), dtype=torch.float16)
+        C_torch = torch.zeros((M, N), dtype=torch.float32)
+        A = tvm.runtime.tensor(A_torch, device=dev)
+        B = tvm.runtime.tensor(B_torch, device=dev)
+        C = tvm.runtime.tensor(C_torch, device=dev)
+        mod(A, B, C)
+        ref = torch.matmul(A_torch, B_torch.T)
+        np.testing.assert_allclose(C.numpy(), ref.numpy(), rtol=1e-3, atol=1e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/codegen/test_codegen_cuda.py

@@ -0,0 +1,686 @@
+# 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
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def _get_source(func: tvm.tirx.PrimFunc) -> str:
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source()
+    return src, mod
+
+
+def _helper_source(src: str, helper_name: str) -> str:
+    start = src.index(helper_name)
+    next_helper = src.find("__device__", start + len(helper_name))
+    if next_helper == -1:
+        return src[start:]
+    return src[start:next_helper]
+
+
+def test_tirx_launch_bounds_omits_min_blocks_without_persistent_schedule():
+    @T.prim_func
+    def main(A: T.Buffer((4,), "int32")):
+        T.device_entry()
+        bx = T.cta_id([4])
+        tx = T.thread_id([128])
+        if tx == 0:
+            A[bx] = A[bx] + 1
+
+    src, _ = _get_source(main)
+    assert 'extern "C" __global__ void __launch_bounds__(128) main_kernel' in src
+    assert "__launch_bounds__(128, 1)" not in src
+
+
+def test_tirx_launch_bounds_min_blocks_attr_sets_one_block_per_sm():
+    @T.prim_func
+    def main(A: T.Buffer((4,), "int32")):
+        T.device_entry()
+        T.attr({"tirx.launch_bounds_min_blocks_per_sm": 1})
+        bx = T.cta_id([4])
+        tx = T.thread_id([128])
+        if tx == 0:
+            A[bx] = A[bx] + 1
+
+    src, _ = _get_source(main)
+    assert 'extern "C" __global__ void __launch_bounds__(128, 1) main_kernel' in src
+    assert "tirx.launch_bounds_min_blocks_per_sm" not in src
+
+
+def test_serial_pragma_unroll_codegen():
+    @T.prim_func
+    def main(A: T.Buffer((4,), "int32")):
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            for i in T.serial(4, unroll=True):
+                if i == 2:
+                    break
+                A[i] = A[i] + 1
+
+    src, _ = _get_source(main)
+    assert "#pragma unroll\n" in src
+    assert "for (" in src
+    assert "break;" in src
+
+
+def test_cluster_cta_id_codegen_uses_coordinate_sregs():
+    @T.prim_func
+    def main(A: T.Buffer((1,), "int32")):
+        T.device_entry()
+        cbx, cby = T.cta_id_in_cluster([2, 2])
+        tx = T.thread_id([32])
+        if tx == 0:
+            A[0] = cbx + cby
+
+    src, _ = _get_source(main)
+    assert "%cluster_ctaid.x" in src
+    assert "%cluster_ctaid.y" in src
+    assert "%cluster_ctarank" not in src
+    assert "cooperative_groups::cluster_group::block_index" not in src
+
+
+def test_cuda_handle_uint64_reinterpret_codegen():
+    @T.prim_func
+    def main(A: T.Buffer((1,), "uint64")):
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            ptr = T.reinterpret("handle", A[0])
+            A[0] = T.reinterpret("uint64", ptr)
+
+    src, _ = _get_source(main)
+    assert "reinterpret_cast<void*>" in src
+    assert "reinterpret_cast<uint64_t>" in src
+    assert "*(void* *)" not in src
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cuda_atomic_add():
+    @T.prim_func
+    def main(A: T.Buffer((1,), "int32"), B: T.Buffer((1,), "float32")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        if tx == 0:
+            T.cuda.atomic_add(A.data, T.int32(1))
+            T.cuda.atomic_add(B.data, T.float32(1.0))
+
+    src, mod = _get_source(main)
+    assert "tvm_builtin_cuda_atomic_add" in src
+    A_np = np.zeros(1, dtype="int32")
+    B_np = np.zeros(1, dtype="float32")
+
+    def run_and_check():
+        dev = tvm.device("cuda")
+        A_tvm = tvm.runtime.tensor(A_np, device=dev)
+        B_tvm = tvm.runtime.tensor(B_np, device=dev)
+        mod["main"](A_tvm, B_tvm)
+        np.testing.assert_allclose(A_tvm.numpy(), 1)
+        np.testing.assert_allclose(B_tvm.numpy(), 1.0)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_ptx_ld_acquire_and_volatile_codegen():
+    @T.prim_func
+    def main(A: T.Buffer((1,), "uint64"), B: T.Buffer((1,), "int32"), C: T.Buffer((1,), "uint32")):
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            A[0] = T.ptx.ld_acquire(A.data, "uint64", "u64", scope="gpu", space="global")
+            B[0] = T.ptx.ld_acquire(B.data, "int32", "s32", scope="sys", space="global")
+            C[0] = T.ptx.ld_acquire(C.data, "uint32", "b32", scope="gpu", space="global")
+            T.ptx.ld_global_acquire(B[0], B.data)
+            A[0] = T.ptx.ld_volatile(A.data, "uint64", "u64", space="global")
+
+    src, _ = _get_source(main)
+    assert "ld.acquire.gpu.global.u64" in src
+    assert "ld.acquire.sys.global.s32" in src
+    assert "ld.acquire.gpu.global.b32" in src
+    assert "ptx_ld_global_acquire_int32" in src
+    assert "ptx_ld_global_acquire_b32" not in src
+    assert "ld.volatile.global.u64" in src
+
+
+def test_megamoe_extracted_intrinsics_codegen():
+    @T.prim_func
+    def main(
+        U32: T.Buffer((4,), "uint32"),
+        I32: T.Buffer((1,), "int32"),
+        U64: T.Buffer((1,), "uint64"),
+        F32: T.Buffer((4,), "float32"),
+    ):
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            T.ptx.red_scalar(
+                U64.data,
+                U64[0],
+                sem="release",
+                scope="gpu",
+                space="global",
+                op="or",
+                ptx_type="b64",
+            )
+            T.ptx.red_scalar(
+                I32.data,
+                I32[0],
+                sem="release",
+                scope="sys",
+                space="global",
+                op="add",
+                ptx_type="s32",
+            )
+            U32[0] = T.ptx.atom_scalar(
+                U32.data,
+                U32[0],
+                sem="release",
+                scope="gpu",
+                space="global",
+                op="add",
+                ptx_type="u32",
+            )
+            U64[0] = T.ptx.atom_scalar(
+                U64.data, U64[0], scope="sys", space="global", op="add", ptx_type="u64"
+            )
+            T.ptx.red_scalar(
+                U32.data, U32[0], scope="gpu", space="global", op="add", ptx_type="u32"
+            )
+            T.ptx.st(U32.data, U32[0], space="shared", ptx_type="u32")
+            T.ptx.st(
+                U32.data,
+                U32[0],
+                U32[1],
+                U32[2],
+                U32[3],
+                space="shared",
+                vec="v4",
+                ptx_type="b32",
+            )
+            T.ptx.st_bulk(U32.data, T.uint32(16), weak=True, space="shared::cta")
+            U32[0] = T.ptx.fns_b32(U32[0], U32[1], I32[0])
+            T.ptx.stmatrix(
+                True,  # trans
+                1,  # num
+                ".b8",  # dtype
+                U32.data,  # smem_ptr
+                U32.data,  # src0
+                shape="m16n8",
+                space="shared",
+            )
+
+            F32[1] = T.cuda.uint_as_float(U32[0])
+            F32[2] = T.ptx.ld(F32.data, "float32", "f32", space="global")
+            U32[3] = T.cuda.float_as_uint(F32[1])
+            F32[0] = T.ptx.add_rn_f32_bf16(F32[0], T.cast(U32[0], "uint16"))
+            U64[0] = T.reinterpret("uint64", U32.data)
+            U32[0] = T.cuda.ballot_sync(T.uint32(0xFFFFFFFF), I32[0])
+            I32[0] = T.cuda.ffs_u32(U32[0])
+            U32[0] = T.cuda.reduce_add_sync_u32(T.uint32(0xFFFFFFFF), U32[0])
+            U32[0] = T.cuda.reduce_min_sync_u32(T.uint32(0xFFFFFFFF), U32[0])
+            U64[0] = T.cuda.clock64()
+            U32[0] = T.cuda.float22bfloat162_rn(F32[0], F32[1])
+
+    src, _ = _get_source(main)
+    for snippet in [
+        "red.release.gpu.global.or.b64",
+        "red.release.sys.global.add.s32",
+        "atom.release.gpu.global.add.u32",
+        "atom.sys.global.add.u64",
+        "red.gpu.global.add.u32",
+        "st.shared.u32",
+        "st.shared.v4.b32",
+        "st.bulk.weak.shared::cta",
+        "fns.b32",
+        "stmatrix.sync.aligned.m16n8.x1.trans.shared.b8",
+        "ld.global.f32",
+        "add.rn.f32.bf16",
+        "__uint_as_float",
+        "__float_as_uint",
+        "__ballot_sync",
+        "__ffs",
+        "__reduce_add_sync",
+        "__reduce_min_sync",
+        "clock64()",
+        "__float22bfloat162_rn",
+    ]:
+        assert snippet in src
+
+
+def test_ptx_cp_async_bulk_non_tma_form_codegen():
+    @T.prim_func
+    def main(
+        A: T.Buffer((128,), "float32"),
+        B: T.Buffer((128,), "float32"),
+        C: T.Buffer((1,), "uint64"),
+    ):
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            smem = T.alloc_shared([128], "float32")
+            T.ptx.cp_async_bulk_g2s_cta(
+                smem.ptr_to([0]), A.data, T.uint32(64), smem.ptr_to([0]), cache_policy=C[0]
+            )
+            T.ptx.cp_async_bulk_g2s_cluster(
+                smem.ptr_to([0]), A.data, T.uint32(64), smem.ptr_to([0]), cache_policy=C[0]
+            )
+            T.ptx.cp_async_bulk_s2g(B.data, smem.ptr_to([0]), T.uint32(64), cache_policy=C[0])
+
+    src, _ = _get_source(main)
+    assert "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" in src
+    assert "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" in src
+    assert "cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint" in src
+    assert "unsigned long long cache_policy" in src
+
+
+def test_tensor_map_param_codegen():
+    @T.prim_func
+    def main(A_map: T.TensorMap()):
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            T.evaluate(T.address_of(A_map))
+
+    src, _ = _get_source(main)
+    assert "const __grid_constant__ CUtensorMap A_map" in src
+    assert "((unsigned long long)(&(A_map)))" in src
+
+
+def test_tma_cache_policy_operand_codegen():
+    @T.prim_func
+    def main(Cache: T.Buffer((1,), "uint64")):
+        A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
+        B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
+
+        T.device_entry()
+        tx = T.thread_id([32])
+        if tx == 0:
+            smem = T.alloc_buffer((128,), "float32", scope="shared", align=128)
+            bar = T.shared_scalar("uint64")
+            T.ptx.cp_async.bulk.tensor.g2c(
+                2,
+                smem.data,
+                T.address_of(bar),
+                T.address_of(A_map),
+                1,
+                2,
+                "",
+                0,
+                0,
+                cache_policy=Cache[0],
+            )
+            T.ptx.cp_async.bulk.tensor.g2c(
+                2,
+                smem.data,
+                T.address_of(bar),
+                T.address_of(A_map),
+                3,
+                2,
+                "",
+                0,
+                0,
+                cache_policy=Cache[0],
+            )
+            T.ptx.cp_async.bulk.tensor.s2g(
+                2, smem.data, T.address_of(A_map), "", 0, 0, cache_policy=Cache[0]
+            )
+            masked_bar = T.cuda.sm100_tma_2sm_mbarrier_addr(T.address_of(bar))
+            T.ptx.cp_async.bulk.tensor.g2c_bar_addr(
+                2,
+                smem.data,
+                masked_bar,
+                T.address_of(A_map),
+                1,
+                2,
+                "",
+                0,
+                0,
+                cache_policy=Cache[0],
+            )
+            if tx == 0:
+                T.ptx.cp_async.bulk.tensor.g2c_bar_addr(
+                    2,
+                    smem.data,
+                    masked_bar,
+                    T.address_of(A_map),
+                    1,
+                    2,
+                    "",
+                    0,
+                    0,
+                    cache_policy=Cache[0],
+                )
+            else:
+                T.ptx.cp_async.bulk.tensor.g2c_bar_addr(
+                    2,
+                    smem.data,
+                    masked_bar,
+                    T.address_of(B_map),
+                    1,
+                    2,
+                    "",
+                    0,
+                    0,
+                    cache_policy=Cache[0],
+                )
+
+    src, _ = _get_source(main)
+    assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_cache_hint" in src
+    assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_multicast_cache_hint" in src
+    assert "g2cluster_unicast" not in src
+    assert "ptx_cp_async_bulk_tensor_g2cta" not in src
+    assert (
+        "cp.async.bulk.tensor.2d.shared::cluster.global"
+        ".mbarrier::complete_tx::bytes.cta_group::2.L2::cache_hint"
+    ) in src
+    assert (
+        "cp.async.bulk.tensor.2d.shared::cluster.global"
+        ".mbarrier::complete_tx::bytes.multicast::cluster"
+        ".cta_group::2.L2::cache_hint"
+    ) in src
+    assert "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group.L2::cache_hint" in src
+    assert "tvm_builtin_cp_async_bulk_tensor_2d_g2c_cta_group2" not in src
+    assert "tvm_builtin_cuda_cvta_generic_to_shared((&(bar_ptr[0]))) & (uint)4278190079" in src
+    assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_cache_hint_bar_addr" in src
+    assert "unsigned long long cache_policy" in src
+
+
+def test_cuda_thread_fence():
+    @T.prim_func
+    def main(A: T.Buffer((16, 16), "int32")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        if tx == 0:
+            T.cuda.thread_fence()
+
+    src, mod = _get_source(main)
+    assert "tvm_builtin_cuda_thread_fence" in src
+
+
+def test_cuda_nano_sleep():
+    @T.prim_func
+    def main(A: T.Buffer((16, 16), "int32")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        if tx == 0:
+            T.cuda.nano_sleep(1)
+
+    src, mod = _get_source(main)
+    assert "tvm_builtin_cuda_nano_sleep" in src
+
+
+def test_cuda_atomic_cas():
+    @T.prim_func
+    def main(A: T.Buffer((16, 16), "int32")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        if tx == 0:
+            T.cuda.atomic_cas(A.data, T.int32(1), T.int32(2))
+
+    src, mod = _get_source(main)
+    assert "tvm_builtin_cuda_atomic_cas" in src
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cuda_func_call():
+    def test_add_one():
+        add_one = """
+__device__ int32_t add_one(int32_t a) {
+    return a + 1;
+}
+"""
+
+        @T.prim_func
+        def main(a: T.Buffer((16, 16), "int32"), b: T.Buffer((16, 16), "int32")):
+            T.device_entry()
+            cta_id = T.cta_id([1])
+            tx = T.thread_id([32])
+            if tx == 0:
+                for i, j in T.grid(16, 16):
+                    b[i, j] = T.cuda.func_call(
+                        "add_one", a[i, j], source_code=add_one, return_type="int32"
+                    )
+
+        src, mod = _get_source(main)
+        A = np.random.randint(0, 10, (16, 16)).astype("int32")
+        B = np.zeros((16, 16), dtype="int32")
+
+        def run_and_check():
+            dev = tvm.device("cuda")
+            A_tvm = tvm.runtime.tensor(A, device=dev)
+            B_tvm = tvm.runtime.tensor(B, device=dev)
+            mod["main"](A_tvm, B_tvm)
+            np.testing.assert_allclose(B_tvm.numpy(), A + 1)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+        print(src)
+
+    test_add_one()
+
+    def test_print():
+        print_func = """
+__device__ void print(int32_t a) {
+    printf("%d\\n", a);
+}
+"""
+
+        @T.prim_func
+        def main(a: T.Buffer((16, 16), "int32")):
+            T.device_entry()
+            cta_id = T.cta_id([1])
+            tx = T.thread_id([32])
+            if tx == 0:
+                for i, j in T.grid(16, 16):
+                    T.cuda.func_call("print", a[i, j], source_code=print_func)
+
+        src, mod = _get_source(main)
+        A = np.random.randint(0, 10, (16, 16)).astype("int32")
+
+        def run_and_check():
+            dev = tvm.device("cuda")
+            A_tvm = tvm.runtime.tensor(A, device=dev)
+            mod["main"](A_tvm)
+            dev.sync()
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+        print(src)
+
+    test_print()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_warp_shuffle_xor_sync():
+    # fmt: off
+    @T.prim_func
+    def func(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (32,), dtype="float32", align=16)
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+
+        A_local = T.alloc_buffer([1], "float32", scope="local")
+        i = T.alloc_buffer([1], "int32", scope="local")
+
+        A_local[0] = T.float32(31 - lane_id)
+        i[0] = 16
+        while i[0] >= 1:
+            A_local[0] += T.tvm_warp_shuffle_xor(0xFFFFFFFF, A_local[0], i[0], 32, 32)
+            i[0] = i[0] // 2
+
+        A[lane_id] = A_local[0]
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    A_np = np.zeros(32, dtype="float32")
+    assert "__shfl_xor_sync" in mod.mod.imports[0].inspect_source()
+    A_ref = np.ones(32, dtype="float32") * 496
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_np, device=dev)
+        mod(A)
+        np.testing.assert_allclose(A.numpy(), A_ref)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("cp_size", [4, 8, 16])
+@pytest.mark.parametrize("cache_hint", ["", "evict_last"])
+@pytest.mark.parametrize("prefetch_size", [-1, 64, 128, 256])
+@pytest.mark.parametrize("predicate", [-1, T.int32(0), T.int32(1)])
+@pytest.mark.parametrize("fill_mode", ["", "zero"])
+def test_ptx_cp_async(cp_size, cache_hint, prefetch_size, predicate, fill_mode):
+    if fill_mode != "" and predicate == -1:
+        return
+
+    N = cp_size // 2
+
+    # fmt: off
+    @T.prim_func
+    def main(A: T.Buffer((N), "float16")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([32])
+        A_shared = T.alloc_shared([N], "float16")
+        for i in T.vectorized(N):
+            A_shared[i] = 5.0
+        T.ptx.fence.proxy_async("shared::cta")
+        T.ptx.cp_async(A_shared.ptr_to([0]), A.ptr_to([0]), cp_size, cache_hint=cache_hint, prefetch_size=prefetch_size, predicate=predicate, fill_mode=fill_mode)  # noqa: E501
+        T.ptx.cp_async.commit_group()
+        T.ptx.cp_async.wait_group(0)
+        for i in T.serial(N):
+            A[i] = A_shared[i] + 1.0
+        # fmt: on
+
+    src, mod = _get_source(main)
+    A_np = np.ones(N, dtype="float16")
+    A_ref = np.ones(N, dtype="float16") * 2
+    if int(predicate) == 0:
+        if fill_mode == "zero":
+            A_ref = np.ones(N, dtype="float16")
+        else:
+            A_ref = np.ones(N, dtype="float16") * 6
+
+    def run_and_check():
+        dev = tvm.device("cuda")
+        A = tvm.runtime.tensor(A_np, device=dev)
+        mod(A)
+        np.testing.assert_allclose(A.numpy(), A_ref)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+    print(src)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("trans", [False, True])
+@pytest.mark.parametrize("num", [1, 2, 4])
+def test_ptx_ldmatrix(trans, num):
+    dtype = ".b16"
+
+    # fmt: off
+    @T.prim_func
+    def main(A: T.Buffer((16, 16), "float16"), B: T.Buffer((16, 16), "float16")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tx = T.thread_id([32])
+        A_shared = T.alloc_shared([16, 16], "float16")
+        if tx == 0:
+            for i, j in T.grid(16, 16):
+                A_shared[i, j] = A[i, j]
+        T.cuda.cta_sync()
+        A_local = T.alloc_local([8], "float16")
+        A_local[0] = -1.0
+                # ldmatrix .x{num}.b16 writes `num` 32-bit registers; A_local
+                # is a contiguous fp16[8] buffer, so consecutive register
+                # destinations land 2 fp16 elements apart.
+        if num == 1:
+            T.ptx.ldmatrix(
+                trans, num, dtype,
+                A_shared.ptr_to([tx % 16, tx // 16 * 8]),
+                T.address_of(A_local[0]),
+            )
+        elif num == 2:
+            T.ptx.ldmatrix(
+                trans, num, dtype,
+                A_shared.ptr_to([tx % 16, tx // 16 * 8]),
+                T.address_of(A_local[0]),
+                T.address_of(A_local[2]),
+            )
+        else:
+            T.ptx.ldmatrix(
+                trans, num, dtype,
+                A_shared.ptr_to([tx % 16, tx // 16 * 8]),
+                T.address_of(A_local[0]),
+                T.address_of(A_local[2]),
+                T.address_of(A_local[4]),
+                T.address_of(A_local[6]),
+            )
+        for i in range(8):
+            row: T.let = (i // 2) % 2 * 8
+            col: T.let = (i // 4) * 8
+            B[row + tx // 4, col + tx % 4 * 2 + i % 2] = A_local[i]
+        # fmt: on
+
+    src, mod = _get_source(main)
+    A_np = np.arange(16 * 16, dtype="float16").reshape((16, 16))
+    B_np = np.zeros((16, 16), dtype="float16")
+    B_ref = np.zeros((16, 16), dtype="float16")
+    if num == 1:
+        B_ref[0:8, 0:8] = A_np[0:8, 0:8] if not trans else A_np[0:8, 0:8].T
+    elif num == 2:
+        B_ref[0:8, 0:8] = A_np[0:8, 0:8] if not trans else A_np[0:8, 0:8].T
+        B_ref[8:16, 0:8] = A_np[8:16, 0:8] if not trans else A_np[8:16, 0:8].T
+    elif num == 4:
+        B_ref[0:8, 0:8] = A_np[0:8, 0:8] if not trans else A_np[0:8, 0:8].T
+        B_ref[0:8, 8:16] = A_np[0:8, 8:16] if not trans else A_np[0:8, 8:16].T
+        B_ref[8:16, 0:8] = A_np[8:16, 0:8] if not trans else A_np[8:16, 0:8].T
+        B_ref[8:16, 8:16] = A_np[8:16, 8:16] if not trans else A_np[8:16, 8:16].T
+
+    def run_and_check():
+        dev = tvm.device("cuda")
+        A = tvm.runtime.tensor(A_np, device=dev)
+        B = tvm.runtime.tensor(B_np, device=dev)
+        mod(A, B)
+        np.testing.assert_allclose(B.numpy(), B_ref)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/codegen/test_codegen_dsmem.py

@@ -0,0 +1,92 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-function-docstring
+"""Tests for cp.async.bulk.shared::cluster.shared::cta PTX instruction codegen."""
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+
+def _get_source(func: tvm.tirx.PrimFunc) -> str:
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source()
+    return src
+
+
+def test_ptx_cp_async_bulk_s2c_codegen():
+    """Test that T.ptx.cp_async.bulk.s2c emits the correct PTX instruction."""
+
+    # fmt: off
+    @T.prim_func
+    def main(A: T.Buffer((128,), "float16")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([1])
+        A_smem = T.alloc_shared([128], "float16")
+        for i in T.serial(128):
+            A_smem[i] = A[i]
+                # Use the raw PTX instruction directly
+        dst_ptr = T.ptx.map_shared_rank(A_smem.ptr_to([0]), T.int32(1))
+        mbar_ptr = T.ptx.map_shared_rank(A_smem.ptr_to([0]), T.int32(1))
+        T.ptx.cp_async.bulk.s2c(
+            dst_ptr,
+            A_smem.ptr_to([0]),
+            T.int32(256),  # 128 elements * 2 bytes
+            mbar_ptr,
+        )
+        # fmt: on
+
+    src = _get_source(main)
+    assert "tvm_builtin_ptx_cp_async_bulk_s2s_cluster" in src
+    assert "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" in src
+
+
+def test_ptx_cp_async_bulk_s2c_codegen_address_conversion():
+    """Test that the codegen correctly converts addresses to shared space."""
+
+    # fmt: off
+    @T.prim_func
+    def main(A: T.Buffer((64,), "float32")):
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([1])
+        A_smem = T.alloc_shared([64], "float32")
+        for i in T.serial(64):
+            A_smem[i] = A[i]
+        dst_ptr = T.ptx.map_shared_rank(A_smem.ptr_to([0]), T.int32(0))
+        mbar_ptr = T.ptx.map_shared_rank(A_smem.ptr_to([0]), T.int32(0))
+        T.ptx.cp_async.bulk.s2c(
+            dst_ptr,
+            A_smem.ptr_to([0]),
+            T.int32(256),  # 64 * 4 bytes
+            mbar_ptr,
+        )
+        # fmt: on
+
+    src = _get_source(main)
+    # Verify address conversion to shared space
+    assert "__cvta_generic_to_shared" in src
+    assert "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" in src
+
+
+if __name__ == "__main__":
+    test_ptx_cp_async_bulk_s2c_codegen()
+    test_ptx_cp_async_bulk_s2c_codegen_address_conversion()
+    print("All codegen tests passed!")

tests/python/tirx/codegen/test_codegen_nki.py

@@ -0,0 +1,330 @@
+# 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.
+
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def lower_and_get_source(func):
+    with target:
+        mod = tvm.IRModule({"main": func})
+        mod = tvm.compile(mod, tir_pipeline="trn")
+        src = mod.mod.imports[0].inspect_source()
+        return src
+
+
+def compare_strings_ignore_whitespace(s1, s2):
+    # Remove all whitespace by splitting and joining the string back together
+    return "".join(s1.split()) == "".join(s2.split())
+
+
+def test_nki_add_1():
+    # fmt: off
+    @T.prim_func
+    def func(A: T.Buffer((128, 512)), B: T.Buffer((128, 512))):
+        T.func_attr({"num_inputs": 1})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 512), "float32", scope="trn.sbuf",)
+        B_sbuf = T.alloc_buffer((128, 512), "float32", scope="trn.sbuf",)
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for i in range(0, 128):
+                for j in range(0, 512):
+                    T.nki.load(A_sbuf[i, j], A[i, j])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for i in range(0, 128):
+                for j in range(0, 512):
+                    T.nki.tensorscalar(B_sbuf[i, j], A_sbuf[i, j], T.float32(1.0), "add")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for i in range(0, 128):
+                for j in range(0, 512):
+                    T.nki.store(B[i, j], B_sbuf[i, j])
+        # fmt: on
+    src = lower_and_get_source(func)
+    print(src)
+    expected = """# Function: func_kernel
+import neuronxcc.nki.language as nl
+from neuronxcc.nki import baremetal, benchmark, simulate_kernel, trace
+import numpy as np
+import neuronxcc.nki.isa as nisa
+import math
+import neuronxcc.nki as nki
+import neuronxcc.nki.typing as nt
+import neuronxcc.nki.compiler as ncc
+@nki.compiler.enable_stack_allocator
+@nki.compiler.skip_middle_end_transformations
+@baremetal(experimental_flags='enable-mutable-parameter', additional_compile_opt='--internal-skip-backend-allocation-opt-nki')
+def func_kernel(A_ptr, B_ptr: nt.mutable_tensor, ):
+  B_ptr_buffer = B_ptr.reshape([65536])
+  A_ptr_buffer = A_ptr.reshape([65536])
+  A_sbuf_ptr = nl.ndarray(shape=[128, 512], dtype=np.float32, buffer=ncc.sbuf.mod_alloc(base_addr=0))
+  B_sbuf_ptr = nl.ndarray(shape=[128, 512], dtype=np.float32, buffer=ncc.sbuf.mod_alloc(base_addr=2048))
+  i = nl.arange(128)
+  j = nl.arange(512)
+  A_sbuf_ptr[i[:, None, ], j[None, :, ]] = nl.load(A_ptr_buffer[((i[:, None, ] * 512) + j[None, :, ])])
+  i_1 = nl.arange(128)
+  j_1 = nl.arange(512)
+  B_sbuf_ptr[i_1[:, None, ], j_1[None, :, ]] = nisa.tensor_scalar(A_sbuf_ptr[i_1[:, None, ], j_1[None, :, ]], operand0=1.000000e+00, op0=nki.language.add, reverse0=False)
+  i_2 = nl.arange(128)
+  j_2 = nl.arange(512)
+  nl.store(B_ptr_buffer[((i_2[:, None, ] * 512) + j_2[None, :, ])], B_sbuf_ptr[i_2[:, None, ], j_2[None, :, ]])
+  return B_ptr
+  """  # noqa: E501
+    assert compare_strings_ignore_whitespace(src, expected)
+
+
+def test_nki_add_2():
+    # fmt: off
+    @T.prim_func
+    def func(A: T.Buffer((128, 2048)), B: T.Buffer((128, 2048))):
+        T.func_attr({"num_inputs": 1})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 512), "float32", scope="trn.sbuf",)
+        B_sbuf = T.alloc_buffer((128, 512), "float32", scope="trn.sbuf",)
+        for k in range(0, 4):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for i in range(0, 128):
+                    for j in range(0, 512):
+                        T.nki.load(A_sbuf[i, j], A[i, 512*k+j])
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for i in range(0, 128):
+                    for j in range(0, 512):
+                        T.nki.tensorscalar(B_sbuf[i, j], A_sbuf[i, j], T.float32(1.0), "add")
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for i in range(0, 128):
+                    for j in range(0, 512):
+                        T.nki.store(B[i, 512*k+j], B_sbuf[i, j])
+
+        # fmt: on
+    src = lower_and_get_source(func)
+    print(src)
+    expected = """# Function: func_kernel
+import neuronxcc.nki.language as nl
+from neuronxcc.nki import baremetal, benchmark, simulate_kernel, trace
+import numpy as np
+import neuronxcc.nki.isa as nisa
+import math
+import neuronxcc.nki as nki
+import neuronxcc.nki.typing as nt
+import neuronxcc.nki.compiler as ncc
+@nki.compiler.enable_stack_allocator
+@nki.compiler.skip_middle_end_transformations
+@baremetal(experimental_flags='enable-mutable-parameter', additional_compile_opt='--internal-skip-backend-allocation-opt-nki')
+def func_kernel(A_ptr, B_ptr: nt.mutable_tensor, ):
+  B_ptr_buffer = B_ptr.reshape([262144])
+  A_ptr_buffer = A_ptr.reshape([262144])
+  A_sbuf_ptr = nl.ndarray(shape=[128, 512], dtype=np.float32, buffer=ncc.sbuf.mod_alloc(base_addr=0))
+  B_sbuf_ptr = nl.ndarray(shape=[128, 512], dtype=np.float32, buffer=ncc.sbuf.mod_alloc(base_addr=2048))
+  for k in nl.sequential_range(4, body_no_reorder=True):
+    i = nl.arange(128)
+    j = nl.arange(512)
+    A_sbuf_ptr[i[:, None, ], j[None, :, ]] = nl.load(A_ptr_buffer[(((i[:, None, ] * 2048) + (k * 512)) + j[None, :, ])])
+    i_1 = nl.arange(128)
+    j_1 = nl.arange(512)
+    B_sbuf_ptr[i_1[:, None, ], j_1[None, :, ]] = nisa.tensor_scalar(A_sbuf_ptr[i_1[:, None, ], j_1[None, :, ]], operand0=1.000000e+00, op0=nki.language.add, reverse0=False)
+    i_2 = nl.arange(128)
+    j_2 = nl.arange(512)
+    nl.store(B_ptr_buffer[(((i_2[:, None, ] * 2048) + (k * 512)) + j_2[None, :, ])], B_sbuf_ptr[i_2[:, None, ], j_2[None, :, ]])
+  return B_ptr"""  # noqa: E501
+    assert compare_strings_ignore_whitespace(src, expected)
+
+
+def test_nki_matmul_1():
+    TILES_IN_BLOCK_M = 16
+    TILES_IN_BLOCK_N = 1
+    TILES_IN_BLOCK_K = 8
+    TILE_M = 128
+    TILE_K = 128
+    TILE_N = 512
+    K = 1024
+    M = 4096
+    N = 2048
+    BLOCK_M = TILE_M * TILES_IN_BLOCK_M
+    BLOCK_N = TILE_N * TILES_IN_BLOCK_N
+    BLOCK_K = TILE_K * TILES_IN_BLOCK_K
+    # the size has to be multiple of block size
+    assert M % BLOCK_M == 0
+    assert N % BLOCK_N == 0
+    assert K % BLOCK_K == 0
+
+    NUM_BLOCK_M = M // BLOCK_M
+    NUM_BLOCK_N = N // BLOCK_N
+    NUM_BLOCK_K = K // BLOCK_K
+
+    @T.prim_func
+    def func(
+        lhsT: T.Buffer((K, M), "float16"),
+        rhs: T.Buffer((K, N), "float16"),
+        result: T.buffer((M, N), "float16"),
+    ):
+        T.func_attr({"num_inputs": 2})
+        result_tiles = T.alloc_buffer(
+            (TILE_M, NUM_BLOCK_M, TILES_IN_BLOCK_M, TILES_IN_BLOCK_N, TILE_N),
+            "float32",
+            scope="trn.sbuf",
+        )
+        rhs_tiles = T.alloc_buffer((TILE_K, TILES_IN_BLOCK_K, BLOCK_N), "float16", scope="trn.sbuf")
+        lhsT_tiles = T.alloc_buffer(
+            (TILE_K, TILES_IN_BLOCK_K, BLOCK_M), "float16", scope="trn.sbuf"
+        )
+        res_tile = T.alloc_buffer((1, TILE_M, TILE_N), "float32", scope="trn.psum")
+        result_packed = T.alloc_buffer((TILE_K, BLOCK_N), "float32", scope="trn.sbuf")
+        for n in range(NUM_BLOCK_N):
+            with T.attr(0, "tensorized_nki_instruction", 1):
+                for i0 in range(TILE_M):
+                    for i1 in range(NUM_BLOCK_M):
+                        for i2 in range(TILES_IN_BLOCK_M):
+                            for i3 in range(TILES_IN_BLOCK_N):
+                                for i4 in range(TILE_N):
+                                    T.nki.memset(result_tiles[i0, i1, i2, i3, i4], T.float32(0.0))
+            for k in range(NUM_BLOCK_K):
+                for bk_r in range(TILES_IN_BLOCK_K):
+                    with T.attr(0, "tensorized_nki_instruction", 1):
+                        for i in range(TILE_K):
+                            for j in range(BLOCK_N):
+                                T.nki.load(
+                                    rhs_tiles[i, bk_r, j],
+                                    rhs[
+                                        (TILES_IN_BLOCK_K * k + bk_r) * TILE_K + i,
+                                        n * BLOCK_N + j,
+                                    ],
+                                )
+                for m in range(NUM_BLOCK_M):
+                    for bk_l in range(TILES_IN_BLOCK_K):
+                        with T.attr(0, "tensorized_nki_instruction", 1):
+                            for i in range(TILE_K):
+                                for j in range(BLOCK_M):
+                                    T.nki.load(
+                                        lhsT_tiles[i, bk_l, j],
+                                        lhsT[
+                                            (TILES_IN_BLOCK_K * k + bk_l) * TILE_K + i,
+                                            m * BLOCK_M + j,
+                                        ],
+                                    )
+                    for bn in range(TILES_IN_BLOCK_N):
+                        for bm in range(TILES_IN_BLOCK_M):
+                            with T.attr(0, "tensorized_nki_instruction", 1):
+                                for i in range(TILE_M):
+                                    for j in range(TILE_N):
+                                        T.nki.memset(res_tile[0, i, j], T.float32(0.0))
+                            for bk in range(TILES_IN_BLOCK_K):
+                                with T.attr(0, "tensorized_nki_instruction", 1):
+                                    for i in range(TILE_M):
+                                        for j in range(TILE_N):
+                                            for k in range(TILE_K):
+                                                T.nki.matmul(
+                                                    res_tile[0, i, j],
+                                                    lhsT_tiles[k, bk, bm * TILE_M + i],
+                                                    rhs_tiles[k, bk, bn * TILE_N + j],
+                                                    1,
+                                                )
+                            with T.attr(0, "tensorized_nki_instruction", 1):
+                                for i in range(TILE_M):
+                                    for j in range(TILE_N):
+                                        T.nki.tensortensor(
+                                            result_tiles[i, m, bm, bn, j],
+                                            result_tiles[i, m, bm, bn, j],
+                                            res_tile[0, i, j],
+                                            "add",
+                                        )
+            for m in range(NUM_BLOCK_M):
+                for bm in range(TILES_IN_BLOCK_M):
+                    for bn in range(TILES_IN_BLOCK_N):
+                        with T.attr(0, "tensorized_nki_instruction", 1):
+                            for i in range(TILE_K):
+                                for j in range(TILE_N):
+                                    T.nki.tensor_copy(
+                                        result_packed[i, bn * TILE_N + j],
+                                        result_tiles[i, m, bm, bn, j],
+                                    )
+                    with T.attr(0, "tensorized_nki_instruction", 1):
+                        for i in range(TILE_K):
+                            for j in range(BLOCK_N):
+                                T.nki.store(
+                                    result[m * BLOCK_M + bm * TILE_M + i, n * BLOCK_N + j],
+                                    result_packed[i, j],
+                                )
+
+    # fmt: on
+
+    src = lower_and_get_source(func)
+    print(src)
+    expected = """# Function: func_kernel
+import neuronxcc.nki.language as nl
+from neuronxcc.nki import baremetal, benchmark, simulate_kernel, trace
+import numpy as np
+import neuronxcc.nki.isa as nisa
+import math
+import neuronxcc.nki as nki
+import neuronxcc.nki.typing as nt
+import neuronxcc.nki.compiler as ncc
+@nki.compiler.enable_stack_allocator
+@nki.compiler.skip_middle_end_transformations
+@baremetal(experimental_flags='enable-mutable-parameter', additional_compile_opt='--internal-skip-backend-allocation-opt-nki')
+def func_kernel(lhsT_ptr, rhs_ptr, result_ptr: nt.mutable_tensor, ):
+  result_ptr_buffer = result_ptr.reshape([8388608])
+  rhs_ptr_buffer = rhs_ptr.reshape([2097152])
+  lhsT_ptr_buffer = lhsT_ptr.reshape([4194304])
+  result_tiles_ptr = nl.ndarray(shape=[128, 2, 16, 1, 512], dtype=np.float32, buffer=ncc.sbuf.mod_alloc(base_addr=0))
+  rhs_tiles_ptr = nl.ndarray(shape=[128, 8, 512], dtype=np.float16, buffer=ncc.sbuf.mod_alloc(base_addr=65536))
+  lhsT_tiles_ptr = nl.ndarray(shape=[128, 8, 2048], dtype=np.float16, buffer=ncc.sbuf.mod_alloc(base_addr=73728))
+  res_tile_ptr = nl.ndarray(shape=[1, nl.par_dim(128), 512], dtype=np.float32, buffer=nl.psum)
+  result_packed_ptr = nl.ndarray(shape=[128, 512], dtype=np.float32, buffer=ncc.sbuf.mod_alloc(base_addr=106496))
+  for n in nl.sequential_range(4, body_no_reorder=True):
+    i0 = nl.arange(128)
+    i1 = nl.arange(2)
+    i2 = nl.arange(16)
+    i4 = nl.arange(512)
+    result_tiles_ptr[i0[:, None, None, None, ], i1[None, :, None, None, ], i2[None, None, :, None, ], 0, i4[None, None, None, :, ]] = 0.000000e+00
+    for bk_r in nl.sequential_range(8):
+      i = nl.arange(128)
+      j = nl.arange(512)
+      rhs_tiles_ptr[i[:, None, ], bk_r, j[None, :, ]] = nl.load(rhs_ptr_buffer[((((bk_r * 262144) + (i[:, None, ] * 2048)) + (n * 512)) + j[None, :, ])])
+    for m in nl.sequential_range(2):
+      for bk_l in nl.sequential_range(8):
+        i_1 = nl.arange(128)
+        j_1 = nl.arange(2048)
+        lhsT_tiles_ptr[i_1[:, None, ], bk_l, j_1[None, :, ]] = nl.load(lhsT_ptr_buffer[((((bk_l * 524288) + (i_1[:, None, ] * 4096)) + (m * 2048)) + j_1[None, :, ])])
+      for bm in nl.sequential_range(16):
+        i_2 = nl.arange(128)
+        j_2 = nl.arange(512)
+        res_tile_ptr[0, i_2[:, None, ], j_2[None, :, ]] = 0.000000e+00
+        for bk in nl.sequential_range(8):
+          i_3 = nl.arange(128)
+          j_3 = nl.arange(512)
+          k = nl.arange(128)
+          res_tile_ptr[0, i_3[:, None, ], j_3[None, :, ]] += nisa.nc_matmul(lhsT_tiles_ptr[k[:, None, ], bk, ((bm * 128) + i_3[None, :, ])],rhs_tiles_ptr[k[:, None, ], bk, j_3[None, :, ]])
+        i_4 = nl.arange(128)
+        j_4 = nl.arange(512)
+        result_tiles_ptr[i_4[:, None, ], m, bm, 0, j_4[None, :, ]] = nisa.tensor_tensor(result_tiles_ptr[i_4[:, None, ], m, bm, 0, j_4[None, :, ]], res_tile_ptr[0, i_4[:, None, ], j_4[None, :, ]], op=nki.language.add)
+    for m_1 in nl.sequential_range(2):
+      for bm_1 in nl.sequential_range(16):
+        i_5 = nl.arange(128)
+        j_5 = nl.arange(512)
+        result_packed_ptr[i_5[:, None, ], j_5[None, :, ]] = nisa.tensor_copy(result_tiles_ptr[i_5[:, None, ], m_1, bm_1, 0, j_5[None, :, ]])
+        i_6 = nl.arange(128)
+        j_6 = nl.arange(512)
+        nl.store(result_ptr_buffer[(((((m_1 * 4194304) + (bm_1 * 262144)) + (i_6[:, None, ] * 2048)) + (n * 512)) + j_6[None, :, ])], result_packed_ptr[i_6[:, None, ], j_6[None, :, ]])
+  return result_ptr"""  # noqa: E501
+    assert compare_strings_ignore_whitespace(src, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/codegen/test_codegen_nvshmem.py

@@ -0,0 +1,312 @@
+# 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.
+"""Basic tests for a Disco nvshmem support"""
+
+# pylint: disable=missing-docstring
+import tempfile
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.runtime import ShapeTuple
+from tvm.runtime import disco as di
+from tvm.script import tirx as T
+from tvm.support.popen_pool import PopenWorker
+from tvm.testing import env
+
+NUM_WORKERS = 4
+
+
+def run_prim_func(sess, prim_func, *args):
+    """Compile, export, load, and run a PrimFunc in the shared disco session."""
+    target = tvm.target.Target("cuda")
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = f"{tmpdir}/test.so"
+        mod = tvm.compile(prim_func, target=target, tir_pipeline="tirx")
+        print(mod.mod.imports[0].inspect_source())
+        mod.export_library(path)
+        rt_mod = sess.load_vm_module(path)
+        rt_mod["main"](*args)
+        sess._sync_all()
+
+
+def create_nvshmem_array(sess, shape, dtype, init_data_fn=None, zero_out=True):
+    """Create and optionally initialize an nvshmem-accessible DNDArray."""
+    nvshmem_empty = sess.get_global_func("runtime.disco.nvshmem.empty")
+    arr = nvshmem_empty(ShapeTuple(shape), dtype, None)
+
+    if init_data_fn:
+        for i in range(NUM_WORKERS):
+            arr.debug_copy_from(i, init_data_fn(i, shape, dtype))
+    elif zero_out:
+        zero_data = np.zeros(shape, dtype=dtype)
+        for i in range(NUM_WORKERS):
+            arr.debug_copy_from(i, zero_data)
+
+    return arr
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.skip(reason="nvshmem doesn't work with pytest")
+def test_codegen_nvshmem():
+    def _test_func():
+        ############ setup ############
+        sess = di.ProcessSession(num_workers=NUM_WORKERS)
+        f_init_nvshmem_uid = tvm.get_global_func("runtime.disco.nvshmem.init_nvshmem_uid")
+        uid = f_init_nvshmem_uid()
+        init_dfunc = sess.get_global_func("runtime.disco.nvshmem.init_nvshmem")
+        init_dfunc(uid, NUM_WORKERS, 0)
+        sess.sync_worker_0()
+
+        def test_thread_info(sess):
+            @T.prim_func
+            def main(res: T.Buffer((2,), "int32")):
+                T.device_entry()
+                cta_id = T.cta_id([1])
+                tid = T.thread_id([nwarps * 32])
+                res[0] = T.nvshmem.my_pe()
+                res[1] = T.nvshmem.n_pes()
+
+            res_array = sess.empty((2,), "int32")
+            run_prim_func(sess, main, res_array)
+
+        def test_transfer(sess, scope, shape, nwarps, nelems, op_name):
+            """Tests data transfer operations (get/put) at thread, warp, and block scopes."""
+            dtype = "float32"
+            is_get = "get" in op_name
+            op_func = getattr(T.nvshmem, op_name)
+            if scope != "thread":
+                op_func = getattr(op_func, scope)
+
+            # fmt: off
+            @T.prim_func
+            def main(A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype)):
+                T.device_entry()
+                cta_id = T.cta_id([1])
+                warp_id = T.warp_id([nwarps])
+                lane_id = T.lane_id([32])
+                tid = T.thread_id([nwarps * 32])
+
+                my_pe = T.nvshmem.my_pe()
+                n_pes = T.nvshmem.n_pes()
+                offset = T.if_then_else(
+                    scope == "block", 0, T.if_then_else(scope == "thread", tid, warp_id * 32)
+                )
+                op_func(dst=B.ptr_to([offset]), src=A.ptr_to([offset]), nelems=nelems, pe=(my_pe + 1) % n_pes)  # noqa: E501
+                T.nvshmem.quiet()
+                # fmt: on
+
+            def init_fn(i, s, d):
+                return np.arange(s[0], dtype=d) + i * 100
+
+            A_array = create_nvshmem_array(sess, shape, dtype, init_fn)
+            B_array = create_nvshmem_array(sess, shape, dtype)
+            sess.sync_worker_0()
+            run_prim_func(sess, main, A_array, B_array)
+
+            for i in range(NUM_WORKERS):
+                if is_get:
+                    expected_B = A_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy()
+                    actual_B = B_array.debug_get_from_remote(i).numpy()
+                else:  # put
+                    expected_B = A_array.debug_get_from_remote(i).numpy()
+                    actual_B = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy()
+                np.testing.assert_equal(actual_B, expected_B)
+
+        def test_signal_op(sess, sig_op):
+            """Tests signal_op and wait_until to implement a barrier-like pattern."""
+            cmp_value = 1 if sig_op == "set" else 2
+
+            # fmt: off
+            @T.prim_func
+            def main(res: T.Buffer((1,), "uint64")):
+                T.device_entry()
+                cta_id = T.cta_id([1])
+                tid = T.thread_id([nwarps * 32])
+                my_pe = T.nvshmem.my_pe()
+                n_pes = T.nvshmem.n_pes()
+                dst_pe = (my_pe + 1) % n_pes
+                if sig_op == "add":
+                    res[0] = 1
+                T.nvshmem.barrier_all()
+                T.nvshmem.signal_op(sig_addr=res.ptr_to([0]), signal=1, sig_op=sig_op, pe=dst_pe)
+                T.nvshmem.wait_until(ivar=res.ptr_to([0]), cmp="eq", cmp_value=cmp_value)
+                # fmt: on
+
+            res_array = create_nvshmem_array(sess, (1,), "uint64")
+            sess.sync_worker_0()
+            run_prim_func(sess, main, res_array)
+
+            for i in range(NUM_WORKERS):
+                res = res_array.debug_get_from_remote(i).numpy()
+                if sig_op == "set":
+                    np.testing.assert_equal(res[0], 1)
+                elif sig_op == "add":
+                    np.testing.assert_equal(res[0], 2)
+
+        def test_put_signal(sess, scope, shape, nwarps, nelems, cmp_value):
+            """Tests combined data transfer and signal operations at thread/warp/block scopes."""
+            dtype = "float32"
+            op_func = getattr(T.nvshmem, "putmem_signal_nbi")
+            if scope != "thread":
+                op_func = getattr(op_func, scope)
+
+            @T.prim_func
+            def main(
+                A: T.Buffer(shape, dtype),
+                B: T.Buffer(shape, dtype),
+                signal_array: T.Buffer((1,), "uint64"),
+            ):
+                T.device_entry()
+                cta_id = T.cta_id([1])
+                warp_id = T.warp_id([nwarps])
+                lane_id = T.lane_id([32])
+                tid = T.thread_id([nwarps * 32])
+                my_pe = T.nvshmem.my_pe()
+                n_pes = T.nvshmem.n_pes()
+                dst_pe = (my_pe + 1) % n_pes
+                offset = T.if_then_else(
+                    scope == "block",
+                    0,
+                    T.if_then_else(scope == "thread", tid, warp_id * 32),
+                )
+                op_func(
+                    dst=B.access_ptr("w", offset=offset),
+                    src=A.access_ptr("r", offset=offset),
+                    nelems=nelems,
+                    sig_addr=signal_array.access_ptr("w", offset=0),
+                    signal=1,
+                    sig_op="set",
+                    pe=dst_pe,
+                )
+                T.nvshmem.wait_until(
+                    ivar=signal_array.access_ptr("r", offset=0),
+                    cmp="eq",
+                    cmp_value=cmp_value,
+                )
+
+            def init_A(i, s, d):
+                return np.arange(s[0], dtype=d) + i * 100
+
+            A_array = create_nvshmem_array(sess, shape, dtype, init_A)
+            B_array = create_nvshmem_array(sess, shape, dtype)
+            signal_array = create_nvshmem_array(sess, (1,), "uint64")
+
+            sess.sync_worker_0()
+            run_prim_func(sess, main, A_array, B_array, signal_array)
+
+            for i in range(NUM_WORKERS):
+                expected = A_array.debug_get_from_remote(i).numpy()
+                actual = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy()
+                signal_np = signal_array.debug_get_from_remote(i).numpy()
+                np.testing.assert_equal(actual, expected)
+                np.testing.assert_equal(signal_np[0], cmp_value)
+
+        def test_fence_barrier(sess):
+            shape = (64,)
+            dtype = "float32"
+
+            # fmt: off
+            @T.prim_func
+            def main(A: T.Buffer(shape, dtype), B: T.Buffer(shape, dtype), res: T.Buffer((1,), "uint64")):  # noqa: E501
+                T.device_entry()
+                cta_id = T.cta_id([1])
+                warp_id = T.warp_id([nwarps])
+                lane_id = T.lane_id([32])
+                tid = T.thread_id([2 * 32])
+                my_pe = T.nvshmem.my_pe()
+                n_pes = T.nvshmem.n_pes()
+                dst_pe = (my_pe + 1) % n_pes
+                T.nvshmem.barrier_all()
+                T.nvshmem.putmem_nbi.block(dst=B.ptr_to([0]), src=A.ptr_to([0]), nelems=4 * 64, pe=(my_pe + 1) % n_pes)  # noqa: E501
+                T.nvshmem.fence()
+                if tid == 0:
+                    T.nvshmem.signal_op(sig_addr=res.ptr_to([0]), signal=1, sig_op="set", pe=dst_pe)
+                T.nvshmem.wait_until(ivar=res.ptr_to([0]), cmp="eq", cmp_value=1)
+                # fmt: on
+            def init_fn(i, s, d):
+                return np.arange(s[0], dtype=d) + i * 100
+
+            A_array = create_nvshmem_array(sess, shape, dtype, init_fn)
+            B_array = create_nvshmem_array(sess, shape, dtype)
+            res_array = create_nvshmem_array(sess, (1,), "uint64")
+            run_prim_func(sess, main, A_array, B_array, res_array)
+
+            for i in range(NUM_WORKERS):
+                expected_B = A_array.debug_get_from_remote(i).numpy()
+                actual_B = B_array.debug_get_from_remote((i + 1) % NUM_WORKERS).numpy()
+                np.testing.assert_equal(actual_B, expected_B)
+
+        # test thread info
+        test_thread_info(sess)
+        print("\n\ntest_thread_info done\n\n")
+
+        # test transfer
+        for scope, shape, nwarps, nelems, op_name in [
+            ("thread", (32,), 1, 4, "getmem_nbi"),
+            ("thread", (32,), 1, 4, "putmem_nbi"),
+            ("warp", (64,), 2, 4 * 32, "getmem_nbi"),
+            ("warp", (64,), 2, 4 * 32, "putmem_nbi"),
+            ("block", (64,), 2, 4 * 64, "getmem_nbi"),
+            ("block", (64,), 2, 4 * 64, "putmem_nbi"),
+        ]:
+            test_transfer(sess, scope, shape, nwarps, nelems, op_name)
+            print(f"\n\ntest_transfer done for {scope}, {shape}, {nwarps}, {nelems}, {op_name}\n\n")
+
+        # test signal op
+        for sig_op in ["set", "add"]:
+            test_signal_op(sess, sig_op)
+            print(f"\n\ntest_signal_op done for {sig_op}\n\n")
+
+        # test put signal
+        for scope, shape, nwarps, nelems, cmp_value in [
+            ("thread", (32,), 1, 4, 32),
+            ("warp", (64,), 2, 4 * 32, 2),
+            ("block", (64,), 2, 4 * 64, 1),
+        ]:
+            test_put_signal(sess, scope, shape, nwarps, nelems, cmp_value)
+            print(
+                f"\n\ntest_put_signal done for {scope}, {shape}, {nwarps}, {nelems}, {cmp_value}\n\n"  # noqa: E501
+            )
+
+        # test fence barrier
+        test_fence_barrier(sess)
+        print("\n\ntest_fence_barrier done\n\n")
+
+        ############ cleanup ############
+        finalize_dfunc = sess.get_global_func("runtime.disco.nvshmem.finalize_nvshmem")
+        finalize_dfunc()
+        sess.sync_worker_0()
+        sess.shutdown()
+        return True
+
+    def run_and_check():
+        worker = PopenWorker()
+        try:
+            worker.send(_test_func)
+            assert worker.recv()
+        finally:
+            worker.kill()
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/codegen/test_cuda_cta_reduce.py

@@ -0,0 +1,201 @@
+# 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 T.cuda.cta_reduce / cta_sum / cta_max / cta_min intrinsics."""
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm.script import tirx as T
+from tvm.testing import env
+
+TARGET = tvm.target.Target("cuda")
+
+
+def _build_and_run(func, n):
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=TARGET, tir_pipeline="tirx")
+    out_np = np.zeros(n, dtype="float32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        out = tvm.runtime.tensor(out_np, device=dev)
+        mod(out)
+        return out.numpy()
+
+    return tvm.testing.run_with_gpu_lock(run_and_check), mod
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cta_sum_4_warps():
+    """CTA sum with 4 warps (128 threads): all threads get the same sum."""
+    NUM_WARPS = 4
+    N = NUM_WARPS * 32
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (N,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([NUM_WARPS])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([N])
+        scratch = T.alloc_buffer((NUM_WARPS,), "float32", scope="shared")
+        val: T.f32 = T.float32(tid + 1)
+        val = T.cuda.cta_sum(val, NUM_WARPS, scratch.ptr_to([0]))
+        out[tid] = val
+        # fmt: on
+
+    result, mod = _build_and_run(func, N)
+    expected = np.float32(N * (N + 1) / 2)  # sum(1..128)
+    np.testing.assert_allclose(result, np.full(N, expected))
+    assert "cta_reduce_sum_4" in mod.mod.imports[0].inspect_source()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cta_sum_8_warps():
+    """CTA sum with 8 warps (256 threads)."""
+    NUM_WARPS = 8
+    N = NUM_WARPS * 32
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (N,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([NUM_WARPS])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([N])
+        scratch = T.alloc_buffer((NUM_WARPS,), "float32", scope="shared")
+        val: T.f32 = T.float32(tid + 1)
+        val = T.cuda.cta_sum(val, NUM_WARPS, scratch.ptr_to([0]))
+        out[tid] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func, N)
+    expected = np.float32(N * (N + 1) / 2)
+    np.testing.assert_allclose(result, np.full(N, expected))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cta_max_4_warps():
+    """CTA max with 4 warps: all threads get the maximum value."""
+    NUM_WARPS = 4
+    N = NUM_WARPS * 32
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (N,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([NUM_WARPS])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([N])
+        scratch = T.alloc_buffer((NUM_WARPS,), "float32", scope="shared")
+        val: T.f32 = T.float32(tid + 1)
+        val = T.cuda.cta_max(val, NUM_WARPS, scratch.ptr_to([0]))
+        out[tid] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func, N)
+    np.testing.assert_allclose(result, np.full(N, float(N)))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cta_min_4_warps():
+    """CTA min with 4 warps: all threads get the minimum value."""
+    NUM_WARPS = 4
+    N = NUM_WARPS * 32
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (N,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([NUM_WARPS])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([N])
+        scratch = T.alloc_buffer((NUM_WARPS,), "float32", scope="shared")
+        val: T.f32 = T.float32(tid + 1)
+        val = T.cuda.cta_min(val, NUM_WARPS, scratch.ptr_to([0]))
+        out[tid] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func, N)
+    np.testing.assert_allclose(result, np.full(N, 1.0))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_cta_sum_1_warp():
+    """CTA sum with 1 warp: degenerates to a pure warp reduce."""
+    NUM_WARPS = 1
+    N = 32
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (N,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([NUM_WARPS])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([N])
+        scratch = T.alloc_buffer((NUM_WARPS,), "float32", scope="shared")
+        val: T.f32 = T.float32(tid + 1)
+        val = T.cuda.cta_sum(val, NUM_WARPS, scratch.ptr_to([0]))
+        out[tid] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func, N)
+    expected = np.float32(32 * 33 / 2)
+    np.testing.assert_allclose(result, np.full(N, expected))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("num_warps", [1, 2, 4, 8, 16])
+def test_cta_sum_all_warp_counts(num_warps):
+    """Parametric test: cta_sum with various warp counts."""
+    N = num_warps * 32
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (N,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([num_warps])
+        lane_id = T.lane_id([32])
+        tid = T.thread_id([N])
+        scratch = T.alloc_buffer((num_warps,), "float32", scope="shared")
+        val: T.f32 = T.float32(tid + 1)
+        val = T.cuda.cta_sum(val, num_warps, scratch.ptr_to([0]))
+        out[tid] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func, N)
+    expected = np.float32(N * (N + 1) / 2)
+    np.testing.assert_allclose(result, np.full(N, expected))

tests/python/tirx/codegen/test_cuda_warp_reduce.py

@@ -0,0 +1,198 @@
+# 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 T.cuda.warp_reduce / warp_sum / warp_max / warp_min intrinsics."""
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm.script import tirx as T
+from tvm.testing import env
+
+TARGET = tvm.target.Target("cuda")
+
+
+def _build_and_run(func, n=32):
+    mod = tvm.IRModule({"main": func})
+    mod = tvm.compile(mod, target=TARGET, tir_pipeline="tirx")
+    out_np = np.zeros(n, dtype="float32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        out = tvm.runtime.tensor(out_np, device=dev)
+        mod(out)
+        return out.numpy()
+
+    return tvm.testing.run_with_gpu_lock(run_and_check), mod
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_warp_sum_full():
+    """Full warp sum (width=32): each lane gets the sum of all 32 values."""
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (32,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane = T.lane_id([32])
+        val: T.f32 = T.float32(lane + 1)
+        val = T.cuda.warp_sum(val)
+        out[lane] = val
+        # fmt: on
+
+    result, mod = _build_and_run(func)
+    expected = np.float32(32 * 33 / 2)  # sum(1..32)
+    np.testing.assert_allclose(result, np.full(32, expected))
+    assert "warp_reduce_sum_32" in mod.mod.imports[0].inspect_source()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_warp_sum_partial_8():
+    """Partial warp sum (width=8): 4 groups of 8 lanes, each group sums independently."""
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (32,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane = T.lane_id([32])
+        val: T.f32 = T.float32(lane + 1)
+        val = T.cuda.warp_sum(val, width=8)
+        out[lane] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func)
+    # Group 0: lanes 0-7 → sum(1..8) = 36
+    # Group 1: lanes 8-15 → sum(9..16) = 100
+    # Group 2: lanes 16-23 → sum(17..24) = 164
+    # Group 3: lanes 24-31 → sum(25..32) = 228
+    expected = np.zeros(32, dtype="float32")
+    for g in range(4):
+        group_sum = sum(range(g * 8 + 1, g * 8 + 9))
+        expected[g * 8 : (g + 1) * 8] = group_sum
+    np.testing.assert_allclose(result, expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_warp_max_partial_4():
+    """Partial warp max (width=4): 8 groups of 4 lanes."""
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (32,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane = T.lane_id([32])
+        val: T.f32 = T.float32(lane + 1)
+        val = T.cuda.warp_max(val, width=4)
+        out[lane] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func)
+    expected = np.zeros(32, dtype="float32")
+    for g in range(8):
+        group_max = float(g * 4 + 4)
+        expected[g * 4 : (g + 1) * 4] = group_max
+    np.testing.assert_allclose(result, expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_warp_min_full():
+    """Full warp min (width=32)."""
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (32,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane = T.lane_id([32])
+        val: T.f32 = T.float32(lane + 1)
+        val = T.cuda.warp_min(val)
+        out[lane] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func)
+    np.testing.assert_allclose(result, np.full(32, 1.0))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_warp_sum_partial_2():
+    """Smallest partial warp sum (width=2): 16 pairs of adjacent lanes."""
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (32,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane = T.lane_id([32])
+        val: T.f32 = T.float32(lane)
+        val = T.cuda.warp_sum(val, width=2)
+        out[lane] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func)
+    # Pairs: (0,1)→1, (2,3)→5, (4,5)→9, ...
+    expected = np.zeros(32, dtype="float32")
+    for i in range(16):
+        pair_sum = float(2 * i + 2 * i + 1)
+        expected[2 * i] = pair_sum
+        expected[2 * i + 1] = pair_sum
+    np.testing.assert_allclose(result, expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("width", [2, 4, 8, 16, 32])
+def test_warp_sum_all_widths(width):
+    """Parametric test: warp_sum with every valid width."""
+
+    # fmt: off
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (32,), "float32")
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        warp_id = T.warp_id([1])
+        lane = T.lane_id([32])
+        val: T.f32 = T.float32(lane)
+        val = T.cuda.warp_sum(val, width=width)
+        out[lane] = val
+        # fmt: on
+
+    result, _ = _build_and_run(func)
+    expected = np.zeros(32, dtype="float32")
+    num_groups = 32 // width
+    for g in range(num_groups):
+        group_sum = sum(range(g * width, (g + 1) * width))
+        expected[g * width : (g + 1) * width] = float(group_sum)
+    np.testing.assert_allclose(result, expected)

tests/python/tirx/codegen/test_ptx_ld_st_ops.py

@@ -0,0 +1,206 @@
+# 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.
+"""Unit tests for generic PTX ``T.ptx.ld`` / ``T.ptx.st`` vector copy ops."""
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm.ir import Op
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.testing import env
+from tvm.tirx.cuda.operator.tile_primitive.copy._common import (
+    copy_ptx_form,
+    copy_ptx_ld_return_type,
+)
+
+TARGET = tvm.target.Target("cuda")
+
+# num_bytes → kernel layout. ``fill_offset`` fills lane i with ``i + fill_offset``.
+_SHARED_COPY_CASES = {
+    16: {"nelems": 4, "smem_dtype": "uint32", "tmp_dtype": "uint32", "fill_offset": 1},
+    8: {"nelems": 2, "smem_dtype": "uint32", "tmp_dtype": "uint32", "fill_offset": 10},
+    4: {"nelems": 1, "smem_dtype": "uint32", "tmp_dtype": "uint32", "fill_value": 42},
+    2: {"nelems": 1, "smem_dtype": "float16", "tmp_dtype": "uint16", "fill_fp16": 7.0},
+    1: {"nelems": 1, "smem_dtype": "uint8", "tmp_dtype": "uint32", "fill_u8": 255},
+}
+
+
+def _build_and_run(func, *np_args):
+    mod = tvm.compile(tvm.IRModule({"main": func}), target=TARGET, tir_pipeline="tirx")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        rt_args = [tvm.runtime.tensor(a, device=dev) for a in np_args]
+        mod(*rt_args)
+        return tuple(a.numpy() for a in rt_args)
+
+    return (*tvm.testing.run_with_gpu_lock(run_and_check), mod)
+
+
+def _expected_values(num_bytes: int) -> np.ndarray:
+    spec = _SHARED_COPY_CASES[num_bytes]
+    if "fill_offset" in spec:
+        off, nelems = spec["fill_offset"], spec["nelems"]
+        return np.array([off + i for i in range(nelems)], dtype=np.uint32)
+    if "fill_fp16" in spec:
+        return np.array([spec["fill_fp16"]], dtype=np.float16)
+    if "fill_u8" in spec:
+        return np.array([spec["fill_u8"]], dtype=np.uint8)
+    return np.array([spec["fill_value"]], dtype=np.uint32)
+
+
+def _shared_scratch_copy_kernel(num_bytes: int):
+    """Build shared → local scratch → shared copy kernel for ``num_bytes`` width."""
+    spec = _SHARED_COPY_CASES[num_bytes]
+    smem_dtype = spec["smem_dtype"]
+    tmp_dtype = spec["tmp_dtype"]
+    nelems = spec["nelems"]
+    fill_offset = spec.get("fill_offset")
+    fill_value = spec.get("fill_value")
+    fill_fp16 = spec.get("fill_fp16")
+    fill_u8 = spec.get("fill_u8")
+    vec, ptx_type = copy_ptx_form(num_bytes)
+    return_type = copy_ptx_ld_return_type(ptx_type)
+
+    @T.prim_func
+    def func(out_ptr: T.handle):
+        out = T.match_buffer(out_ptr, (nelems,), smem_dtype)
+        T.device_entry()
+        T.cta_id([1])
+        T.warp_id([1])
+        lane = T.lane_id([32])
+        src_buf = T.alloc_buffer((nelems,), smem_dtype, scope="shared")
+        dst_buf = T.alloc_buffer((nelems,), smem_dtype, scope="shared")
+        tmp = T.alloc_local((nelems,), tmp_dtype)
+        if fill_offset is not None:
+            if lane < nelems:
+                src_buf[lane] = T.uint32(lane + fill_offset)
+        elif fill_fp16 is not None:
+            if lane == 0:
+                src_buf[0] = T.float16(fill_fp16)
+        elif fill_u8 is not None:
+            if lane == 0:
+                src_buf[0] = T.uint8(fill_u8)
+        elif lane == 0:
+            src_buf[0] = T.uint32(fill_value)
+        T.cuda.cta_sync()
+        if lane == 0:
+            T.ptx.ld(
+                src_buf.ptr_to([0]),
+                return_type,
+                ptx_type,
+                dst=tmp.ptr_to([0]),
+                space="shared",
+                vec=vec,
+            )
+            T.ptx.st(
+                dst_buf.ptr_to([0]),
+                src=tmp.ptr_to([0]),
+                space="shared",
+                vec=vec,
+                ptx_type=ptx_type,
+            )
+        T.cuda.cta_sync()
+        if lane < nelems:
+            out[lane] = dst_buf[lane]
+
+    return func
+
+
+def test_ptx_ld_st_ops_registered():
+    """PTX ld/st must be registered TIR ops and exposed on the T.ptx namespace."""
+    for name in ("tirx.ptx.ld", "tirx.ptx.st"):
+        Op.get(name)  # raises if unregistered
+
+    for attr in (
+        "ld",
+        "st",
+        "ld_acquire",
+        "st_release",
+        "ld_volatile",
+        "st_volatile",
+    ):
+        assert hasattr(T.ptx, attr), attr
+
+
+def test_ptx_ld_st_codegen_emits_shared_asm():
+    """Shared ↔ register typed copies must codegen to ``ld.shared`` / ``st.shared``."""
+
+    # fmt: off
+    @T.prim_func
+    def copy_kernel(d_ptr: T.handle) -> None:
+        D = T.match_buffer(d_ptr, (4,), "uint32")
+        T.device_entry()
+        T.warp_id([4])
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        tid_in_wg = T.thread_id_in_wg([128])
+        smem = T.alloc_buffer((4,), "uint32", scope="shared")
+        reg = T.alloc_local((4,), "uint32")
+        if tid_in_wg == 0:
+            T.ptx.st(
+                smem.ptr_to([0]), src=reg.ptr_to([0]), space="shared", vec="v4", ptx_type="u32"
+            )
+        T.cuda.cta_sync()
+        if tid_in_wg == 0:
+            T.ptx.ld(
+                smem.ptr_to([0]), "uint32", "u32", dst=reg.ptr_to([0]), space="shared", vec="v4"
+            )
+        Tx.copy(D[0:4], reg[:])
+    # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source("cuda")
+    assert "ld.shared" in src, "PTX ld did not emit ld.shared"
+    assert "st.shared" in src, "PTX st did not emit st.shared"
+    assert "tvm_builtin_ptx_ld" in src
+    assert "tvm_builtin_ptx_st" in src
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize(
+    "num_bytes",
+    [16, 8, 4, 2, 1],
+    ids=["128b", "64b", "32b", "16b", "8b"],
+)
+def test_ptx_ld_st_shared_copy_gpu(num_bytes):
+    """GPU roundtrip for each supported PTX ld/st copy width (shared → scratch → shared)."""
+    expected = _expected_values(num_bytes)
+    kernel = _shared_scratch_copy_kernel(num_bytes)
+    out_np = np.zeros_like(expected)
+    result, mod = _build_and_run(kernel, out_np)
+    if expected.dtype == np.uint8:
+        np.testing.assert_array_equal(result, expected)
+    elif expected.dtype == np.float16:
+        np.testing.assert_allclose(result, expected)
+    else:
+        np.testing.assert_array_equal(result, expected)
+    src = mod.mod.imports[0].inspect_source("cuda")
+    assert "tvm_builtin_ptx_ld" in src
+    assert "tvm_builtin_ptx_st" in src
+    vec, _ptx_type = copy_ptx_form(num_bytes)
+    if vec == "v4":
+        assert "ld.shared.v4" in src
+        assert "st.shared.v4" in src
+    elif vec == "v2":
+        assert "ld.shared.v2" in src
+        assert "st.shared.v2" in src

tests/python/tirx/conftest.py

@@ -0,0 +1,51 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Suite-level hardware gate for the tirx tests.
+
+The tirx kernels and codegen paths target Blackwell (sm_100a) — they emit
+PTX/SASS (tcgen05, tmem, cp.async ``.async`` modifiers, fp8 conversions, ...)
+that ptxas/NVRTC reject for older targets, and many tests execute on the
+device. Running the suite on a CPU-only node or a pre-sm_100 GPU therefore
+fails at compile/run time rather than skipping. Gate the whole directory on a
+real sm_100a device so it skips cleanly where the hardware is absent and runs
+in full where it is present.
+"""
+
+from pathlib import Path
+
+import pytest
+
+from tvm.testing import env
+
+
+def pytest_collection_modifyitems(config, items):
+    if env.has_cuda_compute(10):
+        return
+    suite_root = Path(__file__).resolve().parent
+    skip = pytest.mark.skip(
+        reason="tirx suite requires a CUDA compute capability 10.0 (sm_100a) device"
+    )
+    for item in items:
+        path = getattr(item, "path", None)
+        if path is None:
+            continue
+        try:
+            path = Path(path).resolve()
+        except TypeError:
+            continue
+        if path.is_relative_to(suite_root):
+            item.add_marker(skip)

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()

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py

@@ -0,0 +1,251 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for the DSMEM (shared::cta → shared::cluster) copy_async variant.
+
+Split out from ``test_copy_async.py`` so the TMA-focused file stays focused
+on the g2s/s2g TMA family. Any cross-cutting copy_async helper that both
+files need should live in a shared module, not be duplicated.
+"""
+
+import functools
+
+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 import IntImm, Var
+from tvm.tirx.cuda.operator.tile_primitive.copy_async.dsmem import copy_dsmem_impl
+from tvm.tirx.exec_scope import ExecScope
+from tvm.tirx.layout import S, TileLayout
+from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext
+from tvm.tirx.operator.tile_primitive.dispatcher import DispatchFail
+from tvm.tirx.operator.tile_primitive.ops import CopyAsync
+from tvm.tirx.stmt_functor import StmtExprVisitor
+
+
+def _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout):
+    """Call copy_dsmem_impl directly. Returns impl or raises DispatchFail."""
+    from tvm.ir import Range
+    from tvm.tirx.stmt import BufferRegion
+
+    src_buf = tvm.tirx.decl_buffer(shape, dtype, "A", scope="shared.dyn", layout=src_layout)
+    dst_buf = tvm.tirx.decl_buffer(shape, dtype, "B", scope="shared.dyn", layout=dst_layout)
+    ranges = [Range.from_min_extent(0, s) for s in shape]
+    config = {"mbar": Var("mbar", "handle"), "remote_cta_id": IntImm("int32", 1)}
+    op_call = CopyAsync(BufferRegion(dst_buf, ranges), BufferRegion(src_buf, ranges), config=config)
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
+    sctx = DispatchContext(target, ExecScope("thread"), {}, {})
+    return copy_dsmem_impl(op_call, sctx)
+
+
+class _S2CCounter(StmtExprVisitor):
+    """Count cp.async.bulk.shared_to_cluster calls including loop iterations."""
+
+    def __init__(self):
+        super().__init__()
+        self._loop_extents = []
+        self.total = 0
+
+    def visit_for_(self, op):
+        self._loop_extents.append(op.extent)
+        self.visit_stmt(op.body)
+        self._loop_extents.pop()
+
+    def visit_evaluate_(self, op):
+        if isinstance(op.value, tvm.ir.Call):
+            if op.value.op.name == "tirx.ptx.cp_async_bulk_shared_to_cluster":
+                n = 1
+                for e in self._loop_extents:
+                    n *= e
+                self.total += n
+
+
+def _count_s2c_ops(impl):
+    c = _S2CCounter()
+    c.visit_stmt(impl.body)
+    return c.total
+
+
+# ---------------------------------------------------------------------------
+# Parametrized DSMEM test: dispatch assertion + GPU correctness
+# ---------------------------------------------------------------------------
+
+# (shape, dtype, src_spec, dst_spec, expected_s2c_ops | "fail")
+# Dispatch assertion uses src_spec/dst_spec as given.
+# GPU correctness (all non-fail cases) uses src_spec as the layout for both CTAs.
+DSMEM_CONFIGS = [
+    pytest.param((128, 64), "float16", S[128, 64], S[128, 64], 1, id="contiguous-2d"),
+    pytest.param((256,), "float16", S[256], S[256], 1, id="contiguous-1d"),
+    # Stride gap: inner 128 contiguous, outer stride=256 (gap) → 8 bulk copies
+    pytest.param(
+        (8, 128), "float16", S[(8, 128) : (256, 1)], S[(8, 128) : (256, 1)], 8, id="stride-gap"
+    ),
+    # Different outer strides → 8 bulk copies in dispatch
+    pytest.param(
+        (8, 128),
+        "float16",
+        S[(8, 128) : (256, 1)],
+        S[(8, 128) : (512, 1)],
+        8,
+        id="partial-contiguity-diff-stride",
+    ),
+    # Incompatible: row-major vs column-major → DispatchFail
+    pytest.param(
+        (4, 64), "float16", S[4, 64], S[(4, 64) : (1, 4)], "fail", id="incompatible-row-vs-col"
+    ),
+]
+
+
+def _layout_physical_elements(layout):
+    """Compute number of physical elements needed for a TileLayout."""
+    max_offset = 0
+    for shard in layout.shard:
+        if shard.axis.is_memory():
+            max_offset += int(shard.stride) * (int(shard.extent) - 1)
+    return max_offset + 1
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
+@pytest.mark.parametrize("shape,dtype,src_spec,dst_spec,expected", DSMEM_CONFIGS)
+def test_dsmem(shape, dtype, src_spec, dst_spec, expected):
+    """Dispatch assertion + GPU correctness for DSMEM copy.
+
+    Always tests dispatch (s2c op count or DispatchFail).
+    For non-fail cases: also runs a 2-CTA cluster kernel via T.copy_async
+    dispatch (using src_spec as layout for both CTAs) and verifies correctness.
+    """
+    from tvm.tirx.lang.pipeline import MBarrier
+
+    src_layout = TileLayout(src_spec)
+    dst_layout = TileLayout(dst_spec)
+
+    # --- Dispatch assertion ---
+    if expected == "fail":
+        with pytest.raises(DispatchFail):
+            _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout)
+        return
+
+    impl = _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout)
+    assert _count_s2c_ops(impl) == expected
+
+    # --- GPU correctness ---
+    # Allocate two separate smem buffers: src_smem (src_layout) and dst_smem
+    # (dst_layout). CTA 0 loads global→src_smem, copy_async copies src_smem→
+    # dst_smem on CTA 1. CTA 1 reads dst_smem and writes to global output.
+
+    CLUSTER_N = 2
+    n_elements = functools.reduce(lambda a, b: a * b, shape, 1)
+    copy_bytes = n_elements * tvm.DataType(dtype).bits // 8
+    src_phys = _layout_physical_elements(src_layout)
+    dst_phys = _layout_physical_elements(dst_layout)
+    r = tuple(slice(0, s) for s in shape)
+
+    # fmt: off
+    @T.prim_func
+    def dsmem_copy(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()
+        cbx = T.cta_id_in_cluster([CLUSTER_N])
+        T.cta_id([CLUSTER_N])
+        tid = T.thread_id([1])
+        pool = T.SMEMPool()
+                # src_smem: CTA 0 writes here, dispatch reads from here
+        src_raw = pool.alloc([src_phys], dtype, align=128)
+        src_smem = T.decl_buffer(
+            list(shape), dtype, src_raw.data,
+            elem_offset=0, scope="shared.dyn", layout=src_layout,
+        )
+                # dst_smem: dispatch writes here (on remote CTA), CTA 1 reads
+        dst_raw = pool.alloc([dst_phys], dtype, align=128)
+        dst_smem = T.decl_buffer(
+            list(shape), dtype, dst_raw.data,
+            elem_offset=0, scope="shared.dyn", layout=dst_layout,
+        )
+        mbar = MBarrier(pool, 1)
+        pool.commit()
+
+        mbar.init(1)
+        T.ptx.fence.mbarrier_init()
+        T.cuda.cluster_sync()
+
+        if tid == 0:
+            if cbx == 0:
+                Tx.copy(src_smem[r], A[r])
+                T.ptx.fence.proxy_async("shared::cta")
+
+                Tx.copy_async(
+                    dst_smem[r], src_smem[r],
+                    dispatch="dsmem",
+                    mbar=mbar.ptr_to([0]),
+                    remote_cta_id=T.int32(1),
+                )
+            else:
+                T.ptx.mbarrier.arrive.expect_tx(mbar.ptr_to([0]), copy_bytes)
+                mbar.wait(0, 0)
+
+                Tx.copy(B[r], dst_smem[r])
+        # fmt: on
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": dsmem_copy})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        cuda_src = mod.mod.imports[0].inspect_source()
+        assert "cp.async.bulk.shared::cluster.shared::cta" in cuda_src
+
+        np.random.seed(0)
+        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 = tvm.runtime.tensor(A_np, dev)
+        B_tvm = tvm.runtime.tensor(B_np, dev)
+        mod(A_tvm, B_tvm)
+        np.testing.assert_allclose(A_np, B_tvm.numpy())
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_dsmem_dispatch_missing_config():
+    """Dispatch fails when required config keys are missing."""
+    from tvm.ir import Range
+    from tvm.tirx.stmt import BufferRegion
+
+    layout = TileLayout(S[64])
+    buf = tvm.tirx.decl_buffer((64,), "float16", "A", scope="shared.dyn", layout=layout)
+    br = BufferRegion(buf, [Range.from_min_extent(0, 64)])
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
+    sctx = DispatchContext(target, ExecScope("thread"), {}, {})
+
+    with pytest.raises(DispatchFail, match="remote_cta_id"):
+        copy_dsmem_impl(CopyAsync(br, br, config={"mbar": Var("m", "handle")}), sctx)
+    with pytest.raises(DispatchFail, match="mbar"):
+        copy_dsmem_impl(CopyAsync(br, br, config={"remote_cta_id": IntImm("int32", 1)}), sctx)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_ldgsts.py

@@ -0,0 +1,123 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, missing-function-docstring
+"""Tests for the non-bulk CTA-level copy_async dispatch (vectorized load)."""
+
+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
+
+
+@pytest.mark.parametrize(
+    "task",
+    [
+        ################ A[0:128, 0:32] -> A_smem[0:128, 0:32] -> B[0:128, 0:32] ################
+        (
+            (128, 32),  # g_shape
+            (128, 32),  # s_shape
+            (0, 0),  # g_st
+            (128, 32),  # g_extent
+            32,  # thread_cnt
+            TileLayout(S[128, 32]),  # layoutA
+            TileLayout(S[128, 32]),  # layoutB
+            TileLayout(S[128, 32]),  # layoutS
+        ),
+        ################ A[32:64, 32:64] -> A_smem[0:32, 0:32] -> B[32:64, 32:64] ################
+        (
+            (64, 64),  # g_shape
+            (32, 32),  # s_shape
+            (32, 0),  # g_st
+            (32, 32),  # g_extent
+            32,  # thread_cnt
+            TileLayout(S[64, 64]),  # layoutA
+            TileLayout(S[64, 64]),  # layoutB
+            TileLayout(S[32, 32]),  # layoutS
+        ),
+        ################ A[0:1, 0:32, 0:32] -> A_smem[0:32, 0:32] -> B[0:1, 0:32, 0:32] ################  # noqa: E501
+        (
+            (4, 32, 32),  # g_shape
+            (32, 32),  # s_shape
+            (0, 0, 0),  # g_st
+            (1, 32, 32),  # g_extent
+            32,  # thread_cnt
+            TileLayout(S[4, 32, 32]),  # layoutA
+            TileLayout(S[4, 32, 32]),  # layoutB
+            TileLayout(S[32, 32]),  # layoutS
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize(
+    "dtype", ["int8", "float8_e4m3fn", "float8_e5m2", "float16", "bfloat16", "float32"]
+)
+def test_copy_g2s_s2g_cta_vec_load(task, dtype):
+    g_shape, s_shape, g_st, g_extent, thread_cnt, layoutA, layoutB, layoutS = task
+
+    r_smem = list(slice(None) for i in range(len(s_shape)))
+    r_gmem = list(slice(g_st[i], g_st[i] + g_extent[i]) for i in range(len(g_shape)))
+
+    # fmt: off
+    @T.prim_func
+    def copy_async(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()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([thread_cnt])
+        A_smem = T.alloc_buffer(s_shape, dtype, scope="shared", layout=layoutS)
+
+        Tx.cta.copy_async(A_smem[tuple(r_smem)], A[tuple(r_gmem)], dispatch="ldgsts")
+        T.ptx.cp_async.commit_group()
+        T.ptx.cp_async.wait_group()
+        T.cuda.cta_sync()
+        Tx.cta.copy(B[tuple(r_gmem)], A_smem[tuple(r_smem)])
+        # fmt: on
+
+    np_dtype = tvm.testing.np_dtype_from_str(dtype)
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": copy_async})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        np.random.seed(0)
+        A_np = np.random.rand(*g_shape).astype(np_dtype)
+        B_np = np.zeros(g_shape, dtype=np_dtype)
+
+        B_ref = B_np.copy()
+        B_ref[tuple(r_gmem)] = A_np[tuple(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)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py

@@ -0,0 +1,480 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, missing-function-docstring
+"""End-to-end tests for the smem->tmem (tcgen05.cp.32x128b.warpx4) dispatch.
+
+The new dispatch requires the user to declare the t buffer with an
+explicit ``R[4 : 32@TLane]`` indicating warpx4 broadcast — i.e., t.shape[lane] = 32
+with replica 4 → 128 physical lanes.
+
+Run with: pytest test_smem_tmem_dispatch.py -n 8 -v
+"""
+
+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.cuda.operator.tile_primitive.tma_utils import SwizzleMode, mma_shared_layout
+from tvm.tirx.layout import R, S, TCol, TileLayout, TLane
+
+T_LAY_BASIC = TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane])
+
+
+def _make_2d_kernel(
+    s_full,
+    t_full,
+    s_full_shape,
+    t_full_shape,
+    s_r0,
+    s_r1,
+    s_c0,
+    s_c1,
+    t_r0,
+    t_r1,
+    t_c0,
+    t_c1,
+    dtype,
+    cta_group=1,
+):
+    """2D variant: SMEM/TMEM are both 2D; copy a rectangular sub-region."""
+    n_tmem_cols_total = max(32, t_full_shape[-1])
+    OUT_LANES = 32
+    OUT_BYTES = 16
+
+    @T.prim_func(check_well_formed=False)
+    def kernel(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, s_full_shape, dtype)
+        B = T.match_buffer(B_ptr, (OUT_LANES, OUT_BYTES), dtype)
+        T.device_entry()
+        warp_id = T.warp_id([4])
+        wg_id = T.warpgroup_id([1])
+        tid_in_wg = T.thread_id_in_wg([128])
+        lane_id = T.lane_id([32])
+        A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024)
+        tmem_addr = T.alloc_shared([1], "uint32")
+        cp_mbar = T.alloc_shared([1], "uint64")
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(
+                    T.address_of(tmem_addr),
+                    n_cols=n_tmem_cols_total,
+                    cta_group=cta_group,
+                )
+            if tid_in_wg == 0:
+                T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1)
+            T.ptx.fence.proxy_async("shared::cta")
+            T.cuda.cta_sync()
+            Tx.cta.copy(A_smem[:, :], A[:, :])
+            T.cuda.cta_sync()
+            tmem = T.decl_buffer(
+                t_full_shape,
+                dtype,
+                scope="tmem",
+                allocated_addr=tmem_addr[0],
+                layout=t_full,
+            )
+            if tid_in_wg == 0:
+                Tx.copy_async(
+                    tmem[t_r0:t_r1, t_c0:t_c1],
+                    A_smem[s_r0:s_r1, s_c0:s_c1],
+                    cta_group=cta_group,
+                )
+                T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group)
+            T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
+            T.cuda.cta_sync()
+            T.ptx.tcgen05.fence.after_thread_sync()
+            if warp_id == 0:
+                reg = T.alloc_buffer((4,), "uint32", scope="local")
+                for i in range(4):
+                    T.ptx.tcgen05.ld(
+                        tmem.allocated_addr[0],
+                        reg[i],
+                        shape="32x32b",
+                        num=1,
+                        row=0,
+                        col=i,
+                    )
+                T.ptx.tcgen05.wait.ld()
+                B_bytes = reg.view(dtype)
+                for i in range(OUT_BYTES):
+                    B[lane_id, i] = B_bytes[i]
+            if warp_id == 0:
+                T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+                T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group)
+
+    return kernel
+
+
+def _make_3d_4tile_kernel(s_full, t_full, s_full_shape, t_full_shape, dtype, cta_group=1):
+    """3D variant: 4 stacked tiles (NVFP4-style multi-cp test)."""
+    n_tmem_cols_total = max(32, t_full_shape[-1])
+
+    @T.prim_func(check_well_formed=False)
+    def kernel(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, s_full_shape, dtype)
+        B = T.match_buffer(B_ptr, (32, 16), dtype)
+        T.device_entry()
+        warp_id = T.warp_id([4])
+        wg_id = T.warpgroup_id([1])
+        tid_in_wg = T.thread_id_in_wg([128])
+        lane_id = T.lane_id([32])
+        A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024)
+        tmem_addr = T.alloc_shared([1], "uint32")
+        cp_mbar = T.alloc_shared([1], "uint64")
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(
+                    T.address_of(tmem_addr),
+                    n_cols=n_tmem_cols_total,
+                    cta_group=cta_group,
+                )
+            if tid_in_wg == 0:
+                T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1)
+            T.ptx.fence.proxy_async("shared::cta")
+            T.cuda.cta_sync()
+            Tx.cta.copy(A_smem[:, :, :], A[:, :, :])
+            T.cuda.cta_sync()
+            tmem = T.decl_buffer(
+                t_full_shape,
+                dtype,
+                scope="tmem",
+                allocated_addr=tmem_addr[0],
+                layout=t_full,
+            )
+            if tid_in_wg == 0:
+                Tx.copy_async(
+                    tmem[:, :, :],
+                    A_smem[:, :, :],
+                    cta_group=cta_group,
+                )
+                T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group)
+            T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
+            T.cuda.cta_sync()
+            T.ptx.tcgen05.fence.after_thread_sync()
+            if warp_id == 0:
+                reg = T.alloc_buffer((4,), "uint32", scope="local")
+                for i in range(4):
+                    T.ptx.tcgen05.ld(
+                        tmem.allocated_addr[0],
+                        reg[i],
+                        shape="32x32b",
+                        num=1,
+                        row=0,
+                        col=i,
+                    )
+                T.ptx.tcgen05.wait.ld()
+                B_bytes = reg.view(dtype)
+                for i in range(16):
+                    B[lane_id, i] = B_bytes[i]
+            if warp_id == 0:
+                T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group)
+                T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group)
+
+    return kernel
+
+
+def _run_2d(s_full, t_full, s_full_shape, s_region, dtype, A_init, expected):
+    s_r0, s_r1 = s_region[0]
+    s_c0, s_c1 = s_region[1]
+    kernel = _make_2d_kernel(
+        s_full, t_full, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, dtype
+    )
+    return _execute(kernel, A_init, expected)
+
+
+def _run_3d_4tile(s_full, t_full, s_full_shape, dtype, A_init, expected):
+    kernel = _make_3d_4tile_kernel(s_full, t_full, s_full_shape, s_full_shape, dtype)
+    return _execute(kernel, A_init, expected)
+
+
+def _execute(kernel, A_init, expected):
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+    B_np = np.zeros((32, 16), dtype=A_init.dtype)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A = tvm.runtime.tensor(A_init, dev)
+        B = tvm.runtime.tensor(B_np, dev)
+        mod(A, B)
+        B_out = B.numpy()
+        assert np.array_equal(B_out, expected), (
+            f"mismatch:\nlane 0 expected={expected[0].tolist()}\n"
+            f"        got     ={B_out[0].tolist()}"
+        )
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+@pytest.mark.parametrize(
+    "name,s_full,s_full_shape,s_region",
+    [
+        ("sw0_plain_atom_aligned", TileLayout(S[(32, 16) : (16, 1)]), [32, 16], [(0, 32), (0, 16)]),
+        (
+            "sw1_32B_atom",
+            mma_shared_layout("uint8", SwizzleMode.SWIZZLE_32B_ATOM, [32, 32]),
+            [32, 32],
+            [(0, 32), (0, 16)],
+        ),
+        (
+            "sw2_64B_atom",
+            mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 64]),
+            [32, 64],
+            [(0, 32), (0, 16)],
+        ),
+        (
+            "sw3_128B_atom",
+            mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [32, 128]),
+            [32, 128],
+            [(0, 32), (0, 16)],
+        ),
+        (
+            "sw3_64x128_corner",
+            mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]),
+            [64, 128],
+            [(0, 32), (0, 16)],
+        ),
+        (
+            "sw3_64x128_atom_row_8",
+            mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]),
+            [64, 128],
+            [(8, 40), (0, 16)],
+        ),
+        (
+            "sw2_32x256_col_64",
+            mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 256]),
+            [32, 256],
+            [(0, 32), (64, 80)],
+        ),
+        (
+            "sw0_M_atom_major_4_0",
+            TileLayout(S[(8, 8, 2, 16) : (128, 16, 1024, 1)]),
+            [64, 32],
+            [(4, 36), (0, 16)],
+        ),
+    ],
+)
+def test_single_cp(name, s_full, s_full_shape, s_region):
+    A_np = np.arange(int(np.prod(s_full_shape)), dtype=np.uint8).reshape(s_full_shape)
+    r0, r1 = s_region[0]
+    c0, c1 = s_region[1]
+    expected = A_np[r0:r1, c0:c1]
+    _run_2d(s_full, T_LAY_BASIC, s_full_shape, s_region, "uint8", A_np, expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_multi_cp_sw0_4tiles():
+    s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)])
+    t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane])
+    A_np = (np.arange(4 * 32 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(4, 32, 16)
+    expected = A_np[0]
+    _run_3d_4tile(s_full, t_full, [4, 32, 16], "uint8", A_np, expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+def test_align_middle_2_to_1_nvfp4_sfb():
+    """SFB-style nvfp4 case: TMEM mid canonicalizes to single iter
+    (16@TCol + 4@TCol merge), but SMEM mid stays as 2 iters
+    (stride 512 + stride 2048 — outer/inner reversed so canon can't merge).
+    Exercises ``_align_middles`` union-cut algorithm.
+
+    Layout shapes mirror SFB nvfp4 with PIPE=1, SFB_n_chunks=2,
+    MMA_K_BLOCKS=4, sf_mma_k=4.
+    """
+    # SMEM: (2, 4, 32, 4, 4) extents, strides (2048, 4, 16, 512, 1)
+    # — N_chunk outer (stride 2048), then sub-warp tile (4, stride 4), lane
+    # (32, stride 16), K_block (4, stride 512), sf_mma_k (4, stride 1).
+    # Mid post-canon = [(4, 512), (2, 2048)] — non-mergeable in this order.
+    s_full = TileLayout(S[(2, 4, 32, 4, 4) : (2048, 4, 16, 512, 1)])
+    # TMEM: SFB-style 5-axis layout. K_outer (4, 4@TCol) and N_chunk
+    # (2, 16@TCol) merge into single mid iter (8, 4@TCol).
+    t_full = TileLayout(
+        S[(2, 4, 32, 4, 4) : (16 @ TCol, 4 @ TCol, 1 @ TLane, 32 @ TCol, 1 @ TCol)]
+        + R[4 : 32 @ TLane]
+    )
+    s_full_shape = [256, 16]
+    t_full_shape = [256, 16]
+    n_tmem_cols_total = max(32, 32)  # SFB occupies 32 cols total (8*4 elements / 4 epc)
+
+    @T.prim_func(check_well_formed=False)
+    def kernel(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, s_full_shape, "uint8")
+        B = T.match_buffer(B_ptr, (32, 16), "uint8")
+        T.device_entry()
+        warp_id = T.warp_id([4])
+        wg_id = T.warpgroup_id([1])
+        tid_in_wg = T.thread_id_in_wg([128])
+        lane_id = T.lane_id([32])
+        A_smem = T.alloc_buffer(s_full_shape, "uint8", scope="shared", layout=s_full, align=1024)
+        tmem_addr = T.alloc_shared([1], "uint32")
+        cp_mbar = T.alloc_shared([1], "uint64")
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=n_tmem_cols_total, cta_group=1)
+            if tid_in_wg == 0:
+                T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1)
+            T.ptx.fence.proxy_async("shared::cta")
+            T.cuda.cta_sync()
+            Tx.cta.copy(A_smem[:, :], A[:, :])
+            T.cuda.cta_sync()
+            tmem = T.decl_buffer(
+                t_full_shape,
+                "uint8",
+                scope="tmem",
+                allocated_addr=tmem_addr[0],
+                layout=t_full,
+            )
+            if tid_in_wg == 0:
+                Tx.copy_async(tmem[:, :], A_smem[:, :], cta_group=1)
+                T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=1)
+            T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0)
+            T.cuda.cta_sync()
+            T.ptx.tcgen05.fence.after_thread_sync()
+            if warp_id == 0:
+                reg = T.alloc_buffer((4,), "uint32", scope="local")
+                for i in range(4):
+                    T.ptx.tcgen05.ld(
+                        tmem.allocated_addr[0],
+                        reg[i],
+                        shape="32x32b",
+                        num=1,
+                        row=0,
+                        col=i,
+                    )
+                T.ptx.tcgen05.wait.ld()
+                B_bytes = reg.view("uint8")
+                for i in range(16):
+                    B[lane_id, i] = B_bytes[i]
+            if warp_id == 0:
+                T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+                T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=1)
+
+    A_np = (np.arange(256 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(256, 16)
+
+    # Compute expected: for each (lane=L in 0..32, byte b in 0..15), the
+    # tcgen05.ld reads physical (TLane=L, TCol=b). We must invert the TMEM
+    # layout to find which logical (m, k) is at that physical position, then
+    # expected[L, b] = A[m, k].
+    # Layout shard iters (i0..i4) with extents (2, 4, 32, 4, 4) and TMEM
+    # strides (16, 4, 1@TLane, 32, 1) — only TLane and TCol contribute.
+    # For (TLane=L, TCol=p) with L in 0..32, replica r=0:
+    #   i2 = L; remaining iters (i0, i1, i3, i4) contribute to TCol:
+    #   p = 16*i0 + 4*i1 + 32*i3 + i4
+    # For p in 0..15 only i1 and i4 vary (i0 = i3 = 0):
+    #   i1 = p // 4, i4 = p % 4
+    # Logical buffer index: rev row-major over iter coords following shard order.
+    # Shard order outer→inner: (i0, i1, i2, i3, i4) with extents (2, 4, 32, 4, 4).
+    # Logical buffer index = i0*(4*32*4*4) + i1*(32*4*4) + i2*(4*4) + i3*4 + i4
+    expected = np.zeros((32, 16), dtype=np.uint8)
+    for L in range(32):
+        for p in range(16):
+            i0 = 0
+            i3 = 0
+            i1 = p // 4
+            i4 = p % 4
+            i2 = L
+            logical = i0 * (4 * 32 * 4 * 4) + i1 * (32 * 4 * 4) + i2 * (4 * 4) + i3 * 4 + i4
+            m, k = divmod(logical, 16)
+            expected[L, p] = A_np[m, k]
+
+    _execute(kernel, A_np, expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+@pytest.mark.parametrize(
+    "bad",
+    [
+        pytest.param(
+            (
+                "sw3_mid_atom_row",
+                mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]),
+                [64, 128],
+                [(4, 36), (0, 16)],
+            ),
+            id="sw3_mid_atom_row",
+        ),
+        pytest.param(
+            (
+                "sw2_mid_atom_col",
+                mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 128]),
+                [32, 128],
+                [(0, 32), (32, 48)],
+            ),
+            id="sw2_mid_atom_col",
+        ),
+        pytest.param(
+            ("sw0_row_stride_64", TileLayout(S[(64, 64) : (64, 1)]), [64, 64], [(4, 36), (0, 16)]),
+            id="sw0_row_stride_64",
+        ),
+    ],
+)
+def test_dispatch_rejects_bad_inputs(bad):
+    """Configurations where cp 32x128b cannot read the user's intended sub-tile.
+    Compilation should fail with a clear ValueError from the dispatch."""
+    name, s_full, s_full_shape, s_region = bad
+    s_r0, s_r1 = s_region[0]
+    s_c0, s_c1 = s_region[1]
+    kernel = _make_2d_kernel(
+        s_full, T_LAY_BASIC, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, "uint8"
+    )
+    with pytest.raises(Exception):
+        target = tvm.target.Target("cuda")
+        with target:
+            tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+
+
+def test_multi_cp_encodes_descriptor_once_and_patches_addr():
+    """Compile-only regression for the shared-descriptor cp path.
+
+    A multi-tile smem->tmem copy encodes ONE SMEM matrix descriptor template at
+    SMEM base 0 (so the cache key no longer depends on the buffer identity) and
+    patches its 14-bit address field per cp via ``cvta(addr) >> 4 & 0x3FFF``,
+    instead of re-encoding a descriptor per tile. Verifies the 4-tile copy emits
+    a single ``encode_matrix_descriptor`` reused across four
+    ``tcgen05.cp.32x128b.warpx4`` issues, each with the address-field patch.
+    """
+    s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)])
+    t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane])
+    kernel = _make_3d_4tile_kernel(s_full, t_full, [4, 32, 16], [4, 32, 16], "uint8")
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source()
+
+    assert "tcgen05.cp.cta_group::1.32x128b.warpx4" in src, f"cp not emitted; src=\n{src}"
+    # Descriptor encoded once (single matrix-descriptor encode call), then
+    # reused with a per-cp 14-bit SMEM address patch (0x3FFF == 16383 mask).
+    assert "16383" in src, "expected 14-bit SMEM address-field patch (0x3FFF mask)"
+    assert src.count("cp_desc[0] &") == 4, (
+        f"expected 4 address-patched cp's reusing one cp_desc; got "
+        f"{src.count('cp_desc[0] &')}\nsrc=\n{src}"
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem.py

@@ -0,0 +1,344 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, missing-function-docstring
+"""Tests for the TMEM copy_async dispatch (tcgen05-based tmem<->reg and smem<->tmem)."""
+
+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, TCol, TileLayout, TLane
+from tvm.tirx.layout import tid_in_wg as axis_tid_in_wg
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+@pytest.mark.parametrize("dtype", ["float16", "float32"])
+@pytest.mark.parametrize("width_32b", [4, 8, 16, 32])
+def test_copy_tmem2reg_async(dtype, width_32b):
+    """Test async tmem<->local copy using copy_async instead of copy.
+
+    This tests the new copy_async dispatch for tmem<->local that doesn't
+    immediately wait after the operation, allowing for pipelining.
+    """
+
+    def next_power_of_2(x):
+        """Return the smallest power of 2 greater than or equal to x."""
+        if x <= 1:
+            return 1
+        return 1 << (x - 1).bit_length()
+
+    bits = tvm.runtime.DataType(dtype).bits
+    if 128 % bits != 0 or 32 % bits != 0:
+        pytest.skip(f"dtype {dtype} is not supported")
+
+    WIDTH = width_32b * (32 // bits)
+    VEC_LEN = 128 // bits
+    if WIDTH % VEC_LEN != 0:
+        pytest.skip(f"dtype {dtype} + width {width_32b} is not supported")
+
+    g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)])
+    local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)])
+
+    # fmt: off
+    @T.prim_func
+    def copy_async_test(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128, WIDTH), dtype)
+        B = T.match_buffer(B_ptr, (128, WIDTH), dtype)
+
+        A_flat = A.view(-1)
+        B_flat = B.view(-1)
+
+        T.device_entry()
+        warp_id = T.warp_id([(128) // 32])
+        cta_id = T.cta_id([2])
+        wg_id = T.warpgroup_id([1])
+        warp_id_in_wg = T.warp_id_in_wg([4])
+        lane_id = T.lane_id([32])
+        tid_in_wg = T.thread_id([128])
+
+        tmem_addr = T.alloc_shared([1], "uint32")
+
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1)  # noqa: E501
+
+            T.tvm_storage_sync("shared")
+
+            tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0],
+                                 layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)]))
+
+            A_reg = T.alloc_local((WIDTH), dtype)
+            B_reg = T.alloc_local((WIDTH), dtype)
+            A_local = A_reg.view(128, WIDTH, layout=local_view)
+            B_local = B_reg.view(128, WIDTH, layout=local_view)
+            for i in range(WIDTH // VEC_LEN):
+                g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
+                Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN])  # noqa: E501
+            for i in range(WIDTH):
+                B_reg[i] = T.cast(0, dtype)
+            T.cuda.cta_sync()
+
+                    # A_local -> tmem (async)
+            Tx.wg.copy_async(tmem[:, :], A_local[:, :])
+            T.ptx.tcgen05.wait.st()  # explicit wait
+            T.cuda.cta_sync()
+
+                    # tmem -> B_local (async)
+            Tx.wg.copy_async(B_local[:, :], tmem[:, :])
+            T.ptx.tcgen05.wait.ld()  # explicit wait
+            T.cuda.cta_sync()
+            for i in range(WIDTH // VEC_LEN):
+                g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
+                Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN])  # noqa: E501
+
+            if warp_id == 0:
+                T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+                T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1)  # noqa: E501
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": copy_async_test})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH))
+        B_np = np.zeros((128, WIDTH), dtype=dtype)
+
+        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)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ----------------------------------------------------------------------------
+# Migrated from test_copy_sync.py: tmem<->reg round-trip via T.copy_async
+# (the kernels themselves are the actual async tmem dispatch tests; the
+# G↔L copies bookending them just stage data).
+# ----------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+@pytest.mark.parametrize("dtype", ["uint8", "float16", "float32"])
+@pytest.mark.parametrize("width_32b", [2, 4, 8, 16, 32, 64, 128])
+@pytest.mark.parametrize("offset_32b", [0, 3, 10])
+def test_copy_tmem2reg(dtype, width_32b, offset_32b):
+    def next_power_of_2(x):
+        if x <= 1:
+            return 1
+        return 1 << (x - 1).bit_length()
+
+    bits = tvm.runtime.DataType(dtype).bits
+    if 128 % bits != 0 or 32 % bits != 0:
+        pytest.skip(f"dtype {dtype} is not supported")
+
+    WIDTH = width_32b * (32 // bits)
+    OFFSET = offset_32b * (32 // bits)
+    VEC_LEN = 128 // bits
+    if WIDTH % VEC_LEN != 0:
+        pytest.skip(f"dtype {dtype} + width {width_32b} is not supported")
+
+    g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)])
+    local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)])
+
+    # fmt: off
+    @T.prim_func
+    def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128, WIDTH), dtype)
+        B = T.match_buffer(B_ptr, (128, WIDTH), dtype)
+
+        A_flat = A.view(-1)
+        B_flat = B.view(-1)
+
+        T.device_entry()
+        warp_id = T.warp_id([(128) // 32])
+        T.cta_id([2])
+        wg_id = T.warpgroup_id([1])
+        T.warp_id_in_wg([4])
+        T.lane_id([32])
+        tid_in_wg = T.thread_id([128])
+
+        tmem_addr = T.alloc_shared([1], "uint32")
+
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1)  # noqa: E501
+
+            T.tvm_storage_sync("shared")
+
+            tmem = T.decl_buffer((128, OFFSET + WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0],  # noqa: E501
+                                 layout=TileLayout(S[(128, OFFSET + WIDTH) : (1 @ TLane, 1 @ TCol)]))  # noqa: E501
+
+            A_reg = T.alloc_local((WIDTH), dtype)
+            B_reg = T.alloc_local((WIDTH), dtype)
+            A_local = A_reg.view(128, WIDTH, layout=local_view)
+            B_local = B_reg.view(128, WIDTH, layout=local_view)
+            for i in range(WIDTH // VEC_LEN):
+                g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
+                Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN])  # noqa: E501
+            for i in range(WIDTH):
+                B_reg[i] = T.cast(0, dtype)
+            T.cuda.cta_sync()
+
+                    # A_local -> tmem
+            Tx.wg.copy_async(tmem[:, OFFSET: OFFSET + WIDTH], A_local[:, :])
+            T.ptx.tcgen05.wait.st()
+            T.cuda.cta_sync()
+
+                    # tmem -> B_local
+            Tx.wg.copy_async(B_local[:, :], tmem[:, OFFSET: OFFSET + WIDTH])
+            T.ptx.tcgen05.wait.ld()
+            T.cuda.cta_sync()
+            for i in range(WIDTH // VEC_LEN):
+                g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
+                Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN])  # noqa: E501
+
+            if warp_id == 0:
+                T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+                T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1)  # noqa: E501
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": copy_sync})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH))
+        B_np = np.zeros((128, WIDTH), dtype=dtype)
+
+        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)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0")
+@pytest.mark.parametrize("dtype", ["float16", "float32"])
+@pytest.mark.parametrize("width_32b", [4, 8, 16, 32])
+@pytest.mark.parametrize("local_offset_32b", [0, 2, 4])
+def test_copy_tmem2reg_sliced_local(dtype, width_32b, local_offset_32b):
+    """tmem<->local copy with a sliced local buffer region."""
+
+    def next_power_of_2(x):
+        if x <= 1:
+            return 1
+        return 1 << (x - 1).bit_length()
+
+    bits = tvm.runtime.DataType(dtype).bits
+    if 128 % bits != 0 or 32 % bits != 0:
+        pytest.skip(f"dtype {dtype} is not supported")
+
+    WIDTH = width_32b * (32 // bits)
+    LOCAL_OFFSET = local_offset_32b * (32 // bits)
+    TOTAL_LOCAL_WIDTH = WIDTH + LOCAL_OFFSET
+    VEC_LEN = 128 // bits
+    if WIDTH % VEC_LEN != 0 or TOTAL_LOCAL_WIDTH % VEC_LEN != 0:
+        pytest.skip(
+            f"dtype {dtype} + width {width_32b} + offset {local_offset_32b} is not supported"
+        )
+
+    g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)])
+    local_view = TileLayout(S[(128, TOTAL_LOCAL_WIDTH) : (1 @ axis_tid_in_wg, 1)])
+
+    # fmt: off
+    @T.prim_func
+    def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128, WIDTH), dtype)
+        B = T.match_buffer(B_ptr, (128, WIDTH), dtype)
+
+        A_flat = A.view(-1)
+        B_flat = B.view(-1)
+
+        T.device_entry()
+        warp_id = T.warp_id([(128) // 32])
+        T.cta_id([2])
+        wg_id = T.warpgroup_id([1])
+        T.warp_id_in_wg([4])
+        T.lane_id([32])
+        tid_in_wg = T.thread_id([128])
+
+        tmem_addr = T.alloc_shared([1], "uint32")
+
+        if wg_id == 0:
+            if warp_id == 0:
+                T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1)  # noqa: E501
+
+            T.tvm_storage_sync("shared")
+
+            tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0],
+                                 layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)]))
+
+            A_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype)
+            B_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype)
+            A_local = A_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view)
+            B_local = B_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view)
+            for i in range(WIDTH // VEC_LEN):
+                g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
+                Tx.copy(A_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN])  # noqa: E501
+            for i in range(TOTAL_LOCAL_WIDTH):
+                B_reg[i] = T.cast(0, dtype)
+            T.cuda.cta_sync()
+
+                    # A_local[sliced] -> tmem (use sliced region)
+            Tx.wg.copy_async(tmem[:, 0:WIDTH], A_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH])
+            T.ptx.tcgen05.wait.st()
+            T.cuda.cta_sync()
+
+                    # tmem -> B_local[sliced] (use sliced region)
+            Tx.wg.copy_async(B_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH], tmem[:, 0:WIDTH])
+            T.ptx.tcgen05.wait.ld()
+            T.cuda.cta_sync()
+            for i in range(WIDTH // VEC_LEN):
+                g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"])
+                Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN])  # noqa: E501
+
+            if warp_id == 0:
+                T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1)
+                T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1)  # noqa: E501
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": copy_sync})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH))
+        B_np = np.zeros((128, WIDTH), dtype=dtype)
+
+        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)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py

@@ -0,0 +1,870 @@
+# 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.
+import re
+
+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, wg_local_layout
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        ######### basic test #########
+        (
+            (32, 32),  # g_shape
+            (0, 0),  # st_a
+            (0, 0),  # st_b
+            (0, 0),  # st_res
+            (32, 32),  # extent_a
+            (32, 32),  # extent_b
+            (32, 32),  # extent_res
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+        ######### offset test #########
+        (
+            (32, 8, 12),  # g_shape
+            (10, 0, 3),  # st_a
+            (14, 1, 4),  # st_b
+            (20, 0, 2),  # st_res
+            (5, 6, 7),  # extent_a
+            (5, 6, 7),  # extent_b
+            (5, 6, 7),  # extent_res
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+        ######### broadcast test #########
+        (
+            (32, 8, 12),  # g_shape
+            (10, 0, 3),  # st_a
+            (14, 1, 4),  # st_b
+            (20, 0, 2),  # st_res
+            (5, 6, 7),  # extent_a
+            (1, 6, 1),  # extent_b
+            (5, 6, 7),  # extent_res
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
+@pytest.mark.parametrize("operands_type", ["region_region", "region_const", "const_region"])
+@pytest.mark.parametrize("dtype", ["float16"])
+def test_binary_op_shared(input, op_type, operands_type, dtype):
+    # skip test
+    if op_type in ["sub", "fdiv"] and operands_type == "const_region":
+        return
+
+    g_shape, st_a, st_b, st_res, ext_a, ext_b, ext_res, thread_cnt, device_id = input
+    g_layout = s_layout = TileLayout(S[g_shape])
+
+    copy_slice = list(slice(None) for i in range(len(g_shape)))
+    map_slice_a = list(slice(st_a[i], st_a[i] + ext_a[i]) for i in range(len(g_shape)))
+    map_slice_b = list(slice(st_b[i], st_b[i] + ext_b[i]) for i in range(len(g_shape)))
+    map_slice_res = list(slice(st_res[i], st_res[i] + ext_res[i]) for i in range(len(g_shape)))
+
+    const = T.float16(3.0) if dtype == "float16" else T.float32(3.0)
+
+    # fmt: off
+    @T.prim_func
+    def binary_op_region_region(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=g_layout)
+        B = T.match_buffer(B_ptr, g_shape, dtype, layout=g_layout)
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([thread_cnt])
+        A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
+        B_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
+
+        Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
+        Tx.cta.copy(B_smem[tuple(copy_slice)], B[tuple(copy_slice)])
+        T.cuda.cta_sync()
+        if op_type == "add":
+            Tx.cta.add(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        elif op_type == "sub":
+            Tx.cta.sub(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        elif op_type == "mul":
+            Tx.cta.mul(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        elif op_type == "fdiv":
+            Tx.cta.fdiv(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], B_smem[tuple(map_slice_b)])  # noqa: E501
+        T.cuda.cta_sync()
+        Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
+
+    @T.prim_func
+    def binary_op_const_region_or_region_const(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=g_layout)
+        _B = T.match_buffer(B_ptr, g_shape, dtype, layout=g_layout)
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([thread_cnt])
+        A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=s_layout)
+
+        Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
+        T.cuda.cta_sync()
+        if op_type == "add":
+            if operands_type == "const_region":
+                Tx.cta.add(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.add(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        elif op_type == "sub":
+            if operands_type == "const_region":
+                Tx.cta.sub(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.sub(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        elif op_type == "mul":
+            if operands_type == "const_region":
+                Tx.cta.mul(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.mul(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        elif op_type == "fdiv":
+            if operands_type == "const_region":
+                Tx.cta.fdiv(A_smem[tuple(map_slice_res)], const, A_smem[tuple(map_slice_a)])
+            elif operands_type == "region_const":
+                Tx.cta.fdiv(A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)], const)
+        T.cuda.cta_sync()
+        Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
+        # fmt: on
+
+    def get_prim_func(operands_type):
+        if operands_type == "region_region":
+            return binary_op_region_region
+        elif operands_type in ["const_region", "region_const"]:
+            return binary_op_const_region_or_region_const
+        raise ValueError(f"operands_type={operands_type} is not supported")
+
+    def get_ref(A_np, B_np):
+        A_ref = A_np.copy()
+        if op_type == "add":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] + B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] + 3.0
+        elif op_type == "sub":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] - B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] - 3.0
+        elif op_type == "mul":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] * B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] * 3.0
+        elif op_type == "fdiv":
+            if operands_type == "region_region":
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] / B_np[tuple(map_slice_b)]
+            elif operands_type in ["const_region", "region_const"]:
+                A_ref[tuple(map_slice_res)] = A_np[tuple(map_slice_a)] / 3.0
+
+        return A_ref
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*g_shape).astype(dtype)
+        B_np = np.random.rand(*g_shape).astype(dtype)
+        mod = tvm.IRModule({"main": get_prim_func(operands_type)})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
+        A_ref = get_ref(A_np, B_np)
+        atol = 1e-3
+
+        def run_and_check():
+            dev = tvm.cuda(device_id)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=atol)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("op_type", ["sub", "fdiv"])
+def test_binary_non_commutative_const_lhs_rejected(op_type):
+    dtype = "float16"
+    shape = (16, 16)
+    layout = TileLayout(S[shape])
+    const = T.float16(3.0)
+
+    with pytest.raises(Exception):
+
+        @T.prim_func
+        def bad_kernel() -> None:
+            T.device_entry()
+            _bx = T.cta_id([1])
+            _tid = T.thread_id([64])
+            A_smem = T.alloc_buffer(shape, dtype, scope="shared", layout=layout)
+            if op_type == "sub":
+                Tx.cta.sub(A_smem, const, A_smem)
+            elif op_type == "fdiv":
+                Tx.cta.fdiv(A_smem, const, A_smem)
+
+        target = tvm.target.Target("cuda")
+        with target:
+            mod = tvm.IRModule({"main": bad_kernel})
+            tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("exec_scope", ["warp", "warpgroup"])
+@pytest.mark.parametrize("op_type", ["add", "mul"])
+def test_binary_op_shared_subcta_scope(exec_scope, op_type):
+    """Test binary ops in warp/warpgroup scope with shared memory."""
+    dtype = "float16"
+    n_warps = 4 if exec_scope == "warpgroup" else 1
+    g_shape = (n_warps * 32, 8)
+    tx_op = {
+        ("warp", "add"): Tx.warp.add,
+        ("warp", "mul"): Tx.warp.mul,
+        ("warpgroup", "add"): Tx.wg.add,
+        ("warpgroup", "mul"): Tx.wg.mul,
+    }[(exec_scope, op_type)]
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, g_shape, dtype, layout=TileLayout(S[g_shape]))
+        B = T.match_buffer(B_ptr, g_shape, dtype, layout=TileLayout(S[g_shape]))
+        T.device_entry()
+        warp_id = T.warp_id([(256) // 32])
+        wg_id = T.warpgroup_id([(256) // 128])
+        _bx = T.cta_id([1])
+        _tid = T.thread_id([256])
+        A_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=TileLayout(S[g_shape]))
+        B_smem = T.alloc_buffer(g_shape, dtype, scope="shared", layout=TileLayout(S[g_shape]))
+        Tx.cta.copy(A_smem, A)
+        Tx.cta.copy(B_smem, B)
+        T.cuda.cta_sync()
+        if exec_scope == "warp":
+            if warp_id == 5:
+                tx_op(A_smem, A_smem, B_smem)
+        elif exec_scope == "warpgroup":
+            if wg_id == 1:
+                tx_op(A_smem, A_smem, B_smem)
+        T.cuda.cta_sync()
+        Tx.cta.copy(A, A_smem)
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*g_shape).astype(dtype)
+        B_np = np.random.rand(*g_shape).astype(dtype)
+        mod = tvm.IRModule({"main": kernel})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        np_op = {"add": np.add, "mul": np.multiply}[op_type]
+        A_ref = np_op(A_np, B_np).astype(dtype)
+
+        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)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("exec_scope", ["cta", "warpgroup", "warp"])
+@pytest.mark.parametrize("rhs_kind", ["region", "broadcast", "const"])
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
+def test_binary_op_local_subcta_trivial(exec_scope, rhs_kind, op_type):
+    dtype = "float16"
+    m, n = 4, 8
+    n_threads = 256 if exec_scope == "cta" else (128 if exec_scope == "warpgroup" else 32)
+    # in this test, use warp3/warpgroup1 to test
+    thr_str = 0 if exec_scope == "cta" else (128 if exec_scope == "warpgroup" else 32 * 3)
+    a_shape = (n_threads, m, n)
+    b_shape = (n_threads, m, n if rhs_kind == "region" else 1)
+    c_shape = a_shape
+    const = T.float16(1.25)
+    tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv": Tx.fdiv}[op_type]
+    tid_in_scope_fn = {"cta": T.thread_id, "warpgroup": T.thread_id_in_wg, "warp": T.lane_id}[
+        exec_scope
+    ]
+
+    @T.prim_func
+    def kernel(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+        C = T.match_buffer(C_ptr, c_shape, dtype, layout=TileLayout(S[c_shape]))
+
+        T.device_entry()
+        wg_id = T.warpgroup_id([(256) // 128])
+        warp_id = T.warp_id([(256) // 32])
+        _bx = T.cta_id([1])
+        _tid = T.thread_id([256])
+        tid_in_scope = tid_in_scope_fn([n_threads])
+        b_n = T.meta_var(n if rhs_kind == "region" else 1)
+        A_local = T.alloc_buffer((m, n), dtype, scope="local", layout=TileLayout(S[(m, n)]))
+        C_local = T.alloc_buffer((m, n), dtype, scope="local", layout=TileLayout(S[(m, n)]))
+        B_local = T.alloc_buffer((m, b_n), dtype, scope="local", layout=TileLayout(S[(m, b_n)]))
+
+        if thr_str <= _tid and _tid < thr_str + n_threads:
+            for i in T.serial(m):
+                for j in T.serial(n):
+                    A_local[i, j] = A[tid_in_scope, i, j]
+            if rhs_kind != "const":
+                for i in T.serial(m):
+                    for j in T.serial(b_n):
+                        B_local[i, j] = B[tid_in_scope, i, j]
+
+        if exec_scope == "cta":
+            if rhs_kind == "const":
+                tx_op(C_local, A_local, const)
+            else:
+                tx_op(C_local, A_local, B_local)
+        elif exec_scope == "warpgroup":
+            if wg_id == 1:
+                if rhs_kind == "const":
+                    tx_op(C_local, A_local, const)
+                else:
+                    tx_op(C_local, A_local, B_local)
+        else:
+            if warp_id == 3:
+                if rhs_kind == "const":
+                    tx_op(C_local, A_local, const)
+                else:
+                    tx_op(C_local, A_local, B_local)
+                # T.cuda.cta_sync()
+
+        if thr_str <= _tid and _tid < thr_str + n_threads:
+            for i in T.serial(m):
+                for j in T.serial(n):
+                    C[tid_in_scope, i, j] = C_local[i, j]
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*a_shape).astype(dtype)
+        B_np = np.random.rand(*b_shape).astype(dtype)
+        C_np = np.zeros(c_shape, dtype=dtype)
+        mod = tvm.IRModule({"main": kernel})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
+        np_op = {"add": np.add, "sub": np.subtract, "mul": np.multiply, "fdiv": np.divide}[op_type]
+        if rhs_kind == "region":
+            C_ref = np_op(A_np, B_np)
+        elif rhs_kind == "broadcast":
+            C_ref = np_op(A_np, np.repeat(B_np, n, axis=2))
+        else:
+            C_ref = np_op(A_np, const.value)
+        atol = 1e-2 if op_type == "fdiv" else 1e-3
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            C = tvm.runtime.tensor(C_np, dev)
+            mod(A, B, C)
+            tvm.testing.assert_allclose(C_ref, C.numpy(), atol=atol)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize(
+    "input",
+    [
+        ######### basic test #########
+        (
+            (64, 32),  # a_shape
+            (64, 32),  # b_shape
+            (64, 32),  # res_shape
+            64,  # thread_cnt
+            0,  # device_id
+        ),
+        ######### broadcast test #########
+        (
+            (32, 5, 4),  # a_shape
+            (32, 1, 4),  # b_shape
+            (32, 5, 4),  # res_shape
+            32,  # thread_cnt (≥ warp size so sctx.intra at cta scope models cleanly)
+            0,  # device_id
+        ),
+    ],
+)
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("storage_scope", ["shared", "local"])
+@pytest.mark.parametrize("exec_scope", ["cta", "thread"])
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
+@pytest.mark.parametrize("dtype", ["float16"])
+def test_binary_op_vectorized(input, storage_scope, exec_scope, op_type, dtype):
+    a_shape, b_shape, res_shape, thread_cnt, device_id = input
+    tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv": Tx.fdiv}[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def test_binary_cta(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([thread_cnt])
+        if storage_scope == "shared":
+            A_smem = T.alloc_buffer(
+                a_shape, dtype, scope="shared", layout=TileLayout(S[a_shape])
+            )
+            B_smem = T.alloc_buffer(
+                b_shape, dtype, scope="shared", layout=TileLayout(S[b_shape])
+            )
+            Tx.cta.copy(A_smem, A)
+            Tx.cta.copy(B_smem, B)
+            T.cuda.cta_sync()
+            tx_op(A_smem, A_smem, B_smem)
+            T.cuda.cta_sync()
+            Tx.cta.copy(A, A_smem)
+        if storage_scope == "local":
+            A_local = T.alloc_buffer(
+                a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
+            )
+            B_local = T.alloc_buffer(
+                b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
+            )
+            Tx.copy(A_local, A[tx])
+            Tx.copy(B_local, B[tx])
+            tx_op(A_local, A_local, B_local)
+            Tx.copy(A[tx], A_local)
+
+    @T.prim_func
+    def test_binary_thread(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([thread_cnt])
+        if storage_scope == "shared":
+            A_smem = T.alloc_buffer(
+                a_shape, dtype, scope="shared", layout=TileLayout(S[a_shape])
+            )
+            B_smem = T.alloc_buffer(
+                b_shape, dtype, scope="shared", layout=TileLayout(S[b_shape])
+            )
+            Tx.copy(A_smem, A)
+            Tx.copy(B_smem, B)
+            T.cuda.cta_sync()
+            tx_op(A_smem, A_smem, B_smem)
+            T.cuda.cta_sync()
+            Tx.copy(A, A_smem)
+        elif storage_scope == "local":
+            A_local = T.alloc_buffer(
+                a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
+            )
+            B_local = T.alloc_buffer(
+                b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
+            )
+            Tx.copy(A_local, A[tx])
+            Tx.copy(B_local, B[tx])
+            tx_op(A_local, A_local, B_local)
+            Tx.copy(A[tx], A_local)
+        # fmt: on
+
+    def get_prim_func():
+        if exec_scope == "cta":
+            return test_binary_cta
+        elif exec_scope == "thread":
+            return test_binary_thread
+        else:
+            raise ValueError(f"exec_scope={exec_scope} is not supported")
+
+    target = tvm.target.Target("cuda")
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*a_shape).astype(dtype)
+        B_np = np.random.rand(*b_shape).astype(dtype)
+        mod = tvm.IRModule({"main": get_prim_func()})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        print(f"compiled source code: {mod.mod.imports[0].inspect_source()}")
+        np_op = {"add": np.add, "sub": np.subtract, "mul": np.multiply, "fdiv": np.divide}[op_type]
+        A_ref = np_op(A_np, B_np)
+        atol = 1e-2 if op_type == "fdiv" else 1e-3
+
+        def run_and_check():
+            dev = tvm.cuda(device_id)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            mod(A, B)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=atol)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul"])
+def test_binary_op_packed_f32x2_auto_dispatch(op_type):
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        pytest.skip(f"unknown target arch: {arch}")
+    sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    if not sm_digits:
+        pytest.skip(f"cannot parse target arch: {arch}")
+    sm_version = int(sm_digits)
+    if sm_version < 100:
+        pytest.skip(f"packed_f32x2 auto-dispatch requires sm_100+, got {arch}")
+
+    a_shape, b_shape = (64, 32), (64, 32)
+    dtype = "float32"
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, a_shape, dtype, layout=TileLayout(S[a_shape]))
+        B = T.match_buffer(B_ptr, b_shape, dtype, layout=TileLayout(S[b_shape]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        A_local = T.alloc_buffer(
+            a_shape[1:], dtype, scope="local", layout=TileLayout(S[a_shape[1:]])
+        )
+        B_local = T.alloc_buffer(
+            b_shape[1:], dtype, scope="local", layout=TileLayout(S[b_shape[1:]])
+        )
+        Tx.copy(A_local, A[tx])
+        Tx.copy(B_local, B[tx])
+        if op_type == "add":
+            Tx.add(A_local, A_local, B_local)
+        elif op_type == "sub":
+            Tx.sub(A_local, A_local, B_local)
+        elif op_type == "mul":
+            Tx.mul(A_local, A_local, B_local)
+        Tx.copy(A[tx], A_local)
+
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(*a_shape).astype(dtype)
+        B_np = np.random.rand(*b_shape).astype(dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+        ptx_pat = {
+            "add": r"add\.[a-z]+\.ftz\.f32x2",
+            "sub": r"sub\.[a-z]+\.ftz\.f32x2",
+            "mul": r"mul\.[a-z]+\.ftz\.f32x2",
+        }[op_type]
+        builtin_pat = {
+            "add": r"tvm_builtin_ptx_add_packed_",
+            "sub": r"tvm_builtin_ptx_sub_packed_",
+            "mul": r"tvm_builtin_ptx_mul_packed_",
+        }[op_type]
+        assert re.search(ptx_pat, src) or re.search(builtin_pat, src), src
+        if op_type == "add":
+            A_ref = A_np + B_np
+        elif op_type == "sub":
+            A_ref = A_np - B_np
+        elif op_type == "mul":
+            A_ref = A_np * B_np
+
+        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)
+            tvm.testing.assert_allclose(A_ref, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op_name", ["add", "sub", "mul"])
+def test_binary_op_warpgroup_wg_local_layout(op_name):
+    dtype = "float32"
+    rows, cols = 128, 16
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        lhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        rhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        out = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        lhs_row = lhs.local(cols)
+        rhs_row = rhs.local(cols)
+        out_row = out.local(cols)
+        for i in T.serial(cols):
+            lhs_row[i] = A[tid, i]
+            rhs_row[i] = B[tid, i]
+            out_row[i] = T.float32(0)
+        if op_name == "add":
+            Tx.wg.add(out, lhs, rhs)
+        elif op_name == "sub":
+            Tx.wg.sub(out, lhs, rhs)
+        elif op_name == "mul":
+            Tx.wg.mul(out, lhs, rhs)
+        out_row_1 = out.local(cols)
+        for i in T.serial(cols):
+            C[tid, i] = out_row_1[i]
+
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(rows, cols).astype(dtype)
+        B_np = np.random.rand(rows, cols).astype(dtype)
+        C_np = np.zeros((rows, cols), dtype=dtype)
+
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        if op_name == "add":
+            C_ref = A_np + B_np
+        elif op_name == "sub":
+            C_ref = A_np - B_np
+        else:
+            C_ref = A_np * B_np
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            B = tvm.runtime.tensor(B_np, dev)
+            C = tvm.runtime.tensor(C_np, dev)
+            mod(A, B, C)
+            tvm.testing.assert_allclose(C_ref, C.numpy(), atol=1e-5)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize("op_name,ptx_op", [("add", "add"), ("sub", "sub"), ("mul", "mul")])
+def test_binary_op_warpgroup_wg_local_emits_packed_f32x2(op_name, ptx_op):
+    """Warpgroup-scope binary on a wg-local fp32 view must lower to packed
+    f32x2 PTX on SM100+, mirroring the thread-scope packed dispatch.
+
+    Regression test for the fa4 perf path: rescale-style ``T.{add,sub,mul}``
+    calls in warpgroup scope used to fall through to scalar codegen because
+    ``_emit_binary_local_view`` only emitted ``op_func(...)`` per element.
+    """
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        pytest.skip(f"unknown target arch: {arch}")
+    sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    if not sm_digits or int(sm_digits) < 100:
+        pytest.skip(f"packed_f32x2 wg-local path requires sm_100+, got {arch}")
+
+    dtype = "float32"
+    rows, cols = 128, 16
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        lhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        rhs = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        out = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        lhs_row = lhs.local(cols)
+        rhs_row = rhs.local(cols)
+        out_row = out.local(cols)
+        for i in T.serial(cols):
+            lhs_row[i] = A[tid, i]
+            rhs_row[i] = B[tid, i]
+            out_row[i] = T.float32(0)
+        if op_name == "add":
+            Tx.wg.add(out, lhs, rhs)
+        elif op_name == "sub":
+            Tx.wg.sub(out, lhs, rhs)
+        else:
+            Tx.wg.mul(out, lhs, rhs)
+        out_row_1 = out.local(cols)
+        for i in T.serial(cols):
+            C[tid, i] = out_row_1[i]
+
+    with target:
+        mod = tvm.IRModule({"main": test_func})
+        ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = ex.mod.imports[0].inspect_source()
+
+    # Codegen must use the packed f32x2 path, not scalar fallback.
+    assert re.search(rf"{ptx_op}\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        rf"tvm_builtin_ptx_{ptx_op}_packed_[a-z]+_f32x2", src
+    ), f"expected packed f32x2 PTX for op={op_name}, source preview:\n{src[:2000]}"
+
+
+def test_fma_warpgroup_wg_local_emits_packed_f32x2():
+    """Same regression coverage as the binary case but for ``T.fma``."""
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        pytest.skip(f"unknown target arch: {arch}")
+    sm_digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    if not sm_digits or int(sm_digits) < 100:
+        pytest.skip(f"packed_f32x2 wg-local path requires sm_100+, got {arch}")
+
+    dtype = "float32"
+    rows, cols = 128, 16
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, C_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        C = T.match_buffer(C_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        buf = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        buf_row = buf.local(cols)
+        for i in T.serial(cols):
+            buf_row[i] = A[tid, i]
+        Tx.wg.fma(buf, buf, T.float32(2.0), T.float32(0.5))
+        buf_row_1 = buf.local(cols)
+        for i in T.serial(cols):
+            C[tid, i] = buf_row_1[i]
+
+    with target:
+        mod = tvm.IRModule({"main": test_func})
+        ex = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = ex.mod.imports[0].inspect_source()
+
+    assert re.search(r"fma\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        r"tvm_builtin_ptx_fma_packed_[a-z]+_f32x2", src
+    ), f"expected packed f32x2 fma PTX, source preview:\n{src[:2000]}"
+
+
+# -----------------------------------------------------------------------------
+# Dispatch codegen checks (no GPU runtime — explicit target arch).
+# These complement the existing `*_warpgroup_wg_local_layout` / `*_auto_dispatch`
+# variants by forcing the arch in the Target dict, so the codegen path runs
+# even on hosts where ``Target("cuda")`` cannot detect the GPU.
+# -----------------------------------------------------------------------------
+def test_binary_add_f32_sm100_packed_f32x2_dispatch():
+    """add f32 + all-local → reg.py + add_f32x2 packed (no T.vectorized)."""
+    shape = (64, 32)
+    lay = TileLayout(S[shape])
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float32", layout=lay)
+        B = T.match_buffer(B_ptr, shape, "float32", layout=lay)
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        ra = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rb = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        Tx.copy(ra, A[tx])
+        Tx.copy(rb, B[tx])
+        Tx.add(ra, ra, rb)
+        Tx.copy(A[tx], ra)
+
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+    with target:
+        mod = tvm.IRModule({"main": k})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+    assert re.search(r"add\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        r"tvm_builtin_ptx_add_packed_", src
+    ), f"expected packed add_f32x2; got:\n{src[:2000]}"
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_binary_maximum_reg():
+    N = 128
+
+    @T.prim_func
+    def relu_max(a_ptr: T.handle, b_ptr: T.handle, c_ptr: T.handle) -> None:
+        A = T.match_buffer(a_ptr, (N,), "float32")
+        B = T.match_buffer(b_ptr, (N,), "float32")
+        C = T.match_buffer(c_ptr, (N,), "float32")
+        T.device_entry()
+        T.warp_id([4])
+        T.cta_id([1])
+        T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([N])
+        a_reg = T.alloc_local((1,), "float32")
+        b_reg = T.alloc_local((1,), "float32")
+        Tx.copy(a_reg[:], A[tid : tid + 1])
+        Tx.copy(b_reg[:], B[tid : tid + 1])
+        Tx.maximum(a_reg[:], a_reg[:], b_reg[:])
+        Tx.copy(C[tid : tid + 1], a_reg[:])
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(tvm.IRModule({"main": relu_max}), target=target, tir_pipeline="tirx")
+    np.random.seed(0)
+    a = np.random.randn(N).astype("float32")
+    b = np.random.randn(N).astype("float32")
+    c = np.zeros(N, dtype="float32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        a_t, b_t, c_t = (tvm.runtime.tensor(x, dev) for x in (a, b, c))
+        mod["main"](a_t, b_t, c_t)
+        np.testing.assert_allclose(c_t.numpy(), np.maximum(a, b), atol=0, rtol=0)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+def test_binary_add_f16_scalar_fallback_dispatch():
+    """add f16 has no packed VecImpl → reg.py scalar fallback (T.vectorized)."""
+    shape = (64, 32)
+    lay = TileLayout(S[shape])
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float16", layout=lay)
+        B = T.match_buffer(B_ptr, shape, "float16", layout=lay)
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        ra = T.alloc_buffer(shape[1:], "float16", scope="local", layout=TileLayout(S[shape[1:]]))
+        rb = T.alloc_buffer(shape[1:], "float16", scope="local", layout=TileLayout(S[shape[1:]]))
+        Tx.copy(ra, A[tx])
+        Tx.copy(rb, B[tx])
+        Tx.add(ra, ra, rb)
+        Tx.copy(A[tx], ra)
+
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+    with target:
+        mod = tvm.IRModule({"main": k})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+    assert "half" in src or "__half" in src, f"expected scalar half add; got:\n{src[:2000]}"
+    assert not re.search(r"add\.[a-z]+\.ftz\.f(32|16)x2", src), (
+        f"unexpected packed f32x2/f16x2 add in scalar-fallback path; got:\n{src[:2000]}"
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_fma.py

@@ -0,0 +1,403 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for FMA op dispatch, layout=None local dispatch, scalar broadcast,
+and rounding mode support."""
+
+import re
+
+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, wg_local_layout
+
+
+def _get_sm_version():
+    target = tvm.target.Target("cuda")
+    arch = target.arch if hasattr(target, "arch") else ""
+    if not arch.startswith("sm_"):
+        return 0
+    digits = "".join(ch for ch in arch.split("_", 1)[1] if ch.isdigit())
+    return int(digits) if digits else 0
+
+
+# ---------------------------------------------------------------------------
+# FMA op: scalar scale + scalar bias
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_scalar_scalar():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
+
+    N = 128
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    scale_val = 0.5
+    bias_val = -1.0
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([N])
+        buf = T.alloc_buffer((1,), dtype, scope="local", layout=TileLayout(S[1]))
+        Tx.copy(buf, A[tx : tx + 1])
+        Tx.fma(buf, buf, T.float32(scale_val), T.float32(bias_val))
+        Tx.copy(A[tx : tx + 1], buf)
+
+    with target:
+        A_np = np.random.rand(N).astype(dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A = tvm.runtime.tensor(A_np, dev)
+            mod(A)
+            expected = A_np * scale_val + bias_val
+            tvm.testing.assert_allclose(expected, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# FMA op: buffer scale + scalar bias (Horner pattern)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_buffer_scale_scalar_bias():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
+
+    N = 2
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    coeff = 0.695
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        B = T.match_buffer(B_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        acc = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        frac = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        Tx.copy(acc, A[0:N])
+        Tx.copy(frac, B[0:N])
+        Tx.fma(acc, acc, frac, T.float32(coeff))
+        Tx.copy(A[0:N], acc)
+
+    with target:
+        A_np = np.random.rand(N).astype(dtype)
+        B_np = np.random.rand(N).astype(dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        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)
+            expected = A_np * B_np + coeff
+            tvm.testing.assert_allclose(expected, A.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Binary op with scalar broadcast (Expr scalar, e.g. BufferLoad)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_mul_scalar_broadcast():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed mul requires sm_100+, got sm_{sm}")
+
+    N = 16
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, S_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        Scale = T.match_buffer(S_ptr, (1,), dtype, layout=TileLayout(S[1]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        a_local = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        s_local = T.alloc_buffer((1,), dtype, scope="local", layout=TileLayout(S[1]))
+        Tx.copy(a_local, A[0:N])
+        Tx.copy(s_local, Scale[0:1])
+        Tx.mul(a_local, a_local, s_local[0])
+        Tx.copy(A[0:N], a_local)
+
+    with target:
+        A_np = np.random.rand(N).astype(dtype)
+        S_np = np.array([2.5], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            S_dev = tvm.runtime.tensor(S_np, dev)
+            mod(A_dev, S_dev)
+            expected = A_np * S_np[0]
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Binary add with rounding mode
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_add_rounding_mode():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed add with rounding requires sm_100+, got sm_{sm}")
+
+    N = 2
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    round_const = float(2**23 + 2**22)
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        Tx.copy(buf, A[0:N])
+        Tx.add(buf, buf, T.float32(round_const), rounding_mode="rm")
+        Tx.copy(A[0:N], buf)
+
+    with target:
+        A_np = np.array([1.3, 2.7], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        # Check that the PTX uses the rounding mode
+        src = mod.mod.imports[0].inspect_source()
+        assert re.search(r"add\.rm\.ftz\.f32x2", src) or re.search(
+            r"tvm_builtin_ptx_add_packed_", src
+        ), f"Expected packed add with rm rounding in PTX:\n{src}"
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            mod(A_dev)
+            expected = A_np + round_const
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1.0)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# FMA op: layout=None local buffer (no TileLayout)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_no_layout():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed fma requires sm_100+, got sm_{sm}")
+
+    N = 4
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    scale_val = 2.0
+    bias_val = 1.0
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        buf = T.alloc_local([N], dtype)
+        for i in T.serial(N):
+            buf[i] = A[i]
+        Tx.fma(buf[0:N], buf[0:N], T.float32(scale_val), T.float32(bias_val))
+        for i in T.serial(N):
+            A[i] = buf[i]
+
+    with target:
+        A_np = np.array([1.0, 2.0, 3.0, 4.0], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            mod(A_dev)
+            expected = A_np * scale_val + bias_val
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-3)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ---------------------------------------------------------------------------
+# Binary sub with rounding mode (buffer-buffer)
+# ---------------------------------------------------------------------------
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_sub_buffer_buffer_rounding():
+    sm = _get_sm_version()
+    if sm < 100:
+        pytest.skip(f"packed sub with rounding requires sm_100+, got sm_{sm}")
+
+    N = 2
+    dtype = "float32"
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        B = T.match_buffer(B_ptr, (N,), dtype, layout=TileLayout(S[N]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        _tx = T.thread_id([1])
+        a_buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        b_buf = T.alloc_buffer((N,), dtype, scope="local", layout=TileLayout(S[N]))
+        Tx.copy(a_buf, A[0:N])
+        Tx.copy(b_buf, B[0:N])
+        Tx.sub(a_buf, a_buf, b_buf, rounding_mode="rn")
+        Tx.copy(A[0:N], a_buf)
+
+    with target:
+        A_np = np.array([3.14, 2.71], dtype=dtype)
+        B_np = np.array([1.41, 0.57], dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+        assert re.search(r"sub\.rn\.ftz\.f32x2", src) or re.search(
+            r"tvm_builtin_ptx_sub_packed_", src
+        ), f"Expected packed sub with rn rounding in PTX:\n{src}"
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            B_dev = tvm.runtime.tensor(B_np, dev)
+            mod(A_dev, B_dev)
+            expected = A_np - B_np
+            tvm.testing.assert_allclose(expected, A_dev.numpy(), atol=1e-6)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_fma_warpgroup_wg_local_layout():
+    rows, cols = 128, 8
+    dtype = "float32"
+    scale_val = 1.5
+    bias_val = -0.25
+    target = tvm.target.Target("cuda")
+
+    @T.prim_func
+    def test_func(A_ptr: T.handle, B_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        B = T.match_buffer(B_ptr, (rows, cols), dtype, layout=TileLayout(S[(rows, cols)]))
+        T.device_entry()
+        _bx = T.cta_id([1])
+        wg_id = T.warpgroup_id([1])
+        tid = T.thread_id_in_wg([rows])
+
+        reg = T.alloc_buffer((rows, cols), dtype, scope="local", layout=wg_local_layout(cols))
+        reg_row = reg.local(cols)
+        for i in T.serial(cols):
+            reg_row[i] = A[tid, i]
+        Tx.wg.fma(reg, reg, T.float32(scale_val), T.float32(bias_val))
+        reg_row_1 = reg.local(cols)
+        for i in T.serial(cols):
+            B[tid, i] = reg_row_1[i]
+
+    with target:
+        np.random.seed(0)
+        A_np = np.random.rand(rows, cols).astype(dtype)
+        B_np = np.zeros((rows, cols), dtype=dtype)
+        mod = tvm.IRModule({"main": test_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+        def run_and_check():
+            dev = tvm.cuda(0)
+            A_dev = tvm.runtime.tensor(A_np, dev)
+            B_dev = tvm.runtime.tensor(B_np, dev)
+            mod(A_dev, B_dev)
+            expected = A_np * scale_val + bias_val
+            tvm.testing.assert_allclose(expected, B_dev.numpy(), atol=1e-5)
+
+        tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# -----------------------------------------------------------------------------
+# Dispatch codegen check (no GPU runtime — explicit target arch).
+# Complements ``test_fma_warpgroup_wg_local_emits_packed_f32x2`` (which uses
+# the host-detected ``Target("cuda")`` and skips when arch < sm_100).
+# -----------------------------------------------------------------------------
+def test_fma_f32_sm100_packed_f32x2_dispatch():
+    """fma f32 + all-local → reg.py + fma_f32x2 packed (no T.vectorized)."""
+    shape = (64, 32)
+    lay = TileLayout(S[shape])
+
+    @T.prim_func
+    def k(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, D_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, shape, "float32", layout=lay)
+        B = T.match_buffer(B_ptr, shape, "float32", layout=lay)
+        C = T.match_buffer(C_ptr, shape, "float32", layout=lay)
+        D = T.match_buffer(D_ptr, shape, "float32", layout=lay)
+        T.device_entry()
+        _bx = T.cta_id([1])
+        tx = T.thread_id([64])
+        ra = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rb = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rc = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        rd = T.alloc_buffer(shape[1:], "float32", scope="local", layout=TileLayout(S[shape[1:]]))
+        Tx.copy(ra, A[tx])
+        Tx.copy(rb, B[tx])
+        Tx.copy(rc, C[tx])
+        Tx.fma(rd, ra, rb, rc)
+        Tx.copy(D[tx], rd)
+
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+    with target:
+        mod = tvm.IRModule({"main": k})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+        src = mod.mod.imports[0].inspect_source()
+    assert re.search(r"fma\.[a-z]+\.ftz\.f32x2", src) or re.search(
+        r"tvm_builtin_ptx_fma_packed_", src
+    ), f"expected packed fma_f32x2; got:\n{src[:2000]}"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py

@@ -0,0 +1,690 @@
+# 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 CUDA synchronous ``gemm`` (mma.sync) tensor-core dispatch.
+
+The dispatch lowers ``tirx.tile.gemm`` over pure-register fragments to warp-level
+``mma.sync.aligned.m16n8k16/k8`` for bf16/f16 inputs with f32 accumulation.
+
+The fragment layouts below are the standard m16n8 register maps (PTX ISA
+§9.7.13; see ``tests/python/tirx-base/test_tir_ptx_mma.py``):
+
+    lane = 4*g + t        (g = lane >> 2 in [0, 8),  t = lane & 3 in [0, 4))
+    D/C[M, N]: M = g + 8*rM,  N = 2*t + rN,          c_id = 2*rM + rN
+    A[M, K]:   M = g + 8*rM,  K = 2*t + p + 8*kHi,   ma   = p + 2*rM + 4*kHi
+    B[K, N]:   N = g,         K = 2*t + p + 8*kHi,   mb   = p + 2*kHi
+
+Most assertions run the CPU-only ``LowerTIRx`` transform; the numerical check
+is guarded by ``requires_cuda`` since it needs a real device.
+"""
+
+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
+from tvm.tirx.operator.tile_primitive import list_registered_schedules
+
+# Single-tile m16n8k8 fragment layouts -- the smallest unit everything else is
+# built from. A is 16x8, B is 8x8 as [K, N], D/C is 16x8 (the accumulator does
+# not depend on K). Every other layout (the k16 single tile, all tilings, and
+# the transposed orientations) is derived from these three via ``tile_to`` /
+# ``group`` + ``permute_by_groups``.
+D_FRAG = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
+A_FRAG_K8 = TileLayout(S[(2, 8, 4, 2) : (2, 4 @ laneid, 1 @ laneid, 1)])
+B_FRAG_K8 = TileLayout(S[(4, 2, 8) : (1 @ laneid, 1, 4 @ laneid)])
+# m16n8k16 single tile = two k8 tiles stacked along K.
+A_FRAG = A_FRAG_K8.tile_to([16, 16], [16, 8])
+B_FRAG = B_FRAG_K8.tile_to([16, 8], [8, 8])
+
+
+def _transpose_frag(layout, shape):
+    """Swap the two logical axes of a 2D fragment layout.
+
+    The transposed input orientations (A as [K, M], B as [N, K]) hold the exact
+    same per-lane/per-register element distribution as the K-major fragments --
+    only the buffer's logical axes are swapped. So instead of writing them out
+    by hand, derive them: ``group`` the shard into the logical dims, then
+    ``permute_by_groups`` to exchange the two groups.
+    """
+    grouped, seps = layout.group(shape)
+    return grouped.permute_by_groups(seps, [1, 0])
+
+
+# Transposed input orientations of the same single tile: A as [K, M], B as
+# [N, K]. The dispatch swaps axes per the transpose flags; the .row.col mma is
+# unchanged.
+A_KM_FRAG = _transpose_frag(A_FRAG, [16, 16])
+B_NK_FRAG = _transpose_frag(B_FRAG, [16, 8])
+
+
+def _frag(Mt, Nt, Kt, kinst):
+    """Fragment layouts for an Mt x Nt x Kt tiling of m16n8k{8,16}.
+
+    Logical shapes: A = (16*Mt, kinst*Kt), B = (kinst*Kt, 8*Nt), D/C = (16*Mt, 8*Nt).
+    Each operand's tiled layout is the single-tile base ``tile_to`` the full
+    logical shape -- ``tile_to`` repeats the base's per-lane/per-register element
+    map over the tile grid, so a tiling is just a grid of the single-tile
+    fragments (the base is the single source of truth, k8 and k16 alike).
+    """
+    A_base = A_FRAG if kinst == 16 else A_FRAG_K8
+    B_base = B_FRAG if kinst == 16 else B_FRAG_K8
+    D = D_FRAG.tile_to([16 * Mt, 8 * Nt], [16, 8])
+    A = A_base.tile_to([16 * Mt, kinst * Kt], [16, kinst])
+    B = B_base.tile_to([kinst * Kt, 8 * Nt], [kinst, 8])
+    return D, A, B
+
+
+def _build_tiled(Mt, Nt, Kt, kinst, *, beta=0.0, dtype="float16", store=False):
+    """A single-warp kernel issuing one ``T.gemm`` over an Mt x Nt x Kt tiling.
+
+    With ``store=True`` the result is written back to a global buffer (a full
+    kernel for codegen); otherwise only the ``T.gemm`` is emitted (for
+    ``LowerTIRx`` dispatch checks).
+    """
+    Dl, Al, Bl = _frag(Mt, Nt, Kt, kinst)
+    M, N, K = 16 * Mt, 8 * Nt, kinst * Kt
+
+    if not store:
+
+        @T.prim_func
+        def gemm():
+            T.device_entry()
+            _cta = T.cta_id([1])
+            _warp = T.warp_id([1])
+            _lane = T.lane_id([32])
+            A = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
+            B = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
+            C = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
+            D = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
+            Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
+
+        return gemm
+
+    @T.prim_func
+    def gemm(D_ptr: T.handle):
+        D_g = T.match_buffer(D_ptr, (M, N), "float32")
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _warp = T.warp_id([1])
+        lane = T.lane_id([32])
+        A = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
+        B = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
+        C = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
+        D = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
+        Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
+        # Decode D's per-thread registers (c = ((mt*Nt + nt)*2 + rM)*2 + rN)
+        # back to logical (M, N) and store, exercising the whole tiling.
+        D_reg = D.local(Mt * Nt * 4)
+        for c in T.unroll(Mt * Nt * 4):
+            rN = c % 2
+            rM = (c // 2) % 2
+            nt = (c // 4) % Nt
+            mt = c // (4 * Nt)
+            D_g[mt * 16 + lane // 4 + rM * 8, nt * 8 + (lane % 4) * 2 + rN] = D_reg[c]
+
+    return gemm
+
+
+def _build_gemm(alpha=1.0, beta=0.0, dtype="bfloat16"):
+    """A single-warp kernel issuing one ``T.gemm`` over register fragments."""
+
+    @T.prim_func
+    def gemm_min():
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _tid = T.thread_id([32])
+        D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+        C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+        A = T.alloc_buffer((16, 16), dtype, scope="local", layout=A_FRAG)
+        B = T.alloc_buffer((16, 8), dtype, scope="local", layout=B_FRAG)
+        Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=alpha, beta=beta)
+
+    return gemm_min
+
+
+def _build_transpose(transpose_A, transpose_B, *, store=False):
+    """Single m16n8k16 tile with the requested A/B input orientations."""
+    Al = A_KM_FRAG if transpose_A else A_FRAG
+    Bl = B_NK_FRAG if transpose_B else B_FRAG
+    A_shape = (16, 16)  # [K, M] or [M, K] -- both 16x16 for one tile
+    B_shape = (8, 16) if transpose_B else (16, 8)
+
+    if not store:
+
+        @T.prim_func
+        def gemm():
+            T.device_entry()
+            _cta = T.cta_id([1])
+            _warp = T.warp_id([1])
+            _lane = T.lane_id([32])
+            A = T.alloc_buffer(A_shape, "float16", scope="local", layout=Al)
+            B = T.alloc_buffer(B_shape, "float16", scope="local", layout=Bl)
+            C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+            D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+            Tx.warp.gemm(
+                D,
+                A,
+                B,
+                C,
+                transpose_A=transpose_A,
+                transpose_B=transpose_B,
+                alpha=1.0,
+                beta=0.0,
+            )
+
+        return gemm
+
+    @T.prim_func
+    def gemm(D_ptr: T.handle):
+        D_g = T.match_buffer(D_ptr, (16, 8), "float32")
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _warp = T.warp_id([1])
+        lane = T.lane_id([32])
+        A = T.alloc_buffer(A_shape, "float16", scope="local", layout=Al)
+        B = T.alloc_buffer(B_shape, "float16", scope="local", layout=Bl)
+        C = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+        D = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+        Tx.warp.gemm(
+            D,
+            A,
+            B,
+            C,
+            transpose_A=transpose_A,
+            transpose_B=transpose_B,
+            alpha=1.0,
+            beta=0.0,
+        )
+        D_reg = D.local(4)
+        for c in T.unroll(4):
+            D_g[lane // 4 + (c // 2) * 8, (lane % 4) * 2 + c % 2] = D_reg[c]
+
+    return gemm
+
+
+def _build_dtypes(a_dtype, b_dtype, c_dtype, d_dtype):
+    """Single tile with explicit per-operand dtypes (for decline checks)."""
+
+    @T.prim_func
+    def gemm_min():
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _tid = T.thread_id([32])
+        D = T.alloc_buffer((16, 8), d_dtype, scope="local", layout=D_FRAG)
+        C = T.alloc_buffer((16, 8), c_dtype, scope="local", layout=D_FRAG)
+        A = T.alloc_buffer((16, 16), a_dtype, scope="local", layout=A_FRAG)
+        B = T.alloc_buffer((16, 8), b_dtype, scope="local", layout=B_FRAG)
+        Tx.warp.gemm(D, A, B, C, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0)
+
+    return gemm_min
+
+
+def _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype):
+    """End-to-end ``T.gemm`` over an Mt x Nt x Kt tiling, with the A/B inputs
+    loaded and the D output stored register-by-register.
+
+    Fragments are indexed through their per-register multi-dim ``.local()`` views
+    (the shard's non-lane dims, in shard order): A = [Mt, rM(2), Kt, kHi, kp],
+    B = [Kt, kHi, kp, Nt], D/C = [Mt, rM(2), Nt, rN(2)]. The lane owns g = lane>>2
+    and t = lane&3; within a tile M = mt*16 + rM*8 + g, N = nt*8 + t*2 + rN,
+    K = kt*kinst + kHi*8 + t*2 + kp.
+    """
+    Dl, Al, Bl = _frag(Mt, Nt, Kt, kinst)
+    M, N, K = 16 * Mt, 8 * Nt, kinst * Kt
+    KP = 2
+    kHi_n = kinst // (4 * KP)
+
+    @T.prim_func
+    def gemm(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, D_ptr: T.handle):
+        A_g = T.match_buffer(A_ptr, (M, K), dtype)
+        B_g = T.match_buffer(B_ptr, (K, N), dtype)
+        C_g = T.match_buffer(C_ptr, (M, N), "float32")
+        D_g = T.match_buffer(D_ptr, (M, N), "float32")
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _warp = T.warp_id([1])
+        lane = T.lane_id([32])
+        A_f = T.alloc_buffer((M, K), dtype, scope="local", layout=Al)
+        B_f = T.alloc_buffer((K, N), dtype, scope="local", layout=Bl)
+        C_f = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
+        D_f = T.alloc_buffer((M, N), "float32", scope="local", layout=Dl)
+        A_reg = A_f.local(Mt, 2, Kt, kHi_n, KP)
+        for mt, rM, kt, kHi, kp in T.grid(Mt, 2, Kt, kHi_n, KP):
+            A_reg[mt, rM, kt, kHi, kp] = A_g[
+                mt * 16 + lane // 4 + 8 * rM,
+                kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
+            ]
+        B_reg = B_f.local(Kt, kHi_n, KP, Nt)
+        for kt, kHi, kp, nt in T.grid(Kt, kHi_n, KP, Nt):
+            B_reg[kt, kHi, kp, nt] = B_g[
+                kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
+                nt * 8 + lane // 4,
+            ]
+        if beta == 1.0:
+            C_reg = C_f.local(Mt, 2, Nt, 2)
+            for mt, rM, nt, rN in T.grid(Mt, 2, Nt, 2):
+                C_reg[mt, rM, nt, rN] = C_g[
+                    mt * 16 + lane // 4 + 8 * rM, nt * 8 + 2 * (lane % 4) + rN
+                ]
+        Tx.warp.gemm(D_f, A_f, B_f, C_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=beta)
+        D_reg = D_f.local(Mt, 2, Nt, 2)
+        for mt, rM, nt, rN in T.grid(Mt, 2, Nt, 2):
+            D_g[mt * 16 + lane // 4 + 8 * rM, nt * 8 + 2 * (lane % 4) + rN] = D_reg[mt, rM, nt, rN]
+
+    return gemm, M, N, K
+
+
+def _build_transpose_numeric(transpose_A, transpose_B, dtype="float16"):
+    """End-to-end single-tile ``T.gemm`` for one A/B input orientation.
+
+    The transposed A fragment (``A_KM_FRAG``) carries its registers in the
+    [kHi, kp, rM] shard order (vs [rM, kHi, kp] for the K-major ``A_FRAG``); B's
+    register order ([kHi, kp]) is the same for both orientations. The buffer
+    index axes swap with the orientation, but each register still holds the same
+    logical (M, K) / (K, N) element.
+    """
+    Al = A_KM_FRAG if transpose_A else A_FRAG
+    Bl = B_NK_FRAG if transpose_B else B_FRAG
+    A_shape = (16, 16)
+    B_shape = (8, 16) if transpose_B else (16, 8)
+
+    @T.prim_func
+    def gemm(A_ptr: T.handle, B_ptr: T.handle, D_ptr: T.handle):
+        A_g = T.match_buffer(A_ptr, A_shape, dtype)
+        B_g = T.match_buffer(B_ptr, B_shape, dtype)
+        D_g = T.match_buffer(D_ptr, (16, 8), "float32")
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _warp = T.warp_id([1])
+        lane = T.lane_id([32])
+        A_f = T.alloc_buffer(A_shape, dtype, scope="local", layout=Al)
+        B_f = T.alloc_buffer(B_shape, dtype, scope="local", layout=Bl)
+        D_f = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+        A_reg = A_f.local(2, 2, 2)
+        if transpose_A:
+            # A_KM_FRAG register order is [kHi, kp, rM]; buffer is [K, M].
+            for kHi, kp, rM in T.grid(2, 2, 2):
+                A_reg[kHi, kp, rM] = A_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4 + 8 * rM]
+        else:
+            # A_FRAG register order is [rM, kHi, kp]; buffer is [M, K].
+            for rM, kHi, kp in T.grid(2, 2, 2):
+                A_reg[rM, kHi, kp] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
+        B_reg = B_f.local(2, 2)
+        if transpose_B:
+            # B_NK_FRAG buffer is [N, K].
+            for kHi, kp in T.grid(2, 2):
+                B_reg[kHi, kp] = B_g[lane // 4, 2 * (lane % 4) + kp + 8 * kHi]
+        else:
+            for kHi, kp in T.grid(2, 2):
+                B_reg[kHi, kp] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
+        Tx.warp.gemm(
+            D_f,
+            A_f,
+            B_f,
+            D_f,
+            transpose_A=transpose_A,
+            transpose_B=transpose_B,
+            alpha=1.0,
+            beta=0.0,
+        )
+        D_reg = D_f.local(2, 2)
+        for rM, rN in T.grid(2, 2):
+            D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[rM, rN]
+
+    return gemm
+
+
+def _lower(func):
+    with tvm.target.Target("cuda"):
+        return tvm.tirx.transform.LowerTIRx()(tvm.IRModule({"main": func}))
+
+
+def test_cuda_gemm_mma_variant_is_registered():
+    # Importing tvm.tirx registers all per-target schedule variants. The new
+    # synchronous CUDA mma path must show up for ("gemm", "cuda"). The registry
+    # keys ops by their full name (``op.name`` == "tirx.tile.gemm").
+    schedules = list_registered_schedules()
+    cuda_gemm = schedules.get("tirx.tile.gemm", {}).get("cuda", [])
+    assert "mma.m16n8k*" in cuda_gemm, (
+        f"mma.m16n8k* not registered; tirx.tile.gemm schedules = {schedules.get('tirx.tile.gemm')}"
+    )
+
+
+@pytest.mark.parametrize("dtype", ["bfloat16", "float16"])
+def test_cuda_gemm_mma_lowers_to_mma_sync(dtype):
+    """beta=0: the dispatch clears D, then issues a single accumulating mma with
+    the registers laid out in the fixed PTX fragment order."""
+    script = _lower(_build_gemm(alpha=1.0, beta=0.0, dtype=dtype))["main"].script()
+
+    assert "T.ptx.mma(" in script
+    assert "m16n8k16" in script
+    # beta == 0 clears the accumulator before the K loop.
+    assert "T.float32(0" in script
+    # D accumulator: c_id = 2*rM + rN -> regs 0..3.
+    for r in range(4):
+        assert f"d_local[{r}]" in script
+    # A multiplicand: b32 = rM + 2*kHi (kHi outer) -> ma in {0, 2, 4, 6}.
+    for r in (0, 2, 4, 6):
+        assert f"a_local[{r}]" in script
+    # B multiplicand: b32 = kHi -> mb in {0, 2}.
+    for r in (0, 2):
+        assert f"b_local[{r}]" in script
+
+
+def test_cuda_gemm_mma_accumulates_c_when_beta_one():
+    """beta=1: the accumulator is initialized by copying C instead of zeroing."""
+    script = _lower(_build_gemm(alpha=1.0, beta=1.0))["main"].script()
+
+    assert "T.ptx.mma(" in script
+    assert "m16n8k16" in script
+    # The init reads C into D; nothing is zeroed.
+    assert "c_local[" in script
+    assert "T.float32(0" not in script
+
+
+def test_cuda_gemm_mma_rejects_nonunit_alpha():
+    """alpha != 1 is unsupported (ptx mma has no scale); dispatch must fail."""
+    with pytest.raises(RuntimeError, match="dispatch failed"):
+        _lower(_build_gemm(alpha=2.0, beta=0.0))
+
+
+def test_cuda_gemm_mma_rejects_fractional_beta():
+    """beta must be 0 or 1 (mma only accumulates 1*C); other values must fail."""
+    with pytest.raises(RuntimeError, match="dispatch failed"):
+        _lower(_build_gemm(alpha=1.0, beta=0.5))
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
+def test_cuda_gemm_mma_numerical(dtype):
+    """End-to-end D = A @ B on a single m16n8k16 tile (one warp).
+
+    A is [M, K] = [16, 16], B is [K, N] = [16, 8], D is [M, N] = [16, 8].
+
+    The lane-distributed register fragments cannot be filled with a whole-tile
+    ``T.copy`` (the per-thread axis can't be matched coordinate-wise), so each
+    of a lane's registers is loaded/stored by decoding the m16n8k16 register map
+    with ``g = lane >> 2`` and ``t = lane & 3``. The per-register *slot* order
+    matches the dispatch's fragment register layout:
+
+        A reg slot = 4*rM + 2*kHi + kp  -> M = g + 8*rM, K = 2*t + kp + 8*kHi
+        B reg slot = 2*kHi + kp         -> K = 2*t + kp + 8*kHi, N = g
+        D reg slot = 2*rM + rN          -> M = g + 8*rM, N = 2*t + rN
+    """
+    if dtype == "bfloat16":
+        ml_dtypes = pytest.importorskip("ml_dtypes")
+        np_dtype = ml_dtypes.bfloat16
+    else:
+        np_dtype = np.float16
+
+    @T.prim_func
+    def gemm(A_ptr: T.handle, B_ptr: T.handle, D_ptr: T.handle):
+        A_g = T.match_buffer(A_ptr, (16, 16), dtype)
+        B_g = T.match_buffer(B_ptr, (16, 8), dtype)
+        D_g = T.match_buffer(D_ptr, (16, 8), "float32")
+        T.device_entry()
+        _cta = T.cta_id([1])
+        _warp = T.warp_id([1])
+        lane = T.lane_id([32])
+        A_f = T.alloc_buffer((16, 16), dtype, scope="local", layout=A_FRAG)
+        B_f = T.alloc_buffer((16, 8), dtype, scope="local", layout=B_FRAG)
+        D_f = T.alloc_buffer((16, 8), "float32", scope="local", layout=D_FRAG)
+        A_reg = A_f.local(8)
+        for s in T.unroll(8):
+            kp = s % 2
+            kHi = (s // 2) % 2
+            rM = s // 4
+            A_reg[s] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
+        B_reg = B_f.local(4)
+        for s in T.unroll(4):
+            kp = s % 2
+            kHi = s // 2
+            B_reg[s] = B_g[2 * (lane % 4) + kp + 8 * kHi, lane // 4]
+        Tx.warp.gemm(D_f, A_f, B_f, D_f, transpose_A=False, transpose_B=False, alpha=1.0, beta=0.0)
+        D_reg = D_f.local(4)
+        for s in T.unroll(4):
+            rN = s % 2
+            rM = s // 2
+            D_g[lane // 4 + 8 * rM, 2 * (lane % 4) + rN] = D_reg[s]
+
+    with tvm.target.Target("cuda"):
+        mod = tvm.compile(tvm.IRModule({"main": gemm}), target="cuda", tir_pipeline="tirx")
+
+    np.random.seed(0)
+    A_np = np.random.uniform(-1, 1, (16, 16)).astype(np.float32)
+    B_np = np.random.uniform(-1, 1, (16, 8)).astype(np.float32)
+    golden = A_np @ B_np
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_dev = tvm.runtime.tensor(A_np.astype(np_dtype), dev)
+        B_dev = tvm.runtime.tensor(B_np.astype(np_dtype), dev)
+        D_dev = tvm.runtime.tensor(np.zeros((16, 8), np.float32), dev)
+        mod(A_dev, B_dev, D_dev)
+        tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=1e-2, rtol=1e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# (Mt, Nt, Kt, kinst) tilings: single tile, each dim multi-tiled, fully tiled,
+# M = 64, and the m16n8k8 (kHi == 1) variants including a non-16-divisible K.
+_TILED_SHAPES = [
+    (1, 1, 1, 16),  # single m16n8k16 tile
+    (2, 1, 1, 16),  # two M-tiles
+    (1, 2, 1, 16),  # two N-tiles
+    (1, 1, 2, 16),  # two K-tiles (accumulated in place)
+    (2, 2, 2, 16),  # every dim tiled
+    (4, 1, 1, 16),  # M = 64
+    (1, 1, 1, 8),  # m16n8k8 single tile (kHi == 1)
+    (1, 1, 3, 8),  # K = 24 -> three k8 tiles
+    (2, 2, 3, 8),  # k8, every dim tiled
+]
+# (dtype, beta) input modes crossed against every shape: f16/bf16 inputs, with
+# beta = 0 (D = A @ B) and beta = 1 (D = A @ B + C, accumulating C in place).
+_TILED_MODES = [
+    ("float16", 0.0),
+    ("bfloat16", 0.0),
+    ("float16", 1.0),
+    ("bfloat16", 1.0),
+]
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("Mt, Nt, Kt, kinst", _TILED_SHAPES)
+@pytest.mark.parametrize("dtype, beta", _TILED_MODES)
+def test_cuda_gemm_mma_numerical_tiled(dtype, beta, Mt, Nt, Kt, kinst):
+    """End-to-end D = A @ B (+ C when beta==1) over an Mt x Nt x Kt tiling.
+
+    The two stacked ``parametrize`` decorators form the cartesian product of
+    every tiling shape with every (dtype, beta) input mode, so each combination
+    is an independent pytest item (pytest-xdist runs them in parallel)."""
+    if dtype == "bfloat16":
+        ml_dtypes = pytest.importorskip("ml_dtypes")
+        np_dtype = ml_dtypes.bfloat16
+    else:
+        np_dtype = np.float16
+
+    func, M, N, K = _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype)
+    with tvm.target.Target("cuda"):
+        mod = tvm.compile(tvm.IRModule({"main": func}), target="cuda", tir_pipeline="tirx")
+
+    np.random.seed(0)
+    A_np = np.random.uniform(-1, 1, (M, K)).astype(np.float32)
+    B_np = np.random.uniform(-1, 1, (K, N)).astype(np.float32)
+    C_np = np.random.uniform(-1, 1, (M, N)).astype(np.float32)
+    golden = A_np @ B_np + (C_np if beta == 1.0 else 0.0)
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_dev = tvm.runtime.tensor(A_np.astype(np_dtype), dev)
+        B_dev = tvm.runtime.tensor(B_np.astype(np_dtype), dev)
+        C_dev = tvm.runtime.tensor(C_np, dev)
+        D_dev = tvm.runtime.tensor(np.zeros((M, N), np.float32), dev)
+        mod(A_dev, B_dev, C_dev, D_dev)
+        tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=2e-2, rtol=2e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("dtype", ["float16", "bfloat16"])
+@pytest.mark.parametrize(
+    "transpose_A, transpose_B",
+    [(False, False), (True, False), (False, True), (True, True)],
+)
+def test_cuda_gemm_mma_numerical_transpose(transpose_A, transpose_B, dtype):
+    """End-to-end D = A @ B for every A/B input orientation, crossed with dtype.
+
+    The orientation and dtype decorators form a cartesian product, so each
+    (transpose_A, transpose_B, dtype) is an independent pytest item."""
+    if dtype == "bfloat16":
+        ml_dtypes = pytest.importorskip("ml_dtypes")
+        np_dtype = ml_dtypes.bfloat16
+    else:
+        np_dtype = np.float16
+
+    func = _build_transpose_numeric(transpose_A, transpose_B, dtype)
+    with tvm.target.Target("cuda"):
+        mod = tvm.compile(tvm.IRModule({"main": func}), target="cuda", tir_pipeline="tirx")
+
+    np.random.seed(0)
+    A_log = np.random.uniform(-1, 1, (16, 16)).astype(np.float32)  # logical A[M, K]
+    B_log = np.random.uniform(-1, 1, (16, 8)).astype(np.float32)  # logical B[K, N]
+    A_buf = (A_log.T if transpose_A else A_log).astype(np_dtype)
+    B_buf = (B_log.T if transpose_B else B_log).astype(np_dtype)
+    golden = A_log @ B_log
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        A_dev = tvm.runtime.tensor(A_buf, dev)
+        B_dev = tvm.runtime.tensor(B_buf, dev)
+        D_dev = tvm.runtime.tensor(np.zeros((16, 8), np.float32), dev)
+        mod(A_dev, B_dev, D_dev)
+        tvm.testing.assert_allclose(golden, D_dev.numpy(), atol=2e-2, rtol=2e-2)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.parametrize(
+    "Mt, Nt, Kt, kinst",
+    [
+        (2, 1, 1, 16),  # M = 32 (two M-tiles)
+        (1, 2, 1, 16),  # N = 16 (two N-tiles)
+        (1, 1, 2, 16),  # K = 32 (two K-tiles, accumulated in place)
+        (2, 2, 2, 16),  # 8 tiles
+        (4, 1, 1, 16),  # M = 64
+        (1, 1, 1, 8),  # m16n8k8 single tile (kHi == 1)
+        (1, 1, 3, 8),  # K = 24 -> three k8 tiles (16 does not divide 24)
+        (2, 2, 3, 8),  # k8, every dim tiled
+    ],
+)
+def test_cuda_gemm_mma_lowers_tiled(Mt, Nt, Kt, kinst):
+    """Every tiling we expect to dispatch must lower, selecting the right mma.
+
+    The k8 cases are the regression guard for the kHi == 1 fragment grouping
+    (an extent-1 high-K register group must not be rejected as a thread axis).
+    """
+    script = _lower(_build_tiled(Mt, Nt, Kt, kinst))["main"].script()
+    assert "T.ptx.mma(" in script
+    assert f"m16n8k{kinst}" in script
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize(
+    "Mt, Nt, Kt, kinst",
+    [
+        (1, 1, 1, 16),
+        (2, 2, 2, 16),
+        (4, 1, 1, 16),
+        (1, 1, 1, 8),
+        (1, 1, 3, 8),
+        (2, 2, 3, 8),
+    ],
+)
+def test_cuda_gemm_mma_codegen_issue_count(Mt, Nt, Kt, kinst):
+    """Full pipeline (UnrollLoop + CUDA codegen) emits one mma per (Mt, Nt, Kt)
+    tile; K-tiles accumulate in place, so D is cleared once per output tile."""
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+    with target:
+        mod = tvm.compile(
+            tvm.IRModule({"main": _build_tiled(Mt, Nt, Kt, kinst, store=True)}),
+            target=target,
+            tir_pipeline="tirx",
+        )
+    src = mod.mod.imports[0].inspect_source()
+    assert f"mma.sync.aligned.m16n8k{kinst}" in src
+    # mma is emitted as one __device__ helper, invoked once per tile.
+    helper = f"ptx_mma_m16n8k{kinst}_row_col"
+    assert src.count(helper) - 1 == Mt * Nt * Kt
+
+
+@pytest.mark.parametrize(
+    "transpose_A, transpose_B",
+    [(False, False), (True, False), (False, True), (True, True)],
+)
+def test_cuda_gemm_mma_lowers_transpose(transpose_A, transpose_B):
+    """All four A/B orientations dispatch to the same m16n8k16. transpose only
+    describes the input's logical orientation; the .row.col mma is unchanged."""
+    script = _lower(_build_transpose(transpose_A, transpose_B))["main"].script()
+    assert "T.ptx.mma(" in script
+    assert "m16n8k16" in script
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize(
+    "transpose_A, transpose_B",
+    [(False, False), (True, False), (False, True), (True, True)],
+)
+def test_cuda_gemm_mma_codegen_transpose(transpose_A, transpose_B):
+    """Every orientation codegens to a valid m16n8k16 kernel."""
+    target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+    with target:
+        mod = tvm.compile(
+            tvm.IRModule({"main": _build_transpose(transpose_A, transpose_B, store=True)}),
+            target=target,
+            tir_pipeline="tirx",
+        )
+    assert "mma.sync.aligned.m16n8k16" in mod.mod.imports[0].inspect_source()
+
+
+@pytest.mark.parametrize(
+    "a, b, c, d",
+    [
+        ("float16", "float16", "float16", "float16"),  # f16 accumulate
+        ("bfloat16", "float16", "float32", "float32"),  # mixed A/B inputs
+        ("float32", "float32", "float32", "float32"),  # f32 (tf32) inputs
+        ("int8", "int8", "int32", "int32"),  # integer
+    ],
+)
+def test_cuda_gemm_mma_rejects_unsupported_dtype(a, b, c, d):
+    """The table holds only (bf16|f16, same, f32, f32); any other dtype
+    signature must decline rather than emit a wrong mma."""
+    with pytest.raises(RuntimeError, match="dispatch failed"):
+        _lower(_build_dtypes(a, b, c, d))
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py

@@ -0,0 +1,466 @@
+# 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 ``T.permute_layout``.
+
+Coverage:
+
+- The algorithm helpers (`_bank_free`, `_check_bijection`, `_choose_xor_k`)
+  directly, with a NumPy oracle.
+- End-to-end compiled-kernel byte-for-byte equivalence on CUDA for the SF
+  fp8-blockwise-gemm transpose shapes (BLK_SFA = 128, 256) plus a few
+  generic linear↔stride-permuted layouts and additional dtypes (u8, fp16,
+  i32, u64).
+- Reject cases: non-warp scope, dtype mismatch, shape mismatch, swizzle/
+  compose layouts, layouts whose strides don't form a bijection on the
+  slice.  Each must surface as a ``RuntimeError`` from the dispatcher and
+  NOT silently emit a wrong kernel.
+"""
+
+from __future__ import annotations
+
+import math
+
+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
+
+# Helpers exposed by the dispatcher module for direct algorithm tests.
+from tvm.tirx.cuda.operator.tile_primitive.permute_layout.warp_xor_swizzle import (
+    _bank_free,
+    _check_bijection,
+    _choose_xor_k,
+)
+from tvm.tirx.layout import S, SwizzleLayout, TileLayout
+
+# ---------------------------------------------------------------------------
+# Algorithm-only tests (no CUDA needed).
+# ---------------------------------------------------------------------------
+
+
+def _np_layout_offset(extent, strides, multi_idx):
+    return int(sum(s * i for s, i in zip(strides, multi_idx)))
+
+
+def _expected_permute(src_np, src_strides, dst_strides, extent):
+    """Compute the expected output: dst at byte offset ``L_dst(i)`` holds the
+    value at ``src`` byte offset ``L_src(i)``, for every logical index i.
+    """
+    V = math.prod(extent)
+    dst_np = np.zeros_like(src_np)
+    for flat in range(V):
+        idx = []
+        rem = flat
+        for e in reversed(extent):
+            idx.append(rem % e)
+            rem //= e
+        idx = list(reversed(idx))
+        src_off = _np_layout_offset(extent, src_strides, idx)
+        dst_off = _np_layout_offset(extent, dst_strides, idx)
+        dst_np.reshape(-1)[dst_off] = src_np.reshape(-1)[src_off]
+    return dst_np
+
+
+def test_bank_free_sf_128_u32():
+    """SF BLK_SFA=128: write phase has 4-way conflict at k=0, free at k=2."""
+    extent = [4, 32]
+    src = [32, 1]
+    dst = [1, 4]
+    bytes_per = 4
+    P = 4
+    assert _bank_free(extent, src, bytes_per, P, 0)
+    assert not _bank_free(extent, dst, bytes_per, P, 0)
+    assert _bank_free(extent, dst, bytes_per, P, 2)
+    assert _choose_xor_k(extent, src, dst, bytes_per, P) == 2
+
+
+def test_bank_free_sf_256_u32():
+    """SF BLK_SFA=256: same shift=3 (k=2) handles the high block too."""
+    extent = [2, 4, 32]
+    src = [128, 32, 1]
+    dst = [128, 1, 4]
+    bytes_per = 4
+    P = 8
+    assert _bank_free(extent, src, bytes_per, P, 0)
+    assert not _bank_free(extent, dst, bytes_per, P, 0)
+    assert _bank_free(extent, dst, bytes_per, P, 2)
+    assert _choose_xor_k(extent, src, dst, bytes_per, P) == 2
+
+
+def test_identity_no_xor():
+    """L_src == L_dst => k=0 (no XOR needed and the op is essentially a copy)."""
+    assert _choose_xor_k([4, 32], [32, 1], [32, 1], 4, 4) == 0
+    # A 2D buffer with row-major to row-major is a true no-op.
+    assert _bank_free([4, 32], [32, 1], 4, 4, 0)
+
+
+def test_bijection_check_rejects_aliased():
+    """If two logical indices map to the same physical byte, reject."""
+    # Stride 0 on a non-singleton extent => alias.
+    assert not _check_bijection([4, 32], [0, 1])
+    # Negative or non-contiguous-but-bijective is still fine.
+    assert _check_bijection([4, 32], [1, 4])
+
+
+def test_dtype_widths_choose_xor_k():
+    """Each dtype's outcome:
+
+    The unvectorized algorithm is provably correct only when every per-lane
+    access maps to a single 4-byte bank.  For 4-byte dtypes that always holds
+    (one element per bank), so we expect a valid k.  For sub-4-byte dtypes
+    with stride-1 reads, multiple lanes share a bank no matter how we permute
+    register slots — the dispatcher correctly rejects those (k is None).
+    """
+    extent = [4, 32]
+    src = [32, 1]  # linear
+    dst = [1, 4]  # transposed
+    # u32: this is the SF case; the algorithm must find k=2.
+    assert _choose_xor_k(extent, src, dst, 4, 4) == 2
+    # u16/fp16, u8: stride-1 in bytes < 4 packs >1 lane into the same bank;
+    # register-slot XOR cannot fix that, so the dispatcher rejects.
+    assert _choose_xor_k(extent, src, dst, 2, 4) is None
+    assert _choose_xor_k(extent, src, dst, 1, 4) is None
+
+
+# ---------------------------------------------------------------------------
+# End-to-end compiled-kernel tests on CUDA.
+# ---------------------------------------------------------------------------
+
+
+def _has_cuda():
+    try:
+        return tvm.cuda(0).exist
+    except Exception:
+        return False
+
+
+needs_cuda = pytest.mark.skipif(not _has_cuda(), reason="needs CUDA")
+
+
+def _compile_and_run(prim_func, np_inputs):
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.IRModule({"main": prim_func})
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        tensors = [tvm.runtime.tensor(a, dev) for a in np_inputs]
+        mod(*tensors)
+        return [tensor.numpy() for tensor in tensors]
+
+    outputs = tvm.testing.run_with_gpu_lock(run_and_check)
+    return outputs, mod.mod.imports[0].inspect_source()
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@needs_cuda
+@pytest.mark.parametrize(
+    "name, pipe, blk, dtype",
+    [
+        ("sf_128_u32", 2, 128, "uint32"),
+        ("sf_256_u32", 2, 256, "uint32"),
+        ("sf_128_i32", 2, 128, "int32"),
+        ("sf_128_fp32", 2, 128, "float32"),
+    ],
+)
+def test_sf_blockwise_transpose(name, pipe, blk, dtype):
+    """SF blockwise-GEMM scale-factor transpose, the canonical use case."""
+    high = blk // 128 if blk >= 128 else 1
+    # Use 4D logical shape (PIPE, high, 4, 32) to keep the high-block factored.
+    shape = (pipe, high, 4, 32)
+
+    # Element strides for src (linear) and dst (transposed within each
+    # 128-block).  Stage stride = blk; each 128-block contributes 128 to the
+    # high stride.
+    src_strides = (blk, 128, 32, 1)
+    dst_strides = (blk, 128, 1, 4)
+    pre = TileLayout(S[shape:src_strides])
+    post = TileLayout(S[shape:dst_strides])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, shape, dtype, layout=pre)
+        B_buf = T.match_buffer(B, shape, dtype, layout=post)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        for s in T.serial(0, pipe):
+            Tx.warp.permute_layout(
+                B_buf[s, 0:high, 0:4, 0:32], A_buf[s, 0:high, 0:4, 0:32]
+            )
+        # fmt: on
+
+    np.random.seed(0)
+    A_np = tvm.testing.generate_random_array(dtype, shape)
+    B_np = np.zeros_like(A_np)
+
+    [_, B_out], src = _compile_and_run(f, [A_np, B_np])
+
+    # The dispatcher must have picked the XOR-swizzled variant; check that
+    # the generated CUDA contains the per-lane XOR pattern.  This is the
+    # "no perf regression" smoke test: any future variant that omits the
+    # XOR would re-introduce 4-way bank conflicts.
+    assert ">> 3" in src, f"expected XOR-swizzle (lane>>3) in CUDA for {name}"
+    assert "warp_sync" in src or "syncwarp" in src
+
+    # Byte-for-byte equality via numpy reference.
+    for s in range(pipe):
+        A_flat = A_np[s].reshape(-1)
+        B_flat = B_out[s].reshape(-1)
+        ref = _expected_permute(
+            A_flat,
+            list(src_strides[1:]),
+            list(dst_strides[1:]),
+            list(shape[1:]),
+        )
+        np.testing.assert_array_equal(B_flat, ref, err_msg=f"{name} stage {s}")
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@needs_cuda
+def test_identity_passes_through_as_copy():
+    """L_src == L_dst should still compile and produce a correct (identity) copy."""
+    shape = (4, 32)
+    layout = TileLayout(S[shape : (32, 1)])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, shape, "uint32", layout=layout)
+        B_buf = T.match_buffer(B, shape, "uint32", layout=layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.warp.permute_layout(B_buf, A_buf)
+        # fmt: on
+
+    np.random.seed(0)
+    A_np = tvm.testing.generate_random_array("uint32", shape)
+    B_np = np.zeros_like(A_np)
+
+    [_, B_out], _ = _compile_and_run(f, [A_np, B_np])
+    np.testing.assert_array_equal(B_out, A_np)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@needs_cuda
+@pytest.mark.parametrize("dtype", ["uint32", "int32", "float32"])
+@pytest.mark.parametrize(
+    "shape, src_strides, dst_strides",
+    [
+        # (8, 32) → (8, 32) transposed: src linear, dst column-major.
+        ((8, 32), (32, 1), (1, 8)),
+        # (16, 32): per_thread = 16 — tests P=16 path.
+        ((16, 32), (32, 1), (1, 16)),
+    ],
+)
+def test_generic_transpose(shape, src_strides, dst_strides, dtype):
+    """Generic linear↔transposed pairs at various P values."""
+    pre = TileLayout(S[shape:src_strides])
+    post = TileLayout(S[shape:dst_strides])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, shape, dtype, layout=pre)
+        B_buf = T.match_buffer(B, shape, dtype, layout=post)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.warp.permute_layout(B_buf, A_buf)
+        # fmt: on
+
+    np.random.seed(0)
+    A_np = tvm.testing.generate_random_array(dtype, shape)
+    B_np = np.zeros_like(A_np)
+    [_, B_out], _ = _compile_and_run(f, [A_np, B_np])
+
+    ref = _expected_permute(A_np.reshape(-1), list(src_strides), list(dst_strides), list(shape))
+    np.testing.assert_array_equal(B_out.reshape(-1), ref)
+
+
+# ---------------------------------------------------------------------------
+# Reject cases: the dispatcher must surface a clear error, never silently
+# emit a wrong kernel.
+# ---------------------------------------------------------------------------
+
+
+def _build_and_assert_rejected(shape, src_layout, dst_layout, dtype, msg_substr):
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, shape, dtype, layout=src_layout)
+        B_buf = T.match_buffer(B, shape, dtype, layout=dst_layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.warp.permute_layout(B_buf, A_buf)
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.raises(RuntimeError) as exc_info:
+        mod = tvm.IRModule({"main": f})
+        tvm.compile(mod, target=target, tir_pipeline="tirx")
+    assert msg_substr in str(exc_info.value), (
+        f"expected reject reason to mention {msg_substr!r}, got: {exc_info.value}"
+    )
+
+
+def test_reject_dtype_mismatch():
+    shape = (4, 32)
+    layout = TileLayout(S[shape : (32, 1)])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, shape, "uint32", layout=layout)
+        B_buf = T.match_buffer(B, shape, "uint16", layout=layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.warp.permute_layout(B_buf, A_buf)
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.raises(RuntimeError) as exc_info:
+        tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
+    assert "dtype mismatch" in str(exc_info.value)
+
+
+def test_reject_shape_mismatch():
+    src_layout = TileLayout(S[(4, 32) : (32, 1)])
+    dst_layout = TileLayout(S[(8, 16) : (16, 1)])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, (4, 32), "uint32", layout=src_layout)
+        B_buf = T.match_buffer(B, (8, 16), "uint32", layout=dst_layout)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.warp.permute_layout(B_buf, A_buf)
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.raises(RuntimeError) as exc_info:
+        tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
+    assert "shape mismatch" in str(exc_info.value)
+
+
+def test_reject_swizzle_layout():
+    """ComposeLayout(SwizzleLayout, TileLayout) is not supported by the warp variant."""
+    from tvm.tirx.layout import ComposeLayout
+
+    inner = TileLayout(S[(4, 32) : (32, 1)])
+    sw = SwizzleLayout(per_element=2, swizzle_len=2, atom_len=4)
+    swizzled = ComposeLayout(sw, inner)
+    plain = TileLayout(S[(4, 32) : (1, 4)])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, (4, 32), "uint32", layout=swizzled)
+        B_buf = T.match_buffer(B, (4, 32), "uint32", layout=plain)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.warp.permute_layout(B_buf, A_buf)
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.raises(RuntimeError) as exc_info:
+        tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
+    assert "TileLayout" in str(exc_info.value)
+
+
+def test_reject_non_warp_scope():
+    layout_pre = TileLayout(S[(4, 32) : (32, 1)])
+    layout_post = TileLayout(S[(4, 32) : (1, 4)])
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, (4, 32), "uint32", layout=layout_pre)
+        B_buf = T.match_buffer(B, (4, 32), "uint32", layout=layout_post)
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([32])
+        Tx.cta.permute_layout(B_buf, A_buf)  # cta scope, not warp
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target, pytest.raises(RuntimeError) as exc_info:
+        tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
+    assert "warp" in str(exc_info.value)
+
+
+@pytest.mark.parametrize("dtype", ["uint32", "float32"])
+def test_shared_to_shared_uses_direct_ldst(dtype):
+    """Compile-only: a shared->shared 32b transpose must take the direct
+    base-ptr + byte-offset ``ld.shared`` / ``st.shared`` path.
+
+    For a 4/8-byte dtype with both operands in shared memory, indexing through
+    ``buf[...]`` lowers the swizzled layout to a per-element IMAD flatten. The
+    direct path computes one base ptr (``ptr_to(stride_offset)``) and adds a
+    compile-time ``off * dtype_bytes`` per register slot, then issues
+    ``T.ptx.ld/st(..., space="shared")``. The bits move through a uint
+    container, so ``float32`` (whose ``ld.b32`` cannot return a float) lowers
+    the same way as the ``uint32`` SF case.
+    """
+    shape = (4, 32)
+    pre = TileLayout(S[shape : (32, 1)])  # linear
+    post = TileLayout(S[shape : (1, 4)])  # transposed within the 128-block
+
+    # fmt: off
+    @T.prim_func
+    def f(A: T.handle, B: T.handle):
+        A_buf = T.match_buffer(A, shape, dtype, layout=pre)
+        B_buf = T.match_buffer(B, shape, dtype, layout=post)
+        T.device_entry()
+        T.cta_id([1])
+        tid = T.thread_id([32])
+        sA = T.alloc_buffer(shape, dtype, scope="shared", layout=pre)
+        sB = T.alloc_buffer(shape, dtype, scope="shared", layout=post)
+        Tx.cta.copy(sA[:, :], A_buf[:, :])
+        T.cuda.cta_sync()
+        Tx.warp.permute_layout(sB[:, :], sA[:, :])
+        T.cuda.cta_sync()
+        Tx.cta.copy(B_buf[:, :], sB[:, :])
+        # fmt: on
+
+    target = tvm.target.Target("cuda")
+    with target:
+        mod = tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx")
+    src = mod.mod.imports[0].inspect_source()
+    assert "ld.shared" in src, f"expected direct ld.shared in permute; src=\n{src}"
+    assert "st.shared" in src, f"expected direct st.shared in permute; src=\n{src}"
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/test_dispatcher.py

@@ -0,0 +1,158 @@
+# 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.
+
+import pytest
+
+
+def _import_and_register():
+    # Ensure all schedule registrations (legacy + dispatcher variants) are loaded
+    import tvm.tirx.operator.tile_primitive as _  # noqa: F401
+
+
+class _DummyKind:
+    def __init__(self, name: str):
+        self.name = name
+
+    def __str__(self) -> str:  # used in messages
+        return self.name
+
+
+class _DummyTarget:
+    def __init__(self, kind_name: str):
+        self.kind = _DummyKind(kind_name)
+
+
+class _DummyExecScope:
+    def __init__(self, name: str):
+        self.name = name
+
+
+class _DummySctx:
+    def __init__(self, target_kind: str, exec_scope: str):
+        self.target = _DummyTarget(target_kind)
+        self.exec_scope = _DummyExecScope(exec_scope)
+        self.scope_kind = exec_scope
+
+
+def test_dispatch_prints_predicate_reasons():
+    """Validate TRACE mode prints per-variant predicate failure reasons."""
+    _import_and_register()
+    from tvm.ir import Op
+    from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
+
+    class _OpCall:
+        def __init__(self, op):
+            self.op = op
+            self.args = []  # not used by the tested predicates
+
+    # Use TRN copy; predicate requires exec_scope == "thread".
+    op_call = _OpCall(Op.get("tirx.tile.copy"))
+    sctx = _DummySctx(target_kind="trn", exec_scope="warp")  # intentionally wrong
+
+    with pytest.raises(RuntimeError) as e:
+        run_dispatch(op_call, sctx)
+
+    out = str(e.value)
+    print(out)
+    # Header + per-variant reason must be printed in table format
+    assert "TIRx schedule dispatch failed: op=tirx.tile.copy target=trn" in out
+    assert "Variant" in out  # table header present
+    assert "default" in out  # variant name present
+    assert "rejected: exec_scope" in out
+    # opcall object IR should be printed in the table
+    assert "opcall:" in out
+
+
+def test_dispatch_forced_variant_missing_table_and_message():
+    _import_and_register()
+    from tvm.ir import Op
+    from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
+
+    class _OpCall:
+        def __init__(self, op):
+            self.op = op
+            self.dispatch = "__nonexistent__"
+            self.args = []
+
+    op_call = _OpCall(Op.get("tirx.tile.copy"))
+    sctx = _DummySctx(target_kind="trn", exec_scope="thread")
+
+    with pytest.raises(RuntimeError) as e:
+        run_dispatch(op_call, sctx)
+
+    msg = str(e.value)
+    print(msg)
+    assert "TIRx schedule dispatch failed: op=tirx.tile.copy target=trn" in msg
+    assert "no variant named '__nonexistent__' is registered" in msg
+
+
+def test_dispatch_raises_with_aggregated_reasons():
+    """Validate STRICT mode raises aggregated error message with reasons."""
+    _import_and_register()
+    from tvm.ir import Op
+    from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
+
+    class _OpCall:
+        def __init__(self, op):
+            self.op = op
+            self.args = []
+
+    # Use TRN compose_op; variant implementation raises NotImplementedError
+    op_call = _OpCall(Op.get("tirx.tile.compose_op"))
+    sctx = _DummySctx(target_kind="trn", exec_scope="thread")
+
+    with pytest.raises(RuntimeError) as e:
+        run_dispatch(op_call, sctx)
+
+    msg = str(e.value)
+    print(msg)
+    assert "TIRx schedule dispatch failed: op=tirx.tile.compose_op target=trn" in msg
+    assert "default" in msg
+    assert "exception — NotImplementedError" in msg
+    # opcall content and backtrace should be included inside the table
+    assert "opcall:" in msg
+    assert "Traceback (most recent call last):" in msg
+
+
+def test_dispatch_prints_real_opcall_ir():
+    """Create a real TilePrimitiveCall via BufferRegions and ensure its IR is in the table."""
+    _import_and_register()
+    from tvm.ir import Op
+    from tvm.tirx.buffer import decl_buffer
+    from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch
+    from tvm.tirx.stmt import TilePrimitiveCall
+
+    # Build a real TIRx TilePrimitiveCall: tirx.tile.copy(A[0:64], B[0:64])
+    A = decl_buffer((64,), "float32", scope="global")
+    B = decl_buffer((64,), "float32", scope="shared")
+    real_opcall = TilePrimitiveCall(
+        A[0:64], B[0:64], op=Op.get("tirx.tile.copy"), workspace={}, config={}
+    )
+
+    # Force predicate rejection to trigger formatted error with opcall IR
+    sctx = _DummySctx(target_kind="trn", exec_scope="warp")
+    with pytest.raises(RuntimeError) as e:
+        run_dispatch(real_opcall, sctx)
+
+    out = str(e.value)
+    print(out)
+    # Verify header and that the opcall IR is included in the table
+    assert "TIRx schedule dispatch failed: op=tirx.tile.copy target=trn" in out
+    assert "Variant" in out
+    assert "opcall:" in out
+    # IR should mention the operator name
+    assert "tirx.tile.copy" in out

tests/python/tirx/operator/tile_primitive/trn/test_binary_trn.py

@@ -0,0 +1,354 @@
+# 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.
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+Tx_func_map = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "min": Tx.minimum, "max": Tx.maximum}
+
+
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "min", "max"])
+@pytest.mark.parametrize(
+    "operands_type",
+    [
+        "region_region",
+        "const_region",
+        "region_const",
+        "region_broadcast_lhs",
+        "region_broadcast_rhs",
+    ],
+)
+def test_simple_binary(op_type, operands_type):
+    const = T.float32(3.0)
+    src1_shape = [128, 512] if operands_type != "region_broadcast_lhs" else [128, 1]
+    src1_layout = TileLayout(S[src1_shape : (1 @ P, 1 @ F)])
+    src2_shape = [128, 512] if operands_type != "region_broadcast_rhs" else [128, 1]
+    src2_layout = TileLayout(S[src2_shape : (1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    Tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def binary() ->None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        if operands_type == "region_region" or operands_type.startswith("region_broadcast"):
+            Tx_func(C_sbuf, A_sbuf, B_sbuf)
+        elif operands_type == "const_region":
+            Tx_func(C_sbuf, const, A_sbuf)
+        elif operands_type == "region_const":
+            Tx_func(C_sbuf, A_sbuf, const)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer(src1_shape, scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer(src2_shape, scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer(dst_shape, scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if operands_type == "region_region":
+                        T.nki.tensortensor(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], B_sbuf[p_loop, f_loop], op_type)  # noqa: E501
+                    elif operands_type == "region_const":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], T.float32(3.0), op_type, T.bool(False))  # noqa: E501
+                    elif operands_type == "const_region":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], T.float32(3.0), op_type, T.bool(True))  # noqa: E501
+                    elif operands_type == "region_broadcast_rhs":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop], B_sbuf[p_loop, 0], op_type, T.bool(False))  # noqa: E501
+                    elif operands_type == "region_broadcast_lhs":
+                        T.nki.tensorscalar(C_sbuf[p_loop, f_loop], B_sbuf[p_loop, f_loop], A_sbuf[p_loop, 0], op_type, T.bool(True))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "min", "max"])
+@pytest.mark.parametrize(
+    "operands_type",
+    [
+        "region_region",
+        "const_region",
+        "region_const",
+        "region_broadcast_lhs",
+        "region_broadcast_rhs",
+    ],
+)
+def test_binary_complex(op_type, operands_type):
+    src1_shape = [1024, 512] if operands_type != "region_broadcast_lhs" else [1024, 4]
+    src1_layout_data_iter = (128, 4096) if operands_type != "region_broadcast_lhs" else (128, 32)
+    src1_layout = TileLayout(S[src1_layout_data_iter : (1 @ P, 1 @ F)])
+    src2_shape = [512, 512] if operands_type != "region_broadcast_rhs" else [128, 512]
+    src2_layout_data_iter = (128, 2048) if operands_type != "region_broadcast_rhs" else (128, 512)
+    src2_layout = TileLayout(S[src2_layout_data_iter : (1 @ P, 1 @ F)])
+
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    const = T.float32(3.0)
+    Tx_func = Tx_func_map[op_type]
+
+    src1_view_shape = [128, 8, 512]
+    src2_view_shape = [128, 4, 512] if operands_type != "region_broadcast_rhs" else [128, 1, 512]
+    dst_view_shape = [128, 4, 512]
+    if operands_type == "region_broadcast_lhs":
+        src1_view_shape = [128, 8, 4, 1]
+        src2_view_shape = [128, 4, 4, 128]
+        dst_view_shape = [128, 4, 4, 128]
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        A_sbuf_view = A_sbuf.view(*src1_view_shape)
+        B_sbuf_view = B_sbuf.view(*src2_view_shape)
+        C_sbuf_view = C_sbuf.view(*dst_view_shape)
+        for i in range(4):
+            if operands_type == "region_region":
+                Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], B_sbuf_view[:, i, :])
+            elif operands_type == "region_const":
+                Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], const)
+            elif operands_type == "const_region":
+                Tx_func(C_sbuf_view[:, i, :], const, A_sbuf_view[:, i * 2, :])
+            elif operands_type == "region_broadcast_rhs":
+                Tx_func(C_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :], B_sbuf_view[:, 0, :])
+            elif operands_type == "region_broadcast_lhs":
+                Tx_func(C_sbuf_view[:, i, :, :], A_sbuf_view[:, i*2,:, :], B_sbuf_view[:, i, :, :])
+
+    f_extent = 128 if operands_type == "region_broadcast_lhs" else 512
+    b_extent = 4 if operands_type == "region_broadcast_lhs" else 1
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer(src1_layout_data_iter, scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer(src2_layout_data_iter, scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        A_sbuf_view = T.decl_buffer(src1_layout_data_iter, data=A_sbuf.data, scope="trn.sbuf", layout=None)  # noqa: E501
+        B_sbuf_view = T.decl_buffer(src2_layout_data_iter, data=B_sbuf.data, scope="trn.sbuf", layout=None)  # noqa: E501
+        C_sbuf_view = T.decl_buffer((128, 2048), data=C_sbuf.data, scope="trn.sbuf", layout=None)
+        for i, b_loop in T.grid(4, b_extent):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, f_extent, annotations={"nki_dim":"F"}):
+                    if operands_type == "region_region":
+                        T.nki.tensortensor(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], B_sbuf_view[p_loop, i * 512 + f_loop], op_type)  # noqa: E501
+                    elif operands_type == "const_region":
+                        T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], T.float32(3.0), op_type, T.bool(True))  # noqa: E501
+                    elif operands_type == "region_const":
+                        T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], T.float32(3.0), op_type, T.bool(False))  # noqa: E501
+                    elif operands_type == "region_broadcast_lhs":
+                        T.nki.tensorscalar(C_sbuf_view[p_loop, i * 512 + b_loop * 128 + f_loop], B_sbuf_view[p_loop, i * 512 + b_loop * 128 + f_loop], A_sbuf_view[p_loop, i * 8 + b_loop], op_type, T.bool(True))  # noqa: E501
+                    elif operands_type == "region_broadcast_rhs":
+                        T.nki.tensortensor(C_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop], B_sbuf_view[p_loop, f_loop], op_type)  # noqa: E501
+
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_broadcast1():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.add(C_sbuf, A_sbuf, B_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorscalar(C_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], "add", T.bool(False))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_broadcast2():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.add(C_sbuf, A_sbuf, B_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 128):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    T.nki.tensortensor(C_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 128 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 128 + f_loop], B_sbuf[p_loop, b_loop % 4 * 128 + f_loop], "add")  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_broadcast3():
+    src1_shape = [128, 512]
+    src1_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    src2_shape = [32, 128, 512]
+    src2_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.add(C_sbuf, A_sbuf, B_sbuf[0])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    T.nki.tensortensor(C_sbuf[p_loop, b_loop * 128 + f_loop], A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 4096 + f_loop], "add")  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_with_guard():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            Tx.add(C_sbuf[:, :, 0:j*128], A_sbuf[:, :, 0:j*128], B_sbuf[:, 0:j*128])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for j, b_loop in T.grid(4, 96):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    if b_loop % 3 - j < 0:
+                        T.nki.tensortensor(C_sbuf[p_loop, b_loop % 3 * 4096 + b_loop // 3 * 128 + f_loop], A_sbuf[p_loop, b_loop % 3 * 4096 + b_loop // 3 * 128 + f_loop], B_sbuf[p_loop, b_loop % 3 * 128 + f_loop], "add")  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_compose_op_trn.py

@@ -0,0 +1,772 @@
+# 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.
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_simple_activation_reduce():
+    A_shape = (128, 512)
+    A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    B_shape = (128, 512)
+    B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    C_shape = (128, 1)
+    C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        Tx.unary_reduce(B, C, A, "sqrt", "sum", reduce_axes=1)
+
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, 0], B[p_loop, f_loop], A[p_loop, f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_in_loop():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (16, 128)
+    C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop % 8 // 2 * 2048 + b_loop // 8 * 1024 + b_loop % 2 * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_in_loop2():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(16 * 512, 128) : (1 @ F, 1 @ P)])
+    C_shape = (16, 128)
+    C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_two_stage():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        partial_reduce = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 1):
+            for reduction_b_loop in range(8):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.activation_reduce(partial_reduce[p_loop, reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop // 4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C[p_loop, 0], partial_reduce[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_with_bias_scale():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(16 * 512, 128) : (1 @ F, 1 @ P)])
+    C_shape = (16, 128)
+    C_layout = TileLayout(S[(2, 4, 2, 128) : (2 @ F, 4 @ F, 1 @ F, 1 @ P)])
+    bias_shape = 128
+    bias_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        bias = T.alloc_buffer(bias_shape, dtype="float32", scope="trn.sbuf", layout=bias_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=1, bias=bias, scale=2.0)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        bias = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 + b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop], "sqrt", "add", bias[p_loop, 0], T.float32(2.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_simple_tensor_scalar_reduce():
+    A_shape = (128, 512)
+    A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    B_shape = (128, 512)
+    B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    C_shape = (128, 1)
+    C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        Tx.binary_reduce(B, C, A, 1.0, "add", "sum", reduce_axes=1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        A = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorscalar_reduce(C[p_loop, 0], B[p_loop, f_loop], A[p_loop, f_loop], T.float32(1.0), "add", "add", T.bool(False))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_tensor_tensor_reduce_fail():
+    A_shape = (128, 512)
+    A_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    B_shape = (128, 512)
+    B_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    D_shape = (128, 512)
+    D_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    C_shape = (128, 1)
+    C_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        D = T.alloc_buffer(D_shape, dtype="float32", scope="trn.sbuf", layout=D_layout)
+        Tx.binary_reduce(B, C, A, D, "add", "sum", reduce_axes=1)
+
+        # fmt: off
+    with pytest.raises(Exception):
+        with target:
+            mod = tvm.IRModule({"main": tensor_scalar_reduce})
+            mod = tvm.tirx.transform.LowerTIRx()(mod)
+
+
+def test_tensor_scalar_reduce_complex():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (128 @ F, 1 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+    reduce_dst_shape = [128, 512]
+    reduce_dst_layout = TileLayout(S[(128, 4, 128) : (1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        D_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(C_sbuf, D_sbuf, B_sbuf, A_sbuf, "add", "sum", reduce_axes=0)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        D_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in range(512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorscalar_reduce(D_sbuf[p_loop, b_loop % 4 * 128 + b_loop // 4], C_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], A_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], B_sbuf[p_loop, b_loop % 4 * 128 + b_loop // 4], "add", "add", T.bool(True))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_tensor_scalar_reduce_two_stage():
+    src1_shape = [512, 1024, 4]
+    src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
+    dst1_shape = src1_shape
+    dst1_layout = src1_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        partial_reduce = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.tensorscalar_reduce(partial_reduce[p_loop, reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], T.float32(1.0), "add", "add", T.bool(False))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C_sbuf[p_loop, b_loop], partial_reduce[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_vector_chain():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
+    src3_shape = [512]
+    src3_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        _C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        D_sbuf = T.alloc_buffer(src3_shape, "float32", scope="trn.sbuf", layout=src3_layout)
+        E_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.binary_chain(E_sbuf, A_sbuf, B_sbuf, D_sbuf, "add", "add", reverse1=True)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        _C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        D_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        E_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.scalar_tensor_scalar(E_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4], "add", "add", T.bool(False), T.bool(True))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_vector_chain_2():
+    src1_shape = [32, 128, 512]
+    src1_layout = TileLayout(S[(32, 128, 4, 128) : (1 @ F, 32 @ F, 32 * 128 @ F, 1 @ P)])
+    src2_shape = [128, 512]
+    src2_layout = TileLayout(S[(512, 128) : (1 @ F, 1 @ P)])
+    src3_shape = src1_shape
+    src3_layout = src1_layout
+    dst_shape = src1_shape
+    dst_layout = src1_layout
+
+    # fmt: off
+    @T.prim_func
+    def binary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        _C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        D_sbuf = T.alloc_buffer(src3_shape, "float32", scope="trn.sbuf", layout=src3_layout)
+        E_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.binary_chain(E_sbuf, A_sbuf, B_sbuf, D_sbuf, "add", "add", reverse1=True)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        _C_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        D_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        E_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.scalar_tensor_tensor(E_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32 + f_loop], "add", "add", T.bool(False), T.bool(True))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": binary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduce_negate():
+    src_shape = [128, 512, 4]
+    src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 4 @ F, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.reduce_negate(B_sbuf[:, i], A_sbuf[:, :, i], reduce_op="sum", reduce_axes=-2)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop, f_loop * 4 + i], "add", True, -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_reduce_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def binary_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        for j in range(4):
+            for i in range(4):
+                Tx.binary_reduce(B_sbuf[0:128*(j+1), 0:128*(i+1)], C_sbuf[0:128*(j+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], 0.0, "add", "sum", [-1])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary_reduce"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop - j < 1 and f_loop < i * 128 + 128:
+                        T.nki.tensorscalar_reduce(C_sbuf[p_loop, b_loop], B_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0), "add", "add", T.bool(False))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_reduce_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def unary_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        for j in range(4):
+            for i in range(4):
+                Tx.unary_reduce(B_sbuf[0:128*(j+1), 0:128*(i+1)], C_sbuf[0:128*(j+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], "sqrt", "sum", reduce_axes=[-1])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary_reduce"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    if b_loop - j < 1 and f_loop < i * 128 + 128:
+                        T.nki.activation_reduce(C_sbuf[p_loop, b_loop], B_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0))  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_chain_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    src2_shape = [512, 1]
+    src2_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def binary_chain() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(src2_shape, "float32", scope="trn.sbuf", layout=src2_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            for i in range(4):
+                Tx.binary_chain(C_sbuf[0:128*(j+1), 0:128*(i+1)], A_sbuf[0:128*(j+1), 0:128*(i+1)], B_sbuf[0:128*(j+1), 0], 1.0, "add", "sub", reverse1=True)  # noqa: E501
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "binary_chain"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop - j < 1 and f_loop < i * 128 + 128:
+                        T.nki.scalar_tensor_scalar(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], B_sbuf[p_loop, b_loop], T.float32(1.0), "add", "sub", T.bool(False), T.bool(True))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": binary_chain})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_two_stage_workspace():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1), workspace={"partial_reduce": intermediate_buffer})  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        intermediate_buffer = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        C = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 1):
+            for reduction_b_loop in range(8):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.activation_reduce(intermediate_buffer[p_loop, reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop // 4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop], "sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C[p_loop, 0], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_tensor_scalar_reduce_two_stage_workspace():
+    src1_shape = [512, 1024, 4]
+    src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
+    dst1_shape = src1_shape
+    dst1_layout = src1_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2), workspace={"partial_reduce": intermediate_buffer})  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        intermediate_buffer = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                        T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop], T.float32(1.0), "add", "add", T.bool(False))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
+                    T.nki.tensorreduce(C_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_reduce_complex():
+    # fmt: off
+    @T.prim_func
+    def unary_reduce():
+        T.device_entry()
+        p = T.alloc_buffer((128, 8192), "float16", scope="trn.sbuf", layout="PF")
+        rowsum_p = T.alloc_buffer((2, 128, 1), scope="trn.sbuf", layout="FPF")
+        qk = T.alloc_buffer((2, 128, 8192), scope="trn.sbuf", layout="FPF")
+        running_max = T.alloc_buffer((16384, 1), dtype="float32", scope="trn.sbuf", layout="PF")
+        for i in range(4):
+            Tx.unary_reduce(p[0:128, 0:8192], rowsum_p[i % 2, 0:128, 0], qk[i % 2, 0:128, 0:8192], "exp", "sum", bias=running_max[i * 128:i * 128 + 128, 0])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary_reduce"})
+        p = T.alloc_buffer((128, 8192), "float16", scope="trn.sbuf")
+        rowsum_p = T.alloc_buffer((128, 2), scope="trn.sbuf")
+        qk = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        running_max = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(8192, annotations={"nki_dim": "F"}):
+                    T.nki.activation_reduce(rowsum_p[p_loop, i % 2], p[p_loop, f_loop], qk[p_loop, i % 2 * 8192 + f_loop], "exp", "add", running_max[p_loop, i], T.float32(1.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary_reduce})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py

@@ -0,0 +1,841 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_simple_copy():
+    src_shape = [128, 512]
+    src_layout = T.TileLayout(T.S[(128, 512) : (512, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (128, 512), layout=None)
+        A_1 = T.decl_buffer((65536,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(A_sbuf[p_loop, f_loop], A_1[p_loop * 512 + f_loop])
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_simple_copy_2():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 4, 128) : (512, 128, 1)])
+
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 4, 128) : (4 @ F, 1 @ F, 1 @ P)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (128, 512), layout=None)
+        A_1 = T.decl_buffer((65536,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 512):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 1, annotations={"nki_dim": "F"}):
+                    T.nki.load(A_sbuf[p_loop, b_loop], A_1[b_loop * 128 + p_loop])
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_in_a_loop():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], A[i * 128 : i * 128 + 128, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf[p_loop, i * 512 + f_loop], A_1[i * 65536 + p_loop * 512 + f_loop]
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_in_a_loop_2():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(128, 2048) : (2048, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        A_sbuf_view = A_sbuf.view(128, 4, 512)
+        A_view = A.view(128, 4, 512)
+        for i in range(4):
+            Tx.copy(A_sbuf_view[:, i, :], A_view[:, i, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        _A_flat = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        A_sbuf_view = T.decl_buffer((128, 2048), data=A_sbuf.data, scope="trn.sbuf", layout=None)
+        A_view = T.decl_buffer((262144,), data=A.data, layout=None)
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf_view[p_loop, i * 512 + f_loop],
+                        A_view[p_loop * 2048 + i * 512 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod.show()
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for b_loop in range(16):
+            for extend_b_loop in range(1):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            T.nki.matmul(acc_psum[b_loop % 8, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop], acc_psum[b_loop % 8, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_2():
+    src_shape = [65536]
+    src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [4, 65536]
+    dst_layout = TileLayout(S[(4, 128, 128, 4) : (4 @ F, 16 @ F, 1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(B_sbuf[i, :], A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i in range(4):
+            for b_loop in range(4):
+                for extend_b_loop in range(1):
+                    T.attr(0, "tensorized_nki_instruction", 1)
+                    for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                        for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                            for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                                T.nki.matmul(acc_psum[b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_f_loop * 4 + b_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                        T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + i * 4 + b_loop], acc_psum[b_loop, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_different_f():
+    src_shape = [512, 64]
+    src_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 16 @ F, 4 @ F, 1 @ F)])
+    dst_shape = [512, 64]
+    dst_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 4 @ F, 16 @ F, 1 @ F)])
+
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        for b_loop in T.serial(0, 64):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(
+                        B_sbuf[
+                            p_loop,
+                            b_loop // 16 * 64 + b_loop % 4 * 16 + b_loop % 16 // 4 * 4 + f_loop,
+                        ],
+                        A_sbuf[p_loop, b_loop * 4 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_different_shape():
+    src_shape = [512, 64]
+    src_layout = TileLayout(S[(4, 128, 4, 4, 4) : (64 @ F, 1 @ P, 16 @ F, 4 @ F, 1 @ F)])
+    dst_shape = [4, 128, 4]
+    dst_layout = TileLayout(S[(4, 128, 4) : (4 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        B_sbuf_view = B_sbuf.view(512, 4)
+        Tx.copy(B_sbuf_view, A_sbuf[:, 0:4])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        _B_sbuf_view = T.decl_buffer((128, 16), data=B_sbuf.data, scope="trn.sbuf", layout=None)
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(
+                        B_sbuf[p_loop, b_loop * 4 + f_loop],
+                        A_sbuf[p_loop, b_loop * 64 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_irregular_shape():
+    src_shape = [128, 10000]
+    src_layout = TileLayout(S[(128, 10000) : (10000, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A[:, i * 512 : i * 512 + 512], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (128, 10000), layout=None)
+        A_1 = T.decl_buffer((1280000,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.store(A_1[p_loop * 10000 + i * 512 + f_loop], A_sbuf[p_loop, f_loop])
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_different_shape_dim():
+    src_shape = [32, 128, 512]
+    src_layout = TileLayout(S[(32, 128, 512) : (128 * 512, 128, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(32):
+            Tx.copy(A_sbuf, A[i, :, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (32, 128, 512), layout=None)
+        A_1 = T.decl_buffer((2097152,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for i, b_loop in T.grid(32, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.load(A_sbuf[p_loop, f_loop], A_1[i * 65536 + p_loop * 128 + f_loop])
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_offset():
+    src_shape = [256, 512]
+    src_layout = TileLayout(S[(256, 512) : (512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(2):
+            Tx.copy(A_sbuf[i * 256 : i * 256 + 256, :], A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (256, 512), layout=None)
+        A_1 = T.decl_buffer((131072,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 2):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf[p_loop, i * 1024 + b_loop * 512 + f_loop],
+                        A_1[b_loop * 65536 + p_loop * 512 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_large_dma_copy():
+    src_shape = [512, 4096]
+    src_layout = T.TileLayout(T.S[(4, 128, 4096) : (4096 * 128, 4096, 1)])
+    dst_shape = [512, 4096]
+    dst_layout = TileLayout(S[(4, 128, 4096) : (4096 @ F, 1 @ P, 1 @ F)])
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], A[i * 128 : i * 128 + 128, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 4096), layout=None)
+        A_1 = T.decl_buffer((2097152,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 4096, annotations={"nki_dim": "F"}):
+                    T.nki.load(
+                        A_sbuf[p_loop, i * 4096 + f_loop],
+                        A_1[i * 524288 + p_loop * 4096 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_inst_size_limit():
+    src_shape = [512, 4096]
+    src_layout = dst_layout = TileLayout(S[(4, 128, 4096) : (4096 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        B_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.copy(A_sbuf[i * 128 : i * 128 + 128, :], B_sbuf[i * 128 : i * 128 + 128, :])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+        B_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(
+                        A_sbuf[p_loop, i * 4096 + b_loop * 512 + f_loop],
+                        B_sbuf[p_loop, i * 4096 + b_loop * 512 + f_loop],
+                    )
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_complex_index():
+    A_shape = [4096, 4096]
+    A_layout = T.TileLayout(T.S[(4096, 4096) : (1, 4096)])
+    A_sbuf_shape = (2, 2048, 1024)
+    A_sbuf_layout = TileLayout(S[(2, 2048, 8, 128) : (16384 @ F, 1 @ F, 2048 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle, ) -> None:
+        A = T.match_buffer(A_ptr, A_shape, "float32", layout=A_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(A_sbuf_shape, "float32", scope="trn.sbuf", layout=A_sbuf_layout)
+        Tx.copy(A_sbuf[1, 0:2048, 0:1024], A[2048: 4096, 3072:4096])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (4096, 4096), layout=None)
+        A_1 = T.decl_buffer((16777216,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 32768), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
+                    T.nki.load(A_sbuf[p_loop, b_loop * 2048 + f_loop + 16384], A_1[b_loop * 524288 + p_loop * 4096 + f_loop + 12584960])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_complex_index_2():
+    A_sbuf_shape = [4096, 4096]
+    A_sbuf_layout = T.TileLayout(T.S[(4096, 32, 128) : (1 @ F, 4096 @ F, 1 @ P)])
+    A_shape = (2, 2048, 1024)
+    A_layout = T.TileLayout(T.S[(2, 2048, 1024) : (2048 * 1024, 1, 2048)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle, ) -> None:
+        A = T.match_buffer(A_ptr, A_shape, "float32", layout=A_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(A_sbuf_shape, "float32", scope="trn.sbuf", layout=A_sbuf_layout)
+        Tx.copy(A_sbuf[2048: 4096, 3072:4096], A[1, 0:2048, 0:1024])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (2, 2048, 1024), layout=None)
+        A_1 = T.decl_buffer((4194304,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 131072), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
+                    T.nki.load(A_sbuf[p_loop, b_loop * 4096 + f_loop + 100352], A_1[b_loop * 262144 + p_loop * 2048 + f_loop + 2097152])  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_with_workspace():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        identity = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((1, 128, 512), "float32", scope="trn.psum", allocated_addr=(0, 0))
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        Tx.copy(B_sbuf, A_sbuf, workspace={"identity": identity, "acc_psum": acc_psum})
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        for b_loop in range(16):
+            for extend_b_loop in range(1):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            T.nki.matmul(acc_psum[0, lhs_f_loop, extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop], acc_psum[0, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_guard():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            for i in range(4):
+                Tx.copy(A_sbuf[i * 128 : i * 128 + 128, 0:128*j], A[i * 128 : i * 128 + 128, 0:128*j])  # noqa: E501
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
+                    if f_loop < j * 128:
+                        T.nki.load(A_sbuf[p_loop, i * 512 + f_loop], A_1[i * 65536 + p_loop * 512 + f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_guard_2():
+    src_shape = [512, 512]
+    src_layout = T.TileLayout(T.S[(4, 128, 512) : (512 * 128, 512, 1)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for j in range(4):
+            for i in range(4):
+                Tx.copy(A_sbuf[0:128*j, 0:128*i], A[0:128*j, 0:128*i])
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+
+        A = T.match_buffer(A_ptr, (512, 512), layout=None)
+        A_1 = T.decl_buffer((262144,), data=A.data, layout=None)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for j, i, b_loop in T.grid(4, 4, 3):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
+                    if b_loop - j < 0 and f_loop < i * 128:
+                        T.nki.load(A_sbuf[p_loop, b_loop * 512 + f_loop], A_1[b_loop * 65536 + p_loop * 512 + f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_with_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.copy(B_sbuf[i * 128 : i * 128 + 128, 0:128*j], A_sbuf[i * 128 : i * 128 + 128, 0:128*j])  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, j, b_loop in T.grid(4, 4, 3):
+            for extend_b_loop in range(1):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            if b_loop - j < 0:
+                                T.nki.matmul(acc_psum[b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, i * 512 + b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    if b_loop - j < 0:
+                        T.nki.tensor_copy(B_sbuf[p_loop, i * 512 + f_loop * 4 + b_loop], acc_psum[b_loop, p_loop, f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_with_specified_max_inst_size():
+    src_shape = [128, 512]
+    src_layout = "PF"
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, B_sbuf, max_inst_size=128)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf", layout=None)
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf", layout=None)
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 128 + f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_copy_transpose_with_extended_f():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 2048), "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer((128, 2048), "float32", scope="trn.sbuf", layout="FP")
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle):
+        T.func_attr({"global_symbol": "copy"})
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for b_loop in range(4):
+            for extend_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for lhs_f_loop in T.serial(128, annotations={"nki_dim": "lhs_F"}):
+                        for rhs_f_loop in T.serial(128, annotations={"nki_dim": "rhs_F"}):
+                            T.nki.matmul(acc_psum[b_loop, lhs_f_loop, extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 512 + extend_b_loop * 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True))  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.tensor_copy(B_sbuf[p_loop, b_loop * 512 + f_loop], acc_psum[b_loop, p_loop, f_loop])  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_gemm_trn.py

@@ -0,0 +1,583 @@
+# 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.
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_simple_gemm():
+    A_layout = TileLayout(S[(128, 128) : (1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
+
+    C_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((128, 128), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((1, 128, 128), scope="trn.psum")
+        for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(1, 1, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_f_loop], B_sbuf[p_loop, rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_larger_gemm():
+    A_layout = TileLayout(S[(2, 128, 4, 128) : (512 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(2, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((256, 512), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((256, 256), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum")
+        for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[0, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, lhs_b_loop * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_in_a_loop():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_with_stride():
+    A_layout = TileLayout(S[(4, 128, 128, 8) : (1024 @ F, 1 @ F, 1 @ P, 128 @ F)])
+    B_layout = TileLayout(S[(128, 8, 2, 128) : (1 @ P, 512 @ F, 256 @ F, 2 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 512, 2), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((512, 2, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, :, k],
+                    B_sbuf[:, k, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4095), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + reduction_b_loop * 256 + k * 128 + lhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 1024 + k * 512 + rhs_f_loop * 2], True)  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_swap_lhs_rhs():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 2, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_with_sbuf_output():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        buffer = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                  for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                    for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                        T.nki.matmul(buffer[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], buffer[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_different_shape():
+    A_layout = TileLayout(S[(2, 4, 128, 8, 128) : (4096 @ F, 1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((2, 512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[1, 256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 2, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop + 4096], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_too_large_f_size():
+    A_layout = TileLayout(S[(256, 128) : (1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(128, 1024) : (1 @ P, 1 @ F)])
+
+    C_layout = TileLayout(S[(2, 128, 1024) : (1024 @ F, 1 @ P, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((256, 128), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((128, 1024), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((256, 1024), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 256), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((4, 128, 512), scope="trn.psum")
+        for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 512, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], A_sbuf[p_loop, lhs_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, rhs_b_loop * 512 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_sbuf_output_with_workspace():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        C_psum = T.alloc_buffer((1, 128, 512), "float32", scope="trn.psum", allocated_addr=(0, 0))
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    workspace={"acc_psum": C_psum}
+                )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((1, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
+            for reduction_b_loop in range(4):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                  for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                    for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                        T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], C_psum[0, lhs_f_loop, rhs_f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_pf_mismatch_fail():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(2, 128, 8, 128) : (128 @ F, 1 @ F, 256 @ F, 1 @ P)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((256, 1024), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[:, 512 * k : 512 * k + 512],
+                    C_psum[256 * i : 256 * i + 256, :],
+                )
+        # fmt: on
+    with pytest.raises(Exception):
+        with target:
+            mod = tvm.IRModule({"main": gemm})
+            mod = tvm.tirx.transform.LowerTIRx()(mod)
+
+
+def test_gemm_transpose_AB():
+    A_layout = TileLayout(S[(8, 128, 4, 128) : (128 @ F, 1 @ P, 1024 @ F, 1 @ F)])
+    B_layout = TileLayout(S[(2, 128, 8, 128) : (128 @ F, 1 @ F, 256 @ F, 1 @ P)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((1024, 512), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((256, 1024), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_psum[256 * i : 256 * i + 256, :],
+                    A_sbuf[512 * k : 512 * k + 512, 256 * i : 256 * i + 256],
+                    B_sbuf[:, 512 * k : 512 * k + 512],
+                    C_psum[256 * i : 256 * i + 256, :],
+                    transpose_A=True,
+                    transpose_B=True,
+                )
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+
+            #fmt: off
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_guard():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            for j in range(2):
+                for k in range(2):
+                    Tx.gemm(
+                        C_sbuf[0: 256 * i, 0: 128 * (j + 1)],
+                        A_sbuf[0: 256 * i, 0: 512 * (k + 1)],
+                        B_sbuf[0: 512 * (k + 1), 0: 128 * (j + 1)],
+                        C_sbuf[0: 256 * i, 0: 128 * (j + 1)],
+                    )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        for i, j, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2, 2):
+            for reduction_b_loop in range(8):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                  for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                    for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"rhs_F"}):
+                        if reduction_b_loop - k * 4 < 4 and lhs_b_loop - j < 1 and 0 < i and reduction_b_loop - k * 4 < 4:  # noqa: E501
+                            T.nki.matmul(acc_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, rhs_b_loop * 1024 + reduction_b_loop * 128 + rhs_f_loop], True)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
+                if 0 < i and lhs_b_loop - j < 1:
+                    T.nki.tensor_copy(C_sbuf[lhs_f_loop, rhs_b_loop * 256 + lhs_b_loop * 128 + rhs_f_loop], acc_psum[lhs_b_loop * 2 + rhs_b_loop, lhs_f_loop, rhs_f_loop])  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm_guard2():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 128 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)]).to_psum()
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((512, 256), "float32", scope="trn.psum", layout=C_layout)
+        for j in range(4):
+            for i in range(2):
+                for k in range(2):
+                    Tx.gemm(
+                        C_psum[256 * i : 256 * i + 256, :],
+                        A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + (j+1) * 128],
+                        B_sbuf[512 * k : 512 * k + (j+1) * 128, :],
+                        C_psum[256 * i : 256 * i + 256, :],
+                    )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        C_psum = T.alloc_buffer((2, 128, 512), scope="trn.psum")
+        for j, i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(4, 2, 2, 2, 1, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+              for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"lhs_F"}):
+                for rhs_f_loop in T.serial(0, 256, annotations={"nki_dim":"rhs_F"}):
+                    if reduction_b_loop - j < 1 and reduction_b_loop - j < 1:
+                        T.nki.matmul(C_psum[i, lhs_f_loop, lhs_b_loop * 256 + rhs_f_loop], A_sbuf[p_loop, i * 2048 + lhs_b_loop * 1024 + k * 512 + reduction_b_loop * 128 + lhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + rhs_f_loop], True)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_private_alloc_trn.py

@@ -0,0 +1,402 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.trn.transform import TrnPrivateBufferAlloc
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def test_copy_transpose():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(128, 2048) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = TileLayout(S[(2048, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def copy() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        identity = T.alloc_buffer((128, 128), scope="trn.sbuf")
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
+                    T.nki.identity(identity[p_loop, rhs_f_loop], 128)
+        A_sbuf = T.alloc_buffer((512, 512), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 2048) : (1 @ P, 1@F)]))
+        B_sbuf = T.alloc_buffer((512, 512), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(2048, 128) : (1@F, 1@P)]))
+        Tx.copy(B_sbuf[0:512, 0:512], A_sbuf[0:512, 0:512], workspace={"acc_psum": acc_psum, "identity": identity})  # noqa: E501
+
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_normal_copy():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (512, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": copy})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], copy)
+
+
+def test_unary_with_bias_scale():
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias = T.float32(1.0)
+    scale = T.float32(2.0)
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.exp(C_sbuf, A_sbuf, bias=bias, scale=scale)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        T.device_entry()
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
+        A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1@P, 1@F)]))
+        C_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1@P, 1@F)]))
+        Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(1.0), T.float32(2.0), workspace={"const_bias": const_bias})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_two_stage():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 32), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 32, 4, 32), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 32 * 32 * 4) : (1@P, 1@F)]))
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4) : (1@P, 1@F)]))
+        Tx.sum(B_sbuf[0:128, 0:4], A_sbuf[0:128, 0:32, 0:4, 0:32], [1, 3], False, workspace={"partial_reduce": partial_reduce})  # noqa: E501
+
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_gemm():
+    A_layout = TileLayout(S[(4, 128, 8, 128) : (1024 @ F, 1 @ F, 1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(8, 128, 2, 128) : (256 @ F, 1 @ P, 128 @ F, 1 @ F)])
+
+    C_layout = TileLayout(S[(4, 128, 2, 128) : (256 @ F, 1 @ F, 128 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((512, 1024), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((1024, 256), "float32", scope="trn.sbuf", layout=B_layout)
+        C_sbuf = T.alloc_buffer((512, 256), "float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            for k in range(2):
+                Tx.gemm(
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                    A_sbuf[256 * i : 256 * i + 256, 512 * k : 512 * k + 512],
+                    B_sbuf[512 * k : 512 * k + 512, :],
+                    C_sbuf[256 * i : 256 * i + 256, :],
+                )
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "gemm"})
+        T.device_entry()
+        acc_psum = T.alloc_buffer((8, 128, 512), scope="trn.psum", allocated_addr=[0, 0])
+        A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(4, 128, 8, 128) : (1024@F, 1@F, 1@F, 1@P)]))  # noqa: E501
+        B_sbuf = T.alloc_buffer((1024, 256), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(8, 128, 2, 128) : (256@F, 1@P, 128@F, 1@F)]))  # noqa: E501
+        C_sbuf = T.alloc_buffer((512, 256), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(4, 128, 2, 128) : (256@F, 1@F, 128@F, 1@P)]))  # noqa: E501
+        for i, k in T.grid(2, 2):
+            Tx.gemm(C_sbuf[256 * i:256 * i + 256, 0:256], A_sbuf[256 * i:256 * i + 256, 512 * k:512 * k + 512], B_sbuf[512 * k:512 * k + 512, 0:256], C_sbuf[256 * i:256 * i + 256, 0:256], False, False, T.float32(1.0), T.float32(0.0), workspace={"acc_psum": acc_psum})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_binary_reduce_two_stage():
+    src1_shape = [512, 1024, 4]
+    src1_layout = TileLayout(S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)])
+    dst1_shape = src1_shape
+    dst1_layout = src1_layout
+    reduce_dst_shape = [512]
+    reduce_dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def tensor_scalar_reduce() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src1_shape, "float32", scope="trn.sbuf", layout=src1_layout)
+        B_sbuf = T.alloc_buffer(dst1_shape, "float32", scope="trn.sbuf", layout=dst1_layout)
+        C_sbuf = T.alloc_buffer(reduce_dst_shape, "float32", scope="trn.sbuf", layout=reduce_dst_layout)  # noqa: E501
+        Tx.binary_reduce(B_sbuf, C_sbuf, A_sbuf, 1.0, "add", "sum", reduce_axes=(1, 2))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "tensor_scalar_reduce"})
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((512, 1024, 4), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)]))
+        B_sbuf = T.alloc_buffer((512, 1024, 4), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096, 4) : (1 @ P, 1 @ F, 4096 @ F)]))
+        C_sbuf = T.alloc_buffer((512,), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4) : (1 @ P, 1 @ F)]))
+        Tx.binary_reduce(B_sbuf[0:512, 0:1024, 0:4], C_sbuf[0:512], A_sbuf[0:512, 0:1024, 0:4], T.float32(1.0), "add", "sum", [1, 2], workspace={"partial_reduce": partial_reduce})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": tensor_scalar_reduce})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_activation_reduce_two_stage():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 8), scope="trn.sbuf")
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A = T.alloc_buffer((32, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(16 * 1024, 128) : (1@F, 1@P)]))
+        B = T.alloc_buffer((16, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(2, 4, 1024, 128) : (1024@F, 2048@F, 1@F, 1@P)]))
+        C = T.alloc_buffer((1, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
+        for i in range(2):
+            Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i * 16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1], workspace={"const_bias": const_bias, "partial_reduce": partial_reduce})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_partial_workspace_specify():
+    A_shape = (32, 512, 128)
+    A_layout = TileLayout(S[(16 * 1024, 128) : (1 @ F, 1 @ P)])
+    B_shape = (16, 512, 128)
+    B_layout = TileLayout(S[(2, 4, 1024, 128) : (1024 @ F, 2048 @ F, 1 @ F, 1 @ P)])
+    C_shape = (1, 128)
+    C_layout = TileLayout(S[(1, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def activation_reduce():
+        T.device_entry()
+        partial_reduce = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer(A_shape, dtype="float32", scope="trn.sbuf", layout=A_layout)
+        B = T.alloc_buffer(B_shape, dtype="float32", scope="trn.sbuf", layout=B_layout)
+        C = T.alloc_buffer(C_shape, dtype="float32", scope="trn.sbuf", layout=C_layout)
+        for i in range(2):
+            Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum", reduce_axes=(0,1), workspace={"partial_reduce": partial_reduce})  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "activation_reduce"})
+        T.device_entry()
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        partial_reduce = T.alloc_buffer((128, 16), scope="trn.sbuf")
+        A = T.alloc_buffer((32, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(16 * 1024, 128) : (1@F, 1@P)]))
+        B = T.alloc_buffer((16, 512, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(2, 4, 1024, 128) : (1024@F, 2048@F, 1@F, 1@P)]))
+        C = T.alloc_buffer((1, 128), scope="trn.sbuf",
+                           layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
+        for i in range(2):
+            Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i * 16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1], workspace={"const_bias": const_bias, "partial_reduce": partial_reduce})  # noqa: E501
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": activation_reduce})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_workspace_reuse():
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    scale = T.float32(2.0)
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.exp(C_sbuf, A_sbuf, bias=0.0, scale=scale, max_inst_size=1024)
+        Tx.exp(C_sbuf, C_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        T.device_entry()
+        const_bias = T.alloc_buffer((128, 1024), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
+        A_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)]))
+        C_sbuf = T.alloc_buffer((512, 1024), scope="trn.sbuf",
+                                layout=T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)]))
+        Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(0.0), T.float32(2.0), workspace={"const_bias": const_bias}, max_inst_size=1024)  # noqa: E501
+        Tx.exp(C_sbuf[0:512, 0:1024], C_sbuf[0:512, 0:1024], None, None, workspace={"const_bias": const_bias})  # noqa: E501
+
+        # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_no_rewrite_with_existing_workspace():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3), workspace={"partial_reduce": intermediate_buffer})
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], reduction)
+
+
+def test_no_rewrite_with_psum_output():
+    A_layout = TileLayout(S[(128, 128) : (1 @ F, 1 @ P)])
+    B_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
+
+    C_layout = TileLayout(S[(128, 128) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def gemm() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=A_layout)
+        B_sbuf = T.alloc_buffer((128, 128), "float32", scope="trn.sbuf", layout=B_layout)
+        C_psum = T.alloc_buffer((128, 128), "float32", scope="trn.psum", layout=C_layout)
+        Tx.gemm(C_psum, A_sbuf, B_sbuf, C_psum)
+        # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": gemm})
+        mod = TrnPrivateBufferAlloc()(mod)
+        assert_structural_equal(mod["main"], gemm)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_reduction_trn.py

@@ -0,0 +1,283 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+opcode_map = {"sum": "add", "max": "max", "min": "min"}
+
+Tx_func_map = {"sum": Tx.sum, "max": Tx.max, "min": Tx.min}
+
+
+@pytest.mark.parametrize("op_type", ["sum", "max", "min"])
+def test_simple_reduction(op_type):
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 1]
+    dst_layout = TileLayout(S[(128, 1) : (1 @ P, 1 @ F)])
+
+    opcode = opcode_map[op_type]
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def reduction() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        tx_func(B_sbuf, A_sbuf, axes=-1)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop, f_loop], opcode, False, -1)
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_with_multiple_axes():
+    src_shape = [128, 512, 4]
+    src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 1 @ F, 512 @ F)])
+    dst_shape = [128]
+    dst_layout = TileLayout(S[128 : 1 @ P])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 2), max_inst_size=2048)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 1), scope="trn.sbuf")
+        for b_loop in range(1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop, f_loop], "add", False, -1)
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_in_loop():
+    src_shape = [128, 512, 4]
+    src_layout = TileLayout(S[(128, 512, 4) : (1 @ P, 4 @ F, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            Tx.sum(B_sbuf[:, i], A_sbuf[:, :, i], axes=-2)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop, f_loop * 4 + i], "add", False, -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_two_stage():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        intermediate_buffer = T.alloc_buffer((128, 32), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(32):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                    for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                        T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 + f_loop], "add", False, -1)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", False, -1)  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_with_guard():
+    src_shape = [512, 2048]
+    src_layout = TileLayout(S[(4, 128, 2048) : (2048 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [512, 1]
+    dst_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.sum(B_sbuf[0: (i+1) * 128, 0], A_sbuf[0: (i+1) * 128, 0: (j+1) * 256], max_inst_size=512)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        intermediate_buffer = T.alloc_buffer((128, 2), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for i, j in T.grid(4, 4):
+            for b_loop in range(4):
+                for reduction_b_loop in range(2):
+                    T.attr(0, "tensorized_nki_instruction", 1)
+                    for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                        for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                            if (
+                                b_loop - i < 1
+                                and reduction_b_loop * 512 + f_loop < j * 256 + 256
+                            ):
+                                T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, b_loop * 2048 + reduction_b_loop * 512 + f_loop], "add", T.bool(False), -1)  # noqa: E501
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                    for f_loop in T.serial(2, annotations={"nki_dim": "F"}):
+                        if b_loop - i < 1 and f_loop * 2 - j < 1:
+                            T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1)  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_reduction_two_stage_workspace():
+    src_shape = [128, 32, 4, 32]
+    src_layout = TileLayout(S[(128, 32 * 32 * 4) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 4]
+    dst_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def reduction():
+        T.device_entry()
+        intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.sum(B_sbuf, A_sbuf, axes=(1, 3), workspace={"partial_reduce": intermediate_buffer})
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "reduction"})
+        intermediate_buffer = T.alloc_buffer((128, 64), scope="trn.sbuf")
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        for b_loop in range(4):
+            for reduction_b_loop in range(32):
+                T.attr(0, "tensorized_nki_instruction", 1)
+                for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                    for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                        T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 + f_loop], "add", False, -1)  # noqa: E501
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
+                    T.nki.tensorreduce(B_sbuf[p_loop, b_loop], intermediate_buffer[p_loop, f_loop], "add", False, -1)  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": reduction})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_select_trn.py

@@ -0,0 +1,186 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+def test_select():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.select(B_sbuf, A_sbuf, 0.0, lambda i, j: i < j)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.affine_select(B_sbuf[p_loop, f_loop], p_loop < f_loop, A_sbuf[p_loop, f_loop], T.float32(0.0))  # noqa: E501
+            # fmt: on
+
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_select_in_loop():
+    src_shape = [32, 128, 512]
+    src_layout = TileLayout(S[(32, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(2):
+            Tx.select(B_sbuf, A_sbuf[i*16, :, :], 0.0, lambda a, b: (i+1)* a < b)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 16384), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for i, b_loop in T.grid(2, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.affine_select(B_sbuf[p_loop, f_loop], (i + 1) * p_loop < f_loop, A_sbuf[p_loop, i * 8192 + f_loop], T.float32(0.0))  # noqa: E501
+
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_select_expr_affine():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.select(B_sbuf, A_sbuf, 0.0, lambda i, j: i < j)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for b_loop in T.serial(0, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop], b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_select_with_guard():
+    src_shape = [512, 512]
+    src_layout = TileLayout(S[(4, 128, 512) : (512 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+
+    # fmt: off
+    @T.prim_func
+    def select() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.select(B_sbuf[0: (i+1) * 128, 0: (j+1) * 128], A_sbuf[0: (i+1) * 128, 0: (j+1) * 128], 0.0, lambda a, b: a < b)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "select"})
+        A_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        for i, j, b_loop in T.grid(4, 4, 4):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop - i < 1 and f_loop < j * 128 + 128:
+                        T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop], b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))  # noqa: E501
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": select})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py

@@ -0,0 +1,288 @@
+# 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.
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal as _assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.stmt_functor import ir_transform
+
+target = tvm.target.Target("aws/trn1/trn1.2xlarge")
+
+
+def _strip_exec_scope_stmt(stmt):
+    def _postorder(node):
+        if isinstance(node, tvm.tirx.AttrStmt) and node.attr_key == "tirx.device_entry":
+            return node.body
+        return node
+
+    return ir_transform(
+        stmt,
+        preorder=lambda _node: None,
+        postorder=_postorder,
+        only_enable=["tirx.AttrStmt"],
+    )
+
+
+def assert_structural_equal(lhs, rhs, *args, **kwargs):
+    if isinstance(lhs, tvm.tirx.PrimFunc):
+        lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+    if isinstance(rhs, tvm.tirx.PrimFunc):
+        rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+    _assert_structural_equal(lhs, rhs, *args, **kwargs)
+
+
+Tx_func_map = {"reciprocal": Tx.reciprocal, "sqrt": Tx.sqrt, "memset": Tx.memset, "exp": Tx.exp}
+
+
+@pytest.mark.parametrize("op_type", ["reciprocal", "memset"])
+def test_simple_unary(op_type):
+    src_shape = [128, 512]
+    src_layout = T.TileLayout(T.S[(128, 512) : (1 @ P, 1 @ F)])
+    dst_shape = [128, 512]
+    dst_layout = T.TileLayout(T.S[(128, 512) : (1 @ P, 1 @ F)])
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        if op_type == "memset":
+            tx_func(B_sbuf, T.float32(0.0))
+        else:
+            tx_func(B_sbuf, A_sbuf)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        for b_loop in T.serial(0, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if op_type == "reciprocal":
+                        T.nki.reciprocal(
+                            B_sbuf[p_loop, f_loop], A_sbuf[p_loop, f_loop]
+                        )
+                    elif op_type == "memset":
+                        T.nki.memset(B_sbuf[p_loop, f_loop], 0.0)
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["reciprocal", "memset"])
+def test_unary_in_a_loop(op_type):
+    src_shape = [1024, 512]
+    src_layout = T.TileLayout(T.S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = [512, 512]
+    dst_layout = T.TileLayout(T.S[(128, 2048) : (1 @ P, 1 @ F)])
+
+    Tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        A_sbuf_view = A_sbuf.view(128, 8, 512)
+        B_sbuf_view = B_sbuf.view(128, 4, 512)
+        for i in range(4):
+            if op_type == "memset":
+                Tx_func(B_sbuf_view[:, i, :], T.float32(0.0))
+            else:
+                Tx_func(B_sbuf_view[:, i, :], A_sbuf_view[:, i * 2, :])
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 2048), scope="trn.sbuf")
+        A_sbuf_view = T.decl_buffer((128, 4096), data=A_sbuf.data, scope="trn.sbuf", layout=None)
+        B_sbuf_view = T.decl_buffer((128, 2048), data=B_sbuf.data, scope="trn.sbuf", layout=None)
+        for i, b_loop in T.grid(4, 1):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if op_type == "reciprocal":
+                        T.nki.reciprocal(B_sbuf_view[p_loop, i * 512 + f_loop], A_sbuf_view[p_loop, i * 1024 + f_loop])  # noqa: E501
+                    elif op_type == "memset":
+                        T.nki.memset(B_sbuf[p_loop, i * 512 + f_loop], 0.0)
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_complex1():
+    dst_layout = TileLayout(S[(32, 128, 256) : (256 @ F, 1 @ P, 1 @ F)])
+    dst_shape = [4096, 256]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.memset(A_sbuf, T.float32(0.0))
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 8192), scope="trn.sbuf")
+        for b_loop in T.serial(0, 16):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.memset(A_sbuf[p_loop, b_loop * 512 + f_loop], T.float32(0.0))
+            # fmt: on
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["sqrt", "exp"])
+def test_unary_with_bias_scale(op_type):
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias_shape = [512, 1]
+    bias_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
+    scale = T.float32(2.0)
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(bias_shape, "float32", scope="trn.sbuf", layout=bias_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        tx_func(C_sbuf, A_sbuf, bias=B_sbuf, scale=scale)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    T.nki.activation(C_sbuf[p_loop, b_loop * 512  + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, B_sbuf[p_loop, b_loop//2], T.float32(2.0))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+@pytest.mark.parametrize("op_type", ["sqrt", "exp"])
+def test_unary_with_bias_scale_2(op_type):
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(128, 4096) : (1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias = T.float32(1.0)
+    scale = T.float32(2.0)
+    tx_func = Tx_func_map[op_type]
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        tx_func(C_sbuf, A_sbuf, bias=bias, scale=scale)
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        const_bias = T.alloc_buffer((128, 512), scope="trn.sbuf")
+        with T.attr(0, "tensorized_nki_instruction", 1):
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        for b_loop in T.serial(0, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
+                for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
+                    T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, const_bias[p_loop, f_loop], T.float32(2.0))  # noqa: E501
+            # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.trn.transform.TrnPrivateBufferAlloc()(mod)
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+def test_unary_with_guard():
+    src_shape = [512, 1024]
+    src_layout = TileLayout(S[(4, 128, 1024) : (1024 @ F, 1 @ P, 1 @ F)])
+    dst_shape = src_shape
+    dst_layout = src_layout
+    bias_shape = [512, 1]
+    bias_layout = TileLayout(S[(4, 128) : (1 @ F, 1 @ P)])
+    scale = T.float32(2.0)
+
+    # fmt: off
+    @T.prim_func
+    def unary() -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(src_shape, "float32", scope="trn.sbuf", layout=src_layout)
+        B_sbuf = T.alloc_buffer(bias_shape, "float32", scope="trn.sbuf", layout=bias_layout)
+        C_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        for i in range(4):
+            for j in range(4):
+                Tx.sqrt(C_sbuf[0: (i+1) * 128, 0: (j+1)*256], A_sbuf[0: (i+1) * 128, 0: (j+1)*256], bias=B_sbuf[0: (i+1) * 128, 0], scale=scale)  # noqa: E501
+
+    @T.prim_func
+    def expected():
+        T.func_attr({"global_symbol": "unary"})
+        A_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        B_sbuf = T.alloc_buffer((128, 4), scope="trn.sbuf")
+        C_sbuf = T.alloc_buffer((128, 4096), scope="trn.sbuf")
+        for i, j, b_loop in T.grid(4, 4, 8):
+            T.attr(0, "tensorized_nki_instruction", 1)
+            for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
+                for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
+                    if b_loop // 2 - i < 1 and b_loop % 2 * 512 + f_loop < j * 256 + 256:
+                        T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop], A_sbuf[p_loop, b_loop * 512 + f_loop], "sqrt", B_sbuf[p_loop, b_loop // 2], T.float32(2.0))  # noqa: E501
+             # fmt: off
+    with target:
+        mod = tvm.IRModule({"main": unary})
+        mod = tvm.tirx.transform.LowerTIRx()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
+        assert_structural_equal(mod["main"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/test_alloc_pool.py

@@ -0,0 +1,117 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for CUDA allocation pool validation."""
+
+import pytest
+
+from tvm.tirx.cuda.lang.alloc_pool import _validate_mma_alloc_shape
+from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+
+# ---------------------------------------------------------------------------
+# alloc_mma shape validation: bad inputs raise actionable ValueError instead of
+# the opaque "Divide by zero" diagnostic that ``Layout.tile_to`` would emit.
+# ---------------------------------------------------------------------------
+
+
+class TestAllocMmaValidationRowBytes:
+    """row width (cols * itemsize) must be a positive multiple of swizzle atom bytes."""
+
+    def test_bf16_32cols_128b_swizzle_too_narrow(self):
+        # The exact case that bit gdn-prefill v1_0 / v1_2 (eval R10).
+        # Row = 32 * 2B = 64B < 128B atom.
+        with pytest.raises(ValueError, match=r"64B rows.*128B swizzle atom"):
+            _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+    def test_error_suggests_smaller_swizzle(self):
+        try:
+            _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+        except ValueError as e:
+            assert "SWIZZLE_64B_ATOM" in str(e), f"missing fix-it hint: {e}"
+        else:
+            pytest.fail("should have raised")
+
+    def test_error_suggests_widening_cols(self):
+        try:
+            _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+        except ValueError as e:
+            assert "multiple of 64 elements" in str(e), f"missing widen hint: {e}"
+        else:
+            pytest.fail("should have raised")
+
+    def test_fp32_16cols_128b_swizzle_too_narrow(self):
+        # Row = 16 * 4B = 64B < 128B atom.
+        with pytest.raises(ValueError, match=r"64B rows.*128B swizzle atom"):
+            _validate_mma_alloc_shape((128, 16), "float32", SwizzleMode.SWIZZLE_128B_ATOM)
+
+    def test_3d_shape_validates_last_dim(self):
+        # Validation must consider shape[-1], not shape[0].
+        with pytest.raises(ValueError, match=r"64B rows"):
+            _validate_mma_alloc_shape((2, 128, 32), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+
+class TestAllocMmaValidationRowCount:
+    """rows (shape[-2]) must be a positive multiple of the 8-row atom."""
+
+    def test_rows_below_atom_rejected(self):
+        with pytest.raises(ValueError, match=r"shape\[-2\]=4.*multiple of 8"):
+            _validate_mma_alloc_shape((4, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+    def test_rows_not_multiple_of_8_rejected(self):
+        with pytest.raises(ValueError, match=r"shape\[-2\]=12.*multiple of 8"):
+            _validate_mma_alloc_shape((12, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+
+class TestAllocMmaValidationRank:
+    """rank-1 shapes cannot be tiled with a 2-D swizzle atom."""
+
+    def test_rank_one_rejected(self):
+        with pytest.raises(ValueError, match=r"fewer than 2 dimensions"):
+            _validate_mma_alloc_shape((128,), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM)
+
+
+class TestAllocMmaValidationValid:
+    """combinations that should succeed must not be rejected."""
+
+    @pytest.mark.parametrize(
+        "shape,dtype,mode",
+        [
+            # The fix path the agent should pick when row_bytes >= 128.
+            ((128, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM),
+            ((128, 128), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM),
+            # Or downgrade to a swizzle whose atom matches the row.
+            ((128, 32), "bfloat16", SwizzleMode.SWIZZLE_64B_ATOM),
+            ((128, 16), "bfloat16", SwizzleMode.SWIZZLE_32B_ATOM),
+            # 3-D request validates the last two dims only.
+            ((2, 128, 64), "bfloat16", SwizzleMode.SWIZZLE_128B_ATOM),
+            # fp32 with row width >= atom.
+            ((128, 32), "float32", SwizzleMode.SWIZZLE_128B_ATOM),
+            # fp8 (1B) with row width >= atom.
+            ((128, 128), "float8_e4m3", SwizzleMode.SWIZZLE_128B_ATOM),
+        ],
+    )
+    def test_valid_combinations_accepted(self, shape, dtype, mode):
+        _validate_mma_alloc_shape(shape, dtype, mode)
+
+    def test_swizzle_none_skips_validation(self):
+        # SWIZZLE_NONE has no atom — even otherwise-bad shapes are allowed.
+        _validate_mma_alloc_shape((128, 32), "bfloat16", SwizzleMode.SWIZZLE_NONE)
+        _validate_mma_alloc_shape((3, 5), "bfloat16", SwizzleMode.SWIZZLE_NONE)
+        _validate_mma_alloc_shape((128,), "bfloat16", SwizzleMode.SWIZZLE_NONE)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/python/tirx/test_bench_utils.py

@@ -0,0 +1,245 @@
+# 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 tvm.tirx.bench utilities."""
+
+import pytest
+import torch
+
+import tvm.testing
+
+pytest.importorskip("triton")  # tvm.tirx.bench imports triton.profiler
+
+from tvm.testing import env
+from tvm.tirx.bench import _compute_group_count, _parse_proton_tree, bench, tensor_bytes
+
+# ── _parse_proton_tree ──────────────────────────────────────────────────────
+
+
+SAMPLE_TREE = """\
+├─ 1.500 tir
+│  ├─ 1.500 my_kernel_fn
+│  └─ 0.001 vectorized_elementwise_kernel
+└─ 0.800 cublas
+   └─ 0.800 sm90_xmma_gemm_f16f16
+"""
+
+
+def test_parse_proton_tree_basic():
+    impls, errors = _parse_proton_tree(SAMPLE_TREE)
+    assert impls == {"tir": 1.5, "cublas": 0.8}
+    assert errors == {}
+
+
+def test_parse_proton_tree_filters_elementwise():
+    """vectorized_elementwise_kernel and elementwise_kernel_with_index are skipped."""
+    tree = """\
+├─ 0.500 tir
+│  ├─ 0.500 real_kernel
+│  └─ 0.001 elementwise_kernel_with_index
+"""
+    impls, _ = _parse_proton_tree(tree)
+    assert impls == {"tir": 0.5}
+
+
+def test_parse_proton_tree_slowest_child():
+    """Takes the slowest depth-2 child per impl."""
+    tree = """\
+├─ 2.000 tir
+│  ├─ 0.300 kernel_a
+│  └─ 0.700 kernel_b
+"""
+    impls, _ = _parse_proton_tree(tree)
+    assert impls == {"tir": 0.7}
+
+
+def test_parse_proton_tree_baseline_errors():
+    tree = """\
+BASELINE_ERROR: cublas: CUDA OOM
+├─ 1.000 tir
+│  └─ 1.000 my_kernel
+"""
+    impls, errors = _parse_proton_tree(tree)
+    assert impls == {"tir": 1.0}
+    assert errors == {"cublas": "CUDA OOM"}
+
+
+def test_parse_proton_tree_ansi_stripped():
+    """ANSI color codes are stripped before parsing."""
+    tree = "\x1b[32m├─ 1.000 tir\x1b[0m\n│  └─ 1.000 k\n"
+    impls, _ = _parse_proton_tree(tree)
+    assert impls == {"tir": 1.0}
+
+
+def test_parse_proton_tree_empty():
+    impls, errors = _parse_proton_tree("")
+    assert impls == {}
+    assert errors == {}
+
+
+# ── bench ───────────────────────────────────────────────────────────────────
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_bench_basic():
+    """bench returns positive times for each impl."""
+    M, N = 256, 256
+
+    funcs = {"matmul": lambda case: torch.mm(case[0], case[1])}
+
+    def make_input():
+        A = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        B = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        return (A, B), tensor_bytes(A, B)
+
+    def run_and_check():
+        results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event")
+        assert "matmul" in results["impls"]
+        assert results["impls"]["matmul"] > 0
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_bench_multiple_impls():
+    """Multiple impls each get their own timing."""
+    M, N = 128, 128
+    funcs = {
+        "mm": lambda case: torch.mm(case[0], case[1]),
+        "addmm": lambda case: torch.addmm(
+            torch.zeros(M, N, device="cuda", dtype=torch.float16), case[0], case[1]
+        ),
+    }
+
+    def make_input():
+        A = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        B = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        return (A, B), tensor_bytes(A, B)
+
+    def run_and_check():
+        results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event")
+        assert set(results["impls"].keys()) == {"mm", "addmm"}
+        assert all(v > 0 for v in results["impls"].values())
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_bench_multiple_input_groups():
+    """Multiple input groups cycle correctly (L2 eviction)."""
+    M, N = 128, 128
+    call_count = [0]
+
+    def make_input():
+        call_count[0] += 1
+        A = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        B = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        return (A, B), tensor_bytes(A, B)
+
+    funcs = {"mm": lambda case: torch.mm(case[0], case[1])}
+
+    def run_and_check():
+        results = bench(
+            funcs,
+            make_input,
+            warmup=5,
+            repeat=20,
+            cooldown_s=0.0,
+            timer="event",
+            l2_bytes=64 * 1024,
+        )
+        assert results["impls"]["mm"] > 0
+        assert call_count[0] > 1
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+# ── _compute_group_count ───────────────────────────────────────────────────
+
+
+def test_compute_groups_small_tensors():
+    """Small tensors need many groups to fill 3x L2."""
+    # 128x128 fp16 = 32KB.  3*128MB / 32KB = 12288, +1 = 12289
+    input_bytes = tensor_bytes(torch.empty(128, 128, dtype=torch.float16))
+    n = _compute_group_count(input_bytes, l2_bytes=128 * 1024 * 1024)
+    assert n == 12289
+
+
+def test_compute_groups_large_tensors():
+    """Inputs >= 3x L2 need only 1 group."""
+    # 16384x16384 fp32 = 1GB >> 3*128MB = 384MB
+    input_bytes = tensor_bytes(torch.empty(16384, 16384, dtype=torch.float32))
+    n = _compute_group_count(input_bytes, l2_bytes=128 * 1024 * 1024)
+    assert n == 1
+
+
+def test_compute_groups_moderate_tensors():
+    """Moderate tensors: floor(3*L2 / input) + 1."""
+    # 8192x8192 bf16 = 128MB.  floor(384M / 128M) + 1 = 4
+    input_bytes = tensor_bytes(torch.empty(8192, 8192, dtype=torch.bfloat16))
+    n = _compute_group_count(input_bytes, l2_bytes=128 * 1024 * 1024)
+    assert n == 4
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_bench_legacy_callable_api():
+    """bench still accepts the existing single-callable API used by TIRx tests."""
+    M, N = 128, 128
+
+    def run_and_check():
+        A = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        B = torch.randn(M, N, device="cuda", dtype=torch.float16)
+        result = bench(
+            lambda: torch.mm(A, B), warmup=1, repeat=2, proton_name="legacy", flush_l2_size=1
+        )
+        assert result > 0
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_bench_callable_inputs():
+    """bench accepts a factory callable and auto-computes groups."""
+    M, N = 256, 256
+
+    call_count = [0]
+
+    def make_input():
+        call_count[0] += 1
+        case = (
+            torch.randn(M, N, device="cuda", dtype=torch.float16),
+            torch.randn(M, N, device="cuda", dtype=torch.float16),
+        )
+        return case, tensor_bytes(*case)
+
+    funcs = {"mm": lambda case: torch.mm(case[0], case[1])}
+
+    def run_and_check():
+        results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event")
+        assert "mm" in results["impls"]
+        assert results["impls"]["mm"] > 0
+        assert call_count[0] >= 2  # at least 2 groups created
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/python/tirx/test_buffer_print.py

@@ -0,0 +1,396 @@
+# 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.
+
+import re
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def generate_random_data(shape, dtype):
+    np.random.seed(0)
+    return np.random.randn(*shape).astype(dtype)
+
+
+def create_tvm_arrays(data_np, device):
+    return [tvm.runtime.tensor(data, device=device) for data in data_np]
+
+
+def make_tvm_runner(func, input_data, initial_output, expected_output):
+    def run(device):
+        args = create_tvm_arrays([*input_data, initial_output], device)
+        func(*args)
+        verify_result(args[-1], expected_output)
+
+    return run
+
+
+def from_source(code):
+    return tvm.script.from_source(code, s_tir=True)
+
+
+def verify_result(C_tvm, C_np):
+    tvm.testing.assert_allclose(C_tvm.numpy(), C_np, rtol=1e-5)
+
+
+def verify_tir_code(code):
+    assert from_source(code).script() == code
+
+
+def verify_cuda_code_array(func, dim_num, dtype, *dims):
+    generated_code = func.mod.imports[0].inspect_source()
+
+    match = re.search(r"// print_buffer starts(.*?)// print_buffer ends", generated_code, re.DOTALL)
+    if not match:
+        raise AssertionError("print_buffer section not found in generated code")
+
+    print_buffer_section = match.group(1).strip()
+    loop_pattern = re.compile(r"for \(int i(\d+) = 0; i\1 < (\d+); \+\+i\1\)")
+    loops = loop_pattern.findall(print_buffer_section)
+    if len(loops) != dim_num:
+        raise AssertionError(f"Expected {dim_num} nested loops, but found {len(loops)}")
+
+    loop_limits = [int(limit) for _, limit in loops]
+    if loop_limits != list(dims):
+        raise AssertionError(f"Expected loop limits {dims}, but found {loop_limits}")
+
+    dtype_to_printf = {"float32": "%f", "float16": "%f", "int32": "%d", "uint32": "%u"}
+    expected_printf_specifier = dtype_to_printf.get(dtype)
+    if not expected_printf_specifier:
+        raise AssertionError(f"Unsupported dtype {dtype}")
+    variable_access_pattern = r"\w+\[.*\]"
+
+    if dtype == "float16":
+        # Look for `printf("%f", static_cast<float>(C[...]))`
+        printf_pattern = re.compile(
+            r'printf\s*\(\s*"'
+            + re.escape(expected_printf_specifier)
+            + r'"\s*,\s*static_cast<float>\('
+            + variable_access_pattern
+            + r"\)\s*\)"
+        )
+    else:
+        # Look for `printf("%f", C[...])`
+        printf_pattern = re.compile(
+            r'printf\s*\(\s*"'
+            + re.escape(expected_printf_specifier)
+            + r'"\s*,\s*'
+            + variable_access_pattern
+            + r"\s*\)"
+        )
+
+    if not printf_pattern.search(print_buffer_section):
+        raise AssertionError(
+            f'Expected element printf statement with format "{expected_printf_specifier}" and a buffer access, but not found'  # noqa: E501
+        )
+
+
+def verify_cuda_code_scalar(func, dtype, expected_value_or_varname):
+    generated_code = func.mod.imports[0].inspect_source()
+
+    all_print_blocks = re.findall(
+        r"// print_buffer starts(.*?)// print_buffer ends", generated_code, re.DOTALL
+    )
+    if not all_print_blocks:
+        raise AssertionError("No print_buffer sections found in generated code")
+
+    dtype_to_printf = {"float32": "%f", "float16": "%f", "int32": "%d", "uint32": "%u"}
+    expected_printf = dtype_to_printf.get(dtype)
+    if not expected_printf:
+        raise AssertionError(f"Unsupported dtype for scalar verification: {dtype}")
+
+    value_pattern = ""
+    if isinstance(expected_value_or_varname, int | float):
+        if "float" in dtype:
+            value_pattern = re.escape(str(float(expected_value_or_varname))) + "f?"
+        else:
+            value_pattern = re.escape(str(int(expected_value_or_varname)))
+    elif isinstance(expected_value_or_varname, str):
+        value_pattern = re.escape(expected_value_or_varname)
+    else:
+        raise TypeError(
+            "expected_value_or_varname must be a number (for literals) or a string (for variables)"
+        )
+
+    if dtype == "float16":
+        printf_pattern = re.compile(
+            r'printf\s*\(\s*".*?'
+            + re.escape(expected_printf)
+            + r'.*?",\s*static_cast<float>\(\s*'
+            + value_pattern
+            + r"\s*\)\s*\)"
+        )
+    else:
+        printf_pattern = re.compile(
+            r'printf\s*\(\s*".*?'
+            + re.escape(expected_printf)
+            + r'.*?",\s*'
+            + value_pattern
+            + r"\s*\)"
+        )
+
+    for block in all_print_blocks:
+        if printf_pattern.search(block):
+            return
+
+    raise AssertionError(
+        f'Could not find a scalar printf with format "{expected_printf}" and value/variable '
+        f'"{expected_value_or_varname}" in any print_buffer block.'
+    )
+
+
+def verify_cuda_code_string(func, expected_var_name, expected_string_literal):
+    generated_code = func.mod.imports[0].inspect_source()
+
+    all_print_blocks = re.findall(
+        r"// print_buffer starts(.*?)// print_buffer ends", generated_code, re.DOTALL
+    )
+    if not all_print_blocks:
+        raise AssertionError("No print_buffer sections found in generated code")
+
+    var_printf_pattern = re.compile(
+        r'printf\s*\(\s*".*?%s.*?",\s*\(char\*\)' + re.escape(expected_var_name) + r"\s*\)"
+    )
+    literal_printf_pattern = re.compile(
+        r'printf\s*\(\s*".*?%s.*?",\s*\(char\*\)\s*"'
+        + re.escape(expected_string_literal)
+        + r'"\s*\)'
+    )
+
+    for block in all_print_blocks:
+        if var_printf_pattern.search(block) or literal_printf_pattern.search(block):
+            return
+
+    raise AssertionError(
+        f'Could not find a string printf using variable "{expected_var_name}" or '
+        f'string literal "{expected_string_literal}" in any print_buffer block.'
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_print():
+    target = tvm.target.Target("cuda")
+
+    def test_vector_add_1D(dtype, dtype_str):
+        M = 6
+        M_BLK = 6
+        dim_num = 1
+        A_np, B_np = generate_random_data((M,), dtype), generate_random_data((M,), dtype)
+        C_np = A_np + B_np
+
+        @T.prim_func(s_tir=True)
+        def add_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M,), dtype_str)
+            B = T.match_buffer(B_ptr, (M,), dtype_str)
+            C = T.match_buffer(C_ptr, (M,), dtype_str)
+
+            for i in T.grid(M):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(M, i)
+                    C[vi] = A[vi] + B[vi]
+                T.print_buffer(C.data, dtype_str, False, False, dim_num, (M,))
+
+        sch = tvm.s_tir.Schedule(add_func)
+        blk = sch.get_sblock("C")
+        i = sch.get_loops(blk)[0]
+
+        i0, i1 = sch.split(i, factors=[None, M_BLK])
+
+        sch.bind(i0, "blockIdx.x")
+        sch.bind(i1, "threadIdx.x")
+
+        C_np_tmp = np.zeros((M,), dtype=dtype)
+        func = tvm.compile(sch.mod, target=target)
+        verify_tir_code(add_func.script())
+        verify_cuda_code_array(func, dim_num, dtype_str, M)
+        return make_tvm_runner(func, [A_np, B_np], C_np_tmp, C_np)
+
+    def test_vector_add_2D(dtype, dtype_str):
+        M, N = 6, 6
+        M_BLK, N_BLK = 6, 6
+        dim_num = 2
+        A_np, B_np = generate_random_data((M, N), dtype), generate_random_data((M, N), dtype)
+        C_np = A_np + B_np
+
+        @T.prim_func(s_tir=True)
+        def add_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N), dtype_str)
+            B = T.match_buffer(B_ptr, (M, N), dtype_str)
+            C = T.match_buffer(C_ptr, (M, N), dtype_str)
+
+            for i, j in T.grid(M, N):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(M, i)
+                    vj = T.axis.spatial(N, j)
+                    C[vi, vj] = A[vi, vj] + B[vi, vj]
+                T.print_buffer(C.data, C.dtype, False, False, dim_num, (M, N))
+
+        sch = tvm.s_tir.Schedule(add_func)
+        blk = sch.get_sblock("C")
+        i, j = sch.get_loops(blk)
+
+        i0, i1 = sch.split(i, factors=[None, M_BLK])
+        j0, j1 = sch.split(j, factors=[None, N_BLK])
+
+        sch.bind(i0, "blockIdx.x")
+        sch.bind(j0, "blockIdx.y")
+        sch.bind(i1, "threadIdx.x")
+        sch.bind(j1, "threadIdx.y")
+
+        C_np_tmp = np.zeros((M, N), dtype=dtype)
+        func = tvm.compile(sch.mod, target=target)
+        verify_tir_code(add_func.script())
+        verify_cuda_code_array(func, dim_num, dtype_str, M, N)
+        return make_tvm_runner(func, [A_np, B_np], C_np_tmp, C_np)
+
+    def test_vector_add_3D(dtype, dtype_str):
+        M, N, K = 6, 6, 6
+        M_BLK, N_BLK, K_BLK = 6, 6, 6
+        dim_num = 3
+        A_np, B_np = generate_random_data((M, N, K), dtype), generate_random_data((M, N, K), dtype)
+        C_np = A_np + B_np
+
+        @T.prim_func(s_tir=True)
+        def add_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M, N, K), dtype_str)
+            B = T.match_buffer(B_ptr, (M, N, K), dtype_str)
+            C = T.match_buffer(C_ptr, (M, N, K), dtype_str)
+
+            for i, j, k in T.grid(M, N, K):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(M, i)
+                    vj = T.axis.spatial(N, j)
+                    vk = T.axis.spatial(K, k)
+                    C[vi, vj, vk] = A[vi, vj, vk] + B[vi, vj, vk]
+                T.print_buffer(C.data, C.dtype, False, False, dim_num, (M, N, K))
+
+        sch = tvm.s_tir.Schedule(add_func)
+        blk = sch.get_sblock("C")
+        i, j, k = sch.get_loops(blk)
+
+        i0, i1 = sch.split(i, factors=[None, M_BLK])
+        j0, j1 = sch.split(j, factors=[None, N_BLK])
+        k0, k1 = sch.split(k, factors=[None, K_BLK])
+
+        sch.bind(i0, "blockIdx.x")
+        sch.bind(j0, "blockIdx.y")
+        sch.bind(k0, "blockIdx.z")
+        sch.bind(i1, "threadIdx.x")
+        sch.bind(j1, "threadIdx.y")
+        sch.bind(k1, "threadIdx.z")
+
+        C_np_tmp = np.zeros((M, N, K), dtype=dtype)
+        func = tvm.compile(sch.mod, target=target)
+        verify_tir_code(add_func.script())
+        verify_cuda_code_array(func, dim_num, dtype_str, M, N, K)
+        return make_tvm_runner(func, [A_np, B_np], C_np_tmp, C_np)
+
+    def test_const_scalar(dtype, dtype_str):
+        M = 6
+        M_BLK = 6
+        dim_num = 1
+        A_np, B_np = generate_random_data((M,), dtype), generate_random_data((M,), dtype)
+        C_np = A_np + B_np
+
+        @T.prim_func(s_tir=True)
+        def add_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M,), dtype_str)
+            B = T.match_buffer(B_ptr, (M,), dtype_str)
+            C = T.match_buffer(C_ptr, (M,), dtype_str)
+            Ten: T.let = T.IntImm(dtype_str, 10)
+
+            for i in T.grid(M):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(M, i)
+                    C[vi] = A[vi] + B[vi]
+                T.print_buffer(Ten, "int32", False, True, dim_num, ())
+
+        sch = tvm.s_tir.Schedule(add_func)
+        blk = sch.get_sblock("C")
+        i = sch.get_loops(blk)[0]
+
+        i0, i1 = sch.split(i, factors=[None, M_BLK])
+
+        sch.bind(i0, "blockIdx.x")
+        sch.bind(i1, "threadIdx.x")
+
+        C_np_tmp = np.zeros((M,), dtype=dtype)
+        func = tvm.compile(sch.mod, target=target)
+        verify_tir_code(add_func.script())
+        verify_cuda_code_scalar(func, dtype_str, 10)
+        return make_tvm_runner(func, [A_np, B_np], C_np_tmp, C_np)
+
+    def test_string(dtype, dtype_str, test_string):
+        M = 6
+        M_BLK = 6
+        dim_num = 1
+        A_np, B_np = generate_random_data((M,), dtype), generate_random_data((M,), dtype)
+        C_np = A_np + B_np
+
+        @T.prim_func(s_tir=True)
+        def add_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
+            A = T.match_buffer(A_ptr, (M,), dtype_str)
+            B = T.match_buffer(B_ptr, (M,), dtype_str)
+            C = T.match_buffer(C_ptr, (M,), dtype_str)
+            string_var = T.StringImm(test_string)
+
+            for i in T.grid(M):
+                with T.sblock("C"):
+                    vi = T.axis.spatial(M, i)
+                    C[vi] = A[vi] + B[vi]
+                T.print_buffer(string_var, "int8", True, False, dim_num, ())
+
+        sch = tvm.s_tir.Schedule(add_func)
+        blk = sch.get_sblock("C")
+        i = sch.get_loops(blk)[0]
+
+        i0, i1 = sch.split(i, factors=[None, M_BLK])
+
+        sch.bind(i0, "blockIdx.x")
+        sch.bind(i1, "threadIdx.x")
+
+        C_np_tmp = np.zeros((M,), dtype=dtype)
+        func = tvm.compile(sch.mod, target=target)
+        verify_tir_code(add_func.script())
+        verify_cuda_code_string(func, "string_var", test_string)
+        return make_tvm_runner(func, [A_np, B_np], C_np_tmp, C_np)
+
+    runners = [
+        test_vector_add_1D(np.float32, "float32"),
+        test_vector_add_2D(np.int32, "int32"),
+        test_vector_add_2D(np.float16, "float16"),
+        test_vector_add_3D(np.uint32, "uint32"),
+        test_string(np.float32, "float32", "hello tirx!"),
+        test_const_scalar(np.int32, "int32"),
+    ]
+
+    def run_and_check():
+        device = tvm.cuda()
+        for runner in runners:
+            runner(device)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/test_control_flow.py

@@ -0,0 +1,122 @@
+# 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.
+import numpy as np
+import pytest
+
+import tvm
+from tvm.script import tirx as T
+from tvm.testing import env
+
+
+def run_test_break_continue(func, shape, expected):
+    target = tvm.target.Target("cuda")
+    mod = tvm.IRModule({"main": func})
+    with target:
+        mod = tvm.compile(mod, target=target, tir_pipeline="tirx")
+    arr_np = np.zeros(shape, dtype="int32")
+
+    def run_and_check():
+        dev = tvm.cuda(0)
+        arr = tvm.runtime.tensor(arr_np, device=dev)
+        mod(arr)
+        np.testing.assert_allclose(arr.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_break_continue1():
+    # fmt: off
+    @T.prim_func
+    def func(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (10,), "int32")
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([32])
+        for i in T.serial(10):
+            if i == 2:
+                continue
+            if i == 7:
+                break
+            A[i] = i
+        # fmt: on
+
+    expected = np.array([0, 1, 0, 3, 4, 5, 6, 0, 0, 0], dtype="int32")
+    run_test_break_continue(func, (10,), expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_break_continue2():
+    # fmt: off
+    @T.prim_func
+    def func(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (9,), "int32")
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([32])
+        idx = T.alloc_buffer((1,), "int32", scope="local")
+        idx[0] = 0
+        for i in T.serial(3):
+            if i == 0:
+                idx[0] += 1
+                continue
+            for j in T.serial(3):
+                A[idx[0]] = i * 10 + j
+                idx[0] += 1
+                if j == 1:
+                    break
+        # fmt: on
+
+    expected = np.array([0, 10, 11, 20, 21, 0, 0, 0, 0], dtype="int32")
+    run_test_break_continue(func, (9,), expected)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_break_continue3():
+    # fmt: off
+    @T.prim_func
+    def func(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (10,), "int32")
+
+        T.device_entry()
+        cta_id = T.cta_id([1])
+        tid = T.thread_id([32])
+        i = T.alloc_buffer((1,), "int32", scope="local")
+        i[0] = 0
+        while i[0] < 10:
+            if (i[0] % 2) == 1:
+                i[0] += 1
+                continue
+            A[i[0]] = i[0]
+            i[0] += 1
+            if i[0] == 7:
+                break
+        # fmt: on
+
+    expected = np.array([0, 0, 2, 0, 4, 0, 6, 0, 0, 0], dtype="int32")
+    run_test_break_continue(func, (10,), expected)
+
+
+if __name__ == "__main__":
+    test_break_continue1()
+    test_break_continue2()
+    test_break_continue3()

tests/python/tirx/test_exec_context.py

@@ -0,0 +1,428 @@
+# 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.
+"""Unit tests for ExecContext (RFC v3 §6). Cases mirror RFC §8.1 -- §8.10."""
+
+from __future__ import annotations
+
+import pytest
+
+from tvm.tirx.exec_context import (
+    CLUSTER,
+    CTA,
+    LANE_CTA_THREAD,
+    LANE_FLAT,
+    LANE_W_INNER,
+    LANE_WG_OUTER,
+    LANE_WG_THREAD,
+    THREAD,
+    WARP,
+    WARPGROUP,
+    AxisRange,
+    ExecContext,
+    ExecContextError,
+    LaneBinding,
+    filter_modulo,
+    filter_narrow,
+    initial_A,
+    scope_switch,
+)
+
+# -- canonical bindings declared at kernel entry (see RFC §8 naming conv) --
+WARP_FLAT = LaneBinding(axis="warpid", kind=LANE_FLAT, declared_extent=16)
+WG_OUTER = LaneBinding(axis="warpid", kind=LANE_WG_OUTER, declared_extent=4)
+W_INNER = LaneBinding(axis="warpid", kind=LANE_W_INNER, declared_extent=4)
+LANE_BIND = LaneBinding(axis="laneid", kind=LANE_FLAT, declared_extent=32)
+CTA_BIND = LaneBinding(axis="cta_id", kind=LANE_FLAT, declared_extent=1)
+CTA_THREAD_BIND = LaneBinding(axis="thread", kind=LANE_CTA_THREAD, declared_extent=256)
+WG_THREAD_BIND = LaneBinding(axis="thread", kind=LANE_WG_THREAD, declared_extent=128)
+
+
+# ---------------------------------------------------------------------------
+# §3 scope_switch: split table
+# ---------------------------------------------------------------------------
+
+
+def test_initial_A_single_cta():
+    A = initial_A(warp_ext=16)
+    assert A.laneid == AxisRange(32, 0)
+    assert A.warpid == AxisRange(16, 0)
+    assert A.cta_id == AxisRange(1, 0)
+    assert A.size == 512
+
+
+def test_initial_A_cluster():
+    A = initial_A(warp_ext=16, cta_ext=4)
+    assert A.cta_id == AxisRange(4, 0)
+    assert A.size == 2048
+
+
+def test_axis_modulo_filter_uses_stride():
+    A = initial_A(warp_ext=16, cta_ext=4)
+    A = filter_modulo(A, "cta_id", 2, 0)
+    assert A.cta_id == AxisRange(2, 0, 2)
+    A = filter_narrow(A, CTA_BIND, 1, 4)
+    assert A.cta_id == AxisRange(1, 2, 2)
+
+
+def test_axis_modulo_filter_two_cta_pair_residues():
+    A = initial_A(warp_ext=16, cta_ext=2)
+    assert filter_modulo(A, "cta_id", 2, 0).cta_id == AxisRange(1, 0, 2)
+    assert filter_modulo(A, "cta_id", 2, 1).cta_id == AxisRange(1, 1, 2)
+
+
+@pytest.mark.parametrize(
+    "kappa,expected_inter_axes,expected_intra_axes",
+    [
+        (THREAD, {"laneid", "warpid", "cta_id"}, set()),
+        (WARP, {"warpid", "cta_id"}, {"laneid"}),
+        (CTA, {"cta_id"}, {"laneid", "warpid"}),
+        (CLUSTER, set(), {"laneid", "warpid", "cta_id"}),
+    ],
+)
+def test_scope_switch_trivial(kappa, expected_inter_axes, expected_intra_axes):
+    A = initial_A(warp_ext=16, cta_ext=4)
+    split = scope_switch(A, kappa)
+    assert set(split.inter) == expected_inter_axes
+    assert set(split.intra) == expected_intra_axes
+
+
+def test_scope_switch_warpgroup_aligned():
+    A = initial_A(warp_ext=16)
+    split = scope_switch(A, WARPGROUP)
+    assert split.inter["wgid"] == AxisRange(4, 0)
+    assert split.inter["cta_id"] == AxisRange(1, 0)
+    assert split.intra["laneid"] == AxisRange(32, 0)
+    assert split.intra["wid_in_wg"] == AxisRange(4, 0)
+
+
+# ---------------------------------------------------------------------------
+# §4.2 warpgroup factoring: 3 cases
+# ---------------------------------------------------------------------------
+
+
+def test_factor_case1_aligned():
+    A = initial_A(warp_ext=8)  # ext=8, off=0 -- aligned
+    split = scope_switch(A, WARPGROUP)
+    assert split.inter["wgid"] == AxisRange(2, 0)
+    assert split.intra["wid_in_wg"] == AxisRange(4, 0)
+
+
+def test_factor_case2_fits_in_one_wg():
+    # warpid ext=2, off=0 -- fits in one wg
+    A = initial_A(warp_ext=16)
+    A = filter_narrow(A, WARP_FLAT, 0, 2)
+    split = scope_switch(A, WARPGROUP)
+    assert split.inter["wgid"] == AxisRange(1, 0)
+    assert split.intra["wid_in_wg"] == AxisRange(2, 0)
+
+
+def test_factor_case2_offset():
+    # warpid ext=2, off=6 -> wid_off=2, fits (2 <= 4-2)
+    A = initial_A(warp_ext=16)
+    A = filter_narrow(A, WARP_FLAT, 6, 8)
+    split = scope_switch(A, WARPGROUP)
+    assert split.inter["wgid"] == AxisRange(1, 1)
+    assert split.intra["wid_in_wg"] == AxisRange(2, 2)
+
+
+def test_factor_case3_fails():
+    # RFC §8.6: warpid[2:6] crosses wg boundary unaligned
+    A = initial_A(warp_ext=16)
+    A = filter_narrow(A, WARP_FLAT, 2, 6)
+    assert A.warpid == AxisRange(4, 2)
+    with pytest.raises(ExecContextError, match="crosses warpgroup boundary"):
+        scope_switch(A, WARPGROUP)
+
+
+# ---------------------------------------------------------------------------
+# §8.1 -- Pure narrowing CTA -> WG -> W
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_1_cta_wg_warp():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16)
+    # with T.cta()
+    ctx = ctx.with_scope_switch(CTA)
+    assert ctx.inter == {"cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(16, 0)}
+    # with T.warpgroup()
+    ctx = ctx.with_scope_switch(WARPGROUP)
+    assert ctx.inter == {"wgid": AxisRange(4, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
+    # with T.warp()
+    ctx = ctx.with_scope_switch(WARP)
+    assert ctx.inter == {"warpid": AxisRange(16, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.2 -- Filter + scope_switch
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_2_filter_then_warpgroup():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    ctx = ctx.with_filter(WARP_FLAT, 0, 8)
+    assert ctx.A.warpid == AxisRange(8, 0)
+    # recompute at cta: intra=(lane:32, warp:8)
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(8, 0)}
+    # enter warpgroup: factor(8, 0) -> case 1
+    ctx = ctx.with_scope_switch(WARPGROUP)
+    assert ctx.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.3 -- Sugar form T.warp(warpid[2:4])
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_3_sugar_warp_range():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    # desugar: filter warpid[2:4], then warp
+    ctx = ctx.with_filter(WARP_FLAT, 2, 4).with_scope_switch(WARP)
+    assert ctx.A.warpid == AxisRange(2, 2)
+    assert ctx.inter == {"warpid": AxisRange(2, 2), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.4 -- Widen after filter (warp -> warpgroup)
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_4_widen_warp_to_wg():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    ctx = ctx.with_filter(WARP_FLAT, 0, 4).with_scope_switch(WARP)
+    # widen to warpgroup
+    ctx = ctx.with_scope_switch(WARPGROUP)
+    assert ctx.inter == {"wgid": AxisRange(1, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.5 -- Partial warp selection -> warpgroup (partial intra)
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_5_partial_wg():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    ctx = ctx.with_filter(WARP_FLAT, 0, 2).with_scope_switch(WARPGROUP)
+    # case 2: 2 <= 4-0
+    assert ctx.inter == {"wgid": AxisRange(1, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(2, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.6 -- Cross warpgroup boundary (factor fails)
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_6_factor_fail():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    # with_filter recomputes (inter, intra) for current scope_kind=cta -- still OK
+    ctx2 = ctx.with_filter(WARP_FLAT, 2, 6)
+    assert ctx2.A.warpid == AxisRange(4, 2)
+    # scope_switch to warpgroup is the one that must fail
+    with pytest.raises(ExecContextError, match="crosses warpgroup boundary"):
+        ctx2.with_scope_switch(WARPGROUP)
+
+
+# ---------------------------------------------------------------------------
+# §8.7 -- Deep mixed nesting
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_7_deep_nested():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    ctx = ctx.with_filter(WARP_FLAT, 0, 8).with_scope_switch(WARPGROUP)
+    assert ctx.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
+    ctx = ctx.with_filter(WARP_FLAT, 0, 2)
+    # recompute at warpgroup: factor(2, 0) -> case 2
+    assert ctx.inter == {"wgid": AxisRange(1, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(2, 0)}
+    ctx = ctx.with_scope_switch(WARP)
+    assert ctx.inter == {"warpid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
+    assert ctx.intra == {"laneid": AxisRange(32, 0)}
+    ctx = ctx.with_filter(LANE_BIND, 0, 8)
+    assert ctx.intra == {"laneid": AxisRange(8, 0)}
+    assert ctx.inter == {"warpid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.8 -- FA4 pattern: 3 sibling filter branches
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_8_fa4_pattern():
+    root = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+
+    # Branch 1: warp 12 (single warp, tcgen05 MMA elected)
+    b1 = root.with_filter(WARP_FLAT, 12, 13)
+    assert b1.A.warpid == AxisRange(1, 12)
+    assert b1.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(1, 12)}
+
+    # Branch 2: softmax warpgroups (warps 0-7)
+    b2 = root.with_filter(WARP_FLAT, 0, 8).with_scope_switch(WARPGROUP)
+    assert b2.inter == {"wgid": AxisRange(2, 0), "cta_id": AxisRange(1, 0)}
+    assert b2.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
+
+    # Branch 3: correction warpgroup (warps 8-11 = wg2)
+    b3 = root.with_filter(WARP_FLAT, 8, 12)
+    assert b3.A.warpid == AxisRange(4, 8)
+    assert b3.intra == {"laneid": AxisRange(32, 0), "warpid": AxisRange(4, 8)}
+    # And should factor cleanly when entering warpgroup
+    b3wg = b3.with_scope_switch(WARPGROUP)
+    assert b3wg.inter == {"wgid": AxisRange(1, 2), "cta_id": AxisRange(1, 0)}
+    assert b3wg.intra == {"laneid": AxisRange(32, 0), "wid_in_wg": AxisRange(4, 0)}
+
+
+# ---------------------------------------------------------------------------
+# §8.9 -- Cross-CTA with widening to cluster
+# ---------------------------------------------------------------------------
+
+
+def test_ex_8_9_cross_cta_cluster():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16, cta_ext=4).with_scope_switch(CTA)
+    assert ctx.inter == {"cta_id": AxisRange(4, 0)}
+    # filter to warp 0, then warp
+    w = ctx.with_filter(WARP_FLAT, 0, 1).with_scope_switch(WARP)
+    assert w.inter == {"warpid": AxisRange(1, 0), "cta_id": AxisRange(4, 0)}
+    assert w.intra == {"laneid": AxisRange(32, 0)}
+
+    # back at cta scope, enter warpgroup
+    wg = ctx.with_scope_switch(WARPGROUP)
+    assert wg.inter == {"wgid": AxisRange(4, 0), "cta_id": AxisRange(4, 0)}
+    # widen to cluster
+    cl = wg.with_scope_switch(CLUSTER)
+    assert cl.inter == {}
+    assert cl.intra == {
+        "laneid": AxisRange(32, 0),
+        "warpid": AxisRange(16, 0),
+        "cta_id": AxisRange(4, 0),
+    }
+
+
+# ---------------------------------------------------------------------------
+# §8.10 -- identical to 8.3 modulo prose; covered above
+# ---------------------------------------------------------------------------
+
+# ---------------------------------------------------------------------------
+# Rule 1 & 5: filter can only shrink A; saved/restored across scope exit
+# (Restoration is the caller's (IR walker's) responsibility -- ExecContext
+# is immutable, each with_filter returns a fresh ctx. Test that the parent
+# is untouched.)
+# ---------------------------------------------------------------------------
+
+
+def test_filter_is_pure():
+    ctx = ExecContext.at_kernel_entry(warp_ext=16).with_scope_switch(CTA)
+    child = ctx.with_filter(WARP_FLAT, 0, 8)
+    assert ctx.A.warpid == AxisRange(16, 0)  # parent not mutated
+    assert child.A.warpid == AxisRange(8, 0)
+
+
+def test_filter_empty_range_rejected():
+    A = initial_A(warp_ext=16)
+    with pytest.raises(ExecContextError, match="empty or inverted"):
+        filter_narrow(A, WARP_FLAT, 5, 5)
+
+
+def test_filter_out_of_range_rejected():
+    A = initial_A(warp_ext=16)
+    A = filter_narrow(A, WARP_FLAT, 0, 4)
+    with pytest.raises(ExecContextError, match="empty range"):
+        filter_narrow(A, WARP_FLAT, 8, 12)  # disjoint from [0, 4)
+
+
+def test_filter_flat_cta_thread_full_warp_range():
+    A = initial_A(warp_ext=8)
+    A = filter_narrow(A, CTA_THREAD_BIND, 0, 128)
+    assert A.warpid == AxisRange(4, 0)
+    assert A.laneid == AxisRange(32, 0)
+
+
+def test_filter_flat_cta_thread_single_warp_lane_range():
+    A = initial_A(warp_ext=8)
+    A = filter_narrow(A, CTA_THREAD_BIND, 34, 40)
+    assert A.warpid == AxisRange(1, 1)
+    assert A.laneid == AxisRange(6, 2)
+
+
+def test_filter_flat_cta_thread_nonrectangular_rejected():
+    A = initial_A(warp_ext=8)
+    with pytest.raises(ExecContextError, match="non-rectangular"):
+        filter_narrow(A, CTA_THREAD_BIND, 20, 50)
+
+
+def test_filter_flat_warpgroup_thread_range_inside_one_warpgroup():
+    A = initial_A(warp_ext=8)
+    A = filter_narrow(A, WG_OUTER, 1, 2)
+    A = filter_narrow(A, WG_THREAD_BIND, 32, 64)
+    assert A.warpid == AxisRange(1, 5)
+    assert A.laneid == AxisRange(32, 0)
+
+
+def test_filter_flat_warpgroup_thread_full_range_across_warpgroups_is_noop():
+    A = initial_A(warp_ext=8)
+    A2 = filter_narrow(A, WG_THREAD_BIND, 0, 128)
+    assert A2.warpid == AxisRange(8, 0)
+    assert A2.laneid == AxisRange(32, 0)
+
+
+def test_filter_flat_warpgroup_thread_partial_range_across_warpgroups_rejected():
+    A = initial_A(warp_ext=8)
+    with pytest.raises(ExecContextError, match="multiple warpgroups"):
+        filter_narrow(A, WG_THREAD_BIND, 0, 64)
+
+
+# ---------------------------------------------------------------------------
+# Factor-lane bindings: wg_outer and w_inner
+# ---------------------------------------------------------------------------
+
+
+def test_filter_wg_outer():
+    A = initial_A(warp_ext=16)
+    A2 = filter_narrow(A, WG_OUTER, 1, 3)  # wg 1..2 -> warps 4..11
+    assert A2.warpid == AxisRange(8, 4)
+
+
+def test_filter_wg_outer_unaligned_rejected():
+    A = initial_A(warp_ext=16)
+    A = filter_narrow(A, WARP_FLAT, 2, 6)  # warp offset 2 (not WG-aligned)
+    with pytest.raises(ExecContextError, match="aligned to WG_SIZE"):
+        filter_narrow(A, WG_OUTER, 0, 1)
+
+
+def test_filter_w_inner():
+    A = initial_A(warp_ext=16)
+    # First narrow into a single warpgroup, then inner filter is valid
+    A = filter_narrow(A, WARP_FLAT, 4, 8)  # wg1: warps 4..7
+    A2 = filter_narrow(A, W_INNER, 1, 3)  # pick inner lanes 1..2
+    assert A2.warpid == AxisRange(2, 5)
+
+
+def test_filter_w_inner_spanning_wg_rejected():
+    A = initial_A(warp_ext=16)  # spans all 4 wgs
+    with pytest.raises(ExecContextError, match="spans multiple warpgroups"):
+        filter_narrow(A, W_INNER, 0, 2)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))

tests/python/tirx/test_exec_scope.py

@@ -0,0 +1,43 @@
+# 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.
+import pytest
+
+from tvm.tirx.exec_scope import ExecScope
+
+
+def test_exec_scope_create():
+    def is_trivial_scope(scope, name):
+        return isinstance(scope, ExecScope) and scope.name == name
+
+    thread = ExecScope("thread")
+    warp = ExecScope("warp")
+    wg = ExecScope("warpgroup")
+    cta = ExecScope("cta")
+    cluster = ExecScope("cluster")
+
+    assert is_trivial_scope(thread, "thread")
+    assert is_trivial_scope(warp, "warp")
+    assert is_trivial_scope(wg, "warpgroup")
+    assert is_trivial_scope(cta, "cta")
+    assert is_trivial_scope(cluster, "cluster")
+
+    with pytest.raises(Exception, match="Unknown scope kind name"):
+        ExecScope("aaa")
+
+
+if __name__ == "__main__":
+    test_exec_scope_create()

tests/python/tirx/test_hint.py

@@ -0,0 +1,268 @@
+# 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 T.hint() — universal directive primitive for TIRx sketch language."""
+
+import tvm
+import tvm.script
+import tvm.testing
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.tirx import AttrStmt
+
+
+def from_source(code):
+    return tvm.script.from_source(code)
+
+
+def test_hint_statement():
+    """T.hint("msg") as a bare statement produces an AttrStmt with attr_key=tirx_hint."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("persistent tile scheduler with L2 swizzle")
+        T.evaluate(0)
+
+    # Walk the IR to find the AttrStmt with tirx_hint
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            # node is now a Map with "message" key
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "persistent tile scheduler with L2 swizzle"
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found"
+
+
+def test_hint_context_manager():
+    """with T.hint("msg"): scopes its body inside the AttrStmt."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        with T.hint("software pipeline, depth 4"):
+            T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "software pipeline, depth 4"
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found"
+
+
+def test_hint_with_attrs():
+    """T.hint("msg", key="value") passes structured attrs in Map node."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("scheduler", mode="persistent", depth="4")
+        T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "scheduler"
+            assert str(stmt.node["mode"]) == "persistent"
+            assert str(stmt.node["depth"]) == "4"
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' not found"
+
+
+def test_hint_printer_roundtrip_statement():
+    """Verify T.hint("msg") prints as T.hint("msg") and roundtrips through script/parse."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("persistent tile scheduler with L2 swizzle")
+        T.evaluate(0)
+
+    code = func.script()
+    assert 'hint("persistent tile scheduler with L2 swizzle")' in code
+    reparsed = from_source(code)
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_printer_roundtrip_context_manager():
+    """Verify with T.hint("msg"): prints correctly and roundtrips."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        with T.hint("software pipeline, depth 4"):
+            T.evaluate(0)
+
+    code = func.script()
+    assert 'hint("software pipeline, depth 4")' in code
+    reparsed = from_source(code)
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_printer_roundtrip_with_attrs():
+    """Verify T.hint("msg", key="val") prints with kwargs and roundtrips."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        _A = T.match_buffer(A_ptr, (64,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("scheduler", mode="persistent")
+        T.evaluate(0)
+
+    code = func.script()
+    assert 'hint("scheduler"' in code
+    assert 'mode="persistent"' in code
+    reparsed = from_source(code)
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_keyword_arg_on_tx_op():
+    """Tx.op(..., hint="msg") stores hint in TilePrimitiveCall.config."""
+    from tvm.tirx.buffer import decl_buffer
+    from tvm.tirx.stmt import TilePrimitiveCall
+
+    A = decl_buffer((64, 64), "float32", scope="global")
+    A_sm = decl_buffer((64, 64), "float32", scope="shared")
+
+    op_call = TilePrimitiveCall(
+        A[0:64, 0:64],
+        A_sm[0:64, 0:64],
+        op=tvm.ir.Op.get("tirx.tile.copy"),
+        workspace={},
+        config={"hint": "3-input ptx"},
+    )
+    assert "hint" in op_call.config
+    assert str(op_call.config["hint"]) == "3-input ptx"
+
+
+def test_hint_keyword_arg_on_tx_op_roundtrip():
+    """Tx.op(..., hint="msg") roundtrips through printer/parser."""
+    from tvm.script.tirx import tile as Tx
+
+    @T.prim_func
+    def func(A_ptr: T.handle, B_ptr: T.handle):
+        A = T.match_buffer(A_ptr, [10], "float32", scope="global")
+        B = T.match_buffer(B_ptr, [10], "float32", scope="global")
+        Tx.add(B, A, T.float32(1), hint="use_fast_math")
+
+    code = func.script()
+    assert 'hint="use_fast_math"' in code
+    reparsed = from_source(code)
+    assert reparsed.script() == code
+    assert_structural_equal(func, reparsed)
+
+
+def test_hint_no_message():
+    """T.hint(access=...) with no message string."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128,), "float32", scope="global")
+        bx, by, bz = T.cta_id([1, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint(access=A[0:64])
+        T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            # Should have "access" key but no "message" key
+            assert "access" in stmt.node
+            assert "message" not in stmt.node
+            from tvm.tirx import BufferRegion
+
+            assert isinstance(stmt.node["access"], BufferRegion)
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with attr_key='tirx_hint' containing access not found"
+
+
+def test_hint_access_buffer_region():
+    """T.hint(access=A[region]) stores the BufferRegion structurally in the IR."""
+
+    @T.prim_func
+    def func(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, (128, 64), "float32", scope="global")
+        bx, by, bz = T.cta_id([2, 1, 1])
+        warp_id = T.warp_id([1])
+        lane_id = T.lane_id([32])
+        T.hint("partition", access=A[bx * 64 : (bx + 1) * 64, 0:64])
+        T.evaluate(0)
+
+    found = [False]
+
+    def visit(stmt):
+        if isinstance(stmt, AttrStmt) and stmt.attr_key == "tirx_hint":
+            assert isinstance(stmt.node, tvm.ir.Map)
+            assert str(stmt.node["message"]) == "partition"
+            assert "access" in stmt.node
+            from tvm.tirx import BufferRegion
+
+            assert isinstance(stmt.node["access"], BufferRegion)
+            br = stmt.node["access"]
+            assert br.buffer.name == "A"
+            assert len(br.region) == 2
+            found[0] = True
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    assert found[0], "Expected AttrStmt with structured BufferRegion access not found"
+
+
+if __name__ == "__main__":
+    test_hint_statement()
+    test_hint_context_manager()
+    test_hint_with_attrs()
+    test_hint_printer_roundtrip_statement()
+    test_hint_printer_roundtrip_context_manager()
+    test_hint_printer_roundtrip_with_attrs()
+    test_hint_keyword_arg_on_tx_op()
+    test_hint_keyword_arg_on_tx_op_roundtrip()
+    test_hint_no_message()
+    test_hint_access_buffer_region()

tests/python/tirx/test_inline.py

@@ -0,0 +1,260 @@
+# 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 T.inline / T.inline with Python LEGB scoping semantics."""
+
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+
+# Module-level constant for testing global visibility
+MODULE_CONST = 42
+
+
+def test_local_shadows_enclosing():
+    """A local parameter in the inline shadows a variable from the enclosing scope."""
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        T.int32(10)
+
+        @T.inline
+        def write(x):
+            # x here is the parameter, not the enclosing x=10
+            A[0] = x
+
+        write(T.int32(20))
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        T.int32(10)
+        A[0] = T.int32(20)
+
+    assert_structural_equal(func, expected)
+
+
+def test_enclosing_variable_capture():
+    """Inline captures a variable from its enclosing scope (not a parameter)."""
+    val = 64
+
+    @T.inline
+    def write_val(A):
+        A[0] = val
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        write_val(A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        A[0] = 64
+
+    assert_structural_equal(func, expected)
+
+
+def test_nested_inline():
+    """Inner inline can call outer inline (inline-in-inline)."""
+
+    @T.inline
+    def add_one(A):
+        A[0] = A[0] + 1
+
+    @T.inline
+    def add_two(A):
+        add_one(A)
+        add_one(A)
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        add_two(A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        A[0] = A[0] + 1
+        A[0] = A[0] + 1
+
+    assert_structural_equal(func, expected)
+
+
+def test_module_globals_visible():
+    """Inline can see module-level globals."""
+
+    @T.inline
+    def write_const(A):
+        A[0] = MODULE_CONST
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        write_const(A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        A[0] = 42
+
+    assert_structural_equal(func, expected)
+
+
+def test_shadowing_in_inner_scope():
+    """An inline defined inside a for-loop captures the loop variable."""
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((10,), "int32")) -> None:
+        for i in T.serial(10):
+
+            @T.inline
+            def write_i(A):
+                A[i] = i
+
+            write_i(A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((10,), "int32")) -> None:
+        for i in range(10):
+            A[i] = i
+
+    assert_structural_equal(func, expected)
+
+
+def test_lexical_not_dynamic():
+    """An inline defined outside prim_func does NOT see the caller's locals.
+    Specifically, x_value captured at definition time (128) is used,
+    not the loop variable x_value from the caller."""
+    x_value = 128
+
+    @T.inline
+    def static_capture(A, B):
+        B[()] = A[x_value]
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
+        for x_value in T.serial(10):
+            static_capture(A, B)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((1024,), "int32"), B: T.Buffer((), "int32")) -> None:
+        for x_value in range(10):
+            B[()] = A[128]
+
+    assert_structural_equal(func, expected)
+
+
+def test_callback_pattern():
+    """Inline passed as an argument to another inline."""
+
+    @T.inline
+    def apply_fn(fn, A):
+        fn(A)
+
+    @T.inline
+    def inc(A):
+        A[0] = A[0] + 1
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        apply_fn(inc, A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        A[0] = A[0] + 1
+
+    assert_structural_equal(func, expected)
+
+
+def test_sibling_calls():
+    """Two independent inlines called in sequence."""
+
+    @T.inline
+    def write_a(A):
+        A[0] = 1
+
+    @T.inline
+    def write_b(A):
+        A[1] = 2
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        write_a(A)
+        write_b(A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        A[0] = 1
+        A[1] = 2
+
+    assert_structural_equal(func, expected)
+
+
+def test_recursive_inline():
+    """Recursive inline (defined inside prim_func)."""
+
+    # fmt: off
+    @T.prim_func(private=True)
+    def func():
+        T.device_entry()
+        for x in T.serial(10):
+
+            @T.inline
+            def add(x, c):
+                if c > 0:
+                    add(x, c - 1)
+                T.evaluate(x)
+
+            add(x, 3)
+
+    @T.prim_func(private=True)
+    def expected():
+        T.device_entry()
+        for x in range(10):
+            T.evaluate(x)
+            T.evaluate(x)
+            T.evaluate(x)
+            T.evaluate(x)
+        # fmt: on
+
+    assert_structural_equal(func, expected)
+
+
+def test_late_binding():
+    """Variable defined after inline but before call (inside prim_func)."""
+
+    @T.prim_func(private=True)
+    def func(A: T.Buffer((128,), "int32")) -> None:
+        @T.inline
+        def write(A):
+            A[0] = val
+
+        val = T.int32(99)
+        write(A)
+
+    @T.prim_func(private=True)
+    def expected(A: T.Buffer((128,), "int32")) -> None:
+        val = T.int32(99)
+        A[0] = val
+
+    assert_structural_equal(func, expected)
+
+
+if __name__ == "__main__":
+    test_local_shadows_enclosing()
+    test_enclosing_variable_capture()
+    test_nested_inline()
+    test_module_globals_visible()
+    test_shadowing_in_inner_scope()
+    test_lexical_not_dynamic()
+    test_callback_pattern()
+    test_sibling_calls()
+    test_recursive_inline()
+    test_late_binding()
+    print("All tests passed!")

tests/python/tirx/test_jit.py

@@ -0,0 +1,225 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: F821
+"""Tests for ``@T.jit`` + ``T.constexpr``."""
+
+from __future__ import annotations
+
+import pytest
+
+import tvm
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+
+
+def test_int_constexpr_specializes_loop_bound():
+    @T.jit(private=True)
+    def add(
+        A: T.Buffer((N,), "int32"),
+        B: T.Buffer((N,), "int32"),
+        C: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+    ):
+        for i in range(N):
+            C[i] = A[i] + B[i]
+
+    @T.prim_func(private=True)
+    def expected(
+        A: T.Buffer((128,), "int32"),
+        B: T.Buffer((128,), "int32"),
+        C: T.Buffer((128,), "int32"),
+    ):
+        for i in range(128):
+            C[i] = A[i] + B[i]
+
+    assert_structural_equal(add.specialize(N=128), expected, map_free_vars=True)
+
+
+def test_constexpr_in_2d_buffer_shape():
+    @T.jit(private=True)
+    def matadd(
+        A: T.Buffer((M, K), "int32"),
+        B: T.Buffer((M, K), "int32"),
+        C: T.Buffer((M, K), "int32"),
+        *,
+        M: T.constexpr,
+        K: T.constexpr,
+    ):
+        for m in range(M):
+            for k in range(K):
+                C[m, k] = A[m, k] + B[m, k]
+
+    @T.prim_func(private=True)
+    def expected(
+        A: T.Buffer((4, 8), "int32"),
+        B: T.Buffer((4, 8), "int32"),
+        C: T.Buffer((4, 8), "int32"),
+    ):
+        for m in range(4):
+            for k in range(8):
+                C[m, k] = A[m, k] + B[m, k]
+
+    assert_structural_equal(matadd.specialize(M=4, K=8), expected, map_free_vars=True)
+
+
+def test_constexpr_in_body_expression():
+    @T.jit(private=True)
+    def scaled_copy(
+        A: T.Buffer((N,), "int32"),
+        B: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+        SCALE: T.constexpr,
+    ):
+        for i in range(N):
+            B[i] = A[i] * SCALE
+
+    @T.prim_func(private=True)
+    def expected(
+        A: T.Buffer((16,), "int32"),
+        B: T.Buffer((16,), "int32"),
+    ):
+        for i in range(16):
+            B[i] = A[i] * 3
+
+    assert_structural_equal(scaled_copy.specialize(N=16, SCALE=3), expected, map_free_vars=True)
+
+
+def test_specialize_cache_returns_same_instance():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+    ):
+        for i in range(N):
+            A[i] = 0
+
+    a = k.specialize(N=8)
+    b = k.specialize(N=8)
+    assert a is b
+
+
+def test_specialize_different_args_produce_different_funcs():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+    ):
+        for i in range(N):
+            A[i] = 0
+
+    assert k.specialize(N=8) is not k.specialize(N=16)
+
+
+def test_specialize_missing_constexpr_raises():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+        SCALE: T.constexpr,
+    ):
+        for i in range(N):
+            A[i] = SCALE
+
+    with pytest.raises(TypeError, match="missing"):
+        k.specialize(N=8)
+
+
+def test_specialize_extra_kwarg_raises():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+    ):
+        for i in range(N):
+            A[i] = 0
+
+    with pytest.raises(TypeError, match="unexpected"):
+        k.specialize(N=8, BOGUS=42)
+
+
+def test_jit_kernel_with_nested_inline_helper():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+    ):
+        @T.inline
+        def double(x):
+            A[x] = A[x] * 2
+
+        for i in range(N):
+            double(i)
+
+    @T.prim_func(private=True)
+    def expected(
+        A: T.Buffer((4,), "int32"),
+    ):
+        for i in range(4):
+            A[i] = A[i] * 2
+
+    assert_structural_equal(k.specialize(N=4), expected, map_free_vars=True)
+
+
+def test_constexpr_default_value():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+        SCALE: T.constexpr = 7,
+    ):
+        for i in range(N):
+            A[i] = SCALE
+
+    @T.prim_func(private=True)
+    def expected(
+        A: T.Buffer((8,), "int32"),
+    ):
+        for i in range(8):
+            A[i] = 7
+
+    assert_structural_equal(k.specialize(N=8), expected, map_free_vars=True)
+    # Override the default
+    overridden = k.specialize(N=8, SCALE=99)
+    assert k.specialize(N=8) is not overridden
+
+
+def test_specialize_returns_primfunc():
+    @T.jit(private=True)
+    def k(
+        A: T.Buffer((N,), "int32"),
+        *,
+        N: T.constexpr,
+    ):
+        for i in range(N):
+            A[i] = 0
+
+    spec = k.specialize(N=8)
+    assert isinstance(spec, tvm.tirx.PrimFunc)
+    # Specialized PrimFunc has only the runtime params (constexpr stripped).
+    assert len(spec.params) == 1
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/python/tirx/test_op.py

@@ -0,0 +1,96 @@
+# 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.
+import pytest
+
+from tvm.ir import Op
+from tvm.tirx.buffer import decl_buffer
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def _test(op: str, *args):
+    return TilePrimitiveCall(*args, op=Op.get("tirx.tile." + op), workspace={}, config={})
+
+
+def test_copy():
+    A = decl_buffer((64, 64), "float32", scope="global")
+    A_sm = decl_buffer((64, 64), "float32", scope="shared")
+    _test("copy", A[0:64, 0:64], A_sm[0:64, 0:64])
+
+
+def test_fill():
+    A = decl_buffer((64, 64), "float32", scope="global")
+    _test("fill", A[0:64, 0:64], 1.0)
+
+
+def test_gemm():
+    A = decl_buffer((64, 64), "float32", scope="global")
+    B = decl_buffer((64, 64), "float32", scope="global")
+    C = decl_buffer((64, 64), "float32", scope="global")
+    D = decl_buffer((64, 64), "float32", scope="global")
+    _test("gemm", D[:, :], A[:, :], B[:, :], C[:, :], True, False, 1.0, 0.0)
+
+
+def test_buffer_replacer_no_shared_default():
+    """Regression test for F4: BufferReplacer default dicts must not be shared."""
+    from tvm.tirx.transform.common import BufferReplacer
+
+    r1 = BufferReplacer()
+    r2 = BufferReplacer()
+    A = decl_buffer((64,), "float32")
+    B = decl_buffer((64,), "float32")
+    r1.buffer_map[A] = B
+    # r2 must not see r1's mutation
+    assert len(r2.buffer_map) == 0
+
+
+def test_buffer_replacer_replaces_strides_and_elem_offset():
+    """Vars in buffer strides/elem_offset must be replaced, not passed through."""
+    from tvm.tirx import BufferStore, Var
+    from tvm.tirx.transform.common import BufferReplacer
+
+    n = Var("n", "int32")
+    m = Var("m", "int32")
+    A = decl_buffer((64,), "float32", strides=[n], elem_offset=n)
+    store = BufferStore(A, 1.0, [0])
+
+    new = BufferReplacer(var_map={n: m})(store)
+    assert new.buffer.strides[0].same_as(m)
+    assert new.buffer.elem_offset.same_as(m)
+
+
+def test_gemm_async_partial_scale_factor():
+    """Regression test for F7: gemm_async must reject partial scale factors."""
+    from tvm.tirx.script.builder.tirx import gemm_async
+
+    A = decl_buffer((64, 64), "float16", scope="shared")
+    B = decl_buffer((64, 64), "float16", scope="shared")
+    C = decl_buffer((64, 64), "float16", scope="shared")
+    SF = decl_buffer((64,), "float16", scope="shared")
+
+    with pytest.raises(ValueError, match="SFA and SFB must both be provided or both be None"):
+        gemm_async(C[:, :], A[:, :], B[:, :], SFA=SF[:])
+
+    with pytest.raises(ValueError, match="SFA and SFB must both be provided or both be None"):
+        gemm_async(C[:, :], A[:, :], B[:, :], SFB=SF[:])
+
+
+if __name__ == "__main__":
+    test_copy()
+    test_fill()
+    test_gemm()
+    test_buffer_replacer_no_shared_default()
+    test_gemm_async_partial_scale_factor()

tests/python/tirx/test_op_namespace_cleanup.py

@@ -0,0 +1,394 @@
+# 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 TIRx op namespace split between T, T.tile, and device namespaces."""
+
+import importlib
+import sys
+import types
+
+import pytest
+
+import tvm
+from tvm.ir import Op, assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.stmt import TilePrimitiveCall
+
+
+def _tile_calls(func):
+    calls = []
+
+    def visit(stmt):
+        if isinstance(stmt, TilePrimitiveCall):
+            calls.append(stmt)
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    return calls
+
+
+def _expr_calls(func):
+    calls = []
+
+    def visit(node):
+        if isinstance(node, tvm.ir.Call):
+            calls.append(node)
+
+    tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+    return calls
+
+
+def _op_attr(op_name, attr_name):
+    return Op.get(op_name).get_attr(attr_name)
+
+
+def _has_path(root, path):
+    cur = root
+    for part in path.split("."):
+        if not hasattr(cur, part):
+            return False
+        cur = getattr(cur, part)
+    return True
+
+
+def test_tx_is_tile_shorthand_only():
+    assert T.tile is Tx
+    assert T.tile.copy is Tx.copy
+    assert not hasattr(T, "copy")
+    assert not hasattr(Tx, "SMEMPool")
+    assert not hasattr(Tx, "ScopedOp")
+    assert not hasattr(Tx, "meta_class")
+    assert T.cast is not Tx.cast
+    assert T.sqrt is not Tx.sqrt
+
+
+def test_tx_rejects_expression_overloads():
+    x = tvm.tirx.Var("x", "float32")
+    y = tvm.tirx.Var("y", "int32")
+
+    with pytest.raises(TypeError, match="tile-only"):
+        Tx.sqrt(x)
+    with pytest.raises(TypeError, match="tile-only"):
+        T.tile.sqrt(x)
+    with pytest.raises(TypeError, match="tile-only"):
+        Tx.cast(y, "float32")
+    with pytest.raises(TypeError, match="tile-only"):
+        T.tile.cast(y, "float32")
+
+
+def test_builtin_expression_ops_are_not_tile_primitives():
+    x = tvm.tirx.Var("x", "int32")
+    y = tvm.tirx.Var("y", "float32")
+
+    cast = T.cast(x, "float32")
+    assert isinstance(cast, tvm.tirx.Cast)
+    assert cast.ty.dtype == "float32"
+
+    sqrt = T.sqrt(y)
+    assert sqrt.op.name == "tirx.sqrt"
+
+    fma = T.fma(y, y, y)
+    assert fma.op.name == "tirx.fma"
+
+
+def test_kernel_replace_point_is_builtin_marker_not_tile_primitive():
+    assert _op_attr("tirx.tvm_kernel_replace_point", "TIRxOpCategory") == "builtin"
+    assert "tirx.tile.tvm_kernel_replace_point" not in Op.list_op_names()
+    assert hasattr(T, "tvm_kernel_replace_point")
+    assert not hasattr(Tx, "tvm_kernel_replace_point")
+
+    @T.prim_func(check_well_formed=False)
+    def marker():
+        T.tvm_kernel_replace_point()
+
+    calls = _expr_calls(marker)
+    assert [call.op.name for call in calls] == ["tirx.tvm_kernel_replace_point"]
+    assert _tile_calls(marker) == []
+
+    code = marker.script()
+    assert "T.tvm_kernel_replace_point()" in code
+    assert "tvm_kernel_replace_point" in code
+    assert "T.tile.tvm_kernel_replace_point" not in code
+    assert "Tx.tvm_kernel_replace_point" not in code
+    reparsed = tvm.script.from_source(code)
+    assert_structural_equal(marker, reparsed)
+
+
+def test_tile_shorthand_and_scoped_aliases_use_tile_ops():
+    @T.prim_func(check_well_formed=False)
+    def tile_aliases(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.tile.copy(A[0:16], B[0:16])
+        Tx.cast(A[0:16], B[0:16])
+        T.cta.cast(A[0:16], B[0:16])
+        Tx.cta.sqrt(A[0:16], B[0:16])
+
+    calls = _tile_calls(tile_aliases)
+    assert [call.op.name for call in calls] == [
+        "tirx.tile.copy",
+        "tirx.tile.cast",
+        "tirx.tile.cast",
+        "tirx.tile.sqrt",
+    ]
+    assert [call.scope.name for call in calls] == ["thread", "thread", "cta", "cta"]
+
+
+def test_device_intrinsic_namespaces_are_canonical_and_classified():
+    from tvm.backend.cuda.script import (
+        CUDANamespace as BackendCUDANamespace,
+    )
+    from tvm.backend.cuda.script import (
+        NVSHMEMNamespace as BackendNVSHMEMNamespace,
+    )
+    from tvm.backend.cuda.script import (
+        PTXNamespace as BackendPTXNamespace,
+    )
+    from tvm.backend.metal.script import MetalNamespace as BackendMetalNamespace
+    from tvm.backend.trn.script import NKINamespace as BackendNKINamespace
+    from tvm.tirx.script.builder import ir as builder_ir
+
+    assert isinstance(builder_ir.cuda, BackendCUDANamespace)
+    assert isinstance(builder_ir.ptx, BackendPTXNamespace)
+    assert isinstance(builder_ir.nvshmem, BackendNVSHMEMNamespace)
+    assert isinstance(builder_ir.metal, BackendMetalNamespace)
+    assert isinstance(builder_ir.nki, BackendNKINamespace)
+    assert T.cuda is builder_ir.cuda
+    assert T.ptx is builder_ir.ptx
+    assert T.nvshmem is builder_ir.nvshmem
+    assert T.metal is builder_ir.metal
+    assert T.nki is builder_ir.nki
+
+    buffer = tvm.tirx.decl_buffer((1,), "float32")
+    calls = [
+        T.ptx.elect_sync(),
+        T.cuda.thread_fence(),
+        T.nvshmem.fence(),
+        T.nki.identity(buffer[0:1], 1),
+    ]
+
+    expected = [
+        ("tirx.ptx.elect_sync", "ptx"),
+        ("tirx.cuda.thread_fence", "cuda"),
+        ("tirx.nvshmem.fence", "nvshmem"),
+        ("tirx.nki.identity", "nki"),
+    ]
+    assert [
+        (call.op.name, _op_attr(call.op.name, "TDeviceIntrinsicNamespace")) for call in calls
+    ] == expected
+    for op_name, namespace in expected:
+        assert _op_attr(op_name, "TIRxOpCategory") == "device_intrin"
+        assert _op_attr(op_name, "TDeviceIntrinsicNamespace") == namespace
+
+
+def test_backend_specific_wrappers_are_not_root_exports():
+    from tvm.tirx.cuda import op as cuda_op
+    from tvm.tirx.metal import op as metal_op
+    from tvm.tirx.trn import op as trn_op
+
+    backend_only_names = [
+        "ptx_mma",
+        "mma_store",
+        "cuda_thread_fence",
+        "nvshmem_fence",
+        "make_filled_simdgroup_matrix",
+        "simdgroup_load",
+        "nki_load",
+    ]
+    for name in backend_only_names:
+        assert not hasattr(tvm.tirx.op, name)
+        assert not hasattr(tvm.tirx, name)
+        assert not hasattr(T, name)
+
+    assert cuda_op.ptx_mma
+    assert cuda_op.mma_store
+    assert cuda_op.cuda_thread_fence
+    assert cuda_op.nvshmem_fence
+    assert metal_op.make_filled_simdgroup_matrix
+    assert metal_op.simdgroup_load
+    assert trn_op.nki_load
+    assert hasattr(T, "cuda")
+    assert hasattr(T, "ptx")
+    assert hasattr(T, "nvshmem")
+    assert hasattr(T, "metal")
+    assert hasattr(T, "nki")
+
+
+def test_backend_load_updates_tirx_alias_and_script_facades(monkeypatch):
+    from tvm.tirx.script import builder, parser
+    from tvm.tirx.script.builder import ir as builder_ir
+
+    backend_name = "unit_test_backend"
+    backend_module_name = f"tvm.backend.{backend_name}"
+    public_module_name = f"tvm.tirx.{backend_name}"
+    public_op_module_name = f"{public_module_name}.op"
+    namespace_name = "unit_test_backend_ns"
+    register_calls = []
+
+    class UnitTestNamespace:
+        pass
+
+    module = types.ModuleType(backend_module_name)
+    module.__path__ = []
+    module.__package__ = backend_module_name
+    op_module = types.ModuleType(f"{backend_module_name}.op")
+    op_module.marker = object()
+
+    def register_backend():
+        register_calls.append(True)
+        builder_ir.register_script_namespace(namespace_name, UnitTestNamespace())
+
+    module.register_backend = register_backend
+    monkeypatch.setitem(sys.modules, backend_module_name, module)
+    monkeypatch.setitem(sys.modules, op_module.__name__, op_module)
+    sys.modules.pop(public_module_name, None)
+    sys.modules.pop(public_op_module_name, None)
+    tvm.backend.loader._LOADED_BACKENDS.pop(backend_name, None)
+    if hasattr(tvm.tirx, backend_name):
+        delattr(tvm.tirx, backend_name)
+
+    with pytest.raises(ModuleNotFoundError):
+        importlib.import_module(public_op_module_name)
+
+    try:
+        assert tvm.backend.load(backend_name) is None
+        assert tvm.backend.load(backend_name) is None
+        assert register_calls == [True]
+        assert tvm.backend.is_loaded(backend_name)
+        assert getattr(tvm.tirx, backend_name) is module
+        assert sys.modules[public_module_name] is module
+        public_op_module = importlib.import_module(public_op_module_name)
+        assert public_op_module.__tvm_backend_module__ is op_module
+        assert public_op_module.marker is op_module.marker
+
+        namespace = getattr(builder_ir, namespace_name)
+        assert isinstance(namespace, UnitTestNamespace)
+        assert getattr(builder, namespace_name) is namespace
+        assert getattr(parser, namespace_name) is namespace
+        assert getattr(T, namespace_name) is namespace
+    finally:
+        tvm.backend.loader._LOADED_BACKENDS.pop(backend_name, None)
+        if hasattr(tvm.tirx, backend_name):
+            delattr(tvm.tirx, backend_name)
+        sys.modules.pop(public_module_name, None)
+        sys.modules.pop(public_op_module_name, None)
+
+
+def test_device_intrinsic_printer_roundtrips_canonical_namespaces():
+    @T.prim_func
+    def device_namespaces(dst: T.handle, src: T.handle):
+        A = T.match_buffer(src, (1,), "float32")
+        R = T.alloc_buffer((1,), "float32", scope="local")
+        T.cuda.cta_sync()
+        T.ptx.ldg32(R[0], 1, A[0], 0)
+        T.metal.simd_shuffle(A[0], 0)
+        T.metal.simd_shuffle_up(A[0], 1)
+        T.metal.simd_shuffle_down(A[0], 1)
+
+    calls = _expr_calls(device_namespaces)
+    assert [call.op.name for call in calls] == [
+        "tirx.cuda.cta_sync",
+        "tirx.ptx.ldg32",
+        "tirx.metal.simd_shuffle",
+        "tirx.metal.simd_shuffle_up",
+        "tirx.metal.simd_shuffle_down",
+    ]
+    for op_name, namespace in [
+        ("tirx.cuda.cta_sync", "cuda"),
+        ("tirx.ptx.ldg32", "ptx"),
+        ("tirx.metal.simd_shuffle", "metal"),
+        ("tirx.metal.simd_shuffle_up", "metal"),
+        ("tirx.metal.simd_shuffle_down", "metal"),
+    ]:
+        assert _op_attr(op_name, "TIRxOpCategory") == "device_intrin"
+        assert _op_attr(op_name, "TDeviceIntrinsicNamespace") == namespace
+        assert _op_attr(op_name, "TCallEffectKind") in (1, 3)
+
+    code = device_namespaces.script()
+    assert "T.cuda.cta_sync(" in code
+    assert "T.ptx.ldg32(" in code
+    assert "T.metal.simd_shuffle(" in code
+    assert "T.metal.simd_shuffle_up(" in code
+    assert "T.metal.simd_shuffle_down(" in code
+    assert "T.tirx." not in code
+    reparsed = tvm.script.from_source(code)
+    assert reparsed.script() == code
+    assert_structural_equal(device_namespaces, reparsed)
+
+
+def test_registered_tirx_ops_have_exactly_one_category():
+    if _op_attr("tirx.sqrt", "TIRxOpCategory") is None:
+        pytest.skip("TIRx op categories require a rebuilt C++ runtime")
+
+    categories = {"builtin", "tile_primitive", "device_intrin"}
+    device_namespaces = {"cuda", "ptx", "nvshmem", "nki", "metal", "webgpu"}
+    flat_tile_only_names = {
+        "tirx.add",
+        "tirx.binary_chain",
+        "tirx.binary_reduce",
+        "tirx.compose_op",
+        "tirx.copy",
+        "tirx.copy_async",
+        "tirx.fdiv",
+        "tirx.fill",
+        "tirx.gemm",
+        "tirx.gemm_async",
+        "tirx.maximum",
+        "tirx.memset",
+        "tirx.minimum",
+        "tirx.mul",
+        "tirx.permute_layout",
+        "tirx.reduce_negate",
+        "tirx.select",
+        "tirx.sub",
+        "tirx.sum",
+        "tirx.unary_reduce",
+        "tirx.zero",
+    }
+
+    missing = []
+    invalid = []
+    lingering_flat_tile = []
+    for op_name in sorted(name for name in Op.list_op_names() if name.startswith("tirx.")):
+        category = _op_attr(op_name, "TIRxOpCategory")
+        device_namespace = _op_attr(op_name, "TDeviceIntrinsicNamespace")
+
+        if category is None:
+            missing.append(op_name)
+            continue
+        if category not in categories:
+            invalid.append((op_name, category))
+            continue
+        if op_name in flat_tile_only_names:
+            lingering_flat_tile.append(op_name)
+
+        if category == "tile_primitive":
+            if not op_name.startswith("tirx.tile."):
+                lingering_flat_tile.append(op_name)
+            assert device_namespace is None, op_name
+        elif category == "device_intrin":
+            assert device_namespace in device_namespaces, op_name
+            printer_name = _op_attr(op_name, "TScriptPrinterName")
+            assert printer_name is not None, op_name
+            assert printer_name.startswith(device_namespace + "."), op_name
+            assert _has_path(T, printer_name), op_name
+        else:
+            assert category == "builtin"
+            assert device_namespace is None, op_name
+
+    assert not missing
+    assert not invalid
+    assert not lingering_flat_tile

tests/python/tirx/test_printer_tir_namespaces.py

@@ -0,0 +1,488 @@
+# 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.
+
+
+import tvm
+from tvm import tirx as tir
+from tvm.script import tirx as T
+from tvm.tirx.cuda import op as cuda_op
+from tvm.tirx.trn import op as trn_op
+
+
+def _assert_print(obj, expected):
+    # Standalone TIR nodes use the canonical tirx script prefix.
+    out = obj.script(verbose_expr=True, extra_config={"tirx.prefix": "T"}).strip()
+    assert out == expected.strip()
+
+
+def test_printer_cuda_namespace_printf():
+    node = tir.Evaluate(cuda_op.cuda_printf("x=%d", tir.IntImm("int32", 1)))
+    _assert_print(node, 'T.cuda.printf("x=%d", 1)')
+
+
+def test_printer_ptx_namespace_wgmma_commit_group():
+    node = tir.Evaluate(cuda_op.ptx_wgmma_commit_group())
+    _assert_print(node, "T.ptx.wgmma.commit_group()")
+
+
+def test_printer_cuda_cluster_sync():
+    node = tir.Evaluate(cuda_op.cuda_cluster_sync())
+    _assert_print(node, "T.cuda.cluster_sync()")
+
+
+def test_printer_ptx_namespace_cp_async_wait_group():
+    node = tir.Evaluate(cuda_op.ptx_cp_async_wait_group(tir.IntImm("int32", 0)))
+    _assert_print(node, "T.ptx.cp_async.wait_group(0)")
+
+
+def test_printer_nvshmem_namespace():
+    node = tir.Evaluate(cuda_op.nvshmem_fence())
+    _assert_print(node, "T.nvshmem.fence()")
+
+
+def test_printer_ptx_more():
+    r = tir.Var("r", "handle")
+    s = tir.Var("s", "handle")
+    _assert_print(
+        # New API: (trans, num, dtype, smem_ptr, *dst_handles).
+        # .x1.b16 has 1 dst register, so 1 dst handle.
+        cuda_op.ptx_ldmatrix(True, 1, ".b16", s, r),
+        's = T.handle()\nr = T.handle()\nT.ptx.ldmatrix(T.bool(True), 1, ".b16", s, r)',
+    )
+    _assert_print(
+        # New API: (trans, num, dtype, smem_ptr, *src_handles).
+        # .x1.b16 has 1 src register, so 1 src handle.
+        cuda_op.ptx_stmatrix(False, 1, ".b16", s, r),
+        (
+            "s = T.handle()\nr = T.handle()\nT.ptx.stmatrix("
+            'T.bool(False), 1, ".b16", "m8n8", "shared", s, r)'
+        ),
+    )
+    _assert_print(cuda_op.ptx_setmaxnreg(True, 64), "T.ptx.setmaxnreg(T.bool(True), 64)")
+    _assert_print(cuda_op.ptx_fetch_register(32, "laneid"), 'T.ptx.fetch_register(32, "laneid")')
+    _assert_print(cuda_op.ptx_wgmma_fence(), "T.ptx.wgmma.fence()")
+    _assert_print(cuda_op.ptx_wgmma_wait_group(0), "T.ptx.wgmma.wait_group(0)")
+    _assert_print(cuda_op.ptx_cp_async_commit_group(), "T.ptx.cp_async.commit_group()")
+    _assert_print(cuda_op.ptx_cp_async_bulk_commit_group(), "T.ptx.cp_async.bulk.commit_group()")
+    _assert_print(
+        cuda_op.ptx_cp_async_bulk_wait_group(0, True),
+        "T.ptx.cp_async.bulk.wait_group(0, T.bool(True))",
+    )
+    _assert_print(cuda_op.ptx_cp_async_mbarrier_arrive(0), "T.ptx.cp_async.mbarrier.arrive(0)")
+    _assert_print(cuda_op.ptx_fence("acq_rel", "gpu"), 'T.ptx.fence("acq_rel", "gpu")')
+    _assert_print(cuda_op.ptx_fence("sc", "cta"), 'T.ptx.fence("sc", "cta")')
+    _assert_print(
+        cuda_op.ptx_fence_proxy_async("shared::cta"), 'T.ptx.fence.proxy_async("shared::cta")'
+    )
+    _assert_print(cuda_op.ptx_fence_proxy_async("global"), 'T.ptx.fence.proxy_async("global")')
+    _assert_print(cuda_op.ptx_fence_mbarrier_init(), "T.ptx.fence.mbarrier_init()")
+    _assert_print(cuda_op.ptx_elect_sync(), "T.ptx.elect_sync()")
+    lane = tir.Var("lane", "int32")
+    _assert_print(
+        tir.op.selector(lane, cuda_op.ptx_elect_sync()),
+        "lane = T.int32()\nT.selector(lane, T.ptx.elect_sync())",
+    )
+    _assert_print(
+        cuda_op.ptx_ld_global_acquire(r, s),
+        "r = T.handle()\ns = T.handle()\nT.ptx.ld_global_acquire(r, s)",
+    )
+    _assert_print(
+        cuda_op.ptx_map_shared_rank(r, 2), 'r = T.handle()\nT.ptx.mapa(r, 2, "", "u64", "uint64")'
+    )
+    _assert_print(cuda_op.ptx_bar_arrive(0, 128), "T.ptx.bar.arrive(0, 128)")
+    _assert_print(cuda_op.ptx_bar_sync(0, 128), "T.ptx.bar.sync(0, 128)")
+    _assert_print(
+        cuda_op.ptx_tcgen05_alloc(s, 64, 1), "s = T.handle()\nT.ptx.tcgen05.alloc(s, 64, 1)"
+    )
+    _assert_print(
+        cuda_op.ptx_tcgen05_dealloc(s, 64, 1), "s = T.handle()\nT.ptx.tcgen05.dealloc(s, 64, 1)"
+    )
+    d = tir.Var("d", "handle")
+    a = tir.Var("a", "handle")
+    b = tir.Var("b", "handle")
+    _assert_print(
+        cuda_op.ptx_tcgen05_encode_matrix_descriptor(d, a, 1, 2, 0),
+        "d = T.handle()\na = T.handle()\nT.ptx.tcgen05.encode_matrix_descriptor(d, a, 1, 2, 0)",
+    )
+    _assert_print(
+        cuda_op.ptx_tcgen05_encode_instr_descriptor(
+            d,
+            d_dtype="f16",
+            a_dtype="f16",
+            b_dtype="f16",
+            M=16,
+            N=16,
+            K=16,
+            trans_a=True,
+            trans_b=False,
+            n_cta_groups=1,
+            neg_a=False,
+            neg_b=False,
+            sat_d=False,
+            is_sparse=False,
+        ),
+        'd = T.handle()\nT.ptx.tcgen05.encode_instr_descriptor(d, "f16", "f16", "f16", 16, 16, 16, T.bool(True), T.bool(False), 1, T.bool(False), T.bool(False), T.bool(False), T.bool(False))',  # noqa: E501
+    )
+    _assert_print(
+        cuda_op.ptx_tcgen05_encode_instr_descriptor_block_scaled(
+            d,
+            d_dtype="f16",
+            a_dtype="f16",
+            b_dtype="f16",
+            sfa_dtype="f16",
+            sfb_dtype="f16",
+            sfa_tmem_addr=a,
+            sfb_tmem_addr=b,
+            M=16,
+            N=16,
+            K=16,
+            trans_a=True,
+            trans_b=False,
+            is_sparse=True,
+            n_cta_groups=1,
+            neg_a=False,
+            neg_b=False,
+        ),
+        "d = T.handle()\n"
+        "a = T.handle()\n"
+        "b = T.handle()\n"
+        'T.ptx.tcgen05.encode_instr_descriptor_block_scaled(d, "f16", "f16", "f16", "f16", "f16", a, b, 16, 16, 16, T.bool(True), T.bool(False), 1, T.bool(False), T.bool(False), T.bool(True))',  # noqa: E501
+    )
+    _assert_print(
+        cuda_op.ptx_tcgen05_cp(a, d, shape="64x128b", cta_group=1, multicast="warpx2::02_13"),
+        "a = T.handle()\n"
+        "d = T.handle()\n"
+        'T.ptx.tcgen05.cp(a, d, "64x128b", 1, "warpx2::02_13", "", 0, 0)',
+    )
+    _assert_print(cuda_op.ptx_tcgen05_shift(a, 1), "a = T.handle()\nT.ptx.tcgen05.shift(a, 1)")
+    _assert_print(
+        cuda_op.ptx_tcgen05_ld(a, 0, shape="16x64b", num=1, row=0, col=0, pack=False),
+        'a = T.handle()\nT.ptx.tcgen05.ld(a, 0, 0, "16x64b", 1, T.bool(False), 0)',
+    )
+    _assert_print(
+        cuda_op.ptx_tcgen05_st(a, 0, shape="16x64b", num=1, row=0, col=0, unpack=False),
+        'a = T.handle()\nT.ptx.tcgen05.st(a, 0, 0, "16x64b", 1, T.bool(False), 0)',
+    )
+    _assert_print(cuda_op.ptx_tcgen05_wait_ld(), "T.ptx.tcgen05.wait.ld()")
+    _assert_print(cuda_op.ptx_tcgen05_wait_st(), "T.ptx.tcgen05.wait.st()")
+    _assert_print(
+        cuda_op.ptx_tcgen05_commit(a, 1, 0), "a = T.handle()\nT.ptx.tcgen05.commit(a, 1, 0)"
+    )
+    _assert_print(
+        cuda_op.ptx_tcgen05_relinquish_alloc_permit(1), "T.ptx.tcgen05.relinquish_alloc_permit(1)"
+    )
+
+
+def test_printer_ptx_mbarrier():
+    bar = tir.Var("bar", "handle")
+    _assert_print(
+        cuda_op.ptx_mbarrier_init(bar, 32), "bar = T.handle()\nT.ptx.mbarrier.init(bar, 32)"
+    )
+    _assert_print(cuda_op.ptx_mbarrier_arrive(bar), "bar = T.handle()\nT.ptx.mbarrier.arrive(bar)")
+    _assert_print(
+        cuda_op.ptx_mbarrier_arrive_expect_tx(bar, 128),
+        "bar = T.handle()\nT.ptx.mbarrier.arrive.expect_tx(bar, 128)",
+    )
+    _assert_print(
+        cuda_op.ptx_mbarrier_try_wait(bar, 1), "bar = T.handle()\nT.ptx.mbarrier.try_wait(bar, 1)"
+    )
+    _assert_print(cuda_op.cuda_cluster_sync(), "T.cuda.cluster_sync()")
+
+
+def test_printer_cuda_more():
+    p = tir.Var("p", "handle")
+    _assert_print(cuda_op.cuda_thread_fence(), "T.cuda.thread_fence()")
+    _assert_print(cuda_op.cuda_warp_sync(), "T.cuda.warp_sync()")
+    _assert_print(cuda_op.cuda_cta_sync(), "T.cuda.cta_sync()")
+    _assert_print(cuda_op.cuda_grid_sync(), "T.cuda.grid_sync()")
+    _assert_print(cuda_op.cuda_cluster_sync(), "T.cuda.cluster_sync()")
+    _assert_print(cuda_op.cuda_syncthreads_and(1), "T.cuda.syncthreads_and(1)")
+    _assert_print(cuda_op.cuda_syncthreads_or(1), "T.cuda.syncthreads_or(1)")
+    _assert_print(cuda_op.cuda_nano_sleep(100), "T.cuda.nano_sleep(100)")
+    _assert_print(
+        cuda_op.cuda_atomic_add(p, tir.IntImm("int32", 1)),
+        "p = T.handle()\nT.cuda.atomic_add(p, 1)",
+    )
+    _assert_print(cuda_op.cuda_atomic_cas(p, 1, 2), "p = T.handle()\nT.cuda.atomic_cas(p, 1, 2)")
+    _assert_print(cuda_op.cuda_ldg(p, "float32"), 'p = T.handle()\nT.cuda.ldg(p, "float32")')
+    _assert_print(
+        cuda_op.cuda_func_call("f", 1, source_code=""), 'T.cuda.func_call("f", 1, source_code="")'
+    )
+
+
+def test_printer_cuda_low_level_warp_intrinsics_roundtrip():
+    @T.prim_func(check_well_formed=False)
+    def kernel():
+        x = T.int32()
+        mask = T.cuda.__activemask()
+        T.evaluate(T.cuda.__shfl_sync(mask, x, 0, 32))
+        T.evaluate(T.cuda.__shfl_up_sync(mask, x, 1, 32))
+        T.evaluate(T.cuda.__shfl_down_sync(mask, x, 1, 32))
+        T.evaluate(T.cuda.__shfl_xor_sync(mask, x, 1, 32))
+
+    code = kernel.script()
+    assert "T.cuda.__activemask()" in code
+    assert "T.cuda.__shfl_sync(" in code
+    assert "T.cuda.__shfl_up_sync(" in code
+    assert "T.cuda.__shfl_down_sync(" in code
+    assert "T.cuda.__shfl_xor_sync(" in code
+    assert "T.tirx." not in code
+    assert tvm.script.from_source(code).script() == code
+
+
+def test_printer_webgpu_namespace_roundtrip():
+    @T.prim_func(check_well_formed=False)
+    def kernel():
+        x = T.int32()
+        T.evaluate(T.webgpu.subgroup_shuffle(x, 0))
+        T.evaluate(T.webgpu.subgroup_shuffle_up(x, 1))
+        T.evaluate(T.webgpu.subgroup_shuffle_down(x, 1))
+
+    code = kernel.script()
+    assert "T.webgpu.subgroup_shuffle(" in code
+    assert "T.webgpu.subgroup_shuffle_up(" in code
+    assert "T.webgpu.subgroup_shuffle_down(" in code
+    assert "T.tirx." not in code
+    assert tvm.script.from_source(code).script() == code
+
+
+def test_printer_nvshmem_more():
+    p = tir.Var("p", "handle")
+    _assert_print(cuda_op.nvshmem_my_pe(), "T.nvshmem.my_pe()")
+    _assert_print(cuda_op.nvshmem_n_pes(), "T.nvshmem.n_pes()")
+    _assert_print(
+        cuda_op.nvshmem_signal_op(p, 1, "set", 0),
+        'p = T.handle()\nT.nvshmem.signal_op(p, 1, "set", 0)',
+    )
+    _assert_print(
+        cuda_op.nvshmem_wait_until(p, "eq", 0),
+        'p = T.handle()\nT.nvshmem.wait_until(p, "eq", 0, "uint64_t")',
+    )
+    _assert_print(cuda_op.nvshmem_quiet(), "T.nvshmem.quiet()")
+    _assert_print(cuda_op.nvshmem_barrier_all(), "T.nvshmem.barrier_all()")
+    _assert_print(
+        cuda_op.nvshmem_getmem_nbi(p, p, 16, 0),
+        "p = T.handle()\nT.nvshmem.getmem_nbi(p, p, 16, 0)",
+    )
+    _assert_print(
+        cuda_op.nvshmem_getmem_nbi_warp(p, p, 16, 0),
+        "p = T.handle()\nT.nvshmem.getmem_nbi.warp(p, p, 16, 0)",
+    )
+    _assert_print(
+        cuda_op.nvshmem_putmem_nbi_block(p, p, 16, 0),
+        "p = T.handle()\nT.nvshmem.putmem_nbi.block(p, p, 16, 0)",
+    )
+    _assert_print(
+        cuda_op.nvshmem_putmem_nbi(p, p, 16, 0),
+        "p = T.handle()\nT.nvshmem.putmem_nbi(p, p, 16, 0)",
+    )
+    _assert_print(
+        cuda_op.nvshmem_putmem_nbi_warp(p, p, 16, 0),
+        "p = T.handle()\nT.nvshmem.putmem_nbi.warp(p, p, 16, 0)",
+    )
+    _assert_print(
+        cuda_op.nvshmem_putmem_signal_nbi(p, p, 16, p, 1, "set", 0),
+        'p = T.handle()\nT.nvshmem.putmem_signal_nbi(p, p, 16, p, 1, "set", 0)',
+    )
+    _assert_print(
+        cuda_op.nvshmem_putmem_signal_nbi_warp(p, p, 16, p, 1, "set", 0),
+        'p = T.handle()\nT.nvshmem.putmem_signal_nbi.warp(p, p, 16, p, 1, "set", 0)',
+    )
+    _assert_print(
+        cuda_op.nvshmem_putmem_signal_nbi_block(p, p, 16, p, 1, "set", 0),
+        'p = T.handle()\nT.nvshmem.putmem_signal_nbi.block(p, p, 16, p, 1, "set", 0)',
+    )
+
+
+def test_printer_nki_namespace():
+    A = tir.decl_buffer([1], dtype="float16", name="A")
+    B = tir.decl_buffer([1], dtype="float16", name="B")
+    a0 = A[0]
+    b0 = B[0]
+    _assert_print(
+        trn_op.nki_load(a0, b0),
+        'A = T.Buffer((1,), "float16")\nB = T.Buffer((1,), "float16")\nT.nki.load(A, B)',
+    )
+    _assert_print(
+        trn_op.nki_store(a0, b0),
+        'A = T.Buffer((1,), "float16")\nB = T.Buffer((1,), "float16")\nT.nki.store(A, B)',
+    )
+    _assert_print(
+        trn_op.nki_tensor_copy(a0, b0),
+        'A = T.Buffer((1,), "float16")\nB = T.Buffer((1,), "float16")\nT.nki.tensor_copy(A, B)',
+    )
+    _assert_print(
+        trn_op.nki_matmul(a0, a0, b0),
+        'A = T.Buffer((1,), "float16")\n'
+        'B = T.Buffer((1,), "float16")\n'
+        "T.nki.matmul(A, A, B, T.bool(True))",
+    )
+    _assert_print(
+        trn_op.nki_activation(a0, b0, "relu", 0.0, 1.0),
+        'A = T.Buffer((1,), "float16")\n'
+        'B = T.Buffer((1,), "float16")\n'
+        'T.nki.activation(A, B, "relu", T.float32(0.0), T.float32(1.0))',
+    )
+    _assert_print(
+        trn_op.nki_memset(a0, 0),
+        'A = T.Buffer((1,), "float16")\nT.nki.memset(A, 0)',
+    )
+    _assert_print(
+        trn_op.nki_identity(a0, 1),
+        'A = T.Buffer((1,), "float16")\nT.nki.identity(A, 1)',
+    )
+    _assert_print(
+        trn_op.nki_reciprocal(a0, b0),
+        'A = T.Buffer((1,), "float16")\nB = T.Buffer((1,), "float16")\nT.nki.reciprocal(A, B)',
+    )
+    _assert_print(
+        trn_op.nki_tensorreduce(a0, b0, "sum", False, 0),
+        'A = T.Buffer((1,), "float16")\n'
+        'B = T.Buffer((1,), "float16")\n'
+        'T.nki.tensorreduce(A, B, "sum", T.bool(False), 0)',
+    )
+    _assert_print(
+        trn_op.nki_tensortensor(a0, a0, b0, "add"),
+        'A = T.Buffer((1,), "float16")\n'
+        'B = T.Buffer((1,), "float16")\n'
+        'T.nki.tensortensor(A, A, B, "add")',
+    )
+    _assert_print(
+        trn_op.nki_tensorscalar(a0, a0, 1.0, "mul", False),
+        'A = T.Buffer((1,), "float16")\n'
+        'T.nki.tensorscalar(A, A, T.float32(1.0), "mul", T.bool(False))',
+    )
+    _assert_print(
+        trn_op.nki_tensorscalar_reduce(a0, a0, 1.0, "mul", "sum", False),
+        'A = T.Buffer((1,), "float16")\n'
+        'T.nki.tensorscalar_reduce(A, A, T.float32(1.0), "mul", "sum", T.bool(False), T.bool(False))',  # noqa: E501
+    )
+    _assert_print(
+        trn_op.nki_scalar_tensor_tensor(a0, a0, 1.0, a0, "add", "add"),
+        'A = T.Buffer((1,), "float16")\n'
+        'T.nki.scalar_tensor_tensor(A, A, T.float32(1.0), A, "add", "add", T.bool(False), T.bool(False))',  # noqa: E501
+    )
+    _assert_print(
+        trn_op.nki_scalar_tensor_scalar(a0, a0, 1.0, 1.0, "add", "add"),
+        'A = T.Buffer((1,), "float16")\n'
+        'T.nki.scalar_tensor_scalar(A, A, T.float32(1.0), T.float32(1.0), "add", "add", T.bool(False), T.bool(False))',  # noqa: E501
+    )
+    _assert_print(
+        trn_op.nki_activation_reduce(a0, a0, b0, "relu", "sum", 0.0, 1.0),
+        'A = T.Buffer((1,), "float16")\n'
+        'B = T.Buffer((1,), "float16")\n'
+        'T.nki.activation_reduce(A, A, B, "relu", "sum", T.float32(0.0), T.float32(1.0))',
+    )
+    _assert_print(
+        trn_op.nki_affine_select(a0, a0, a0, 1.0),
+        'A = T.Buffer((1,), "float16")\nT.nki.affine_select(A, A, A, T.float32(1.0))',
+    )
+
+
+def test_printer_ptx_mma_and_wgmma():
+    r = tir.Var("r", "handle")
+    d = tir.Var("d", "handle")
+    a = tir.Var("a", "handle")
+    tir.Var("b", "handle")
+    _assert_print(
+        cuda_op.ptx_mma("m8n8k4", "row", "row", "fp16", "fp16", "fp16", "fp16", [r], [r], [r]),
+        'r = T.handle()\nT.ptx.mma("m8n8k4", "row", "row", "fp16", "fp16", "fp16", "fp16", 1, 1, 1, 0, T.bool(True), r, r, r, T.bool(False))',  # noqa: E501
+    )
+    _assert_print(
+        cuda_op.ptx_wgmma_encode_matrix_descriptor(d, a, 1, 1, 0),
+        "d = T.handle()\na = T.handle()\nT.ptx.wgmma.encode_matrix_descriptor(d, a, 1, 1, 0)",
+    )
+    _assert_print(cuda_op.ptx_wgmma_noop_barrier(0), "T.ptx.wgmma.noop_barrier(0)")
+    _assert_print(
+        cuda_op.ptx_wgmma_mma_async_ss(
+            d,
+            d,
+            0,
+            0,
+            M=16,
+            N=16,
+            K=16,
+            in_dtype="f16",
+            out_dtype="f16",
+            transA=True,
+            transB=False,
+            scaleA=1.0,
+            scaleB=1.0,
+            scaleD=True,
+        ),
+        'd = T.handle()\nT.ptx.wgmma.mma_async.ss(16, 16, 16, "f16", "f16", T.bool(True), T.bool(False), T.float32(1.0), T.float32(1.0), T.bool(True), d, d, 0, 0)',  # noqa: E501
+    )
+    _assert_print(
+        cuda_op.ptx_wgmma_mma_async_rs(
+            d,
+            0,
+            0,
+            M=16,
+            N=16,
+            K=16,
+            in_dtype="f16",
+            out_dtype="f16",
+            transA=True,
+            transB=False,
+            scaleA=1.0,
+            scaleB=1.0,
+            scaleD=True,
+        ),
+        'd = T.handle()\nT.ptx.wgmma.mma_async.rs(16, 16, 16, "f16", "f16", T.bool(True), T.bool(False), T.float32(1.0), T.float32(1.0), T.bool(True), d, 0, 0)',  # noqa: E501
+    )
+
+
+def test_printer_ptx_cp_async_tensor():
+    tmap = tir.Var("tm", "handle")
+    _assert_print(
+        cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster(2, tmap, 0, tmap, 0, 1, "", 0, 1, ""),
+        "tm = T.handle()\n"
+        'T.ptx.cp_async.bulk.tensor.g2c(2, tm, 0, tm, 0, 1, T.uint64(0), 0, 0, 1, "")',
+    )
+    _assert_print(
+        cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster(
+            2, tmap, 0, tmap, 0, 1, "", 0, 1, ""
+        ),
+        "tm = T.handle()\n"
+        "T.ptx.cp_async.bulk.tensor.g2c_tile_gather4"
+        '(2, tm, 0, tm, 0, 1, T.uint64(0), 0, 0, 1, "")',
+    )
+    _assert_print(
+        cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch(2, tmap, "", 0, 0, ""),
+        'tm = T.handle()\nT.ptx.cp_async.bulk.tensor.g2c_prefetch(2, tm, T.uint64(0), 0, 0, 0, "")',
+    )
+    _assert_print(
+        cuda_op.ptx_cp_async_bulk_tensor_shared_to_global(2, 0, tmap, "", 0, 0, ""),
+        'tm = T.handle()\nT.ptx.cp_async.bulk.tensor.s2g(2, 0, tm, T.uint64(0), 0, 0, 0, "")',
+    )
+    _assert_print(
+        cuda_op.ptx_cp_async_bulk_tensor_shared_to_global_reduce(2, 0, tmap, "", "add", 0, 0, ""),
+        "tm = T.handle()\n"
+        "T.ptx.cp_async.bulk.tensor.s2g_reduce"
+        '(2, 0, tm, T.uint64(0), 0, "add", 0, 0, "")',
+    )
+
+
+def test_printer_ptx_cp_async_call():
+    sh = tir.Var("sh", "handle")
+    gl = tir.Var("gl", "handle")
+    _assert_print(
+        cuda_op.ptx_cp_async(
+            sh, gl, 16, cache_hint="", prefetch_size=-1, predicate=-1, fill_mode=""
+        ),
+        'sh = T.handle()\ngl = T.handle()\nT.ptx.cp_async(sh, gl, 16, T.uint64(0), 0, -1, -1, "")',
+    )

tests/python/tirx/test_roundtrip_namespaces.py

@@ -0,0 +1,43 @@
+# 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.
+
+import tvm
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+
+
+def from_source(code):
+    return tvm.script.from_source(code)
+
+
+def test_roundtrip_tir_namespaces_minimal():
+    # Exercise a selection of namespace ops and ensure round-trip consistency
+    @T.prim_func
+    def func(a_ptr: T.handle) -> None:
+        A = T.match_buffer(a_ptr, (2, 2), "float16")
+        T.ptx.wgmma.commit_group()
+        T.cuda.cluster_sync()
+        T.ptx.cp_async.wait_group(0)
+        T.ptx.fence.proxy_async("shared::cta")
+        T.cuda.printf("ok")
+        T.nvshmem.quiet()
+        T.nki.identity(A[0, 0], 1)
+
+    code = func.script()
+    roundtripped = from_source(code)
+    assert roundtripped.script() == code
+    assert_structural_equal(func, roundtripped)

tests/python/tirx/test_verifier.py

@@ -0,0 +1,371 @@
+# 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.
+import pytest
+
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.analysis import verify_tirx_well_formed as verify
+
+
+def test_root_scope():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1() -> None:
+        T.device_entry()
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test2() -> None:
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test3() -> None:
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test4() -> None:
+        T.device_entry()
+        pass
+
+        # fmt: on
+
+    verify(test1)
+    verify(test2)
+    verify(test3)
+    verify(test4)
+
+
+def test_nested_scope():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1() -> None:
+        T.device_entry()
+        pass
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test2() -> None:
+        T.device_entry()
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test3() -> None:
+        T.device_entry()
+        pass
+    @T.prim_func(check_well_formed=False)
+    def test4() -> None:
+        T.device_entry()
+        pass
+        pass
+
+        # fmt: on
+
+    verify(test1)
+    verify(test2)
+    verify(test3)
+    verify(test4)
+
+
+def test_scope_id_consistency():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test2():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        T.thread_id([128])
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test3():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([2])
+        T.lane_id([32])
+        T.thread_id([128])
+        pass
+
+    @T.prim_func(check_well_formed=False)
+    def test4():
+        T.device_entry()
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        clx, cly, clz = T.cluster_id([4, 5, 12])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+    @T.prim_func(check_well_formed=False)
+    def test5():
+        T.device_entry()
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        clx, cly, clz = T.cluster_id([3, 5, 12])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+    @T.prim_func(check_well_formed=False)
+    def test6():
+        T.device_entry()
+        clx, cly, clz = T.cluster_id([4, 5, 12])
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+    @T.prim_func(check_well_formed=False)
+    def test7():
+        T.device_entry()
+        clx, cly, clz = T.cluster_id([3, 5, 12])
+        bx, by, bz = T.cta_id([8, 10, 12])
+        cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
+        T.evaluate(bx + by + bz)
+        T.evaluate(cbx + cby + cbz)
+        T.evaluate(clx + cly + clz)
+
+        # fmt: on
+
+    verify(test1)
+    verify(test2)
+    with pytest.raises(Exception, match="Inconsistent extents for scope"):
+        verify(test3)
+    verify(test4)
+    with pytest.raises(Exception, match="Inconsistent extents|non-divisible extents"):
+        verify(test5)
+    verify(test6)
+    with pytest.raises(Exception, match="Inconsistent extents|non-divisible extents"):
+        verify(test7)
+
+
+def test_layout():
+    ### TileLayout
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        A = T.alloc_buffer((2,), layout=T.TileLayout(T.S[2, 1]))
+
+        A[0] = 0
+        # fmt: on
+    verify(test1)
+
+    ### SwizzleLayout
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test2():
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        A = T.alloc_buffer((512,), scope="shared", layout=T.SwizzleLayout(3, 3, 3))
+
+        A[0] = 0
+        # fmt: on
+    verify(test2)
+
+
+def test_host():
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1(A_ptr: T.handle):
+        A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)
+
+        A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
+        T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, "float32", 2, A.data, 16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0)  # noqa: E501
+
+        T.device_entry()
+        for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
+            for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
+                bar = T.alloc_buffer((1,), "uint64", scope="shared", align=8)
+                phase = T.alloc_buffer((1,), "int32", scope="local")
+                A_smem = T.alloc_buffer((16, 16), "float32", scope="shared", align=128)
+
+                phase[0] = 0
+                if threadIdx == 0:
+                    T.ptx.mbarrier.init(bar.data, 1)
+                    T.ptx.fence.proxy_async("shared::cta")
+                    T.ptx.cp_async.bulk.tensor.g2c(2, A_smem.data, bar.data, T.address_of(A_map), 0, 1, "", 0, 0)  # noqa: E501
+                    T.ptx.mbarrier.arrive.expect_tx(bar.data, 16*16*4)
+                T.ptx.mbarrier.try_wait(bar.data, phase[0])
+                phase[0] = phase[0] ^ 1
+                T.print_buffer(A_smem.data, "float32", False, False, 2, 16*16)
+        # fmt: on
+    verify(test1)
+
+
+def test_device_func():
+    # Per-call exec-scope migration: scope is now attached per op via the
+    # ``T.op[scope](...)`` subscription surface instead of a ``with T.cta():``
+    # region. ``test1`` exercises a per-call-scoped op; ``test2`` the plain
+    # (unscoped) op. The old multi-root-scope negative case asserted the removed
+    # "only one root scope" verifier rule and no longer has an equivalent, so it
+    # is dropped.
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def test1(A: T.Buffer((128,), "float32")):
+        T.device_entry()
+        T.cta_id([1])
+        T.thread_id([128])
+        Tx.cta.fill(A, 0.)
+
+    @T.prim_func(check_well_formed=False)
+    def test2(A: T.Buffer((128,), "float32")):
+        T.device_entry()
+        T.cta_id([128])
+        T.thread_id([128])
+        Tx.fill(A, 0.)
+    # fmt: on
+    verify(test1, device_func=True)
+    verify(test2, device_func=True)
+
+
+def test_preferred_cluster_validation():
+    # fmt: off
+    # Valid: cluster→cta with preferred_extents matching size
+    @T.prim_func(check_well_formed=False)
+    def test1() -> None:
+        T.device_entry()
+        cbx, cby = T.cta_id_in_cluster([2, 1], preferred=[2, 2])
+        tx = T.thread_id([128])
+        T.evaluate(cbx + cby + tx)
+
+        # Invalid: preferred size doesn't match extents size (caught at verify time)
+    @T.prim_func(check_well_formed=False)
+    def test2() -> None:
+        T.device_entry()
+        cbx, cby = T.cta_id_in_cluster([2, 1], preferred=[2])
+        tx = T.thread_id([128])
+        T.evaluate(cbx + cby + tx)
+        # fmt: on
+
+    verify(test1)
+    with pytest.raises(Exception, match="preferred_extents must have the same size"):
+        verify(test2)
+
+    # Invalid: preferred on a non-cluster→cta scope (caught at IR build time)
+    with pytest.raises(Exception):
+        # fmt: off
+        @T.prim_func(check_well_formed=False)
+        def test3() -> None:
+            T.device_entry()
+            bx = T.cta_id([128], preferred=[256])
+            tx = T.thread_id([128])
+            T.evaluate(bx + tx)
+            # fmt: on
+
+
+def test_scope_id_deferred_relaxed_at_construction():
+    """Deferred scope_id (no extents) must pass the well-formed check even when
+    no sibling provides enough info to resolve it -- strict resolution is
+    deferred to LowerTIRx."""
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def partial_only_cta():
+        T.device_entry()
+        bx = T.cta_id()           # deferred kernel→cta, no closure source
+        tx = T.thread_id([128])   # explicit
+        T.evaluate(bx + tx)
+
+    @T.prim_func(check_well_formed=False)
+    def all_deferred():
+        T.device_entry()
+        bx = T.cta_id()
+        wg = T.warpgroup_id()
+        warp = T.warp_id_in_wg()
+        lane = T.lane_id()
+        T.evaluate(bx + wg + warp + lane)
+
+    @T.prim_func(check_well_formed=False)
+    def mixed():
+        T.device_entry()
+                # kCtaWarp=4, kWarpThread=32 → kCtaThread=128 derivable.
+        T.warp_id([4])
+        T.lane_id([32])
+        T.thread_id()             # deferred kCtaThread, resolvable via closure
+        pass
+        # fmt: on
+
+        # All three accepted by well-formed: deferred extents are tolerated.
+    verify(partial_only_cta)
+    verify(all_deferred)
+    verify(mixed)
+
+
+def test_scope_id_deferred_consistency_still_enforced():
+    """Even with deferred defs, known-known consistency between sibling defs
+    must still be enforced by the closure check."""
+
+    # fmt: off
+    @T.prim_func(check_well_formed=False)
+    def inconsistent():
+        # 4 warps * 32 lanes = 128 threads, but explicit thread_id says 64 -> error.
+        T.device_entry()
+        T.cta_id([32])
+        T.warp_id([4])
+        T.lane_id([32])
+        T.thread_id()       # deferred (shouldn't shadow the conflict)
+        T.thread_id([64])   # conflicts with derived kCtaThread=128
+        pass
+        # fmt: on
+
+    with pytest.raises(Exception, match="Inconsistent extents for scope"):
+        verify(inconsistent)
+
+
+def test_scope_id_deferred_multi_var_rejected():
+    """Deferred form (no extents) requires exactly one Var. Multi-var defers
+    have no well-defined recovery from fused closure values."""
+
+    # The C++ ScopeIdDef ctor enforces this; constructing such a def from the
+    # parser path is not currently expressible (parser only emits single-Var
+    # deferred), but we exercise the FFI-level guard directly.
+    from tvm.tirx.exec_scope import ScopeIdDef
+    from tvm.tirx.expr import Var
+
+    # Single-Var deferred form is fine.
+    ScopeIdDef([Var("", "int32")], None, "kernel", "cta")
+
+    # Two-Var deferred should be rejected.
+    with pytest.raises(Exception, match="Deferred ScopeIdDef.*must define exactly one Var"):
+        ScopeIdDef([Var("", "int32"), Var("", "int32")], None, "kernel", "cta")
+
+
+if __name__ == "__main__":
+    test_root_scope()
+    test_nested_scope()
+    test_scope_id_consistency()
+    test_layout()
+    test_host()
+    test_device_func()
+    test_scope_id_deferred_relaxed_at_construction()
+    test_scope_id_deferred_consistency_still_enforced()
+    test_scope_id_deferred_multi_var_rejected()

tests/python/tirx/transform/test_tirx_expr_functor.py

@@ -0,0 +1,829 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm import tirx as tir
+from tvm.ir import Call, Op
+from tvm.ir.base import assert_structural_equal
+from tvm.tirx.expr import (
+    EQ,
+    GE,
+    GT,
+    LE,
+    LT,
+    NE,
+    Add,
+    And,
+    Broadcast,
+    BufferLoad,
+    Cast,
+    Div,
+    FloatImm,
+    FloorDiv,
+    FloorMod,
+    IntImm,
+    Let,
+    Max,
+    Min,
+    Mod,
+    Mul,
+    Not,
+    Or,
+    ProducerLoad,
+    Ramp,
+    Reduce,
+    Select,
+    Shuffle,
+    StringImm,
+    Sub,
+    Var,
+)
+from tvm.tirx.expr_functor import ExprMutator, ExprVisitor
+
+# Basic example variables for testing
+n = tir.Var("n", "int32")
+m = tir.Var("m", "int32")
+x = tir.Var("x", "float32")
+y = tir.Var("y", "float32")
+
+
+class BasicVisitor(ExprVisitor):
+    """Default ExprVisitor"""
+
+
+class ASTLog:
+    """Helper class to log AST"""
+
+    def __init__(self) -> None:
+        self.log = []
+        self.indent = "\t"
+        self.level = 0
+
+    def push_scope(self):
+        self.level += 1
+
+    def pop_scope(self):
+        self.level -= 1
+
+    def add(self, s: str):
+        self.log.append(self.indent * self.level + s)
+
+    def __str__(self) -> str:
+        return "\n".join(self.log)
+
+
+class ASTPrinter(ExprVisitor):
+    """Print TIR AST in structured format."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> None:
+        self.log.add("Var")
+
+    def visit_buffer_load_(self, op: BufferLoad) -> None:
+        self.log.add("BufferLoad")
+        self.log.push_scope()
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_producer_load_(self, op: ProducerLoad) -> None:
+        self.log.add("ProducerLoad")
+        self.log.push_scope()
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_let_(self, op: Let) -> None:
+        self.log.add("Let")
+        self.log.push_scope()
+        self.visit_expr(op.var)
+        self.visit_expr(op.value)
+        self.visit_expr(op.body)
+        self.log.pop_scope()
+
+    def visit_call_(self, op: Call) -> None:
+        self.log.add("Call")
+        self.log.push_scope()
+        if isinstance(op.op, Op):
+            self.log.add("Op")
+        else:
+            self.visit_expr(op.op)
+        for arg in op.args:
+            self.visit_expr(arg)
+        self.log.pop_scope()
+
+    def visit_add_(self, op: Add) -> None:
+        self.log.add("Add")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_sub_(self, op: Sub) -> None:
+        self.log.add("Sub")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_mul_(self, op: Mul) -> None:
+        self.log.add("Mul")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_div_(self, op: Div) -> None:
+        self.log.add("Div")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_mod_(self, op: Mod) -> None:
+        self.log.add("Mod")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_floordiv_(self, op: FloorDiv) -> None:
+        self.log.add("FloorDiv")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_floormod_(self, op: FloorMod) -> None:
+        self.log.add("FloorMod")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_min_(self, op: Min) -> None:
+        self.log.add("Min")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_max_(self, op: Max) -> None:
+        self.log.add("Max")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_eq_(self, op: EQ) -> None:
+        self.log.add("EQ")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_ne_(self, op: NE) -> None:
+        self.log.add("NE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_lt_(self, op: LT) -> None:
+        self.log.add("LT")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_le_(self, op: LE) -> None:
+        self.log.add("LE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_gt_(self, op: GT) -> None:
+        self.log.add("GT")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_ge_(self, op: GE) -> None:
+        self.log.add("GE")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_and_(self, op: And) -> None:
+        self.log.add("And")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_or_(self, op: Or) -> None:
+        self.log.add("Or")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.visit_expr(op.b)
+        self.log.pop_scope()
+
+    def visit_reduce_(self, op: Reduce) -> None:
+        self.log.add("Reduce")
+        self.log.push_scope()
+        for source in op.source:
+            self.visit_expr(source)
+        for axis in op.axis:
+            self.visit_expr(axis.var)
+        self.visit_expr(op.condition)
+        self.log.pop_scope()
+
+    def visit_cast_(self, op: Cast) -> None:
+        self.log.add("Cast")
+        self.log.push_scope()
+        self.visit_expr(op.value)
+        self.log.pop_scope()
+
+    def visit_not_(self, op: Not) -> None:
+        self.log.add("Not")
+        self.log.push_scope()
+        self.visit_expr(op.a)
+        self.log.pop_scope()
+
+    def visit_select_(self, op: Select) -> None:
+        self.log.add("Select")
+        self.log.push_scope()
+        self.visit_expr(op.condition)
+        self.visit_expr(op.true_value)
+        self.visit_expr(op.false_value)
+        self.log.pop_scope()
+
+    def visit_ramp_(self, op: Ramp) -> None:
+        self.log.add("Ramp")
+        self.log.push_scope()
+        self.visit_expr(op.base)
+        self.visit_expr(op.stride)
+        self.visit_expr(op.lanes)
+        self.log.pop_scope()
+
+    def visit_broadcast_(self, op: Broadcast) -> None:
+        self.log.add("Broadcast")
+        self.log.push_scope()
+        self.visit_expr(op.value)
+        self.visit_expr(op.lanes)
+        self.log.pop_scope()
+
+    def visit_shuffle_(self, op: Shuffle) -> None:
+        self.log.add("Shuffle")
+        self.log.push_scope()
+        for vec in op.vectors:
+            self.visit_expr(vec)
+        for idx in op.indices:
+            self.visit_expr(idx)
+        self.log.pop_scope()
+
+    def visit_int_imm_(self, op: IntImm) -> None:
+        self.log.add("IntImm")
+
+    def visit_float_imm_(self, op: FloatImm) -> None:
+        self.log.add("FloatImm")
+
+    def visit_string_imm_(self, op: StringImm) -> None:
+        self.log.add("StringImm")
+
+
+class BasicMutator(ExprMutator):
+    """Default ExprMutator"""
+
+
+class ASTPostPrinterMutator(ExprMutator):
+    """Print TIR AST in the post order format."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log = ASTLog()
+
+    def visit_var_(self, op: Var) -> tir.Expr:
+        result = super().visit_var_(op)
+        self.log.add("Var")
+        return result
+
+    def visit_buffer_load_(self, op: BufferLoad) -> tir.Expr:
+        result = super().visit_buffer_load_(op)
+        self.log.add("BufferLoad")
+        return result
+
+    def visit_producer_load_(self, op: ProducerLoad) -> tir.Expr:
+        result = super().visit_producer_load_(op)
+        self.log.add("ProducerLoad")
+        return result
+
+    def visit_let_(self, op: Let) -> tir.Expr:
+        result = super().visit_let_(op)
+        self.log.add("Let")
+        return result
+
+    def visit_call_(self, op: Call) -> tir.Expr:
+        result = super().visit_call_(op)
+        self.log.add("Call")
+        return result
+
+    def visit_add_(self, op: Add) -> tir.Expr:
+        result = super().visit_add_(op)
+        self.log.add("Add")
+        return result
+
+    def visit_sub_(self, op: Sub) -> tir.Expr:
+        result = super().visit_sub_(op)
+        self.log.add("Sub")
+        return result
+
+    def visit_mul_(self, op: Mul) -> tir.Expr:
+        result = super().visit_mul_(op)
+        self.log.add("Mul")
+        return result
+
+    def visit_div_(self, op: Div) -> tir.Expr:
+        result = super().visit_div_(op)
+        self.log.add("Div")
+        return result
+
+    def visit_mod_(self, op: Mod) -> tir.Expr:
+        result = super().visit_mod_(op)
+        self.log.add("Mod")
+        return result
+
+    def visit_floordiv_(self, op: FloorDiv) -> tir.Expr:
+        result = super().visit_floordiv_(op)
+        self.log.add("FloorDiv")
+        return result
+
+    def visit_floormod_(self, op: FloorMod) -> tir.Expr:
+        result = super().visit_floormod_(op)
+        self.log.add("FloorMod")
+        return result
+
+    def visit_min_(self, op: Min) -> tir.Expr:
+        result = super().visit_min_(op)
+        self.log.add("Min")
+        return result
+
+    def visit_max_(self, op: Max) -> tir.Expr:
+        result = super().visit_max_(op)
+        self.log.add("Max")
+        return result
+
+    def visit_eq_(self, op: EQ) -> tir.Expr:
+        result = super().visit_eq_(op)
+        self.log.add("EQ")
+        return result
+
+    def visit_ne_(self, op: NE) -> tir.Expr:
+        result = super().visit_ne_(op)
+        self.log.add("NE")
+        return result
+
+    def visit_lt_(self, op: LT) -> tir.Expr:
+        result = super().visit_lt_(op)
+        self.log.add("LT")
+        return result
+
+    def visit_le_(self, op: LE) -> tir.Expr:
+        result = super().visit_le_(op)
+        self.log.add("LE")
+        return result
+
+    def visit_gt_(self, op: GT) -> tir.Expr:
+        result = super().visit_gt_(op)
+        self.log.add("GT")
+        return result
+
+    def visit_ge_(self, op: GE) -> tir.Expr:
+        result = super().visit_ge_(op)
+        self.log.add("GE")
+        return result
+
+    def visit_and_(self, op: And) -> tir.Expr:
+        result = super().visit_and_(op)
+        self.log.add("And")
+        return result
+
+    def visit_or_(self, op: Or) -> tir.Expr:
+        result = super().visit_or_(op)
+        self.log.add("Or")
+        return result
+
+    def visit_reduce_(self, op: Reduce) -> tir.Expr:
+        result = super().visit_reduce_(op)
+        self.log.add("Reduce")
+        return result
+
+    def visit_cast_(self, op: Cast) -> tir.Expr:
+        result = super().visit_cast_(op)
+        self.log.add("Cast")
+        return result
+
+    def visit_not_(self, op: Not) -> tir.Expr:
+        result = super().visit_not_(op)
+        self.log.add("Not")
+        return result
+
+    def visit_select_(self, op: Select) -> tir.Expr:
+        result = super().visit_select_(op)
+        self.log.add("Select")
+        return result
+
+    def visit_ramp_(self, op: Ramp) -> tir.Expr:
+        result = super().visit_ramp_(op)
+        self.log.add("Ramp")
+        return result
+
+    def visit_broadcast_(self, op: Broadcast) -> tir.Expr:
+        result = super().visit_broadcast_(op)
+        self.log.add("Broadcast")
+        return result
+
+    def visit_shuffle_(self, op: Shuffle) -> tir.Expr:
+        result = super().visit_shuffle_(op)
+        self.log.add("Shuffle")
+        return result
+
+    def visit_int_imm_(self, op: IntImm) -> tir.Expr:
+        result = super().visit_int_imm_(op)
+        self.log.add("IntImm")
+        return result
+
+    def visit_float_imm_(self, op: FloatImm) -> tir.Expr:
+        result = super().visit_float_imm_(op)
+        self.log.add("FloatImm")
+        return result
+
+    def visit_string_imm_(self, op: StringImm) -> tir.Expr:
+        result = super().visit_string_imm_(op)
+        self.log.add("StringImm")
+        return result
+
+
+def basic_check(expr, visitor_str, mutator_str):
+    """Helper function to check visitor and mutator on an expression"""
+
+    # Check visitor
+    basic_visitor = BasicVisitor()
+    basic_visitor.visit_expr(expr)
+    # Check AST printer visitor
+    log_visitor = ASTPrinter()
+    log_visitor.visit_expr(expr)
+    assert str(log_visitor.log) == visitor_str
+
+    # Check basic mutator
+    basic_mutator = BasicMutator()
+    mutated_expr = basic_mutator.visit_expr(expr)
+    assert_structural_equal(mutated_expr, expr)
+
+    # Check post-order printer mutator
+    post_log_mutator = ASTPostPrinterMutator()
+    mutated_expr = post_log_mutator.visit_expr(expr)
+    assert_structural_equal(mutated_expr, expr)
+    assert str(post_log_mutator.log) == mutator_str
+
+
+def test_var():
+    basic_check(n, "Var", "Var")
+
+
+def test_int_imm():
+    basic_check(tir.IntImm("int32", 10), "IntImm", "IntImm")
+
+
+def test_float_imm():
+    basic_check(tir.FloatImm("float32", 1.5), "FloatImm", "FloatImm")
+
+
+def test_string_imm():
+    basic_check(tir.StringImm("hello"), "StringImm", "StringImm")
+
+
+def test_add():
+    add_node = tir.Add(n, m)
+    basic_check(add_node, "\n".join(["Add", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Add"]))
+
+
+def test_sub():
+    sub_node = tir.Sub(n, m)
+    basic_check(sub_node, "\n".join(["Sub", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Sub"]))
+
+
+def test_mul():
+    mul_node = tir.Mul(n, m)
+    basic_check(mul_node, "\n".join(["Mul", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Mul"]))
+
+
+def test_div():
+    div_node = tir.Div(n, m)
+    basic_check(div_node, "\n".join(["Div", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Div"]))
+
+
+def test_floor_div():
+    floor_div_node = tir.FloorDiv(n, m)
+    basic_check(
+        floor_div_node,
+        "\n".join(["FloorDiv", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "FloorDiv"]),
+    )
+
+
+def test_floor_mod():
+    floor_mod_node = tir.FloorMod(n, m)
+    basic_check(
+        floor_mod_node,
+        "\n".join(["FloorMod", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "FloorMod"]),
+    )
+
+
+def test_min():
+    min_node = tir.Min(n, m)
+    basic_check(min_node, "\n".join(["Min", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Min"]))
+
+
+def test_max():
+    max_node = tir.Max(n, m)
+    basic_check(max_node, "\n".join(["Max", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "Max"]))
+
+
+def test_eq():
+    eq_node = tir.EQ(n, m)
+    basic_check(eq_node, "\n".join(["EQ", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "EQ"]))
+
+
+def test_ne():
+    ne_node = tir.NE(n, m)
+    basic_check(ne_node, "\n".join(["NE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "NE"]))
+
+
+def test_lt():
+    lt_node = tir.LT(n, m)
+    basic_check(lt_node, "\n".join(["LT", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "LT"]))
+
+
+def test_le():
+    le_node = tir.LE(n, m)
+    basic_check(le_node, "\n".join(["LE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "LE"]))
+
+
+def test_gt():
+    gt_node = tir.GT(n, m)
+    basic_check(gt_node, "\n".join(["GT", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "GT"]))
+
+
+def test_ge():
+    ge_node = tir.GE(n, m)
+    basic_check(ge_node, "\n".join(["GE", "\tVar", "\tVar"]), "\n".join(["Var", "Var", "GE"]))
+
+
+def test_and():
+    and_node = tir.And(tir.EQ(n, m), tir.LT(n, 10))
+    basic_check(
+        and_node,
+        "\n".join(["And", "\tEQ", "\t\tVar", "\t\tVar", "\tLT", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "IntImm", "LT", "And"]),
+    )
+
+
+def test_or():
+    or_node = tir.Or(tir.EQ(n, m), tir.LT(n, 10))
+    basic_check(
+        or_node,
+        "\n".join(["Or", "\tEQ", "\t\tVar", "\t\tVar", "\tLT", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "IntImm", "LT", "Or"]),
+    )
+
+
+def test_not():
+    not_node = tir.Not(tir.EQ(n, m))
+    basic_check(
+        not_node,
+        "\n".join(["Not", "\tEQ", "\t\tVar", "\t\tVar"]),
+        "\n".join(["Var", "Var", "EQ", "Not"]),
+    )
+
+
+def test_select():
+    select_node = tir.Select(tir.EQ(n, m), n, m)
+    basic_check(
+        select_node,
+        "\n".join(["Select", "\tEQ", "\t\tVar", "\t\tVar", "\tVar", "\tVar"]),
+        "\n".join(["Var", "Var", "EQ", "Var", "Var", "Select"]),
+    )
+
+
+def test_cast():
+    cast_node = tir.Cast("float32", n)
+    basic_check(cast_node, "\n".join(["Cast", "\tVar"]), "\n".join(["Var", "Cast"]))
+
+
+def test_let():
+    let_node = tir.Let(n, tir.IntImm("int32", 10), n + 1)
+    basic_check(
+        let_node,
+        "\n".join(["Let", "\tVar", "\tIntImm", "\tAdd", "\t\tVar", "\t\tIntImm"]),
+        "\n".join(["Var", "IntImm", "Var", "IntImm", "Add", "Let"]),
+    )
+
+
+def test_ramp():
+    ramp_node = tir.Ramp(n, 1, 4)
+    basic_check(
+        ramp_node,
+        "\n".join(["Ramp", "\tVar", "\tIntImm", "\tIntImm"]),
+        "\n".join(["Var", "IntImm", "IntImm", "Ramp"]),
+    )
+
+
+def test_broadcast():
+    broadcast_node = tir.Broadcast(n, 4)
+    basic_check(
+        broadcast_node,
+        "\n".join(["Broadcast", "\tVar", "\tIntImm"]),
+        "\n".join(["Var", "IntImm", "Broadcast"]),
+    )
+
+
+def test_inherit():
+    # The internal class is not instantiated.
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+    add_node = tir.Add(n, m)
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "\tInternalVar", "\tInternalVar"])
+
+
+def test_inherit_with_cls():
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            self.log.push_scope()
+            self.visit_expr(op.a)
+            self.visit_expr(op.b)
+            self.log.pop_scope()
+
+    add_node = tir.Add(n, m)
+    iv = InternalVisitor()
+    iv.visit_expr(add_node)
+    assert str(iv.log) == "\n".join(["InternalAdd", "\tInternalVar", "\tInternalVar"])
+
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "\tInternalVar", "\tInternalVar"])
+
+
+def test_call_visitor_super():
+    class InternalVisitor(ExprVisitor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("InternalAdd")
+            super().visit_add_(op)  # call ExprVisitor.visit_add_
+
+        def visit_var_(self, op: Var) -> None:
+            self.log.add("InternalVar")
+
+        def visit_int_imm_(self, op: IntImm) -> None:
+            self.log.add("InternalIntImm")
+
+    class LeafVisitor(InternalVisitor):
+        def visit_add_(self, op: Add) -> None:
+            self.log.add("LeafAdd")
+            super().visit_add_(op)  # call InternalVisitor.visit_add_
+
+    add_node = tir.Add(n, tir.IntImm("int32", 10))
+    iv = InternalVisitor()
+    iv.visit_expr(add_node)
+    assert str(iv.log) == "\n".join(["InternalAdd", "InternalVar", "InternalIntImm"])
+
+    lv = LeafVisitor()
+    lv.visit_expr(add_node)
+    assert str(lv.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar", "InternalIntImm"])
+
+
+def test_call_mutator_super():
+    class InternalMutator(ExprMutator):
+        def __init__(self) -> None:
+            super().__init__()
+            self.log = ASTLog()
+
+        def visit_add_(self, op: Add) -> tir.Expr:
+            self.log.add("InternalAdd")
+            return super().visit_add_(op)  # call ExprMutator.visit_add_
+
+        def visit_var_(self, op: Var) -> tir.Expr:
+            self.log.add("InternalVar")
+            return super().visit_var_(op)  # call ExprMutator.visit_var_
+
+        def visit_int_imm_(self, op: IntImm) -> tir.Expr:
+            self.log.add("InternalIntImm")
+            return super().visit_int_imm_(op)  # call ExprMutator.visit_int_imm_
+
+    class LeafMutator(InternalMutator):
+        def visit_add_(self, op: Add) -> tir.Expr:
+            self.log.add("LeafAdd")
+            return super().visit_add_(op)  # call InternalMutator.visit_add_
+
+    add_node = tir.Add(n, tir.IntImm("int32", 10))
+    im = InternalMutator()
+    im.visit_expr(add_node)
+    assert str(im.log) == "\n".join(["InternalAdd", "InternalVar", "InternalIntImm"])
+
+    lm = LeafMutator()
+    lm.visit_expr(add_node)
+    assert str(lm.log) == "\n".join(["LeafAdd", "InternalAdd", "InternalVar", "InternalIntImm"])
+
+
+def test_var_mutation():
+    """Test mutating variables in a TIR expression"""
+
+    class VarMutator(ExprMutator):
+        def __init__(self, var_map):
+            super().__init__()
+            self.var_map = var_map
+
+        def visit_var_(self, op: Var) -> tir.Expr:
+            if op.name in self.var_map:
+                return self.var_map[op.name]
+            return op
+
+    # Create a simple expression
+    expr = n + m
+
+    # Create a mutator that replaces 'n' with a constant
+    var_map = {"n": tir.IntImm("int32", 42)}
+    mutator = VarMutator(var_map)
+    result = mutator.visit_expr(expr)
+
+    # The result should be 42 + m
+    expected = tir.Add(tir.IntImm("int32", 42), m)
+    assert_structural_equal(result, expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/transform/test_transform_naive_allocator.py

@@ -0,0 +1,177 @@
+# 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.
+
+import tvm
+import tvm.testing
+from tvm.ir import assert_structural_equal
+from tvm.script import tirx as T
+from tvm.script.tirx import tile as Tx
+from tvm.tirx.layout import F, P, S, TileLayout
+from tvm.tirx.trn.transform import TrnNaiveAllocator
+
+
+def test_one_alloc():
+    src_shape = [128, 512]
+    src_layout = TileLayout(S[(128, 512) : (512, 1)])
+    dst_shape = [128, 512]
+    dst_layout = TileLayout(S[(128, 512) : (1 @ P, 1 @ F)])
+
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout)
+        Tx.copy(A_sbuf, A)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        A = T.match_buffer(A_ptr, src_shape, "float32", layout=src_layout)
+        T.device_entry()
+        A_sbuf = T.alloc_buffer(dst_shape, "float32", scope="trn.sbuf", layout=dst_layout, allocated_addr=[0])  # noqa: E501
+        Tx.copy(A_sbuf, A)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_two_alloc():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_existing_alloc():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[1])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[4*512*4+1])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[1])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_workspace():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        C_sbuf = T.alloc_buffer([128, 1024], "float32", scope="trn.sbuf")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        C_sbuf = T.alloc_buffer([128, 1024], "float32", scope="trn.sbuf", allocated_addr=[2*512*4+4*512*4])  # noqa: E501
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_other_scope_alloc():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        C_sbuf = T.alloc_buffer([8, 128, 512], "float32", scope="global")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        C_sbuf = T.alloc_buffer([8, 128, 512], "float32", scope="global")
+        Tx.copy(B_sbuf[0:256, :], A_sbuf, workspace={"C": C_sbuf})
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+def test_buffer_views():
+    # fmt: off
+    @T.prim_func
+    def copy(A_ptr: T.handle) -> None:
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF")
+        B_view = B_sbuf.view(2, 256, 512)
+        Tx.copy(B_view[0], A_sbuf)
+
+    @T.prim_func
+    def expected(A_ptr: T.handle) -> None:
+        T.func_attr({"global_symbol": "copy"})
+        T.device_entry()
+        A_sbuf = T.alloc_buffer([256, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[0])  # noqa: E501
+        B_sbuf = T.alloc_buffer([512, 512], "float32", scope="trn.sbuf", layout="PF", allocated_addr=[2*512*4])  # noqa: E501
+        B_view = B_sbuf.view(2, 256, 512)
+        Tx.copy(B_view[0], A_sbuf)
+        # fmt: on
+
+    mod = tvm.IRModule({"copy": copy})
+    mod = TrnNaiveAllocator()(mod)
+    assert_structural_equal(mod["copy"], expected)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx/transform/test_transform_static_horizontal_fusion.py

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

tests/python/tirx/utils.py

@@ -0,0 +1,16 @@
+# 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.