mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Lei Wang
1c561f6d46
* [Bugfix] Raise error on zero grid dimension instead of silent clamp (#1993) Fix ThreadWorkLoad::Extract() silently clamping zero grid dims to 1, which caused either CUDA_ERROR_ILLEGAL_ADDRESS crashes (dynamic case) or silent wrong results (static case). Closes #1993 * lint fix
77 lines
2.0 KiB
Python
77 lines
2.0 KiB
Python
import torch
|
|
import pytest
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def _issue1993_dynamic_grid():
|
|
num_tokens = T.dynamic("num_tokens")
|
|
|
|
@T.prim_func
|
|
def kernel(out: T.Tensor[(num_tokens,), T.float32]):
|
|
with T.Kernel(num_tokens, threads=1) as pid:
|
|
out[pid] = T.float32(1.0)
|
|
|
|
return kernel
|
|
|
|
|
|
@tilelang.jit
|
|
def _issue1993_static_grid():
|
|
|
|
@T.prim_func
|
|
def kernel(out: T.Tensor[(4,), T.float32]):
|
|
with T.Kernel(0, threads=1) as pid:
|
|
out[pid] = T.float32(1.0)
|
|
|
|
return kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_issue_1993_dynamic_zero_grid_dim():
|
|
"""Regression test for issue #1993.
|
|
|
|
When a dynamic grid dimension resolves to 0 at runtime, the runtime
|
|
should raise an error instead of silently clamping to 1 and launching
|
|
the kernel (which would write through a NULL pointer and crash with
|
|
CUDA_ERROR_ILLEGAL_ADDRESS).
|
|
"""
|
|
kernel = _issue1993_dynamic_grid()
|
|
|
|
# Positive case: should work correctly
|
|
out = torch.zeros(4, dtype=torch.float32, device="cuda")
|
|
kernel(out)
|
|
torch.cuda.synchronize()
|
|
assert out.eq(1.0).all()
|
|
|
|
# Zero case: should raise an error, not crash with illegal memory access
|
|
out_empty = torch.zeros(0, dtype=torch.float32, device="cuda")
|
|
with pytest.raises(Exception): # noqa: B017
|
|
kernel(out_empty)
|
|
torch.cuda.synchronize()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_issue_1993_static_zero_grid_dim():
|
|
"""Regression test for issue #1993.
|
|
|
|
When T.Kernel(0) is used with a static constant, the runtime should
|
|
raise an error instead of silently clamping to 1 and executing a
|
|
spurious CTA.
|
|
"""
|
|
kernel = _issue1993_static_grid()
|
|
|
|
out = torch.zeros(4, dtype=torch.float32, device="cuda")
|
|
with pytest.raises(Exception): # noqa: B017
|
|
kernel(out)
|
|
torch.cuda.synchronize()
|
|
|
|
# Buffer should remain untouched
|
|
assert out.eq(0.0).all()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|