mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
ChennesandLeiWang1999
a443dde926
* [BugFix] Reject contracting shared-buffer layouts in T.annotate_layout T.annotate_layout could accept a contracting many-to-one forward map on a shared buffer. makeBufferWithLayout then misread the smaller output extent as a replication factor, rewrote the buffer to a higher rank, and eventually failed with an unrelated BufferStore dimensionality ICHECK. Add a pigeonhole check for static layout shapes before the replication branch. When the layout output extent is smaller than its input extent, raise a user-facing ValueError that identifies the buffer and reports both extents. Keep identity, permutation, and output-growing layouts unchanged. This check intentionally targets contracting maps; general injectivity analysis for maps with larger output bounding boxes is outside this change. Fixes #2632 * [Layout] Enforce injective shared-buffer annotations Add a LayoutNode injectivity check with an exact fallback for static domains, and reuse it for FragmentNode thread and replication coordinates. Validate shared-buffer annotations after layout reshaping so sparse many-to-one maps are rejected with a useful collision diagnostic while padded injective maps remain valid. --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
"""Regression test for GitHub issue #2632.
|
|
|
|
A many-to-one forward map passed to ``T.annotate_layout`` on a shared buffer
|
|
can either crash during buffer remapping or silently alias distinct logical
|
|
elements. It must instead be rejected with a ``ValueError`` naming the buffer
|
|
and the non-injectivity, including when its sparse output bounding box is larger
|
|
than the logical input domain.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
from tilelang.layout import Layout
|
|
|
|
N = 128
|
|
|
|
REJECT_MATCH = r"shared buffer `[^`]*`.*must be injective.*both map to physical coordinates"
|
|
|
|
|
|
def _roundtrip_kernel(fwd):
|
|
|
|
@T.prim_func
|
|
def main(A: T.Tensor((N,), "int32"), Out: T.Tensor((N,), "int32")):
|
|
with T.Kernel(1, threads=N):
|
|
s = T.alloc_shared((N,), "int32")
|
|
T.annotate_layout({s: Layout((N,), fwd)})
|
|
for i in T.Parallel(N):
|
|
s[i] = A[i]
|
|
for i in T.Parallel(N):
|
|
Out[i] = s[i]
|
|
|
|
return main
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
"fwd",
|
|
[
|
|
pytest.param(lambda i: i % 64, id="mod64-2-to-1"),
|
|
pytest.param(lambda i: i * 0, id="zero-N-to-1"),
|
|
pytest.param(lambda i: (i % 64) * 4, id="mod64-times4-padded-2-to-1"),
|
|
],
|
|
)
|
|
def test_noninjective_shared_layout_rejected(fwd):
|
|
"""Every many-to-one map must be a clean ValueError, not an ICE or alias."""
|
|
with pytest.raises(ValueError, match=REJECT_MATCH):
|
|
tilelang.compile(_roundtrip_kernel(fwd), out_idx=[1], target="cuda")
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
"fwd",
|
|
[
|
|
pytest.param(lambda i: i, id="identity"),
|
|
pytest.param(lambda i: i ^ 1, id="xor-pair-swap"),
|
|
pytest.param(lambda i: i * 2, id="grow-i-times-2"),
|
|
pytest.param(lambda i: i + 10, id="grow-i-plus-10"),
|
|
],
|
|
)
|
|
def test_injective_shared_layout_still_compiles_and_runs(fwd):
|
|
"""Injective maps (including codomain-growing ones) must keep working."""
|
|
kernel = tilelang.compile(_roundtrip_kernel(fwd), out_idx=[1], target="cuda")
|
|
A = torch.arange(N, dtype=torch.int32, device="cuda")
|
|
torch.testing.assert_close(kernel(A), A)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|