mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
Chennes
e4fba8678a
Some ldmatrix lane maps step beyond the final SM75 MMA micro-tile, and column overflow must carry into the next row. Layout remapping previously linearized the source indices and optional offset before converting them back to multidimensional coordinates, which normalized these addresses as a side effect. Since #2836 preserves the multidimensional indices, the out-of-range coordinates reach the swizzle layout unchanged and the native SM75 MMA GEMM paths fault on illegal shared-memory accesses. Wrap each complete row-major source offset within its operand region in ldmatrix_a and ldmatrix_b before rebuilding the source coordinates. Leading dimensions remain untouched, so the wrap cannot cross into another stage of a multi-version buffer. When the coordinates are provably in range, the simplifier removes the wrap. SM70 and SM80 generated sources remain byte-identical, while the SM75 addresses match the pre-#2836 address mapping. Add multi-stage SM75 GEMM coverage and a lowered-offset bounds regression that cross-compiles explicitly for sm_75 and therefore runs on any CUDA runner.
248 lines
9.5 KiB
Python
248 lines
9.5 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tvm.tirx.stmt_functor import post_order_visit
|
|
|
|
|
|
def _make_gemm_kernel(M, N, K, block_M, block_N, block_K, in_dtype, out_dtype, accum_dtype):
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), in_dtype),
|
|
B: T.Tensor((N, K), in_dtype),
|
|
C: T.Tensor((M, N), out_dtype),
|
|
):
|
|
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_K), in_dtype)
|
|
B_shared = T.alloc_shared((block_N, block_K), in_dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
|
|
T.clear(C_local)
|
|
for ko in T.serial(T.ceildiv(K, block_K)):
|
|
T.copy(A[by * block_M, ko * block_K], A_shared)
|
|
T.copy(B[bx * block_N, ko * block_K], B_shared)
|
|
T.gemm(A_shared, B_shared, C_local, transpose_B=True)
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def _assert_mma_sync_shape(source, a_type, c_type, m, n, k):
|
|
expected = f"tl::mma_sync<tl::DataType::{a_type}, tl::DataType::{a_type}, tl::DataType::{c_type}, {m}, {n}, {k}, false, true>"
|
|
assert expected in source
|
|
assert "mma_sync_sm70" not in source
|
|
|
|
|
|
def _pack_int4(tensor: torch.Tensor) -> torch.Tensor:
|
|
tensor_i16 = tensor.to(torch.int16)
|
|
packed = (tensor_i16[..., ::2] & 0x0F) | ((tensor_i16[..., 1::2] & 0x0F) << 4)
|
|
return packed.to(torch.int8).contiguous()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_f16_gemm_uses_m16n8k8_and_matches_torch():
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_gemm_kernel(M, N, K, 64, 64, 32, T.float16, T.float32, T.float32),
|
|
target="cuda",
|
|
out_idx=[2],
|
|
)
|
|
_assert_mma_sync_shape(kernel.get_kernel_source(), "kFloat16", "kFloat32", 16, 8, 8)
|
|
|
|
a = torch.randn((M, K), device="cuda", dtype=torch.float16)
|
|
b = torch.randn((N, K), device="cuda", dtype=torch.float16)
|
|
c = kernel(a, b)
|
|
ref = a.float() @ b.float().T
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=1e-2, atol=1e-2, max_mismatched_ratio=0.01)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_f16_gemm_f16_accum_uses_m16n8k8_and_matches_torch():
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_gemm_kernel(M, N, K, 64, 64, 32, T.float16, T.float16, T.float16),
|
|
target="cuda",
|
|
out_idx=[2],
|
|
)
|
|
_assert_mma_sync_shape(kernel.get_kernel_source(), "kFloat16", "kFloat16", 16, 8, 8)
|
|
|
|
a = torch.randn((M, K), device="cuda", dtype=torch.float16) * 0.1
|
|
b = torch.randn((N, K), device="cuda", dtype=torch.float16) * 0.1
|
|
c = kernel(a, b)
|
|
ref = (a.float() @ b.float().T).to(torch.float16)
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=1e-2, atol=1e-2, max_mismatched_ratio=0.01)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_uint8_gemm_uses_m8n8k16_and_matches_torch():
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_gemm_kernel(M, N, K, 64, 64, 64, T.uint8, T.int32, T.int32),
|
|
target="cuda",
|
|
out_idx=[2],
|
|
)
|
|
_assert_mma_sync_shape(kernel.get_kernel_source(), "kUInt8", "kInt32", 8, 8, 16)
|
|
|
|
a = torch.randint(0, 256, (M, K), device="cuda", dtype=torch.uint8)
|
|
b = torch.randint(0, 256, (N, K), device="cuda", dtype=torch.uint8)
|
|
c = kernel(a, b)
|
|
ref = (a.cpu().to(torch.int32) @ b.cpu().to(torch.int32).T).to(device="cuda")
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=0, atol=0)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_int8_gemm_uses_m8n8k16_and_matches_torch():
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_gemm_kernel(M, N, K, 64, 64, 64, T.int8, T.int32, T.int32),
|
|
target="cuda",
|
|
out_idx=[2],
|
|
)
|
|
_assert_mma_sync_shape(kernel.get_kernel_source(), "kInt8", "kInt32", 8, 8, 16)
|
|
|
|
a = torch.randint(-8, 8, (M, K), device="cuda", dtype=torch.int8)
|
|
b = torch.randint(-8, 8, (N, K), device="cuda", dtype=torch.int8)
|
|
c = kernel(a, b)
|
|
ref = (a.cpu().to(torch.int32) @ b.cpu().to(torch.int32).T).to(device="cuda")
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=0, atol=0)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
def test_sm75_int4_gemm_uses_m8n8k32_and_matches_torch():
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_gemm_kernel(M, N, K, 64, 64, 64, T.int4, T.int32, T.int32),
|
|
target="cuda",
|
|
out_idx=[2],
|
|
)
|
|
_assert_mma_sync_shape(kernel.get_kernel_source(), "kInt4", "kInt32", 8, 8, 32)
|
|
|
|
a = torch.randint(-8, 8, (M, K), device="cuda", dtype=torch.int8)
|
|
b = torch.randint(-8, 8, (N, K), device="cuda", dtype=torch.int8)
|
|
c = kernel(_pack_int4(a), _pack_int4(b))
|
|
ref = (a.cpu().to(torch.int32) @ b.cpu().to(torch.int32).T).to(device="cuda")
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=0, atol=0)
|
|
|
|
|
|
def _make_pipelined_gemm_kernel(
|
|
M, N, K, block_M, block_N, block_K, num_stages, transpose_B, in_dtype=T.float16, out_dtype=T.float32, accum_dtype=T.float32
|
|
):
|
|
B_shape = (N, K) if transpose_B else (K, N)
|
|
B_shared_shape = (block_N, block_K) if transpose_B else (block_K, block_N)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), 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=128) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_K), 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 ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
T.copy(A[by * block_M, ko * block_K], A_shared)
|
|
if transpose_B:
|
|
T.copy(B[bx * block_N, ko * block_K], B_shared)
|
|
else:
|
|
T.copy(B[ko * block_K, bx * block_N], B_shared)
|
|
T.gemm(A_shared, B_shared, C_local, transpose_B=transpose_B)
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@tilelang.testing.requires_cuda_compute_version_eq(7, 5)
|
|
@pytest.mark.parametrize("transpose_B", [True, False])
|
|
def test_sm75_f16_gemm_pipelined_multi_stage_matches_torch(transpose_B):
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_pipelined_gemm_kernel(M, N, K, 64, 64, 32, 2, transpose_B),
|
|
target="cuda",
|
|
out_idx=[2],
|
|
)
|
|
a = torch.randn((M, K), device="cuda", dtype=torch.float16)
|
|
b_shape = (N, K) if transpose_B else (K, N)
|
|
b = torch.randn(b_shape, device="cuda", dtype=torch.float16)
|
|
c = kernel(a, b)
|
|
ref = a.float() @ (b.float().T if transpose_B else b.float())
|
|
|
|
tilelang.testing.torch_assert_close(c, ref, rtol=1e-2, atol=1e-2, max_mismatched_ratio=0.01)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
"in_dtype, accum_dtype, block_K, transpose_B",
|
|
[
|
|
(T.float16, T.float32, 32, True),
|
|
(T.float16, T.float32, 32, False),
|
|
(T.int8, T.int32, 64, True),
|
|
(T.int4, T.int32, 64, True),
|
|
],
|
|
)
|
|
def test_sm75_ldmatrix_shared_reads_stay_in_bounds(in_dtype, accum_dtype, block_K, transpose_B):
|
|
# Cross-compiles for sm_75, so it guards the Turing address math on any
|
|
# CUDA runner. The ldmatrix lane maps step past the shared tile unless the
|
|
# emitter wraps them (see mma_macro_generator.py).
|
|
M = N = K = 128
|
|
kernel = tilelang.compile(
|
|
_make_pipelined_gemm_kernel(M, N, K, 64, 64, block_K, 0, transpose_B, in_dtype, accum_dtype, accum_dtype),
|
|
target={"kind": "cuda", "arch": "sm_75"},
|
|
out_idx=[2],
|
|
)
|
|
(func,) = kernel.artifact.device_mod.functions.values()
|
|
|
|
analyzer = tvm.arith.Analyzer()
|
|
buffer_sizes = {}
|
|
ldmatrix_calls = []
|
|
|
|
def visit(node):
|
|
if isinstance(node, tvm.tirx.AttrStmt) and str(node.attr_key) == "thread_extent":
|
|
analyzer.update(node.node.var, tvm.arith.ConstIntBound(0, int(node.value) - 1))
|
|
elif isinstance(node, tvm.tirx.For):
|
|
analyzer.update(
|
|
node.loop_var,
|
|
tvm.arith.ConstIntBound(int(node.min), int(node.min) + int(node.extent) - 1),
|
|
)
|
|
elif isinstance(node, (tvm.tirx.AllocBuffer, tvm.tirx.DeclBuffer)):
|
|
size = 1
|
|
for extent in node.buffer.shape:
|
|
size *= int(extent)
|
|
buffer_sizes[node.buffer.data] = size
|
|
elif isinstance(node, tvm.tirx.Call) and str(getattr(node.op, "name", "")) == "tl.ptx_ldmatrix":
|
|
ldmatrix_calls.append(node)
|
|
|
|
post_order_visit(func.body, visit)
|
|
assert ldmatrix_calls, "expected ldmatrix lowering for sm_75"
|
|
|
|
for call in ldmatrix_calls:
|
|
src = call.args[2]
|
|
assert str(src.op.name) == "tirx.tvm_access_ptr"
|
|
data, offset, extent = src.args[1], src.args[2], int(src.args[3])
|
|
assert data in buffer_sizes, f"ldmatrix source {data.name} has no visible allocation"
|
|
bound = analyzer.const_int_bound(offset)
|
|
assert bound.min_value >= 0
|
|
assert bound.max_value + extent <= buffer_sizes[data], (
|
|
f"ldmatrix reads past the shared tile: offset up to {bound.max_value} + extent {extent} exceeds {buffer_sizes[data]} elements"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|