chore: import upstream snapshot with attribution

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