mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
c66cadf87f
* [Fix] Suppress spurious __syncthreads() between T.ws blocks (#2192) and add T.named_barrier_arrive (#2188) Issue #2192 - BugFix: Unexpected CTA-wide __syncthreads() between T.ws producer/consumer - Root cause: FindConflict() in thread_storage_sync.cc did not check for mutually exclusive thread constraints. When T.ws(0) (tx<128) reads and T.ws(1) (tx>=128) writes the same shared buffer, no thread satisfies both constraints, so no __syncthreads() is needed — mbarriers already handle the synchronization. - Fix: add mutual-exclusivity check (prove NOT(P AND C)) early in FindConflict(). If the access constraints are mutually exclusive (disjoint thread sets), return false (no conflict) immediately, before any index-range analysis. - Measured improvement: barrierpipe_stage2 gemm latency 44ms → 25ms on H200. Issue #2188 - Feature: expose native T.named_barrier_arrive(barrier_id, thread_count) - Adds tl.named_barrier_arrive intrinsic op (bar.arrive PTX, one-sided arrive). - Registers op in builtin.cc, declares in builtin.h. - Adds tl::__named_barrier_arrive<id, cnt>() template in common.h. - Adds codegen case in codegen_cuda.cc. - Exposes T.named_barrier_arrive() in tilelang/language/builtin.py. - Adds regression test testing/python/issue/test_tilelang_issue_2188_2192.py. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * lint: clang-format + ruff fixes after #2188/#2192 changes Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * test: simplify issue 2188/2192 test, keep only #2192, add requires_cuda Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * [Fix] Make named_barrier_arrive accept variable barrier_id/thread_count * Update documentation for named barriers and remove obsolete test file - Clarified documentation for `named_barrier_arrive` in `builtin.h` and `builtin.py` to specify the barrier index range. - Removed obsolete regression test for issues #2188 and #2192 as it is no longer needed. * lint fix --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
"""Regression test for GitHub issue #2188.
|
|
|
|
#2188 - Feature: native T.named_barrier_arrive (bar.arrive) API.
|
|
"""
|
|
|
|
import torch
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_named_barrier_arrive_api():
|
|
"""T.named_barrier_arrive emits the CUDA named-barrier arrive helper."""
|
|
|
|
@tilelang.jit(out_idx=[1])
|
|
def kernel_fn():
|
|
@T.prim_func
|
|
def kernel(inp: T.Tensor((1,), T.int32), out: T.Tensor((1,), T.int32)):
|
|
with T.Kernel(1, threads=64) as _:
|
|
tx = T.get_thread_binding()
|
|
slot = T.alloc_shared((1,), T.int32)
|
|
|
|
ready_barrier = 10
|
|
empty_barrier = 11
|
|
participant_threads = 64
|
|
|
|
if tx < 32:
|
|
if tx == 0:
|
|
slot[0] = inp[0] + 1
|
|
T.named_barrier_arrive(ready_barrier, participant_threads)
|
|
T.sync_threads(empty_barrier, participant_threads)
|
|
else:
|
|
T.sync_threads(ready_barrier, participant_threads)
|
|
if tx == 32:
|
|
out[0] = slot[0]
|
|
T.named_barrier_arrive(empty_barrier, participant_threads)
|
|
|
|
return kernel
|
|
|
|
jit_kernel = kernel_fn()
|
|
source = jit_kernel.get_kernel_source()
|
|
|
|
assert "tl::__named_barrier_arrive(10, 64);" in source
|
|
assert "tl::__named_barrier_arrive(11, 64);" in source
|
|
|
|
inp = torch.tensor([41], device="cuda", dtype=torch.int32)
|
|
out = jit_kernel(inp)
|
|
torch.testing.assert_close(out.cpu(), torch.tensor([42], dtype=torch.int32))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_named_barrier_arrive_api()
|