diff --git a/python/tvm/backend/cuda/lang/pipeline.py b/python/tvm/backend/cuda/lang/pipeline.py index 0cf482eedba3..040bdf056f4e 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 @@ 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) @@ -138,19 +158,26 @@ def ptr_to(self, 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): diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tests/python/tirx/codegen/test_codegen_blackwell.py index c40749f977ab..99bd4f57de1d 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_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():