mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Leonard
d863f1d89d
* Add RDNA gfx1151 ROCm target support * Fix optional annotation handling for gfx1100 * Address CodeRabbit ROCm cleanup feedback * Scope MFMA tests to CDNA targets * Scope gfx950 copy async GEMM test * Reject non-divisible warp layouts * Skip CUDA/CDNA GEMM dtype tests on RDNA * Scope RDNA device model to gfx11 * Use RDNA arch helper in block scoring * Make RDNA tensor intrinsics generation-aware * Exclude gfx1010 from SDOT4 guard * Store RDNA tensor instructions as list * Reject unsupported RDNA carver targets
268 lines
8.5 KiB
Python
268 lines
8.5 KiB
Python
"""Tests for gfx950 (MI350) copy.async feature.
|
||
|
||
Two new behaviours introduced in commit dfa63b10:
|
||
1. cp_async_gs<16> on gfx950 lowers to buffer_load_dwordx4 ... lds
|
||
(128-bit direct-to-LDS, bypassing VGPRs) instead of a plain uint4
|
||
scalar store. coalesced_width=8 (8 fp16 = 16 bytes) is required to
|
||
trigger the 16-byte path.
|
||
2. CDNA arch helper reports smem_cap = 160 KB for gfx950, even when
|
||
an older driver reports the conservative 64 KB default.
|
||
"""
|
||
|
||
import pytest
|
||
import tilelang as tl
|
||
import tilelang.language as T
|
||
import tilelang.testing
|
||
from tilelang.testing import _check_is_gfx950 as _is_gfx950
|
||
|
||
|
||
def _matmul_kernel(
|
||
M,
|
||
N,
|
||
K,
|
||
block_M,
|
||
block_N,
|
||
block_K,
|
||
trans_A,
|
||
trans_B,
|
||
in_dtype,
|
||
out_dtype,
|
||
accum_dtype,
|
||
num_stages,
|
||
threads=128,
|
||
k_pack=1,
|
||
# coalesced_width=8 → cp_async_gs<16> (16 bytes, 8×fp16)
|
||
# coalesced_width=4 → cp_async_gs<8> (8 bytes, 4×fp16)
|
||
coalesced_width=4,
|
||
):
|
||
"""Return a prim_func for pipelined GEMM using T.copy (global->shared)."""
|
||
A_shape = (K, M) if trans_A else (M, K)
|
||
B_shape = (N, K) if trans_B else (K, N)
|
||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_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(A_shared_shape, in_dtype)
|
||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||
T.clear(C_local)
|
||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
||
if trans_A:
|
||
T.copy(A[k * block_K, by * block_M], A_shared, coalesced_width=coalesced_width)
|
||
else:
|
||
T.copy(A[by * block_M, k * block_K], A_shared, coalesced_width=coalesced_width)
|
||
if trans_B:
|
||
T.copy(B[bx * block_N, k * block_K], B_shared, coalesced_width=coalesced_width)
|
||
else:
|
||
T.copy(B[k * block_K, bx * block_N], B_shared, coalesced_width=coalesced_width)
|
||
T.gemm(A_shared, B_shared, C_local, trans_A, trans_B, k_pack=k_pack)
|
||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||
|
||
return main
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 1: codegen — cp_async_gs<16> is present in generated HIP source
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _matmul_kernel_async(
|
||
M,
|
||
N,
|
||
K,
|
||
block_M,
|
||
block_N,
|
||
block_K,
|
||
trans_A,
|
||
trans_B,
|
||
in_dtype,
|
||
out_dtype,
|
||
accum_dtype,
|
||
coalesced_width=4,
|
||
threads=128,
|
||
):
|
||
"""Return a prim_func using T.async_copy (explicit async, no pipeline) for codegen tests."""
|
||
A_shape = (K, M) if trans_A else (M, K)
|
||
B_shape = (N, K) if trans_B else (K, N)
|
||
A_shared_shape = (block_K, block_M) if trans_A else (block_M, block_K)
|
||
B_shared_shape = (block_N, block_K) if trans_B else (block_K, block_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(A_shared_shape, in_dtype)
|
||
B_shared = T.alloc_shared(B_shared_shape, in_dtype)
|
||
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
||
T.clear(C_local)
|
||
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=0):
|
||
if trans_A:
|
||
T.async_copy(A[k * block_K, by * block_M], A_shared, coalesced_width=coalesced_width)
|
||
else:
|
||
T.async_copy(A[by * block_M, k * block_K], A_shared, coalesced_width=coalesced_width)
|
||
if trans_B:
|
||
T.async_copy(B[bx * block_N, k * block_K], B_shared, coalesced_width=coalesced_width)
|
||
else:
|
||
T.async_copy(B[k * block_K, bx * block_N], B_shared, coalesced_width=coalesced_width)
|
||
T.gemm(A_shared, B_shared, C_local, trans_A, trans_B)
|
||
T.copy(C_local, C[by * block_M, bx * block_N])
|
||
|
||
return main
|
||
|
||
|
||
@tilelang.testing.requires_gfx950
|
||
def test_gfx950_cp_async_gs_16_in_codegen():
|
||
"""coalesced_width=8 (16 bytes) must emit cp_async_gs<16> in generated HIP source.
|
||
|
||
Uses T.async_copy (explicit async semantics) so that the 16-byte path is
|
||
emitted verbatim without relying on the software pipeline planner.
|
||
Pipelining correctness is covered separately by test_gfx950_copy_async_gemm_pipelined.
|
||
"""
|
||
prog = _matmul_kernel_async(
|
||
256,
|
||
256,
|
||
256,
|
||
128,
|
||
128,
|
||
32,
|
||
False,
|
||
True,
|
||
T.float16,
|
||
T.float32,
|
||
T.float32,
|
||
coalesced_width=8, # 8 fp16 = 16 bytes → cp_async_gs<16>
|
||
)
|
||
kernel = tl.compile(prog, out_idx=[2])
|
||
src = kernel.get_kernel_source()
|
||
assert "cp_async_gs<16>" in src, "Expected cp_async_gs<16> in generated HIP source for 128-bit async copy path"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 2: LDS capacity reported as 160 KB on gfx950
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@tilelang.testing.requires_rocm
|
||
def test_gfx950_smem_cap_160kb():
|
||
"""CDNA arch helper must report 160 KB LDS for gfx950."""
|
||
from tilelang import tvm
|
||
from tilelang.carver.arch.cdna import CDNA, _GFX950_LDS_SIZE
|
||
|
||
target = tvm.target.Target("rocm")
|
||
arch = CDNA(target)
|
||
|
||
if _is_gfx950():
|
||
assert arch.smem_cap == _GFX950_LDS_SIZE, f"Expected smem_cap={_GFX950_LDS_SIZE} for gfx950, got {arch.smem_cap}"
|
||
else:
|
||
# On non-gfx950 devices the override must NOT kick in
|
||
from tilelang import tvm as _tvm
|
||
|
||
dev = _tvm.device("rocm", 0)
|
||
assert arch.smem_cap == dev.max_shared_memory_per_block
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 3: numerical correctness — pipelined copy.async GEMM (num_stages=2)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"trans_A, trans_B, k_pack",
|
||
[
|
||
(False, False, 1),
|
||
(False, True, 1),
|
||
(True, True, 1),
|
||
(True, False, 1),
|
||
],
|
||
)
|
||
@tilelang.testing.requires_gfx950
|
||
def test_gfx950_copy_async_gemm_pipelined(trans_A, trans_B, k_pack):
|
||
"""Pipelined GEMM (num_stages=2) with gfx950 copy.async must be numerically correct."""
|
||
prog = _matmul_kernel(
|
||
512,
|
||
512,
|
||
512,
|
||
128,
|
||
128,
|
||
32,
|
||
trans_A,
|
||
trans_B,
|
||
T.float16,
|
||
T.float32,
|
||
T.float32,
|
||
num_stages=2,
|
||
threads=128,
|
||
k_pack=k_pack,
|
||
coalesced_width=4 * k_pack,
|
||
)
|
||
kernel = tl.compile(prog, out_idx=[2])
|
||
profiler = kernel.get_profiler()
|
||
|
||
def ref_program(A, B):
|
||
import torch
|
||
|
||
a = A.T.float() if trans_A else A.float()
|
||
b = B.T.float() if trans_B else B.float()
|
||
return torch.matmul(a, b)
|
||
|
||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 4: non-pipelined baseline still correct (num_stages=0)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
"trans_A, trans_B",
|
||
[
|
||
(False, False),
|
||
(False, True),
|
||
(True, True),
|
||
(True, False),
|
||
],
|
||
)
|
||
@tilelang.testing.requires_gfx950
|
||
def test_gfx950_copy_async_gemm_no_pipeline(trans_A, trans_B):
|
||
"""Non-pipelined GEMM (num_stages=0) must also produce correct results."""
|
||
prog = _matmul_kernel(
|
||
512,
|
||
512,
|
||
512,
|
||
128,
|
||
128,
|
||
32,
|
||
trans_A,
|
||
trans_B,
|
||
T.float16,
|
||
T.float32,
|
||
T.float32,
|
||
num_stages=0,
|
||
threads=128,
|
||
coalesced_width=4,
|
||
)
|
||
kernel = tl.compile(prog, out_idx=[2])
|
||
profiler = kernel.get_profiler()
|
||
|
||
def ref_program(A, B):
|
||
import torch
|
||
|
||
a = A.T.float() if trans_A else A.float()
|
||
b = B.T.float() if trans_B else B.float()
|
||
return torch.matmul(a, b)
|
||
|
||
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
tilelang.testing.main()
|