This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 62fb780bb0 [FIX][TIRx] Use cluster arrivals for remote mbarrier views
(#20074)
62fb780bb0 is described below
commit 62fb780bb0a8da62e3808f60a2343f6fd1d4b01f
Author: Hongyi Jin <[email protected]>
AuthorDate: Wed Jul 29 13:51:30 2026 -0400
[FIX][TIRx] Use cluster arrivals for remote mbarrier views (#20074)
## Motivation and context
`MBarrier.remote_view(rank)` represents an mbarrier owned by another CTA
in the same cluster. The existing view kept only a buffer whose pointer
had been mapped to the remote CTA with PTX `mapa`. Calling
`remote_bar.arrive(...)` then followed the inherited local-arrive path
and emitted the local `mbarrier.arrive.shared.b64` form against that
mapped remote address.
A remote arrival must instead use the `shared::cluster` instruction
form. TIRx models that form with the owner-local barrier pointer plus
the destination CTA rank and predicate. Using the local instruction with
a remote address is not equivalent and is reported by synccheck as a
local arrival on a remote mbarrier address.
## Changes
- Keep the owner-local buffer and target CTA rank when constructing a
remote mbarrier view.
- Route `MBarrier` remote arrivals through the cluster helper using the
local barrier pointer and stored CTA rank.
- Apply the same routing to `TMABar`, including
`mbarrier.arrive.expect_tx`.
- Keep a typed `PointerType(uint64, shared)` mapped buffer on the view
so `ptr_to` remains available to operations that explicitly consume a
remote shared-memory pointer.
- Reject operations with ambiguous or invalid ownership semantics:
- initializing or waiting on a remote view,
- supplying another `cta_id` to a view that already fixes its target,
- creating a remote view from another remote view.
- Preserve the existing local-CTA behavior for ordinary barriers.
## Testing
- Verify the typed `mapa` binding and remote buffer in TIRx IR.
- Verify CUDA codegen for plain and counted
`mbarrier.arrive.shared::cluster.b64` forms.
- Verify CUDA codegen for remote
`mbarrier.arrive.expect_tx.shared::cluster.b64`.
- Verify that the corresponding local instruction forms are not emitted
for remote views.
- Verify diagnostics for remote init, wait, nested views, and
conflicting `cta_id` arguments.
- Run changed-files pre-commit checks.
Focused result: 3 tests passed.
---
python/tvm/backend/cuda/lang/pipeline.py | 73 +++++++---
.../python/tirx/codegen/test_codegen_blackwell.py | 156 +++++++++++++++++++++
2 files changed, 213 insertions(+), 16 deletions(-)
diff --git a/python/tvm/backend/cuda/lang/pipeline.py
b/python/tvm/backend/cuda/lang/pipeline.py
index 0cf482eedb..040bdf056f 100644
--- a/python/tvm/backend/cuda/lang/pipeline.py
+++ b/python/tvm/backend/cuda/lang/pipeline.py
@@ -93,24 +93,46 @@ class MBarrier:
def __init__(self, pool, depth, phase_offset=0, leader=None):
self.buf = pool.alloc((depth,), "uint64", align=8)
+ self._local_buf = self.buf
+ self._remote_cta_id = None
self.depth = depth
self.phase_offset = phase_offset
self.leader = leader if leader is not None else (T.cuda.thread_rank()
== 0)
- @T.inline
def init(self, count):
+ if self._remote_cta_id is not None:
+ raise ValueError("MBarrier.remote_view() cannot be initialized")
+ self._init(count)
+
+ @T.inline
+ def _init(self, count):
if self.leader:
for i in T.unroll(self.depth):
T.ptx.mbarrier.init(self.buf.ptr_to([i]), count)
- @T.inline
def wait(self, stage, phase):
+ if self._remote_cta_id is not None:
+ raise ValueError("MBarrier.remote_view() cannot be waited on")
+ self._wait(stage, phase)
+
+ @T.inline
+ def _wait(self, stage, phase):
# Blocks: ``mbarrier.try_wait`` loops internally until the phase flips,
# so this returns only once the barrier has completed.
T.ptx.mbarrier.try_wait(self.buf.ptr_to([stage]), phase ^
self.phase_offset)
- @T.inline
def arrive(self, stage, cta_id=None, pred=None, count=None):
+ if self._remote_cta_id is not None:
+ if cta_id is not None:
+ raise ValueError("MBarrier.remote_view().arrive() cannot also
specify cta_id")
+ cta_id = self._remote_cta_id
+ buf = self._local_buf
+ else:
+ buf = self.buf
+ self._arrive(buf.ptr_to([stage]), cta_id, pred, count)
+
+ @T.inline
+ def _arrive(self, bar, cta_id=None, pred=None, count=None):
# Default: local-CTA arrive — emits the simple
# ``mbarrier.arrive.shared.b64`` form. To arrive on a remote
# CTA's mbarrier in a cluster kernel, callers must pass
@@ -125,12 +147,10 @@ class MBarrier:
# When ``None`` the implicit count-of-1 form is emitted. Passing
# ``count=1`` is semantically identical but spells the count
explicitly.
if cta_id is None:
- T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
+ T.ptx.mbarrier.arrive(bar)
else:
actual_pred = True if pred is None else pred
- T.ptx.mbarrier.arrive(
- self.buf.ptr_to([stage]), cta_id=cta_id, pred=actual_pred,
count=count
- )
+ T.ptx.mbarrier.arrive(bar, cta_id=cta_id, pred=actual_pred,
count=count)
def ptr_to(self, idx):
return self.buf.ptr_to(idx)
@@ -138,19 +158,26 @@ class MBarrier:
def remote_view(self, rank):
"""Create a view of this barrier mapped to another CTA's shared memory.
- Arrive-only: the returned view is built with ``object.__new__`` and
- never copies ``self.leader``, so calling ``.init()`` on it would fail.
- Use it solely to ``arrive`` on a remote CTA's mbarrier.
+ The returned view retains the local barrier and target CTA so
+ ``arrive`` emits the cluster form. Its mapped buffer remains available
+ through ``ptr_to`` for operations that consume a remote shared-memory
+ pointer. ``init`` and ``wait`` are local-only and reject remote views.
"""
from tvm.ir import PointerType, PrimType
from tvm.tirx import Var as TIRVar
- expr = T.reinterpret("handle",
T.ptx.map_shared_rank(self.buf.ptr_to([0]), rank))
- ptr = TIRVar("remote_mbar_ptr", PointerType(PrimType("uint64")))
+ if self._remote_cta_id is not None:
+ raise ValueError("MBarrier.remote_view() cannot be applied to a
remote view")
+
+ ptr_ty = PointerType(PrimType("uint64"), "shared")
+ expr = T.reinterpret(ptr_ty,
T.ptx.map_shared_rank(self.buf.ptr_to([0]), rank))
+ ptr = TIRVar("remote_mbar_ptr", ptr_ty)
T.Bind(expr, var=ptr)
buf = T.decl_buffer([self.depth], "uint64", data=ptr, scope="shared")
remote = object.__new__(type(self))
remote.buf = buf
+ remote._local_buf = self._local_buf
+ remote._remote_cta_id = rank
remote.depth = self.depth
remote.phase_offset = self.phase_offset
return remote
@@ -163,8 +190,18 @@ class TMABar(MBarrier):
(matching MBarrier.arrive defaults).
"""
- @T.inline
def arrive(self, stage, tx_count=None, cta_id=None, pred=None):
+ if self._remote_cta_id is not None:
+ if cta_id is not None:
+ raise ValueError("TMABar.remote_view().arrive() cannot also
specify cta_id")
+ cta_id = self._remote_cta_id
+ buf = self._local_buf
+ else:
+ buf = self.buf
+ self._arrive_tma(buf.ptr_to([stage]), tx_count, cta_id, pred)
+
+ @T.inline
+ def _arrive_tma(self, bar, tx_count=None, cta_id=None, pred=None):
# NOTE: this arrive() kwarg set intentionally differs from
# MBarrier.arrive (hardware necessity, LSP-incompatible by design).
# ``tx_count``: TMA byte count for ``mbarrier.arrive.expect_tx``.
@@ -173,12 +210,16 @@ class TMABar(MBarrier):
# arrive is local-CTA only. See ``MBarrier.arrive`` for the
# full default-local rationale.
if tx_count is not None:
- T.ptx.mbarrier.arrive.expect_tx(self.buf.ptr_to([stage]), tx_count)
+ if cta_id is None:
+ T.ptx.mbarrier.arrive.expect_tx(bar, tx_count)
+ else:
+ actual_pred = True if pred is None else pred
+ T.ptx.mbarrier.arrive.expect_tx(bar, tx_count, cta_id=cta_id,
pred=actual_pred)
elif cta_id is None:
- T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]))
+ T.ptx.mbarrier.arrive(bar)
else:
actual_pred = True if pred is None else pred
- T.ptx.mbarrier.arrive(self.buf.ptr_to([stage]), cta_id=cta_id,
pred=actual_pred)
+ T.ptx.mbarrier.arrive(bar, cta_id=cta_id, pred=actual_pred)
class TCGen05Bar(MBarrier):
diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py
b/tests/python/tirx/codegen/test_codegen_blackwell.py
index c40749f977..99bd4f57de 100644
--- a/tests/python/tirx/codegen/test_codegen_blackwell.py
+++ b/tests/python/tirx/codegen/test_codegen_blackwell.py
@@ -33,6 +33,38 @@ def _get_source(func: tvm.tirx.PrimFunc) -> str:
return src, mod
+def _assert_remote_mbarrier_ir(func, arrive_op_name, cta_arg_index):
+ bindings = []
+ buffers = []
+ mapa_calls = []
+ arrive_calls = []
+
+ def visit(node):
+ if isinstance(node, tvm.tirx.Bind) and node.var.name ==
"remote_mbar_ptr":
+ bindings.append(node)
+ if isinstance(node, tvm.tirx.DeclBuffer) and node.buffer.data.name ==
"remote_mbar_ptr":
+ buffers.append(node.buffer)
+ if isinstance(node, tvm.ir.Call) and node.op.name == "tirx.ptx.mapa":
+ mapa_calls.append(node)
+ if isinstance(node, tvm.ir.Call) and node.op.name == arrive_op_name:
+ arrive_calls.append(node)
+
+ tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+ assert len(bindings) == 1
+ assert len(buffers) == 1
+ assert len(mapa_calls) == 1
+ assert arrive_calls
+ assert isinstance(bindings[0].var.ty, tvm.ir.PointerType)
+ assert bindings[0].var.ty.storage_scope == "shared"
+ assert bindings[0].value.ty.storage_scope == "shared"
+ assert buffers[0].data.same_as(bindings[0].var)
+ assert buffers[0].data.ty.storage_scope == "shared"
+ assert buffers[0].scope() == "shared"
+ for arrive in arrive_calls:
+ tvm.ir.assert_structural_equal(arrive.args[0], mapa_calls[0].args[0])
+ tvm.ir.assert_structural_equal(arrive.args[cta_arg_index],
mapa_calls[0].args[1])
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
def test_tmem_alloc_dealloc_relinquish():
@@ -89,6 +121,130 @@ def test_mbarrier_try_wait_once_codegen():
assert "selp.u32" in src
[email protected]
[email protected](not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
+def test_mbarrier_remote_view_codegen():
+ from tvm.tirx.lang.pipeline import MBarrier
+
+ # fmt: off
+ @T.prim_func
+ def test_remote_view():
+ T.device_entry()
+ T.cluster_id([1])
+ T.cta_id_in_cluster([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = MBarrier(pool, 1)
+ pool.commit()
+ remote_bar = bar.remote_view(0)
+ remote_bar.arrive(0)
+ remote_bar.arrive(0, count=2)
+ # fmt: on
+
+ _assert_remote_mbarrier_ir(test_remote_view, "tirx.ptx.mbarrier_arrive", 1)
+ with tvm.target.Target("cuda"):
+ src, _ = _get_source(test_remote_view)
+ assert "tvm_builtin_ptx_mapa_u64" in src
+ assert "tvm_builtin_ptx_mbarrier_arrive_remote" in src
+ assert "tvm_builtin_ptx_mbarrier_arrive_remote_count" in src
+ assert "mbarrier.arrive.shared::cluster.b64" in src
+ 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_tma_mbarrier_remote_view_codegen():
+ from tvm.tirx.lang.pipeline import TMABar
+
+ # fmt: off
+ @T.prim_func
+ def test_remote_view():
+ T.device_entry()
+ T.cluster_id([1])
+ T.cta_id_in_cluster([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = TMABar(pool, 1)
+ pool.commit()
+ remote_bar = bar.remote_view(0)
+ remote_bar.arrive(0, tx_count=128)
+ # fmt: on
+
+ _assert_remote_mbarrier_ir(test_remote_view,
"tirx.ptx.mbarrier_arrive_expect_tx", 2)
+ with tvm.target.Target("cuda"):
+ src, _ = _get_source(test_remote_view)
+ assert "tvm_builtin_ptx_mbarrier_arrive_expect_tx_remote" in src
+ assert "mbarrier.arrive.expect_tx.shared::cluster.b64" in src
+ assert 'asm volatile("mbarrier.arrive.expect_tx.shared.b64' not in src
+
+
+def test_mbarrier_remote_view_rejects_invalid_operations():
+ from tvm.tirx.lang.pipeline import MBarrier, TMABar
+
+ with pytest.raises(tvm.error.DiagnosticError, match=r"remote_view\(\)
cannot be initialized"):
+ # fmt: off
+ @T.prim_func
+ def invalid_init():
+ T.device_entry()
+ T.cta_id([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = MBarrier(pool, 1)
+ bar.remote_view(0).init(1)
+ # fmt: on
+
+ with pytest.raises(tvm.error.DiagnosticError, match=r"remote_view\(\)
cannot be waited on"):
+ # fmt: off
+ @T.prim_func
+ def invalid_wait():
+ T.device_entry()
+ T.cta_id([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = MBarrier(pool, 1)
+ bar.remote_view(0).wait(0, 0)
+ # fmt: on
+
+ with pytest.raises(tvm.error.DiagnosticError, match="cannot also specify
cta_id"):
+ # fmt: off
+ @T.prim_func
+ def ambiguous_mbarrier_arrive():
+ T.device_entry()
+ T.cta_id([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = MBarrier(pool, 1)
+ bar.remote_view(0).arrive(0, cta_id=1)
+ # fmt: on
+
+ with pytest.raises(tvm.error.DiagnosticError, match="cannot also specify
cta_id"):
+ # fmt: off
+ @T.prim_func
+ def ambiguous_tma_arrive():
+ T.device_entry()
+ T.cta_id([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = TMABar(pool, 1)
+ bar.remote_view(0).arrive(0, tx_count=128, cta_id=1)
+ # fmt: on
+
+ with pytest.raises(
+ tvm.error.DiagnosticError,
+ match=r"remote_view\(\) cannot be applied to a remote view",
+ ):
+ # fmt: off
+ @T.prim_func
+ def nested_remote_view():
+ T.device_entry()
+ T.cta_id([2])
+ T.thread_id([128])
+ pool = T.SMEMPool()
+ bar = MBarrier(pool, 1)
+ bar.remote_view(0).remote_view(1)
+ # fmt: on
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
def test_fence_before_after_thread_sync():
