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.
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang import tvm
|
|
from tilelang.engine.lower import lower
|
|
from tilelang.layout import make_linear_layout
|
|
|
|
|
|
def _make_tma_store_kernel(n, block, annotate_linear=False):
|
|
@T.prim_func
|
|
def main(b: T.Tensor((n,), T.float32)):
|
|
with T.Kernel(T.ceildiv(n, block), threads=128) as bx:
|
|
shared = T.alloc_shared((block,), T.float32)
|
|
if annotate_linear:
|
|
T.annotate_layout({shared: make_linear_layout(shared)})
|
|
for i in T.Parallel(block):
|
|
shared[i] = T.cast(bx * block + i, T.float32)
|
|
T.copy(shared, b[bx * block])
|
|
|
|
return main
|
|
|
|
|
|
def _lower_source(n, block, annotate_linear=False):
|
|
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90"})
|
|
with target:
|
|
artifact = lower(
|
|
_make_tma_store_kernel(n, block, annotate_linear),
|
|
target=target,
|
|
enable_device_compile=False,
|
|
)
|
|
return artifact.kernel_source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_tma_store_selection_uses_copy_bounds():
|
|
full_tile_source = _lower_source(1024, 256)
|
|
# The unannotated rank-1 tail exercises descriptor layout inference.
|
|
partial_tile_source = _lower_source(1000, 256)
|
|
# A linear layout cannot itself force descriptor lowering, so this also
|
|
# verifies that Lower re-derives the tail bounds from its analyzer.
|
|
partial_tile_linear_layout_source = _lower_source(1000, 256, annotate_linear=True)
|
|
|
|
assert "tma_store" in full_tile_source
|
|
assert "CUtensorMap" not in full_tile_source
|
|
assert "tma_store" in partial_tile_source
|
|
assert "CUtensorMap" in partial_tile_source
|
|
assert "tma_store" in partial_tile_linear_layout_source
|
|
assert "CUtensorMap" in partial_tile_linear_layout_source
|
|
|
|
|
|
@tilelang.testing.requires_cuda_compute_version(9, 0)
|
|
def test_tma_store_partial_tile_runtime():
|
|
n, block = 1000, 256
|
|
kernel = tilelang.compile(_make_tma_store_kernel(n, block, annotate_linear=True), target="cuda")
|
|
|
|
sentinel = -1.0
|
|
padded_n = (n + block - 1) // block * block
|
|
backing = torch.full((padded_n,), sentinel, device="cuda")
|
|
kernel(backing[:n])
|
|
torch.cuda.synchronize()
|
|
|
|
expected = torch.arange(n, dtype=torch.float32, device="cuda")
|
|
torch.testing.assert_close(backing[:n], expected)
|
|
assert torch.all(backing[n:] == sentinel)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_tma_store_selection_uses_copy_bounds()
|
|
test_tma_store_partial_tile_runtime()
|