Add accumulate method for PortChannel

docs/tutorials/04-port-channel.md

@@ -60,12 +60,15 @@ We need to call `proxyService.startProxy()` before running GPU kernels that use
 
 The device handle of a `PortChannel` provides the following methods. Since the data transfer is offloaded, each method is supposed to be called by a single GPU thread.
 - `put()`: Initiates an asynchronous one-way data transfer from the local memory to the remote memory.
-- `signal()`: Asynchronously signals the completion of all previous `put()`s to the remote side.
+- `accumulate()`: Initiates a 64-bit atomic add on a value in remote memory.
+- `signal()`: Asynchronously signals the completion of previous operations to the remote side.
 - `wait()`: Blocks the calling GPU thread until the corresponding `signal()` is received from the remote side.
 - `poll()`: Non-blocking version of `wait()`. Returns immediately with a boolean indicating whether the signal has been received.
 - `flush()`: Synchronizes the local GPU with the `PortChannel`, ensuring that all previous operations are completed.
 - Fused methods (e.g., `putWithSignal()`): combines multiple sequential operations into a single call for efficiency.
 
+Like `put()`, `accumulate()` is issued asynchronously through the proxy and is ordered with operations issued after it on the same channel. A typical sequence enqueues one or more additions and then uses the application's existing `signal()`/`flush()`/`wait()` protocol before consuming the result. See the `PortChannel::accumulate()` and `Connection::accumulate()` API documentation for transport restrictions, alignment requirements, and concurrency details.
+
 The following diagram illustrates how the `bidirPutKernel()` function in the example code would work when GPU0 is faster than GPU1. The execution order may vary depending on the relative speeds of the GPUs. 
 
 ```{mermaid}

include/mscclpp/atomic_device.hpp

@@ -46,8 +46,13 @@ constexpr auto memoryOrderRelease = __ATOMIC_RELEASE;
 constexpr auto memoryOrderAcqRel = __ATOMIC_ACQ_REL;
 constexpr auto memoryOrderSeqCst = __ATOMIC_SEQ_CST;
 
+#if defined(MSCCLPP_DEVICE_HIP)
+constexpr auto scopeSystem = __HIP_MEMORY_SCOPE_SYSTEM;
+constexpr auto scopeDevice = __HIP_MEMORY_SCOPE_AGENT;
+#else
 constexpr auto scopeSystem = 0;
 constexpr auto scopeDevice = 0;
+#endif  // defined(MSCCLPP_DEVICE_HIP)
 
 template <typename T, int scope = scopeSystem>
 MSCCLPP_HOST_DEVICE_INLINE T atomicLoad(const T* ptr, int memoryOrder) {
@@ -61,7 +66,11 @@ MSCCLPP_HOST_DEVICE_INLINE void atomicStore(T* ptr, const T& val, int memoryOrde
 
 template <typename T, int scope = scopeSystem>
 MSCCLPP_HOST_DEVICE_INLINE T atomicFetchAdd(T* ptr, const T& val, int memoryOrder) {
+#if defined(__HIP_DEVICE_COMPILE__)
+  return __hip_atomic_fetch_add(ptr, val, memoryOrder, scope);
+#else
   return __atomic_fetch_add(ptr, val, memoryOrder);
+#endif  // defined(__HIP_DEVICE_COMPILE__)
 }
 
 #endif  // !defined(MSCCLPP_DEVICE_CUDA)

include/mscclpp/core.hpp

@@ -649,6 +649,36 @@ class Connection {
   /// @param newValue The new value to write.
   void updateAndSync(RegisteredMemory dst, uint64_t dstOffset, uint64_t* src, uint64_t newValue);
 
+  /// Add a value to a 64-bit integer in a destination RegisteredMemory.
+  ///
+  /// The caller supplies only its own contribution, unlike updateAndSync(), which needs the
+  /// destination's current value. Addition is modulo 2^64 (the signed operand contributes its
+  /// two's-complement bit pattern) and commutes, so arrival order does not matter.
+  ///
+  /// The addition must be a real read-modify-write at the destination, so how many concurrent
+  /// writers one address allows depends on the transport:
+  ///
+  /// - IB: any number of writers, via NIC atomic fetch-and-add. Throws in no-atomic mode, where
+  ///   the device has no RDMA atomics.
+  /// - Ethernet: any number of remote writers. The receiving process does the update and
+  ///   serializes its connections. The destination GPU must not write the address concurrently;
+  ///   such a write is lost inside the read-modify-write window.
+  /// - CudaIpc on ROCm: any number of writers. The proxy runs a kernel, which a caller kernel
+  ///   does not block.
+  /// - CudaIpc on CUDA: throws. The host cannot read-modify-write device memory, and a
+  ///   proxy-launched kernel cannot run while the caller's kernel waits. Use a device-side atomic
+  ///   on peer memory reached through a MemoryChannel.
+  ///
+  /// The 8-byte target word must fit within @p dst, and its final address must be naturally
+  /// 8-byte aligned.
+  ///
+  /// @param dst The destination RegisteredMemory.
+  /// @param dstOffset The offset in bytes from the start of the destination RegisteredMemory.
+  /// @param value The 64-bit signed value to add.
+  /// @throws Error with ErrorCode::InvalidUsage if the target is out of bounds or misaligned, or
+  ///   if the transport cannot accumulate.
+  void accumulate(RegisteredMemory dst, uint64_t dstOffset, int64_t value);
+
   /// Flush any pending writes to the remote process.
   /// @param timeoutUsec Timeout in microseconds. Default: -1 (no timeout)
   void flush(int64_t timeoutUsec = -1);

include/mscclpp/env.hpp

@@ -84,9 +84,10 @@ class Env {
   /// enabled and will dump traces to this directory. Unset by default.
   const std::string npkitDumpDir;
 
-  /// Env name: `MSCCLPP_CUDAIPC_USE_DEFAULT_STREAM`. If set to true, the CUDA IPC transport will use the default
-  /// stream for all operations. If set to false, it will use a separate stream for each operation. This is an
-  /// experimental feature and should be false in most cases. Default is false.
+  /// Env name: `MSCCLPP_CUDAIPC_USE_DEFAULT_STREAM`. On CUDA, if set to true, the CUDA IPC transport will use the
+  /// default stream for all operations. If set to false, it will use a separate stream for each operation. ROCm
+  /// always uses a nonblocking stream so proxy-launched kernels can make progress. This is an experimental feature
+  /// and should be false in most cases. Default is false.
   const bool cudaIpcUseDefaultStream;
 
   /// Env name: `MSCCLPP_NCCL_LIB_PATH`. The path to the original NCCL/RCCL shared library. If set, it will be used

include/mscclpp/fifo_device.hpp

@@ -15,10 +15,27 @@
 
 namespace mscclpp {
 
+/// Operation that a trigger asks the proxy to perform.
+///
+/// These are opcodes, not flags: compare one by equality, and never combine two. The encoding
+/// enumerates the combinations the device API can produce rather than composing them, so a
+/// combination nothing emits cannot be expressed, and a trigger whose type field is unset is not
+/// a valid operation.
 using TriggerType = uint64_t;
-constexpr TriggerType TriggerData = 0x1;  // Trigger a data transfer.
-constexpr TriggerType TriggerFlag = 0x2;  // Trigger a signaling.
-constexpr TriggerType TriggerSync = 0x4;  // Trigger a flush.
+constexpr TriggerType TriggerNone = 0;                   // Not an operation; invalid for ProxyService.
+constexpr TriggerType TriggerPut = 1;                    // Transfer data.
+constexpr TriggerType TriggerSignal = 2;                 // Signal the remote semaphore.
+constexpr TriggerType TriggerPutWithSignal = 3;          // Transfer data, then signal.
+constexpr TriggerType TriggerFlush = 4;                  // Flush the connection.
+constexpr TriggerType TriggerAccumulate = 6;             // Add a value to remote memory.
+constexpr TriggerType TriggerPutWithSignalAndFlush = 7;  // Transfer data, signal, then flush.
+// 5 is unassigned.
+
+// Preserve the original flag-combination encodings for existing operations. Triggers are now
+// compared as opcodes, but changing these values would break producers and consumers built from
+// different revisions.
+static_assert(TriggerPut == 1 && TriggerSignal == 2 && TriggerPutWithSignal == 3 && TriggerFlush == 4 &&
+              TriggerPutWithSignalAndFlush == 7);
 
 constexpr unsigned int TriggerBitsSize = 32;
 constexpr unsigned int TriggerBitsOffset = 32;
@@ -29,6 +46,8 @@ constexpr unsigned int TriggerBitsSemaphoreId = 10;
 // there. See FifoDeviceHandle::push().
 constexpr unsigned int TriggerBitsFifoReserved = 1;
 
+static_assert(TriggerAccumulate < (1ULL << TriggerBitsType), "trigger opcodes must fit in the type field");
+
 /// Pair of 64-bit unsigned integers used as a trigger for the proxy.
 /// Used as a work element in the concurrent FIFO.
 /// Most significant bit of snd is reserved.

include/mscclpp/port_channel.hpp

@@ -84,7 +84,7 @@ class ProxyService : public BaseProxyService {
   std::vector<RegisteredMemory> memories_;
   std::shared_ptr<Proxy> proxy_;
   std::unordered_map<std::shared_ptr<BaseConnection>, int> inflightRequests_;
-  // Latest pending TriggerSync FIFO position per connection. Proxy publishes pos+1 to the
+  // Latest pending TriggerFlush FIFO position per connection. Proxy publishes pos+1 to the
   // connection's gpuFlushDonePos_ when the CQ drains, then erases the entry.
   std::unordered_map<std::shared_ptr<BaseConnection>, uint64_t> pendingFlushPos_;
 

include/mscclpp/port_channel_device.hpp

@@ -19,7 +19,7 @@ using MemoryId = uint32_t;
 
 namespace detail {
 #if defined(MSCCLPP_DEVICE_COMPILE)
-/// Wait until the proxy has processed and drained the TriggerSync at FIFO position `fifoPos`.
+/// Wait until the proxy has processed and drained the TriggerFlush at FIFO position `fifoPos`.
 /// The proxy publishes `flushDonePos = latestCompletedPos + 1` when the CQ drains, so the
 /// wait condition `flushDonePos > fifoPos` is satisfied exactly when our own request has
 /// been completed. Using the FIFO push position as the wait target couples the wait to the
@@ -51,18 +51,18 @@ struct BasePortChannelDeviceHandle {
       : semaphoreId_(semaphoreId), semaphore_(semaphore), fifo_(fifo), flushDonePos_(flushDonePos) {}
 
 #if defined(MSCCLPP_DEVICE_COMPILE)
-  /// Push a TriggerData to the FIFO.
+  /// Push a TriggerPut to the FIFO.
   /// @param dstId The ID of destination memory region.
   /// @param dstOffset The offset into the destination memory region.
   /// @param srcId The ID of source memory region.
   /// @param srcOffset The offset into the source memory region.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void put(MemoryId dstId, uint64_t dstOffset, MemoryId srcId, uint64_t srcOffset,
                                  uint64_t size) {
-    fifo_.push({TriggerData, dstId, dstOffset, srcId, srcOffset, size, semaphoreId_});
+    fifo_.push({TriggerPut, dstId, dstOffset, srcId, srcOffset, size, semaphoreId_});
   }
 
-  /// Push a TriggerData to the FIFO.
+  /// Push a TriggerPut to the FIFO.
   /// @param dstId The ID of destination memory region.
   /// @param srcId The ID of source memory region.
   /// @param offset The common offset into the destination and source memory regions.
@@ -71,21 +71,21 @@ struct BasePortChannelDeviceHandle {
     put(dstId, offset, srcId, offset, size);
   }
 
-  /// Push a TriggerFlag to the FIFO.
-  MSCCLPP_DEVICE_INLINE void signal() { fifo_.push({TriggerFlag, 0, 0, 0, 0, 0, semaphoreId_}); }
+  /// Push a TriggerSignal to the FIFO.
+  MSCCLPP_DEVICE_INLINE void signal() { fifo_.push({TriggerSignal, 0, 0, 0, 0, 0, semaphoreId_}); }
 
-  /// Push a TriggerData and a TriggerFlag at the same time to the FIFO.
+  /// Push a TriggerPutWithSignal to the FIFO.
   /// @param dstId The ID of destination memory region.
   /// @param dstOffset The offset into the destination memory region.
   /// @param srcId The ID of source memory region.
   /// @param srcOffset The offset into the source memory region.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void putWithSignal(MemoryId dstId, uint64_t dstOffset, MemoryId srcId, uint64_t srcOffset,
                                            uint64_t size) {
-    fifo_.push({TriggerData | TriggerFlag, dstId, dstOffset, srcId, srcOffset, size, semaphoreId_});
+    fifo_.push({TriggerPutWithSignal, dstId, dstOffset, srcId, srcOffset, size, semaphoreId_});
   }
 
-  /// Push a TriggerData and a TriggerFlag at the same time to the FIFO.
+  /// Push a TriggerPutWithSignal to the FIFO.
   /// @param dstId The ID of destination memory region.
   /// @param srcId The ID of source memory region.
   /// @param offset The common offset into the destination and source memory regions.
@@ -94,7 +94,7 @@ struct BasePortChannelDeviceHandle {
     putWithSignal(dstId, offset, srcId, offset, size);
   }
 
-  /// Push a TriggerData, a TriggerFlag, and a TriggerSync at the same time to the FIFO.
+  /// Push a TriggerPutWithSignalAndFlush to the FIFO.
   /// @param dstId The ID of destination memory region.
   /// @param dstOffset The offset into the destination memory region.
   /// @param srcId The ID of source memory region.
@@ -103,12 +103,11 @@ struct BasePortChannelDeviceHandle {
   /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
   MSCCLPP_DEVICE_INLINE void putWithSignalAndFlush(MemoryId dstId, uint64_t dstOffset, MemoryId srcId,
                                                    uint64_t srcOffset, uint64_t size, int64_t maxSpinCount = 1000000) {
-    uint64_t pos =
-        fifo_.push({TriggerData | TriggerFlag | TriggerSync, dstId, dstOffset, srcId, srcOffset, size, semaphoreId_});
+    uint64_t pos = fifo_.push({TriggerPutWithSignalAndFlush, dstId, dstOffset, srcId, srcOffset, size, semaphoreId_});
     detail::waitFlush(flushDonePos_, pos, maxSpinCount);
   }
 
-  /// Push a TriggerData, a TriggerFlag, and a TriggerSync at the same time to the FIFO.
+  /// Push a TriggerPutWithSignalAndFlush to the FIFO.
   /// @param dstId The ID of destination memory region.
   /// @param srcId The ID of source memory region.
   /// @param offset The common offset into the destination and source memory regions.
@@ -119,13 +118,26 @@ struct BasePortChannelDeviceHandle {
     putWithSignalAndFlush(dstId, offset, srcId, offset, size, maxSpinCount);
   }
 
-  /// Push a TriggerSync to the FIFO.
+  /// Push a TriggerFlush to the FIFO.
   /// @param maxSpinCount The maximum number of spin counts before asserting. Never assert if negative.
   MSCCLPP_DEVICE_INLINE void flush(int64_t maxSpinCount = 1000000) {
-    uint64_t pos = fifo_.push({TriggerSync, 0, 0, 0, 0, 0, semaphoreId_});
+    uint64_t pos = fifo_.push({TriggerFlush, 0, 0, 0, 0, 0, semaphoreId_});
     detail::waitFlush(flushDonePos_, pos, maxSpinCount);
   }
 
+  /// Push an accumulate trigger to the FIFO: add a 64-bit value to remote memory.
+  /// Connection::accumulate() documents how many concurrent writers each transport allows.
+  /// @param dstId The ID of destination memory region.
+  /// @param dstOffset The offset into the destination memory region.
+  /// @param value The 64-bit signed value to add.
+  MSCCLPP_DEVICE_INLINE void accumulate(MemoryId dstId, uint64_t dstOffset, int64_t value) {
+    // The operand occupies fst, spanning the low size and high srcOffset fields.
+    uint64_t operand = static_cast<uint64_t>(value);
+    ProxyTrigger trigger(TriggerAccumulate, dstId, dstOffset, /*srcId=*/0, operand >> TriggerBitsSize,
+                         static_cast<uint32_t>(operand), semaphoreId_);
+    fifo_.push(trigger);
+  }
+
   /// Check if the port channel has been signaled.
   /// @return true if the port channel has been signaled.
   MSCCLPP_DEVICE_INLINE bool poll() { return semaphore_.poll(); }
@@ -149,33 +161,33 @@ struct PortChannelDeviceHandle : public BasePortChannelDeviceHandle {
       : BasePortChannelDeviceHandle(semaphoreId, semaphore, fifo, flushDonePos), dst_(dst), src_(src) {}
 
 #if defined(MSCCLPP_DEVICE_COMPILE)
-  /// Push a TriggerData to the FIFO.
+  /// Push a TriggerPut to the FIFO.
   /// @param dstOffset The offset into the destination memory region.
   /// @param srcOffset The offset into the source memory region.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void put(uint64_t dstOffset, uint64_t srcOffset, uint64_t size) {
     BasePortChannelDeviceHandle::put(dst_, dstOffset, src_, srcOffset, size);
   }
 
-  /// Push a TriggerData to the FIFO.
+  /// Push a TriggerPut to the FIFO.
   /// @param offset The common offset into the destination and source memory regions.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void put(uint64_t offset, uint64_t size) { put(offset, offset, size); }
 
-  /// Push a TriggerData and a TriggerFlag at the same time to the FIFO.
+  /// Push a TriggerPutWithSignal to the FIFO.
   /// @param dstOffset The offset into the destination memory region.
   /// @param srcOffset The offset into the source memory region.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void putWithSignal(uint64_t dstOffset, uint64_t srcOffset, uint64_t size) {
     BasePortChannelDeviceHandle::putWithSignal(dst_, dstOffset, src_, srcOffset, size);
   }
 
-  /// Push a TriggerData and a TriggerFlag at the same time to the FIFO.
+  /// Push a TriggerPutWithSignal to the FIFO.
   /// @param offset The common offset into the destination and source memory regions.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void putWithSignal(uint64_t offset, uint64_t size) { putWithSignal(offset, offset, size); }
 
-  /// Push a TriggerData, a TriggerFlag, and a TriggerSync at the same time to the FIFO.
+  /// Push a TriggerPutWithSignalAndFlush to the FIFO.
   /// @param dstOffset The offset into the destination memory region.
   /// @param srcOffset The offset into the source memory region.
   /// @param size The size of the transfer.
@@ -185,12 +197,19 @@ struct PortChannelDeviceHandle : public BasePortChannelDeviceHandle {
     BasePortChannelDeviceHandle::putWithSignalAndFlush(dst_, dstOffset, src_, srcOffset, size, maxSpinCount);
   }
 
-  /// Push a TriggerData, a TriggerFlag, and a TriggerSync at the same time to the FIFO.
+  /// Push a TriggerPutWithSignalAndFlush to the FIFO.
   /// @param offset The common offset into the destination and source memory regions.
   /// @param size The size of the transfer.
   MSCCLPP_DEVICE_INLINE void putWithSignalAndFlush(uint64_t offset, uint64_t size) {
     putWithSignalAndFlush(offset, offset, size);
   }
+  /// Push an accumulate trigger to the FIFO: add a 64-bit value to the destination memory.
+  /// See Connection::accumulate() for transport support.
+  /// @param dstOffset The offset into the destination memory region.
+  /// @param value The 64-bit signed value to add.
+  MSCCLPP_DEVICE_INLINE void accumulate(uint64_t dstOffset, int64_t value) {
+    BasePortChannelDeviceHandle::accumulate(dst_, dstOffset, value);
+  }
 #endif  // defined(MSCCLPP_DEVICE_COMPILE)
 };
 

src/core/accumulate_kernel.cu

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#include <mscclpp/gpu.hpp>
+
+#if defined(MSCCLPP_USE_ROCM)
+
+#include <mscclpp/atomic_device.hpp>
+#include <mscclpp/gpu_utils.hpp>
+
+#include "context.hpp"
+
+namespace mscclpp {
+
+// System-scope modulo-2^64 atomic add.
+__global__ void accumulateU64Kernel(uint64_t* dst, uint64_t value) {
+  (void)atomicFetchAdd<uint64_t, scopeSystem>(dst, value, memoryOrderRelaxed);
+}
+
+void CudaIpcStream::accumulate(uint64_t* dst, uint64_t value) {
+  CudaDeviceGuard deviceGuard(deviceId_);
+  setStreamIfNeeded();
+  // Submit to this connection's stream, which orders the add ahead of any signal or flush that
+  // follows. On ROCm a kernel runs while the caller's kernel occupies the GPU, so the proxy does
+  // not wait for the caller.
+  accumulateU64Kernel<<<1, 1, 0, *stream_>>>(dst, value);
+  MSCCLPP_CUDATHROW(cudaGetLastError());
+  dirty_ = true;
+}
+
+}  // namespace mscclpp
+
+#endif  // defined(MSCCLPP_USE_ROCM)