Add switch-channel NVLS cross-rank barrier

include/mscclpp/switch_channel.hpp

@@ -4,6 +4,8 @@
 #ifndef MSCCLPP_SWITCH_CHANNEL_HPP_
 #define MSCCLPP_SWITCH_CHANNEL_HPP_
 
+#include <cstdint>
+#include <memory>
 #include <mscclpp/gpu_utils.hpp>
 #include <mscclpp/switch_channel_device.hpp>
 
@@ -16,6 +18,12 @@ struct SwitchChannel {
   void* devicePtr_;
   std::shared_ptr<void> mcPtr_;
   size_t bufferSize_;
+  // Barrier state inherited from the owning NvlsConnection (see NvlsConnection::bindAllocatedMemory).
+  // All are null / zero if the connection was created without barrier support.
+  uint32_t* barrierLocalFlag_ = nullptr;
+  uint32_t* barrierMcFlag_ = nullptr;
+  uint32_t* barrierGen_ = nullptr;
+  int barrierNRanks_ = 0;
 
  public:
   using DeviceHandle = SwitchChannelDeviceHandle;
@@ -41,9 +49,30 @@ class NvlsConnection {
   /// @return SwitchChannel with devicePtr, mcPtr and bufferSize
   SwitchChannel bindAllocatedMemory(CUdeviceptr devicePtr, size_t size);
 
+  /// Attach a device-side barrier resource shared by all SwitchChannels created from this
+  /// connection. After this call, `SwitchChannel::deviceHandle().barrier()` can synchronize all
+  /// ranks in the multicast group without a separate mesh of memory-channel semaphores. This is set
+  /// up automatically by `connectNvlsCollective`; it is an internal setup hook and is not intended
+  /// to be called directly.
+  /// @param barrierConn Auxiliary NVLS connection backing the barrier flag (kept alive).
+  /// @param barrierBuffer Storage for the barrier flag (kept alive); element 0 is the shared arrival
+  /// counter, element 1 is this rank's generation counter.
+  /// @param barrierChannel The bound barrier channel (kept alive for its multicast pointer).
+  /// @param nRanks Number of ranks participating in the multicast group.
+  void attachBarrier(std::shared_ptr<NvlsConnection> barrierConn, std::shared_ptr<void> barrierBuffer,
+                     std::shared_ptr<SwitchChannel> barrierChannel, int nRanks);
+
  private:
   class Impl;
   std::shared_ptr<Impl> pimpl_;
+
+  // Barrier resources, owned by this connection and shared by every SwitchChannel it creates.
+  std::shared_ptr<NvlsConnection> barrierConn_;
+  std::shared_ptr<void> barrierBuffer_;
+  std::shared_ptr<SwitchChannel> barrierChannel_;
+  uint32_t* barrierLocalFlag_ = nullptr;
+  uint32_t* barrierMcFlag_ = nullptr;
+  int barrierNRanks_ = 0;
 };
 
 class Communicator;

include/mscclpp/switch_channel_device.hpp

@@ -11,7 +11,9 @@
 #include <cuda_fp16.h>
 #endif  // defined(MSCCLPP_DEVICE_CUDA)
 
+#include <mscclpp/atomic_device.hpp>
 #include <mscclpp/gpu_data_types.hpp>
+#include <mscclpp/poll_device.hpp>
 
 #include "device.hpp"
 
@@ -25,8 +27,76 @@ struct SwitchChannelDeviceHandle {
   void* devicePtr;
   void* mcPtr;
   size_t bufferSize;
+  /// Multicast pointer to the shared arrival counter used by barrier(). A single multimem add on
+  /// this pointer is reflected into every rank's copy of the counter by the switch. Null if the
+  /// owning connection was created without barrier support.
+  uint32_t* mcBarrierFlag;
+  /// Local (unicast) pointer to this rank's own copy of the arrival counter used by barrier().
+  /// This is the address barrier() spins on. Null if the connection has no barrier support.
+  uint32_t* localBarrierFlag;
+  /// Local pointer to this rank's persistent barrier generation counter. It advances by nRanks on
+  /// every barrier() call and provides the per-rank wait target (see barrier()). Persisting it in
+  /// GPU memory lets barrier() be called repeatedly within and across kernel launches without any
+  /// host-side reset. Null if the connection has no barrier support.
+  uint32_t* barrierGen;
+  /// Number of ranks (devices) participating in the multicast group.
+  int nRanks;
 
 #if defined(MSCCLPP_DEVICE_CUDA)
+  /// Synchronize all ranks in the multicast group using the switch's multimem atomics.
+  ///
+  /// This is a device-side cross-rank barrier: it lets a kernel synchronize all ranks in the NVLS
+  /// group without a separate set of memory-channel semaphores or any host-side barrier.
+  ///
+  /// `memoryOrder` selects how much visibility the barrier carries. The default,
+  /// `cuda::memory_order::relaxed`, is a pure execution barrier: it synchronizes rank arrival but
+  /// makes no cross-rank data-visibility guarantee. Pass a stronger order (release/acq_rel/seq_cst)
+  /// to also publish memory -- the arrival then becomes a `.release` multimem add and the wait an
+  /// `.acquire` load, so writes issued by any rank before its barrier() call are visible to all
+  /// ranks after their barrier() call returns. This ordering is carried by scoped release/acquire
+  /// on the counter itself -- at `.sys` scope only on the counter -- rather than by a pair of
+  /// `__threadfence_system()` calls, which is much cheaper than a full system fence (this matches
+  /// NCCL's LSA switch barrier in `lsa_barrier__funcs.h`).
+  ///
+  /// The protocol: every rank advances its private target by nRanks, performs one multimem add of 1
+  /// on the shared counter (which the switch applies to every rank's copy), then spins on its own
+  /// local copy until the counter reaches the target. Because every rank calls barrier() the same
+  /// number of times and advances its target identically, the targets stay in lock-step and the
+  /// counter is never reset.
+  ///
+  /// @note Must be called by exactly one thread per rank (e.g. block 0, thread 0); the barrier
+  /// counts ranks, not threads. For a grid-wide cross-rank barrier, converge the grid (e.g. via
+  /// `mscclpp::DeviceSyncer::sync`) before and after this call. Requires that the owning
+  /// `NvlsConnection` was created with barrier support, i.e. the barrier pointers are non-null.
+  /// @param memoryOrder Ordering applied to the arrival/wait. `relaxed` (default) gives a pure
+  /// execution barrier; a stronger order (release/acq_rel/seq_cst) additionally publishes each
+  /// rank's pre-barrier writes to all ranks via a release arrival paired with an acquire wait.
+  /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
+  MSCCLPP_DEVICE_INLINE void barrier(cuda::memory_order memoryOrder = cuda::memory_order::relaxed,
+                                     [[maybe_unused]] int64_t maxSpinCount = 100000000) {
+    // Guard against calling barrier() on a channel whose connection has no barrier support. This is
+    // a debug-only diagnostic; in release builds a null pointer here dereferences and crashes, which
+    // is intentionally preferred over a silent no-op barrier (that would hide a cross-rank race).
+    MSCCLPP_ASSERT_DEVICE(barrierGen != nullptr, "SwitchChannel::barrier() called without barrier support");
+    // Advance this rank's private target. All ranks advance identically, so targets stay in lock-step.
+    const uint32_t target = (*barrierGen += static_cast<uint32_t>(nRanks));
+    // Signal arrival: one multimem add increments every rank's copy of the counter through the switch.
+    if (memoryOrder == cuda::memory_order::relaxed) {
+      // Relaxed arrival: pure execution barrier, no data-visibility ordering.
+      asm volatile("multimem.red.relaxed.sys.add.u32 [%0], %1;" ::"l"(mcBarrierFlag), "r"(1U) : "memory");
+    } else {
+      // Release arrival publishes this rank's prior writes before the arrival is observed by peers.
+      asm volatile("multimem.red.release.sys.add.u32 [%0], %1;" ::"l"(mcBarrierFlag), "r"(1U) : "memory");
+    }
+
+    cuda::memory_order waitOrder =
+        (memoryOrder == cuda::memory_order::relaxed) ? cuda::memory_order::relaxed : cuda::memory_order::acquire;
+    // Wait until every rank has arrived. The signed (wrap-safe) compare means "counter is behind target".
+    POLL_MAYBE_JAILBREAK(
+        (static_cast<int32_t>(atomicLoad<uint32_t, scopeSystem>(localBarrierFlag, waitOrder) - target) < 0),
+        maxSpinCount);
+  }
+
   template <typename T>
   MSCCLPP_DEVICE_INLINE T reduce(uint64_t index) {
     return SwitchChannelDeviceHandle::multimemLoadReduce(reinterpret_cast<T*>(mcPtr) + index);

python/mscclpp/language/channel.py

@@ -1007,6 +1007,36 @@ def broadcast_packets(self, rank, src_chunk: Chunk, buffer_offset, size, tb):
         )
         get_program().add_operation(self.src_rank, tb, op)
 
+    def barrier(self, rank, tb_list):
+        """Perform a switch-native cross-rank barrier for this rank.
+
+        Replaces a MemoryChannel signal/wait mesh with a single NVLS switch barrier.
+        All threadblocks in ``tb_list`` (across instances) converge on this rank, one
+        thread issues the single multicast arrival, and all blocks are released once
+        every rank in the group has arrived.
+
+        Args:
+            rank (int): The rank that will execute this barrier operation.
+            tb_list (List[int]): Thread block IDs that participate in the barrier.
+                Must include thread block 0, which acts as the arrival leader.
+
+        Raises:
+            RuntimeError: If tb_list is empty or does not include thread block 0.
+
+        Example:
+            >>> channel.barrier(rank=0, tb_list=[0])
+        """
+        if len(tb_list) == 0:
+            raise RuntimeError("Switch barrier requires at least one thread block.")
+        if 0 not in tb_list:
+            raise RuntimeError("Switch barrier tb_list must include thread block 0 (the arrival leader).")
+
+        self.src_rank = rank
+        for tb in tb_list:
+            tb_channel_ids = get_program().setup_channel(tb, self)
+            op = GroupBarrier(rank, tb_list, tb_channel_ids[0])
+            get_program().add_operation(rank, tb, op)
+
     class SwitchChannelRankView:
         """A rank-specific view of a SwitchChannel for performing operations.
 
@@ -1071,6 +1101,24 @@ def broadcast(self, src_chunk: Chunk, buffer_offset, size, tb):
             """
             return self._channel.broadcast(self._rank, src_chunk, buffer_offset, size, tb)
 
+        def barrier(self, tb_list):
+            """Perform a switch-native barrier from this rank's perspective.
+
+            Convenience method that calls the underlying channel's barrier method
+            with this view's rank automatically provided.
+
+            Args:
+                tb_list (List[int]): Thread block IDs that participate in the barrier.
+                    Must include thread block 0, which acts as the arrival leader.
+
+            Returns:
+                The result of the underlying channel's barrier operation.
+
+            Example:
+                >>> rank_view.barrier(tb_list=[0])
+            """
+            return self._channel.barrier(self._rank, tb_list)
+
         def broadcast_packets(self, src_chunk: Chunk, buffer_offset, size, tb):
             """Perform a packet broadcast operation from this rank's perspective.
 

python/mscclpp/language/internal/buffer_access.py

@@ -19,7 +19,11 @@ def __init__(self):
     def process_operations(self, operations):
         result_operations = []
         for operation in operations:
-            if operation.name == Instruction.nop or operation.name == Instruction.barrier:
+            if (
+                operation.name == Instruction.nop
+                or operation.name == Instruction.barrier
+                or operation.name == Instruction.group_barrier
+            ):
                 self.clear_data_access()
             else:
                 if operation.name == Instruction.pipeline:

python/mscclpp/language/internal/operations.py

@@ -362,19 +362,30 @@ class BarrierOperation(BaseOperation):
     __current_barriers = []
 
     def __init__(self, rank: int, tb_list: List[int]):
-        for _ in range(len(BarrierOperation.__current_barriers), rank + 1):
-            BarrierOperation.__current_barriers.append({})
         barrier_info = BarrierOperation.BarrierInfo(tb_list)
-
-        if barrier_info not in BarrierOperation.__current_barriers[rank]:
-            self.barrier_id = len(BarrierOperation.__current_barriers[rank])
-            BarrierOperation.__current_barriers[rank][barrier_info] = self.barrier_id
-        else:
-            self.barrier_id = BarrierOperation.__current_barriers[rank][barrier_info]
+        self.barrier_id = BarrierOperation.reserve_barrier_id(rank, barrier_info)
 
         super().__init__(Instruction.barrier)
         self.barrier_info = barrier_info
 
+    @classmethod
+    def reserve_barrier_id(cls, rank: int, barrier_info):
+        """Allocate a per-rank DeviceSyncer id, shared across all barrier kinds.
+
+        Both intra-block barriers and switch (NVLS) barriers draw from this single
+        per-rank sequence. Because ``BarrierInfo`` includes a ``kind`` discriminator,
+        different barrier kinds with the same ``tb_list`` receive distinct ids and
+        therefore never alias to the same ``deviceSyncers[]`` slot after instancing.
+        """
+        for _ in range(len(cls.__current_barriers), rank + 1):
+            cls.__current_barriers.append({})
+        if barrier_info not in cls.__current_barriers[rank]:
+            barrier_id = len(cls.__current_barriers[rank])
+            cls.__current_barriers[rank][barrier_info] = barrier_id
+        else:
+            barrier_id = cls.__current_barriers[rank][barrier_info]
+        return barrier_id
+
     def shift_ids(self, instance, num_instances, replication_function):
         self.barrier_id = replication_function(self.barrier_id, instance, num_instances)
 
@@ -395,14 +406,55 @@ def to_dict(self):
         return result
 
     class BarrierInfo:
-        def __init__(self, tb_list):
+        def __init__(self, tb_list, kind="sync"):
             self.tb_list = tb_list
+            self.kind = kind
 
         def __eq__(self, other):
-            return self.tb_list == other.tb_list
+            return self.tb_list == other.tb_list and self.kind == other.kind
 
         def __hash__(self):
-            return hash(tuple(self.tb_list))
+            return hash((tuple(self.tb_list), self.kind))
+
+
+class GroupBarrier(BaseOperation):
+    """Grid-wide cross-rank barrier over an NVLS SwitchChannel.
+
+    Replaces the O(n^2) MemoryChannel signal/wait mesh with a single switch-native
+    barrier. All threadblocks participating on a rank converge via a shared
+    DeviceSyncer, one thread issues the single cross-rank ``multimem`` arrival, and
+    all blocks are then released.
+
+    Instancing semantics (grid-wide collapse): every instance copy maps to the SAME
+    DeviceSyncer slot (leader instance 0), and ``num_threadblocks`` folds in the
+    instance count so the syncer converges every physical block on the rank. Only
+    physical ``blockIdx.x == 0`` issues the arrival, so ``tb_list`` must include tb 0.
+
+    The barrier is fence-free: ordering is carried by scoped release/acquire on the
+    switch counter itself (see ``SwitchChannel::barrier()``), not by system fences.
+    """
+
+    def __init__(self, rank: int, tb_list: List[int], switch_channel_id: int):
+        super().__init__(Instruction.group_barrier)
+        self.barrier_info = BarrierOperation.BarrierInfo(tb_list, kind="switch")
+        self.barrier_id = BarrierOperation.reserve_barrier_id(rank, self.barrier_info)
+        self.switch_channel_id = switch_channel_id
+        self.tb_count = len(tb_list)
+        self.num_threadblocks = self.tb_count
+
+    def shift_ids(self, instance, num_instances, replication_function):
+        # Collapse all instances onto the leader (instance 0) slot so every physical
+        # block shares one DeviceSyncer, and grow num_threadblocks to cover them all.
+        self.barrier_id = replication_function(self.barrier_id, 0, num_instances)
+        self.num_threadblocks = self.tb_count * num_instances
+
+    def to_dict(self):
+        result = {"name": self.name.value}
+        result["switch_channel_id"] = self.switch_channel_id
+        result["barrier_id"] = self.barrier_id
+        result["num_threadblocks"] = self.num_threadblocks
+
+        return result
 
 
 class FlushOperation(BaseOperation):

python/mscclpp/language/internal/types.py

@@ -93,6 +93,7 @@ class Instruction(Enum):
     group_store_packet = "gstorepkt"
     group_load_reduce = "glre"
     group_load_reduce_store = "glres"
+    group_barrier = "gbarrier"
     pipeline = "pipeline"
 
     def __str__(self):

python/mscclpp/language/tests/single_node/allgather_nvls_zero_copy_barrier.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+from mscclpp.language.channel import *
+from mscclpp.language.rank import *
+from mscclpp.language.general import *
+from mscclpp.language.program import *
+from mscclpp.language.collectives import *
+
+
+def allgather_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size, instances):
+    # Defaults instances=8, num_threads_per_block=256 are tuned for 64-GPU (4x GB200) MNNVL NVLS:
+    # they give the best busbw across 1MB-1GB (instances saturate at 8; tpb=256 beats 512/1024).
+    chunksperloop = 1
+    collective = AllGather(gpu_size, chunksperloop, True)
+    with CollectiveProgram(
+        name,
+        collective,
+        gpu_size,
+        instances=instances,
+        protocol="Simple",
+        num_threads_per_block=num_threads_per_block,
+        use_double_scratch_buffer=False,
+        min_message_size=min_message_size,
+        max_message_size=max_message_size,
+    ):
+        # NVLS multicast channel over the output buffer. For Allgather each
+        # rank stores its own chunk to all ranks' output buffers via the switch.
+        nvls_chan = SwitchChannel(rank_list=[gpu for gpu in range(gpu_size)], buffer_type=BufferType.output)
+
+        # Synchronization to ensure all the GPUs are ready. A single switch-native
+        # barrier over the NVLS channel replaces the O(n^2) MemoryChannel signal/wait mesh.
+        for gpu in range(gpu_size):
+            nvls_chan.at_rank(gpu).barrier(tb_list=[0])
+
+        # Broadcasting each rank's chunk to every rank via NVLS multimem store.
+        # Rank `gpu` owns output chunk `gpu` (its input under in-place AllGather) and
+        # stores it to offset `gpu` across all ranks in the switch group.
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            output_buffer = rank.get_output_buffer()
+            nvls_chan.at_rank(gpu).broadcast(src_chunk=output_buffer[gpu : gpu + 1], buffer_offset=gpu, size=1, tb=0)
+
+        # Synchronization to ensure the GPUs finished
+        for gpu in range(gpu_size):
+            nvls_chan.at_rank(gpu).barrier(tb_list=[0])
+
+        print(JSON())
+
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--name", type=str, help="name of the program")
+parser.add_argument("--num_gpus", type=int, help="number of gpus")
+parser.add_argument("--num_threads_per_block", type=int, default=256, help="number of threads per block")
+parser.add_argument("--min_message_size", type=int, default=0, help="minimum message size")
+parser.add_argument("--max_message_size", type=int, default=2**64 - 1, help="maximum message size")
+parser.add_argument("--instances", type=int, default=8, help="number of instances (parallel threadblocks)")
+
+args = parser.parse_args()
+
+allgather_example(
+    args.name,
+    args.num_gpus,
+    args.num_threads_per_block,
+    args.min_message_size,
+    args.max_message_size,
+    args.instances,
+)