"""Test T.tma_copy() with user-managed synchronization. For TMA loads (global -> shared): T.tma_copy() emits only expect_tx + tma_load (no arrive, no wait). The user must explicitly call T.barrier_arrive() and T.mbarrier_wait_parity(). This allows multiple tma_copy operations to share a single barrier arrive. Pipeline buffer versioning expands the barrier to num_stages versions automatically. For TMA stores (shared -> global): T.tma_copy() emits tma_store + tma_store_arrive (no wait). The user must explicitly call T.tma_store_wait() for synchronization. No barrier argument is needed for stores. """ import pytest from tilelang import tvm as tvm import tilelang.testing import tilelang.language as T import tilelang def matmul_tma_copy( M, N, K, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, threads, num_stages, ): A_shape = (M, K) B_shape = (K, N) @T.prim_func def main( A: T.Tensor(A_shape, in_dtype), B: T.Tensor(B_shape, in_dtype), C: T.Tensor((M, N), out_dtype), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by): A_shared = T.alloc_shared((block_M, block_K), in_dtype) B_shared = T.alloc_shared((block_K, block_N), in_dtype) C_local = T.alloc_fragment((block_M, block_N), accum_dtype) mbar_A = T.alloc_barrier(128) mbar_B = T.alloc_barrier(128) T.clear(C_local) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): T.tma_copy(A[by * block_M, k * block_K], A_shared, barrier=mbar_A) T.barrier_arrive(mbar_A) T.tma_copy(B[k * block_K, bx * block_N], B_shared, barrier=mbar_B) T.barrier_arrive(mbar_B) T.barrier_wait(mbar_A, k % 2) T.barrier_wait(mbar_B, k % 2) T.gemm(A_shared, B_shared, C_local) T.copy(C_local, C[by * block_M, bx * block_N]) return main def run_gemm_tma_copy(num_stages): M, N, K = 1024, 1024, 1024 block_M, block_N, block_K = 128, 128, 32 in_dtype = T.float16 out_dtype = T.float16 accum_dtype = T.float32 threads = 128 program = matmul_tma_copy( M, N, K, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, threads, num_stages, ) kernel = tilelang.compile( program, out_idx=[2], pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}, ) print(kernel.get_kernel_source()) profiler = kernel.get_profiler() def ref_program(A, B): import torch C = torch.matmul(A.to(torch.float), B.to(torch.float)) return C.to(torch.__getattribute__(out_dtype)) profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2) @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_pipeline_2_stages(): run_gemm_tma_copy(num_stages=2) @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_pipeline_3_stages(): run_gemm_tma_copy(num_stages=3) def per_warp_tma_copy_kernel(): M = 32 K = 1024 num_threads = 256 warp_size = 32 num_warps = num_threads // warp_size @T.prim_func def main(A: T.Tensor((M, K), T.float32), B: T.Tensor((M,), T.float32)): with T.Kernel(T.ceildiv(M, num_warps), threads=num_threads) as pid: tid = T.get_thread_binding() warp_idx = tid // warp_size row = pid * num_warps + warp_idx a_shared = T.alloc_shared((num_warps, K), dtype=T.float32) mbars = T.alloc_barrier([warp_size] * num_warps) T.tma_copy( A[row, 0:K], a_shared[warp_idx, 0:K], barrier=mbars[warp_idx], leader_scope_threads=warp_size, ) T.mbarrier_arrive(mbarrier=mbars[warp_idx]) T.mbarrier_wait_parity(mbarrier=mbars[warp_idx], parity=0) if tid % warp_size == 0: B[row] = a_shared[warp_idx, 0] return main def block_tma_copy_kernel(): M = 256 K = 1024 num_threads = 256 @T.prim_func def main(A: T.Tensor((M, K), T.float32), B: T.Tensor((1,), T.float32)): with T.Kernel(T.ceildiv(M, num_threads), threads=num_threads) as pid: tid = T.get_thread_binding() a_shared = T.alloc_shared((1, K), dtype=T.float32) mbar = T.alloc_barrier([num_threads]) T.tma_copy(A[pid, 0:K], a_shared[0, 0:K], barrier=mbar) T.mbarrier_arrive(mbarrier=mbar) T.mbarrier_wait_parity(mbarrier=mbar, parity=0) if tid == 0: B[pid] = a_shared[0, 0] return main @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_per_warp_leader_scope_codegen(): kernel = tilelang.compile(per_warp_tma_copy_kernel(), out_idx=[1]) source = kernel.get_kernel_source() assert "tl_shuffle_elect<32>()" in source, "Expected per-warp elect<32>" @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_default_block_leader_scope_codegen(): kernel = tilelang.compile(block_tma_copy_kernel(), out_idx=[1]) source = kernel.get_kernel_source() assert "tl_shuffle_elect<256>()" in source, "Expected block-wide elect<256>" @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_uses_explicit_global_stride_for_row_offset(): rows = 2 cols = 32 row_stride = 40 @T.prim_func def main( A: T.StridedTensor((rows, cols), (row_stride, 1), T.float32), B: T.Tensor((cols,), T.float32), ): with T.Kernel(1, threads=32): a_shared = T.alloc_shared((cols,), T.float32) T.copy(A[1, 0:cols], a_shared, prefer_instruction="tma") T.copy(a_shared, B) import torch kernel = tilelang.compile(main, out_idx=[1]) assert "CUtensorMap" not in kernel.get_kernel_source() padded = torch.arange(rows * row_stride, dtype=torch.float32, device="cuda").reshape(rows, row_stride) source = padded[:, :cols] assert source.stride() == (row_stride, 1) output = kernel(source) torch.testing.assert_close(output, source[1]) @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_uses_descriptor_for_padded_multirow_region(): rows = 2 cols = 32 row_stride = 40 @T.prim_func def load( A: T.StridedTensor((rows, cols), (row_stride, 1), T.float32), B: T.Tensor((rows, cols), T.float32), ): with T.Kernel(1, threads=32): a_shared = T.alloc_shared((rows, cols), T.float32) T.copy(A, a_shared, prefer_instruction="tma") T.copy(a_shared, B, prefer_instruction="sync") @T.prim_func def store( A: T.Tensor((rows, cols), T.float32), B: T.StridedTensor((rows, cols), (row_stride, 1), T.float32), ): with T.Kernel(1, threads=32): a_shared = T.alloc_shared((rows, cols), T.float32) T.copy(A, a_shared, prefer_instruction="sync") T.copy(a_shared, B, prefer_instruction="tma") import torch source = torch.arange(rows * cols, dtype=torch.float32, device="cuda").reshape(rows, cols) load_kernel = tilelang.compile(load, out_idx=[1]) assert "CUtensorMap" in load_kernel.get_kernel_source() padded_source = torch.full((rows, row_stride), -1.0, dtype=torch.float32, device="cuda") strided_source = padded_source[:, :cols] strided_source.copy_(source) output = load_kernel(strided_source) torch.testing.assert_close(output, source) store_kernel = tilelang.compile(store) assert "CUtensorMap" in store_kernel.get_kernel_source() padded_destination = torch.full((rows, row_stride), -1.0, dtype=torch.float32, device="cuda") strided_destination = padded_destination[:, :cols] store_kernel(source, strided_destination) torch.testing.assert_close(strided_destination, source) torch.testing.assert_close(padded_destination[:, cols:], torch.full_like(padded_destination[:, cols:], -1.0)) def matmul_tma_copy_store( M, N, K, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, threads, num_stages, ): """GEMM using T.tma_copy for both load (global->shared) and store (shared->global).""" A_shape = (M, K) B_shape = (K, N) @T.prim_func def main( A: T.Tensor(A_shape, in_dtype), B: T.Tensor(B_shape, in_dtype), C: T.Tensor((M, N), out_dtype), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by): A_shared = T.alloc_shared((block_M, block_K), in_dtype) B_shared = T.alloc_shared((block_K, block_N), in_dtype) C_shared = T.alloc_shared((block_M, block_N), out_dtype) C_local = T.alloc_fragment((block_M, block_N), accum_dtype) mbar_A = T.alloc_barrier(128) mbar_B = T.alloc_barrier(128) T.clear(C_local) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): T.tma_copy(A[by * block_M, k * block_K], A_shared, barrier=mbar_A) T.barrier_arrive(mbar_A) T.tma_copy(B[k * block_K, bx * block_N], B_shared, barrier=mbar_B) T.barrier_arrive(mbar_B) T.mbarrier_wait_parity(mbar_A, k % 2) T.mbarrier_wait_parity(mbar_B, k % 2) T.gemm(A_shared, B_shared, C_local) # Store result: fragment -> shared -> global via T.tma_copy (no barrier needed) T.copy(C_local, C_shared) T.tma_copy(C_shared, C[by * block_M, bx * block_N]) T.tma_store_wait(read=False) return main def run_gemm_tma_copy_store(num_stages): M, N, K = 1024, 1024, 1024 block_M, block_N, block_K = 128, 128, 32 in_dtype = T.float16 out_dtype = T.float16 accum_dtype = T.float32 threads = 128 program = matmul_tma_copy_store( M, N, K, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype, threads, num_stages, ) kernel = tilelang.compile( program, out_idx=[2], pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}, ) print(kernel.get_kernel_source()) profiler = kernel.get_profiler() def ref_program(A, B): import torch C = torch.matmul(A.to(torch.float), B.to(torch.float)) return C.to(torch.__getattribute__(out_dtype)) profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2) def fp4_tma_copy_roundtrip( M=128, N=256, block_M=64, block_N=128, ): @T.prim_func def main( A: T.Tensor((M, N), T.float4_e2m1fn), B: T.Tensor((M, N), T.float4_e2m1fn), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): A_shared = T.alloc_shared((block_M, block_N), T.float4_e2m1fn) mbar = T.alloc_barrier(128) T.tma_copy(A[by * block_M, bx * block_N], A_shared, barrier=mbar) T.barrier_arrive(mbar) T.mbarrier_wait_parity(mbar, 0) T.tma_copy(A_shared, B[by * block_M, bx * block_N]) T.tma_store_wait(read=False) return main def fp4_tma_copy_unpacked_smem_load(M=128, N=256, block_M=64, block_N=128): @T.prim_func def main( A: T.Tensor((M, N), T.float4_e2m1fn), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): A_shared = T.alloc_shared((block_M, block_N), T.float4_e2m1_unpacked) mbar = T.alloc_barrier(128) T.tma_copy( A[by * block_M, bx * block_N], A_shared, barrier=mbar, ) T.barrier_arrive(mbar) T.mbarrier_wait_parity(mbar, 0) return main def fp4_tma_copy_unpacked_smem_store(M=128, N=256, block_M=64, block_N=128): @T.prim_func def main( A: T.Tensor((M, N), T.float4_e2m1fn), B: T.Tensor((M, N), T.float4_e2m1fn), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): A_shared = T.alloc_shared((block_M, block_N), T.float4_e2m1_unpacked) mbar = T.alloc_barrier(128) T.tma_copy( A[by * block_M, bx * block_N], A_shared, barrier=mbar, ) T.barrier_arrive(mbar) T.mbarrier_wait_parity(mbar, 0) T.tma_copy(A_shared, B[by * block_M, bx * block_N]) return main def _fp4_tma_descriptor_init_block(host_source, desc_name): marker = f"[0].v_ptr) = {desc_name};" start = host_source.find(marker) assert start >= 0, f"Missing {desc_name} TensorMap initialization" end = host_source.find("TVMFFIFunctionCall(__tvm_tensormap_create_tiled_packed", start) assert end >= 0, f"Missing {desc_name} TensorMap creation call" return host_source[start:end] def _fp4_tma_stack_int(block, index): import re match = re.search(rf"\[{index}\]\.v_int64\)\s*=\s*\(\(int64_t\)(-?\d+)\);", block) assert match, f"Missing stack[{index}] integer assignment in:\n{block}" return int(match.group(1)) def _assert_fp4_packed_tma_descriptor(host_source, desc_name): expected_tma_args = { 1: 13, # CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B 2: 2, 4: 256, 5: 128, 6: 1, 7: 128, 8: 128, 9: 64, 10: 1, 11: 1, 12: 0, 13: 2, # CU_TENSOR_MAP_SWIZZLE_64B 14: 2, 15: 0, } block = _fp4_tma_descriptor_init_block(host_source, desc_name) for index, expected in expected_tma_args.items(): assert _fp4_tma_stack_int(block, index) == expected def _assert_fp4_unpacked_tma_descriptor(host_source, desc_name, *, expect_swizzle=None): block = _fp4_tma_descriptor_init_block(host_source, desc_name) assert _fp4_tma_stack_int(block, 1) == 14 # CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B assert _fp4_tma_stack_int(block, 8) == 128 # 128-element inner box for 8-bit storage if expect_swizzle is not None: assert _fp4_tma_stack_int(block, 13) == expect_swizzle def run_fp4_tma_copy_roundtrip(): import torch M, N = 128, 256 program = fp4_tma_copy_roundtrip(M=M, N=N) kernel = tilelang.compile( program, out_idx=[1], pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}, ) device_source = kernel.get_kernel_source() host_source = kernel.get_host_source() assert "CUtensorMap" in device_source assert "tl::tma_load" in device_source assert "tl::tma_store" in device_source assert host_source.count("__tvm_tensormap_create_tiled") >= 2 for desc_name in ("A_desc", "B_desc"): _assert_fp4_packed_tma_descriptor(host_source, desc_name) a = torch.randint(-128, 128, (M, N // 2), device="cuda", dtype=torch.int8) b = kernel(a) assert torch.equal(b.view(torch.int8), a) @tilelang.testing.requires_cuda def test_fp4_unpacksmem_tma_descriptor_uses_align16b(): program = fp4_tma_copy_unpacked_smem_load() artifact = tilelang.lower( program, target={"kind": "cuda", "arch": "sm_100"}, enable_device_compile=False, ) host_ir = str(artifact.host_mod) device_ir = str(artifact.device_mod) assert 'T.handle("float4_e2m1fn", "global")' in host_ir assert 'A_shared = T.alloc_buffer((8192,), "custom[float4_e2m1_unpacked]"' in device_ir assert '["__tvm_tensormap_create_tiled", A_desc, 14,' in host_ir assert 'T.call_packed("__tvm_tensormap_create_tiled", A_desc, 14,' in host_ir @tilelang.testing.requires_cuda def test_fp4_unpacksmem_tma_store_is_rejected(): program = fp4_tma_copy_unpacked_smem_store() with pytest.raises(tvm.TVMError, match="only supports float4_e2m1_unpacked as an FP4 unpack load"): tilelang.lower( program, target={"kind": "cuda", "arch": "sm_100"}, enable_device_compile=False, ) @tilelang.testing.requires_cuda def test_copy_prefer_tma_lowers_as_synchronous_tma_load(): @T.prim_func def main(x: T.Tensor((128, 32), T.float32)): with T.Kernel(threads=128): x_shared = T.alloc_shared((128, 32), T.float32) T.copy(x, x_shared, annotations={"prefer_instruction": "tma"}) target = {"kind": "cuda", "arch": "sm_90"} artifact = tilelang.lower( main, target=target, enable_device_compile=False, ) device_source = str(artifact.kernel_source) assert "tl::tma_load" in device_source assert "x_to_x_shared_mbarrier_mem" in device_source assert "x_to_x_shared_mbarrier[0]" in device_source assert "arrive_and_expect_tx" in device_source assert ".wait(0)" in device_source @tilelang.testing.requires_cuda def test_device_bound_pointer_keeps_descriptorless_bulk1d(): @T.prim_func def main( src_ptrs: T.Tensor((1,), T.ptr), out: T.Tensor((256,), T.float16), ): with T.Kernel(threads=128): src = T.make_tensor(src_ptrs[0], (256,), T.float16) src_shared = T.alloc_shared((256,), T.float16) T.copy(src, src_shared, prefer_instruction="tma") T.copy(src_shared, out, prefer_instruction="sync") target = tvm.target.Target({"kind": "cuda", "arch": "sm_90"}) with target: artifact = tilelang.lower( main, target=target, enable_device_compile=False, ) device_source = str(artifact.kernel_source) assert "tl::tma_load" in device_source assert "CUtensorMap" not in device_source @tilelang.testing.requires_cuda def test_device_bound_descriptor_tma_is_rejected_when_ws_disabled(): @T.prim_func def main(src_ptrs: T.Tensor((1,), T.ptr)): with T.Kernel(threads=128): src = T.make_tensor( src_ptrs[0], (16, 32), T.float16, strides=(32, 1), ) src_shared = T.alloc_shared((16, 16), T.float16) T.copy(src[0, 0], src_shared, prefer_instruction="tma") target = tvm.target.Target({"kind": "cuda", "arch": "sm_90"}) pass_configs = { tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED.value: True, } with ( target, tvm.transform.PassContext(config=pass_configs), pytest.raises( tvm.TVMError, match="bound inside the device function body", ), ): tilelang.lower( main, target=target, enable_device_compile=False, ) def run_fp4_tma_copy_unpacked_smem_load(): program = fp4_tma_copy_unpacked_smem_load() kernel = tilelang.compile( program, target="cuda", pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}, ) device_source = kernel.get_kernel_source() host_source = kernel.get_host_source() assert "CUtensorMap" in device_source _assert_fp4_unpacked_tma_descriptor(host_source, "A_desc") # 64 x 128 logical FP4 elems -> 4096 transaction bytes (4b/elem), not 8192. assert "expect_transaction(4096)" in device_source @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(10, 0) def test_fp4_tma_copy_roundtrip_packed_smem(): run_fp4_tma_copy_roundtrip() @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(10, 0) def test_fp4_tma_copy_unpacked_smem_load_descriptor_codegen(): run_fp4_tma_copy_unpacked_smem_load() @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_store_pipeline_2_stages(): run_gemm_tma_copy_store(num_stages=2) @tilelang.testing.requires_cuda @tilelang.testing.requires_cuda_compute_version_ge(9, 0) def test_tma_copy_store_pipeline_3_stages(): run_gemm_tma_copy_store(num_stages=3) if __name__ == "__main__": tilelang.testing.main() # test_tma_copy_pipeline_2_stages()