trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-02 06:34:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix] Guard TMA copies against unaligned box offsets

Browse Source
This commit is contained in:
UnsettingGalaxy
2026-07-21 02:51:47 +00:00
parent f6b23c2721
commit 72fa8d4afc
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
testing/python/issue/test_tilelang_issue_2527.py
+9 -6
View File
@@ -6,23 +6,26 @@ import torch
M, N = 64, 64
m0, n0, mm, nn = 0, 1, 32, 32
@T.prim_func
def main(B: T.Tensor((M, N), "int32")):
with T.Kernel(1, threads=128):
As = T.alloc_shared((mm, nn), "int32")
T.fill(As, 7)
T.copy(As, B[m0:m0+mm, n0:n0+nn]) # -> 2-D bulk TMA store, crd0 = n0 = 1 (unaligned)
T.copy(As, B[m0 : m0 + mm, n0 : n0 + nn]) # -> 2-D bulk TMA store, crd0 = n0 = 1 (unaligned)
m = tilelang.compile(main)
src = m.get_kernel_source()
print("USES_TMA_STORE:", "tma_store" in src) # -> True
print("USES_TMA_STORE:", "tma_store" in src) # -> True
for line in src.splitlines():
if "tl::tma_store(" in line:
print(" ", line.strip()) # -> tl::tma_store(B_desc, (&(As[0])), 1, 0);
print(" ", line.strip()) # -> tl::tma_store(B_desc, (&(As[0])), 1, 0);
B = torch.zeros((M, N), dtype=torch.int32, device="cuda")
try:
m(B); torch.cuda.synchronize()
print("default OK?", bool((B.cpu()[m0:m0+mm, n0:n0+nn] == 7).all())) # -> region == 7
m(B)
torch.cuda.synchronize()
print("default OK?", bool((B.cpu()[m0 : m0 + mm, n0 : n0 + nn] == 7).all())) # -> region == 7
except Exception as e:
print("default CRASH:", repr(e)[:120]) # -> illegal instruction
print("default CRASH:", repr(e)[:120]) # -> illegal instruction