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.
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_runtime_unknown_sign_vector_negative_index_load():
|
|
@T.prim_func
|
|
def main(A: T.Tensor((1024,), T.float32), B: T.Tensor((4, 4), T.float32)):
|
|
with T.Kernel(1, threads=1) as _:
|
|
for t in T.serial(4):
|
|
B[t, T.Ramp(0, 1, 4)] = A[T.Ramp(t - 2, 1, 4)]
|
|
|
|
kernel = tilelang.compile(main, out_idx=[1], target="cuda")
|
|
|
|
a = torch.arange(1024, device="cuda", dtype=torch.float32)
|
|
b = kernel(a)
|
|
expected = torch.tensor(
|
|
[
|
|
[1022.0, 1023.0, 0.0, 1.0],
|
|
[1023.0, 0.0, 1.0, 2.0],
|
|
[0.0, 1.0, 2.0, 3.0],
|
|
[1.0, 2.0, 3.0, 4.0],
|
|
],
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
torch.testing.assert_close(b, expected)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_runtime_unknown_sign_vector_negative_index_store():
|
|
@T.prim_func
|
|
def main(B: T.Tensor((4, 4), T.float32), A: T.Tensor((1024,), T.float32)):
|
|
with T.Kernel(1, threads=1) as _:
|
|
for i in T.serial(1024):
|
|
A[i] = T.float32(-1)
|
|
for t in T.serial(4):
|
|
A[T.Ramp(t - 2, 1, 4)] = B[t, T.Ramp(0, 1, 4)]
|
|
|
|
kernel = tilelang.compile(main, out_idx=[1], target="cuda")
|
|
|
|
b = torch.arange(16, device="cuda", dtype=torch.float32).reshape(4, 4)
|
|
a = kernel(b)
|
|
|
|
expected = torch.full((1024,), -1, device="cuda", dtype=torch.float32)
|
|
expected[0] = b[2, 0]
|
|
expected[1] = b[3, 0]
|
|
expected[2] = b[3, 1]
|
|
expected[3] = b[3, 2]
|
|
expected[4] = b[3, 3]
|
|
expected[1022] = b[0, 0]
|
|
expected[1023] = b[1, 0]
|
|
|
|
torch.testing.assert_close(a, expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|