This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ece32bfaee [TIRx][CUDA] Preserve local mbarrier predicate and count
(#20250)
ece32bfaee is described below
commit ece32bfaee86dc21920ee0e68e0e4a3681afbc49
Author: Hongyi Jin <[email protected]>
AuthorDate: Tue Sep 1 07:46:27 2026 -0400
[TIRx][CUDA] Preserve local mbarrier predicate and count (#20250)
## Motivation
A KDA kernel-evolution experiment introduced a count-1 local handoff
barrier. The producer intentionally selected one lane:
```python
pred = K.cuda.elect_sync()
bar.arrive(0, pred=pred)
```
However, the local branch of `MBarrier.arrive` discarded `pred` and
emitted an unconditional `mbarrier.arrive.shared.b64`. Synccheck and
NumSim consequently observed 32 arrivals from the producer warp where
the barrier expected 1. The kernel had to work around the lowering bug
by placing an unconditional `arrive()` inside an explicit elected-lane
guard.
The same local branch also discarded an explicit `count`, even though
both arguments are part of the public `MBarrier.arrive` contract and are
already preserved by the remote branches.
## Summary
- preserve explicit predicate and arrival count for local
`MBarrier.arrive` calls
- keep local arrivals on the direct shared-memory path while leaving
remote address mapping unchanged
- cover both the TIR intrinsic operands and generated predicated PTX
## Testing
- `python -m pytest -q
tests/python/tirx/codegen/test_codegen_blackwell.py`
- 14 passed
- `pre-commit run --files python/tvm/backend/cuda/lang/pipeline.py
tests/python/tirx/codegen/test_codegen_blackwell.py`
- passed
---
python/tvm/backend/cuda/lang/pipeline.py | 20 ++++++-------
.../python/tirx/codegen/test_codegen_blackwell.py | 33 ++++++++++++++++++++++
2 files changed, 43 insertions(+), 10 deletions(-)
diff --git a/python/tvm/backend/cuda/lang/pipeline.py
b/python/tvm/backend/cuda/lang/pipeline.py
index 69d289e474..2c4ece7c2d 100644
--- a/python/tvm/backend/cuda/lang/pipeline.py
+++ b/python/tvm/backend/cuda/lang/pipeline.py
@@ -205,7 +205,9 @@ class MBarrier:
# reuses the one mapa the view did.
_mbarrier_arrive_remote(self.buf.ptr_to([stage]), pred, count)
elif remote is None:
- self._arrive(self.buf.ptr_to([stage]))
+ # Keep local arrival as the default; remote arrival must pay for
+ # address mapping explicitly.
+ self._arrive(self.buf.ptr_to([stage]), pred, count)
else:
# Split of the legacy fused wrapper: map the address into the
# target CTA, then arrive on it. Plain Python so the mapa scratch
@@ -216,15 +218,13 @@ class MBarrier:
)
@T.inline
- def _arrive(self, bar):
- # Local-CTA arrive. To arrive on a remote CTA's mbarrier in a cluster
- # kernel, callers must pass ``remote=`` explicitly (e.g.
- # ``bar.arrive(stage, remote=0)``) or use
- # ``MBarrier.remote_view(rank).arrive(stage)``. Defaulting the
- # cross-CTA path was both surprising (``bar.arrive(stage)`` silently
- # ``mapa``ed across the cluster) and a per-call cost of ~3 PTX ops on
- # every single-CTA kernel.
- T.ptx.mbarrier.arrive.shared.b64(bar, T.uint32(1))
+ def _arrive(self, bar, pred=None, count=None):
+ if pred is None:
+ T.ptx.mbarrier.arrive.shared.b64(bar, T.uint32(1 if count is None
else count))
+ else:
+ T.ptx.mbarrier.arrive.shared.b64(
+ bar, T.uint32(1 if count is None else count), pred=pred
+ )
def ptr_to(self, idx):
return self.buf.ptr_to(idx)
diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py
b/tests/python/tirx/codegen/test_codegen_blackwell.py
index 41e6c8a26a..f02103a20b 100644
--- a/tests/python/tirx/codegen/test_codegen_blackwell.py
+++ b/tests/python/tirx/codegen/test_codegen_blackwell.py
@@ -179,6 +179,39 @@ def test_mbarrier_remote_view_codegen():
assert 'asm volatile("mbarrier.arrive.shared.b64' not in src
[email protected]
[email protected](not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
+def test_mbarrier_local_arrive_forwards_predicate_and_count():
+ from tvm.tirx.lang.pipeline import MBarrier
+
+ # fmt: off
+ @T.prim_func
+ def test_local_arrive():
+ T.device_entry()
+ thread = T.thread_id([32])
+ pool = T.SMEMPool()
+ bar = MBarrier(pool, 1)
+ pool.commit()
+ bar.arrive(0, pred=(thread == 0), count=2)
+ # fmt: on
+
+ arrive_calls = []
+
+ def visit(node):
+ if isinstance(node, tvm.ir.Call) and node.op.name ==
"tirx.ptx.mbarrier_arrive":
+ arrive_calls.append(node)
+
+ tvm.tirx.stmt_functor.post_order_visit(test_local_arrive.body, visit)
+ assert len(arrive_calls) == 1
+ call = arrive_calls[0]
+ assert call.args[1].value == 2
+ assert call.args[-1].value == "pred"
+
+ with tvm.target.Target("cuda"):
+ src, _ = _get_source(test_local_arrive)
+ assert "@p mbarrier.arrive.shared.b64 _, [%0], %1;" in src
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
def test_tma_mbarrier_remote_view_codegen():
