mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
andyluo7
dd9c20dbf7
A large number of tests import CUDA-only internals directly (for example
`tilelang.cuda.pipeline.CUDAPassPipelineBodyPrologue`) or exercise
NVIDIA-specific features such as TMA, LDGSTG, mbarrier, warp specialization,
tcgen05 and tfloat32. `src/cuda/CMakeLists.txt` only compiles a small
"ALWAYS" subset when `USE_CUDA=OFF`, so those FFI symbols are never
registered and the tests fail with:
AttributeError: module 'tilelang.cuda._ffi_api' has no attribute
'AnnotateDeviceBoundTmaCopies'
This is not ROCm-specific -- the same tests fail on any build without the
CUDA backend, including Metal and CPU-only.
Add `@tilelang.testing.requires_cuda` to those tests, plus three
portability fixes that keep coverage on non-CUDA backends instead of
skipping it:
* `test_tilelang_transform_layout_inference.py`: CUDA emits
`__launch_bounds__(N, 1)` while HIP emits `__launch_bounds__(N)`.
Match both.
* `test_tilelang_transform_pipeline_planning.py`: PipelinePlanning
deliberately strips pipeline annotations on ROCm targets other than
gfx950 (see the guard in src/transform/pipeline_planning.cc), so assert
the documented sequential fallback there.
* `test_tilelang_kernel_gemm.py`: the two f32f32f32 cases use `T.tfloat32`,
which is NVIDIA-only, so narrow `requires_cuda_or_cdna` to
`requires_cuda`.
Also refresh three stale RDNA tests: gfx12 is now a supported generation
(`get_arch` accepts 11 and 12), so they can no longer assert that gfx1200
is rejected. Force an unsupported generation instead, which is the only
way to reach the guard now that `target_is_rdna` reports only gfx11/gfx12.
Verified on 8x MI300X (gfx942, ROCm 7.2 container, torch nightly rocm7.2):
the ROCm suite goes from 105 failures to 0 with no loss of passing tests.
67 lines
2.2 KiB
Python
67 lines
2.2 KiB
Python
import torch
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
def ref_program(x, y):
|
|
return x + y
|
|
|
|
|
|
@tilelang.jit(out_idx=[-1])
|
|
def elementwise_add(M, N, block_M, block_N, in_dtype, out_dtype, threads):
|
|
@T.prim_func
|
|
def elem_add(A: T.Tensor((M, N), in_dtype), B: T.Tensor((M, N), 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_N), in_dtype)
|
|
B_shared = T.alloc_shared((block_M, block_N), in_dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), out_dtype)
|
|
C_shared = T.alloc_shared((block_M, block_N), out_dtype)
|
|
|
|
T.copy(
|
|
A[by * block_M, bx * block_N],
|
|
A_shared,
|
|
annotations={"prefer_instruction": "tma"},
|
|
)
|
|
T.copy(
|
|
B[by * block_M, bx * block_N],
|
|
B_shared,
|
|
annotations={"prefer_instruction": "tma"},
|
|
)
|
|
for local_y, local_x in T.Parallel(block_M, block_N):
|
|
C_local[local_y, local_x] = A_shared[local_y, local_x] + B_shared[local_y, local_x]
|
|
T.copy(C_local, C_shared)
|
|
T.copy(C_shared, C[by * block_M, bx * block_N])
|
|
|
|
return elem_add
|
|
|
|
|
|
def run_elementwise_add(M, N):
|
|
a = torch.randn(M, N, dtype=torch.float32, device="cuda")
|
|
b = torch.randn(M, N, dtype=torch.float32, device="cuda")
|
|
|
|
# Default config
|
|
block_M, block_N = 128, 128
|
|
config = {"block_M": block_M, "block_N": block_N, "threads": 128}
|
|
kernel = elementwise_add(M, N, **config, in_dtype=T.float32, out_dtype=T.float32)
|
|
|
|
out = kernel(a, b)
|
|
torch.testing.assert_close(out, ref_program(a, b), rtol=1e-2, atol=1e-2)
|
|
|
|
code = kernel.get_kernel_source()
|
|
if block_N == N:
|
|
assert "tma_load" in code and "CUtensorMap" not in code
|
|
else:
|
|
assert "tma_load" in code and "CUtensorMap" in code
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_tilelang_language_tma_1d():
|
|
run_elementwise_add(128, 128)
|
|
run_elementwise_add(256, 128)
|
|
run_elementwise_add(256, 256)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_tilelang_language_tma_1d()
|