mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tvm.target import Target
|
|
|
|
|
|
def _get_issue_2307_source() -> str:
|
|
@T.prim_func
|
|
def main(score: T.Tensor((16,), T.float32)):
|
|
with T.Kernel(1, threads=32):
|
|
score_fragment = T.alloc_fragment(32, T.float32)
|
|
for i in T.Parallel(32):
|
|
if i < 16:
|
|
score_fragment[i] = score[i]
|
|
else:
|
|
score_fragment[i] = T.infinity(T.float32)
|
|
for i in T.Parallel(32):
|
|
if i < 16:
|
|
T.device_assert(T.isfinite(score_fragment[i]))
|
|
|
|
target = Target("cuda")
|
|
with target:
|
|
artifact = tilelang.lower(main, target=target)
|
|
return artifact.kernel_source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_fragment_read_assert_parallel_loop_is_partitioned():
|
|
src = _get_issue_2307_source()
|
|
lines = src.splitlines()
|
|
assert "for (int i = 0; i < 32; ++i)" not in src
|
|
|
|
assert_idx = next(i for i, line in enumerate(lines) if "device_assert_with_msg" in line)
|
|
context = "\n".join(lines[max(0, assert_idx - 4) : assert_idx + 1])
|
|
assert "threadIdx.x) < 16" in context
|
|
assert "isfinite(score_fragment[0])" in context
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|