mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
Ftoandcoderabbitai[bot]
ab083cb3c6
* [BugFix] Fix T.Persistent dropping tiles when last dim is not a multiple of group_size (#2433) * Update src/ir.cc Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
"""Regression for tile-ai/tilelang issue #2433."""
|
|
|
|
import torch
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.jit
|
|
def _persistent_visit_kernel(MB, NB, sm_num):
|
|
|
|
@T.prim_func
|
|
def kernel(Visited: T.Tensor((MB, NB), T.int32)):
|
|
with T.Kernel(sm_num, threads=1) as block_id:
|
|
for bx, by in T.Persistent([MB, NB], sm_num, block_id):
|
|
Visited[bx, by] += 1
|
|
|
|
return kernel
|
|
|
|
|
|
@tilelang.jit
|
|
def _persistent_visit_kernel_3d(D0, D1, D2, sm_num):
|
|
|
|
@T.prim_func
|
|
def kernel(Visited: T.Tensor((D0, D1, D2), T.int32)):
|
|
with T.Kernel(sm_num, threads=1) as block_id:
|
|
for i, j, k in T.Persistent([D0, D1, D2], sm_num, block_id):
|
|
Visited[i, j, k] += 1
|
|
|
|
return kernel
|
|
|
|
|
|
def run_persistent_visit(m_blocks: int, n_blocks: int, sm_num: int):
|
|
kernel = _persistent_visit_kernel(m_blocks, n_blocks, sm_num)
|
|
visited = torch.zeros(m_blocks, n_blocks, device="cuda", dtype=torch.int32)
|
|
kernel(visited)
|
|
expected = torch.ones_like(visited)
|
|
assert torch.equal(visited, expected), (
|
|
f"T.Persistent visited grid mismatch for "
|
|
f"m_blocks={m_blocks}, n_blocks={n_blocks}, sm_num={sm_num}: "
|
|
f"missing={(visited == 0).nonzero(as_tuple=False).tolist()}, "
|
|
f"overflow={(visited > 1).nonzero(as_tuple=False).tolist()}"
|
|
)
|
|
|
|
|
|
def run_persistent_visit_3d(d0: int, d1: int, d2: int, sm_num: int):
|
|
kernel = _persistent_visit_kernel_3d(d0, d1, d2, sm_num)
|
|
visited = torch.zeros(d0, d1, d2, device="cuda", dtype=torch.int32)
|
|
kernel(visited)
|
|
expected = torch.ones_like(visited)
|
|
assert torch.equal(visited, expected), (
|
|
f"T.Persistent 3D visited grid mismatch for "
|
|
f"({d0},{d1},{d2}), sm_num={sm_num}: "
|
|
f"missing={(visited == 0).nonzero(as_tuple=False).tolist()}, "
|
|
f"overflow={(visited > 1).nonzero(as_tuple=False).tolist()}"
|
|
)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_issue_2433():
|
|
"""Reproduce the failure table from the issue and the divisor controls.
|
|
|
|
With group_size defaulting to min(8, n_blocks), only the rows where
|
|
n_blocks % group_size != 0 trip the bug; the divisor rows act as
|
|
controls to make sure the fix doesn't regress them.
|
|
"""
|
|
# Divisor controls (must always pass).
|
|
for n_blocks in (1, 2, 4, 8):
|
|
run_persistent_visit(m_blocks=4, n_blocks=n_blocks, sm_num=3)
|
|
|
|
# Non-divisor cases -- exactly the failing rows from the issue.
|
|
for n_blocks in (9, 10, 11, 12):
|
|
run_persistent_visit(m_blocks=4, n_blocks=n_blocks, sm_num=3)
|
|
|
|
# Single-wave control from the issue (sm_num >= total tiles): proves
|
|
# the bug is divisibility, not wave-count related.
|
|
run_persistent_visit(m_blocks=4, n_blocks=9, sm_num=36)
|
|
|
|
# Boundary: leading dim not a multiple of group_size either -- only the
|
|
# last dim is grouped, so this must still pass.
|
|
run_persistent_visit(m_blocks=7, n_blocks=11, sm_num=5)
|
|
|
|
# Boundary: 3D domain with non-divisor last dim, exercising the
|
|
# mixed-radix decode + range guard at higher rank.
|
|
run_persistent_visit_3d(d0=3, d1=5, d2=10, sm_num=4)
|
|
|
|
# Boundary: more SMs than tiles, so several waves end up entirely in
|
|
# the padded tail and must be skipped by the range guard.
|
|
run_persistent_visit(m_blocks=2, n_blocks=3, sm_num=16)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|