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,100 @@ 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)
+  /// Cross-rank switch barrier, split into arrival (signal) and completion (wait) halves.
+  ///
+  /// These four methods implement a device-side cross-rank barrier over the switch's multimem
+  /// atomics, without any memory-channel semaphores or host-side barrier. A full barrier is a
+  /// `signal()`/`wait()` pair (or their relaxed variants): every rank issues one multimem add of 1
+  /// on the shared arrival counter -- which the switch applies to every rank's copy -- and then
+  /// spins on its own local copy until the counter reaches its private target. Splitting arrival
+  /// from completion lets a kernel overlap independent work between the two halves.
+  ///
+  /// The protocol is monotonic and never reset: `wait()`/`relaxedWait()` advance this rank's private
+  /// target by nRanks each call. Because every rank calls the pair the same number of times and
+  /// advances its target identically, the targets stay in lock-step with the shared counter.
+  ///
+  /// Ordering is selected by which pair is used. The relaxed pair (`relaxedSignal`/`relaxedWait`) is
+  /// a pure execution barrier: it synchronizes rank arrival but makes no cross-rank data-visibility
+  /// guarantee. The ordered pair (`signal`/`wait`) additionally publishes memory: the arrival is a
+  /// `.release` multimem add and the wait an `.acquire` load, so writes issued by any rank before its
+  /// `signal()` are visible to all ranks after their `wait()` returns. This ordering is carried by
+  /// scoped release/acquire on the counter itself (at `.sys` scope, on the counter only) 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`).
+  ///
+  /// @note Each method 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`) around the pair. Requires that the owning `NvlsConnection`
+  /// was created with barrier support, i.e. the barrier pointers are non-null.
+
+  /// Issue an ordered cross-rank arrival, publishing this rank's prior writes.
+  ///
+  /// Performs one `multimem.red.release.sys.add` of 1 on the shared counter. The `.release` ordering
+  /// makes writes issued before this call visible to any peer that observes the arrival via an
+  /// acquiring `wait()`. Pair with `wait()`.
+  MSCCLPP_DEVICE_INLINE void signal() {
+    asm volatile("multimem.red.release.sys.add.u32 [%0], %1;" ::"l"(mcBarrierFlag), "r"(1U) : "memory");
+  }
+
+  /// Issue a relaxed cross-rank arrival, without any data-visibility ordering.
+  ///
+  /// Relaxed variant of `signal()`: performs `multimem.red.relaxed.sys.add`, a pure execution arrival
+  /// that synchronizes rank progress but makes no cross-rank memory-visibility guarantee. Pair with
+  /// `relaxedWait()`.
+  MSCCLPP_DEVICE_INLINE void relaxedSignal() {
+    asm volatile("multimem.red.relaxed.sys.add.u32 [%0], %1;" ::"l"(mcBarrierFlag), "r"(1U) : "memory");
+  }
+
+  /// Wait until every rank has arrived, acquiring peers' published writes.
+  ///
+  /// Advances this rank's private target by nRanks, then spins on its local copy of the counter with
+  /// an `.acquire` load until the counter reaches the target (the signed, wrap-safe compare means
+  /// "counter is behind target"). The acquire pairs with peers' `signal()` release so their
+  /// pre-arrival writes are visible after this returns. Pair with `signal()`.
+  ///
+  /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
+  MSCCLPP_DEVICE_INLINE void wait(int64_t maxSpinCount = 10000000) {
+    MSCCLPP_ASSERT_DEVICE(barrierGen != nullptr, "SwitchChannel::wait() called without barrier support");
+    const uint32_t target = (*barrierGen += static_cast<uint32_t>(nRanks));
+    POLL_MAYBE_JAILBREAK(
+        (static_cast<int32_t>(atomicLoad<uint32_t, scopeSystem>(localBarrierFlag, cuda::memory_order::acquire) -
+                              target) < 0),
+        maxSpinCount);
+  }
+
+  /// Wait until every rank has arrived, without any data-visibility ordering.
+  ///
+  /// Relaxed variant of `wait()`: advances this rank's private target by nRanks and spins on its local
+  /// copy of the counter with a relaxed load until the counter reaches the target. Provides rank
+  /// synchronization only (no cross-rank memory ordering). Pair with `relaxedSignal()`.
+  ///
+  /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
+  MSCCLPP_DEVICE_INLINE void relaxedWait(int64_t maxSpinCount = 10000000) {
+    MSCCLPP_ASSERT_DEVICE(barrierGen != nullptr, "SwitchChannel::relaxedWait() called without barrier support");
+    const uint32_t target = (*barrierGen += static_cast<uint32_t>(nRanks));
+    POLL_MAYBE_JAILBREAK(
+        (static_cast<int32_t>(atomicLoad<uint32_t, scopeSystem>(localBarrierFlag, cuda::memory_order::relaxed) -
+                              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,62 @@ def broadcast_packets(self, rank, src_chunk: Chunk, buffer_offset, size, tb):
         )
         get_program().add_operation(self.src_rank, tb, op)
 
+    def signal(self, rank, tb_list, relaxed=False):
+        """Signal all ranks in the group from this rank.
+
+        Sends a signal to all ranks in the rank group, notifying them that
+        an operation has completed or that data is ready.
+
+        Args:
+            rank (int): The rank that will execute this signal operation.
+            tb_list (List[int]): Thread block IDs that participate in the signal.
+            relaxed (bool, optional): Whether to use relaxed signaling. Defaults to False.
+
+        Raises:
+            RuntimeError: If tb_list is empty.
+
+        Example:
+            >>> channel.signal(rank=0, tb_list=[0])
+        """
+        if len(tb_list) == 0:
+            raise RuntimeError("Group signal requires at least one thread block.")
+        if 0 not in tb_list:
+            raise RuntimeError("Group signal 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 = GroupSignal(rank, tb_list, tb_channel_ids[0], relaxed=relaxed)
+            get_program().add_operation(rank, tb, op)
+
+    def wait(self, rank, tb_list, relaxed=False):
+        """Wait for a signal from all ranks in the group.
+
+        Waits for a signal from all ranks in the rank group, ensuring that
+        operations are completed before proceeding.
+
+        Args:
+            rank (int): The rank that will execute this wait operation.
+            tb_list (List[int]): Thread block IDs that participate in the wait.
+            relaxed (bool, optional): Whether to use relaxed waiting. Defaults to False.
+
+        Raises:
+            RuntimeError: If tb_list is empty.
+
+        Example:
+            >>> channel.wait(rank=0, tb_list=[0])
+        """
+        if len(tb_list) == 0:
+            raise RuntimeError("Group wait requires at least one thread block.")
+        if 0 not in tb_list:
+            raise RuntimeError("Group wait 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 = GroupWait(rank, tb_list, tb_channel_ids[0], relaxed=relaxed)
+            get_program().add_operation(rank, tb, op)
+
     class SwitchChannelRankView:
         """A rank-specific view of a SwitchChannel for performing operations.
 
@@ -1071,6 +1127,44 @@ def broadcast(self, src_chunk: Chunk, buffer_offset, size, tb):
             """
             return self._channel.broadcast(self._rank, src_chunk, buffer_offset, size, tb)
 
+        def signal(self, tb_list, relaxed=False):
+            """Signal all ranks in the group from this rank's perspective.
+
+            Convenience method that calls the underlying channel's signal method
+            with this view's rank automatically provided.
+
+            Args:
+                tb_list (List[int]): Thread block IDs that participate in the signal.
+                    Must include thread block 0, which acts as the arrival leader.
+                relaxed (bool, optional): Whether to use relaxed signaling. Defaults to False.
+
+            Returns:
+                The result of the underlying channel's signal operation.
+
+            Example:
+                >>> rank_view.signal(tb_list=[0])
+            """
+            return self._channel.signal(self._rank, tb_list, relaxed)
+
+        def wait(self, tb_list, relaxed=False):
+            """Wait for a signal from all ranks in the group from this rank's perspective.
+
+            Convenience method that calls the underlying channel's wait method
+            with this view's rank automatically provided.
+
+            Args:
+                tb_list (List[int]): Thread block IDs that participate in the wait.
+                    Must include thread block 0, which acts as the arrival leader.
+                relaxed (bool, optional): Whether to use relaxed waiting. Defaults to False.
+
+            Returns:
+                The result of the underlying channel's wait operation.
+
+            Example:
+                >>> rank_view.wait(tb_list=[0])
+            """
+            return self._channel.wait(self._rank, tb_list, relaxed)
+
         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,14 @@ 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_signal
+                or operation.name == Instruction.group_wait
+                or operation.name == Instruction.group_relaxed_signal
+                or operation.name == Instruction.group_relaxed_wait
+            ):
                 self.clear_data_access()
             else:
                 if operation.name == Instruction.pipeline: