mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
Chennes
ed00dfcd7f
* [BugFix] Reject unsupported T.serial access to thread-distributed fragments with clear diagnostics A T.serial loop runs sequentially inside each CUDA thread, but a fragment's elements are distributed across threads, so indexing a fragment by a serial loop variable has no valid thread-ownership mapping. Three fuzzer reports share this root cause: - #2393: serial reduce of a GEMM-output fragment -> unhelpful internal error ("contains inner var j"). - #2395: serial write into a fragment then T.copy out -> nvcc "identifier undefined" (an owner-thread guard built from serial loop vars was emitted outside the loop that binds them). - #2396: serial write of a fragment to global -> silently wrong (threads race-write each cell). AddWrapperForSingleBufStore wraps bare fragment stores in a degenerate T.parallel(1) and only rejected constant non-zero fragment indices, so a dynamic (e.g. serial-loop) index slipped through and lowered to broken code (#2395/#2396). Reject any fragment index that is not the constant all-zero index in this fallback path, with a message pointing to T.Parallel / T.reduce_sum. Rephrase the existing inner-var layout check (#2393) to explain the cause instead of dumping "contains inner var". Also collect every access to a buffer in a statement (not just the last) so the diagnostic is not bypassed when one fragment is accessed at both a constant and a variable index in the same statement; without it such a case falls through to a cryptic downstream StructuralEqual internal error. This is a completeness hardening only, not required by the three issues above. Fixes #2393 Fixes #2395 Fixes #2396 * Generalize fragment diagnostic to non-parallel loops inner_vars_ is populated for every non-parallel For (serial, unroll, vectorized), so the message must not hard-code T.serial. Reword to "non-parallel loop variable" per review feedback. * Narrow fix scope to #2393 only Remove fragment index validation from AddWrapperForSingleBufStore and delete tests for #2395/#2396. These issues cannot be caught at this pass stage because T.grid and T.serial are indistinguishable (both appear as ForKind.SERIAL before LowerTileOp). Keep only the ParallelOpNode diagnostic improvement which successfully catches #2393.
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
# ruff: noqa
|
|
import pytest
|
|
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_issue_2393_serial_reduce_gemm_fragment_is_rejected():
|
|
"""#2393: a serial reduce of a GEMM-output (MMA-layout) fragment used to fail
|
|
with an unhelpful internal error ("contains inner var j"). Now rejected with
|
|
a clear diagnostic suggesting T.reduce_sum."""
|
|
M, N, KD = 64, 64, 64
|
|
|
|
@tilelang.jit
|
|
def k(Q: T.Tensor((M, KD), T.float16), Kk: T.Tensor((N, KD), T.float16)):
|
|
out = T.empty((M,), T.float32)
|
|
with T.Kernel(1, threads=128) as bx:
|
|
Qs = T.alloc_shared((M, KD), T.float16)
|
|
Ks = T.alloc_shared((N, KD), T.float16)
|
|
acc = T.alloc_fragment((M, N), T.float32)
|
|
s = T.alloc_fragment((M,), T.float32)
|
|
T.copy(Q, Qs)
|
|
T.copy(Kk, Ks)
|
|
T.clear(acc)
|
|
T.gemm(Qs, Ks, acc, transpose_B=True)
|
|
T.fill(s, 0.0)
|
|
for i in T.Parallel(M):
|
|
for j in T.serial(N):
|
|
s[i] += acc[i, j]
|
|
T.copy(s, out)
|
|
return out
|
|
|
|
import torch
|
|
|
|
Q = torch.randn(M, KD, dtype=torch.float16, device="cuda")
|
|
Kk = torch.randn(N, KD, dtype=torch.float16, device="cuda")
|
|
with pytest.raises(Exception) as excinfo:
|
|
k(Q, Kk)
|
|
msg = str(excinfo.value)
|
|
assert "thread-distributed fragment" in msg and "T.reduce_sum" in msg
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|