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.
76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
from tilelang import tvm
|
|
from tvm import tirx
|
|
from tvm.tirx import op
|
|
from tilelang.cuda.pipeline import CUDAPassPipelineBodyPrologue
|
|
from tilelang.transform import LowerAccessPtr
|
|
|
|
|
|
def issue_2123_atomic_load_repro(num_tiles, threads=32):
|
|
@T.prim_func
|
|
def kernel(status: T.Tensor((num_tiles,), T.int32), out: T.Tensor((1,), T.int32)):
|
|
with T.Kernel(num_tiles, threads=threads) as tile:
|
|
look = T.alloc_var(T.int32)
|
|
state = T.alloc_var(T.int32)
|
|
done = T.alloc_var(T.bool)
|
|
tx = T.get_thread_binding()
|
|
if tx == 0:
|
|
look = tile - 1
|
|
done = look < 0
|
|
state = 0
|
|
while not done:
|
|
state = T.atomic_load(status[look], memory_order="acquire")
|
|
if state != 0:
|
|
done = True
|
|
else:
|
|
look -= 1
|
|
done = look < 0
|
|
if tile == num_tiles - 1:
|
|
out[0] = state
|
|
|
|
return kernel
|
|
|
|
|
|
def _has_op_call(func, op_name):
|
|
found = False
|
|
target_op = op.Op.get(op_name)
|
|
|
|
def _visit(node):
|
|
nonlocal found
|
|
if isinstance(node, tirx.Call) and node.op.same_as(target_op):
|
|
found = True
|
|
|
|
tirx.stmt_functor.post_order_visit(func.body, _visit)
|
|
return found
|
|
|
|
|
|
def _assert_access_ptr_lowered(mod):
|
|
assert _has_op_call(mod["main"], "tirx.tvm_access_ptr")
|
|
assert not _has_op_call(mod["main"], "tl.access_ptr")
|
|
|
|
|
|
def test_issue_2123_atomic_load_lower_access_ptr_direct():
|
|
func = issue_2123_atomic_load_repro(4).with_attr("global_symbol", "main")
|
|
mod = tvm.IRModule.from_expr(func)
|
|
|
|
lowered = LowerAccessPtr()(mod)
|
|
|
|
_assert_access_ptr_lowered(lowered)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_issue_2123_atomic_load_lower_access_ptr_pipeline():
|
|
target = tvm.target.Target("cuda", host="llvm")
|
|
func = issue_2123_atomic_load_repro(4).with_attr("global_symbol", "main")
|
|
mod = tvm.IRModule.from_expr(func)
|
|
|
|
lowered = CUDAPassPipelineBodyPrologue(mod, target)
|
|
|
|
_assert_access_ptr_lowered(lowered)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|