Add accumulate method for PortChannel

include/mscclpp/core.hpp

@@ -649,6 +649,31 @@ 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 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.
+  ///
+  /// @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 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/fifo_device.hpp

@@ -15,10 +15,21 @@
 
 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 TriggerFlush = 3;                  // Flush the connection.
+constexpr TriggerType TriggerPutWithSignal = 4;          // Transfer data, then signal.
+constexpr TriggerType TriggerPutWithSignalAndFlush = 5;  // Transfer data, signal, then flush.
+constexpr TriggerType TriggerAccumulate = 6;             // Add a value to remote memory.
+// 7 is unassigned.
 
 constexpr unsigned int TriggerBitsSize = 32;
 constexpr unsigned int TriggerBitsOffset = 32;
@@ -29,6 +40,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,31 @@ 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) {
+    ProxyTrigger trigger;
+    // The operand occupies fst, spanning the size and srcOffset fields.
+    trigger.fst = static_cast<uint64_t>(value);
+    // snd carries dstOffset, dstMemoryId, the opcode, and semaphoreId.
+    trigger.snd = 0;
+    trigger.fields.dstOffset = dstOffset;
+    trigger.fields.dstMemoryId = dstId;
+    trigger.fields.type = TriggerAccumulate;
+    trigger.fields.semaphoreId = 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 +166,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 +202,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 atomic add on a signed 64-bit value.
+__global__ void accumulateI64Kernel(int64_t* dst, int64_t value) {
+  (void)atomicFetchAdd<int64_t, scopeSystem>(dst, value, memoryOrderRelaxed);
+}
+
+void CudaIpcStream::accumulate(int64_t* dst, int64_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.
+  accumulateI64Kernel<<<1, 1, 0, *stream_>>>(dst, value);
+  MSCCLPP_CUDATHROW(cudaGetLastError());
+  dirty_ = true;
+}
+
+}  // namespace mscclpp
+
+#endif  // defined(MSCCLPP_USE_ROCM)