[FIX][TIRx] Use cluster arrivals for remote mbarrier views

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,32 +147,37 @@ def arrive(self, stage, cta_id=None, pred=None, count=None):
         # 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)
 
     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 @@ def arrive(self, stage, tx_count=None, cta_id=None, pred=None):
         # 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):

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_try_wait_once(A: T.Buffer((16, 16), "float16")):
         assert "selp.u32" in src
 
 
+@pytest.mark.gpu
+@pytest.mark.skipif(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
+
+
+@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():
+    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():