Change executor to load new execution plan

include/mscclpp/copy_device.hpp

@@ -96,6 +96,16 @@ MSCCLPP_DEVICE_INLINE void copy(void* dst, void* src, uint64_t bytes, uint32_t t
   }
 }
 
+template <typename T>
+MSCCLPP_DEVICE_INLINE void write(void* dst, uint64_t index, const T& v) {
+  *(reinterpret_cast<T*>(dst) + index) = v;
+}
+
+template <typename T>
+MSCCLPP_DEVICE_INLINE T read(void* src, uint64_t index) {
+  return *(reinterpret_cast<T*>(src) + index);
+}
+
 /// Read data from the origin and write packets to the target buffer.
 ///
 /// @param targetPtr The target buffer.

src/executor/executor.cc

@@ -113,11 +113,17 @@ struct ExecutionContext {
   std::shared_ptr<ProxyService> proxyService;
   std::unordered_map<int, std::shared_ptr<Connection>> connections;
   std::vector<std::shared_ptr<NvlsConnection>> nvlsConnections;
-  std::unordered_map<std::pair<BufferType, int>, mscclpp::RegisteredMemory> registeredMemories;
+
+  // For registered memories, registeredMemoryAddresses is used for memoryChannel and registeredMemoryIds is used for
+  // proxy channel
+  std::vector<mscclpp::RegisteredMemory> registeredMemories;
+  std::vector<void*> registeredMemoryAddresses;
+  std::vector<mscclpp::MemoryId> registeredMemoryIds;
+
   std::vector<std::shared_ptr<mscclpp::MemoryDevice2DeviceSemaphore>> memorySemaphores;
   std::vector<mscclpp::SemaphoreId> proxySemaphores;
-  std::vector<mscclpp::MemoryChannel> memoryChannels;
-  std::vector<mscclpp::PortChannel> portChannels;
+  std::vector<mscclpp::BaseMemoryChannel> memoryChannels;
+  std::vector<mscclpp::BasePortChannel> portChannels;
   std::vector<mscclpp::NvlsConnection::DeviceMulticastPointer> nvlsChannels;
   std::unordered_map<DeviceExecutionPlanKey, std::vector<DeviceExecutionPlan>> deviceExecutionPlans;
   std::unordered_map<DeviceExecutionPlanKey, std::shared_ptr<char>> deviceExecutionPlansBuffers;
@@ -165,7 +171,7 @@ struct Executor::Impl {
     }
 
     plan.impl_->reset();
-    plan.impl_->loadExecutionPlan(inputMessageSize, outputMessageSize, constSrcOffset, constDstOffset);
+    plan.impl_->loadExecutionPlan(rank, inputMessageSize, outputMessageSize, constSrcOffset, constDstOffset);
 
     ExecutionContext context;
     size_t maxScratchBufferSize = plan.impl_->getMaxScratchBufferSize(rank);
@@ -177,8 +183,8 @@ struct Executor::Impl {
     context.proxyService = std::make_shared<ProxyService>();
     context.nthreadsPerBlock = plan.impl_->getNThreadsPerBlock();
     this->setupConnections(context, rank, plan, sendMemRange, recvMemRange);
+    this->setupChannels(context, rank, plan);
     this->setupRegisteredMemories(context, sendbuff, recvbuff, sendMemRange, recvMemRange, rank, plan);
-    this->setupChannels(context, sendbuff, recvbuff, sendMemRange, recvMemRange, rank, plan);
     this->setupNvlsChannels(context, sendbuff, recvbuff, sendMemRange, recvMemRange, rank, plan);
     this->setupDeviceExecutionPlan(context, devicePlanKey, rank, plan);
     context.deviceExecutionPlansBuffers[devicePlanKey] =
@@ -192,18 +198,16 @@ struct Executor::Impl {
     return context;
   }
 
-  TransportFlags getTransportFlags(std::vector<ChannelInfo>& infos, int rank) {
+  TransportFlags getTransportFlags(const BufferInfo& info, int rank) {
     TransportFlags flags;
-    for (ChannelInfo& info : infos) {
-      if (info.channelType == ChannelType::MEMORY) {
+    for (const ChannelType& type : info.accessChannelTypes) {
+      if (type == ChannelType::MEMORY) {
         flags |= Transport::CudaIpc;
-      } else if (info.channelType == ChannelType::PORT) {
-        for (int peer : info.connectedPeers) {
-          if (!inSameNode(rank, peer, this->nranksPerNode)) {
-            flags |= IBs[rank % this->nranksPerNode];
-          } else
-            flags |= Transport::CudaIpc;
-        }
+      } else if (type == ChannelType::PORT) {
+        if (!inSameNode(rank, info.accessRank, this->nranksPerNode)) {
+          flags |= IBs[rank % this->nranksPerNode];
+        } else
+          flags |= Transport::CudaIpc;
       }
     }
     return flags;
@@ -244,40 +248,42 @@ struct Executor::Impl {
           throw Error("Invalid buffer type", ErrorCode::ExecutorError);
       }
     };
-    auto getConnectedPeers = [&](std::vector<ChannelInfo>& infos) {
-      std::set<int> peers;
-      for (ChannelInfo& info : infos) {
-        for (int peer : info.connectedPeers) {
-          peers.insert(peer);
-        }
+
+    // Add local src,dst and scratch to registeredMemoryIds
+    for (auto& bufferType : {BufferType::INPUT, BufferType::OUTPUT, BufferType::SCRATCH}) {
+      TransportFlags flags = Transport::CudaIpc;
+#if defined(USE_IBVERBS)
+      flags |= IBs[rank % this->nranksPerNode];
+#endif
+      RegisteredMemory localMemory;
+      auto bufferInfo = getBufferInfo(bufferType);
+      if (bufferInfo.second > 0) {
+        localMemory =
+            this->comm->registerMemory(getBufferInfo(bufferType).first, getBufferInfo(bufferType).second, flags);
       }
-      return std::vector<int>(peers.begin(), peers.end());
-    };
+      context.proxyService->addMemory(localMemory);
+    }
 
-    std::vector<BufferType> bufferTypes = plan.impl_->getConnectedBufferTypes(rank);
-    for (BufferType bufferType : bufferTypes) {
-      std::vector<ChannelInfo> channelInfos = plan.impl_->getChannelInfosByDstRank(rank, bufferType);
-      TransportFlags transportFlags = getTransportFlags(channelInfos, rank);
+    for (const auto& bufferInfo : plan.impl_->getLocalBufferToSend(rank)) {
       RegisteredMemory memory =
-          this->comm->registerMemory(getBufferInfo(bufferType).first, getBufferInfo(bufferType).second, transportFlags);
-      std::vector<int> connectedPeers = getConnectedPeers(channelInfos);
-      std::vector<std::shared_future<mscclpp::RegisteredMemory>> remoteRegMemoryFutures;
-      for (int peer : connectedPeers) {
-        comm->sendMemory(memory, peer, 0);
-      }
-      channelInfos = plan.impl_->getChannelInfos(rank, bufferType);
-      connectedPeers = getConnectedPeers(channelInfos);
-      for (int peer : connectedPeers) {
-        remoteRegMemoryFutures.push_back(comm->recvMemory(peer, 0));
-      }
-      for (size_t i = 0; i < remoteRegMemoryFutures.size(); i++) {
-        context.registeredMemories[{bufferType, connectedPeers[i]}] = std::move(remoteRegMemoryFutures[i].get());
+          this->comm->registerMemory(getBufferInfo(bufferInfo.bufferType).first,
+                                     getBufferInfo(bufferInfo.bufferType).second, getTransportFlags(bufferInfo, rank));
+      comm->sendMemory(memory, bufferInfo.accessRank, 0);
+    }
+    for (const auto& bufferInfo : plan.impl_->getRemoteBufferInfos(rank)) {
+      std::shared_future<RegisteredMemory> remoteRegMemoryFuture = comm->recvMemory(bufferInfo.rank, 0);
+      context.registeredMemories.emplace_back(std::move(remoteRegMemoryFuture.get()));
+      for (ChannelType chanType : bufferInfo.accessChannelTypes) {
+        if (chanType == ChannelType::MEMORY) {
+          context.registeredMemoryAddresses.push_back(context.registeredMemories.back().data());
+        } else if (chanType == ChannelType::PORT) {
+          context.registeredMemoryIds.push_back(context.proxyService->addMemory(context.registeredMemories.back()));
+        }
       }
     }
   }
 
-  void setupChannels(ExecutionContext& context, void* sendbuff, void* recvbuff, size_t sendBufferSize,
-                     size_t recvBufferSize, int rank, const ExecutionPlan& plan) {
+  void setupChannels(ExecutionContext& context, int rank, const ExecutionPlan& plan) {
     const auto channelTypes = {ChannelType::MEMORY, ChannelType::PORT};
     std::vector<std::shared_ptr<MemoryDevice2DeviceSemaphore>> memorySemaphores;
     std::vector<mscclpp::SemaphoreId> proxySemaphores;
@@ -300,44 +306,22 @@ struct Executor::Impl {
       // Current semaphore construction requires two-way communication, e.g., to construct a semaphore signaling from
       // rank 0 to rank 1, both rank 0 and rank 1 need to send a message to each other. This PR fixes an executor bug
       // that fails to conduct two-way communication for constructing such one-way semaphores, and instead hangs
-      // during the semaphore construction. In the future, we may need to change the implementation to construct
-      // semaphore via one-way communication.
+      // during the semaphore construction.
       channelInfos = plan.impl_->getUnpairedChannelInfos(rank, nranks, channelType);
       processChannelInfos(channelInfos);
     }
     context.memorySemaphores = std::move(memorySemaphores);
     context.proxySemaphores = std::move(proxySemaphores);
 
-    auto getBufferSize = [&](BufferType type) {
-      switch (type) {
-        case BufferType::INPUT:
-          return sendBufferSize;
-        case BufferType::OUTPUT:
-          return recvBufferSize;
-        case BufferType::SCRATCH:
-          return context.scratchBufferSize;
-        default:
-          throw Error("Invalid buffer type", ErrorCode::ExecutorError);
-      }
-    };
-
     for (ChannelType channelType : channelTypes) {
       std::vector<ChannelInfo> channelInfos = plan.impl_->getChannelInfos(rank, channelType);
       int index = 0;
       for (ChannelInfo& info : channelInfos) {
-        void* src = getBuffer(info.srcBufferType, sendbuff, recvbuff, context.scratchBuffer.get());
-        size_t bufferSize = getBufferSize(info.srcBufferType);
-        TransportFlags transport = getTransportFlags(channelInfos, rank);
-        RegisteredMemory localMemory = this->comm->registerMemory(src, bufferSize, transport);
-        for (int peer : info.connectedPeers) {
+        for (size_t i = 0; i < info.connectedPeers.size(); i++) {
           if (channelType == ChannelType::MEMORY) {
-            context.memoryChannels.emplace_back(context.memorySemaphores[index++],
-                                                context.registeredMemories[{info.dstBufferType, peer}], src, nullptr);
+            context.memoryChannels.emplace_back(context.memorySemaphores[index++]);
           } else if (channelType == ChannelType::PORT) {
-            context.portChannels.emplace_back(context.proxyService->portChannel(
-                context.proxySemaphores[index++],
-                context.proxyService->addMemory(context.registeredMemories[{info.dstBufferType, peer}]),
-                context.proxyService->addMemory(localMemory)));
+            context.portChannels.emplace_back(context.proxyService->basePortChannel(context.proxySemaphores[index++]));
           }
         }
       }
@@ -367,17 +351,27 @@ struct Executor::Impl {
       deviceExecutionPlan.nMemoryChannels = plan.impl_->threadblockMemoryChannelMap.at(rank).at(threadblock).size();
       deviceExecutionPlan.nPortChannels = plan.impl_->threadblockPortChannelMap.at(rank).at(threadblock).size();
       int chanIndex = 0;
-      for (const auto& [index, _] : plan.impl_->threadblockMemoryChannelMap.at(rank).at(threadblock)) {
+      for (const int index : plan.impl_->threadblockMemoryChannelMap.at(rank).at(threadblock)) {
         deviceExecutionPlan.channels.memoryChannels[chanIndex++] = mscclpp::deviceHandle(context.memoryChannels[index]);
       }
       chanIndex = 0;
-      for (const auto& [index, _] : plan.impl_->threadblockPortChannelMap.at(rank).at(threadblock)) {
+      for (const int index : plan.impl_->threadblockPortChannelMap.at(rank).at(threadblock)) {
         deviceExecutionPlan.channels.portChannels[chanIndex++] = mscclpp::deviceHandle(context.portChannels[index]);
       }
       chanIndex = 0;
-      for (const auto& [index, _] : plan.impl_->threadblockNvlsChannelMap.at(rank).at(threadblock)) {
+      for (const int index : plan.impl_->threadblockNvlsChannelMap.at(rank).at(threadblock)) {
         deviceExecutionPlan.channels.nvlsChannels[chanIndex++] = mscclpp::deviceHandle(context.nvlsChannels[index]);
       }
+      int memIndex = 0;
+      for (const int index : plan.impl_->threadblockMemoryChannelBufferMap.at(rank).at(threadblock)) {
+        deviceExecutionPlan.remoteBuffers.remoteBuffersViaMemoryChannel[memIndex++] =
+            context.registeredMemoryAddresses[index];
+      }
+      memIndex = 0;
+      for (const int index : plan.impl_->threadblockPortChannelBufferMap.at(rank).at(threadblock)) {
+        deviceExecutionPlan.remoteBuffers.remoteBuffersViaPortChannel[memIndex++] = context.registeredMemoryIds[index];
+      }
+
       if (ops.size() > MAX_OPERATION) {
         throw Error("Executor plan launching " + std::to_string(ops.size()) +
                         " operations, exceeding device execution plan support (" + std::to_string(MAX_OPERATION) + ")",

src/include/execution_common.hpp

@@ -13,6 +13,10 @@ namespace mscclpp {
 constexpr int MAX_CHANNEL = 16;
 constexpr int MAX_CHANNEL_PER_OPERATION = 8;
 constexpr int MAX_OPERATION = 64;
+constexpr int INPUT_BUFFER_MEMORY_ID = 0;
+constexpr int OUTPUT_BUFFER_MEMORY_ID = 1;
+constexpr int SCRATCH_BUFFER_MEMORY_ID = 2;
+constexpr int MAX_RESERVED_MEMORY_IDS = 3;
 
 enum class BufferType : uint8_t {
   NONE,
@@ -51,47 +55,60 @@ enum class OperationType : uint8_t {
   READ_REDUCE_COPY,
   READ_REDUCE_COPY_SEND,
   MULTI_LOAD_REDUCE_STORE,
+  RELAXED_SIGNAL,
+  RELAXED_WAIT,
+  PIPELINE,
 };
 
 struct Channels {
-  mscclpp::DeviceHandle<mscclpp::MemoryChannel> memoryChannels[MAX_CHANNEL];
-  mscclpp::DeviceHandle<mscclpp::PortChannel> portChannels[MAX_CHANNEL];
+  mscclpp::DeviceHandle<mscclpp::BaseMemoryChannel> memoryChannels[MAX_CHANNEL];
+  mscclpp::DeviceHandle<mscclpp::BasePortChannel> portChannels[MAX_CHANNEL];
   mscclpp::DeviceHandle<mscclpp::NvlsConnection::DeviceMulticastPointer> nvlsChannels[MAX_CHANNEL];
 };
 
+struct RemoteBuffers {
+  void* remoteBuffersViaMemoryChannel[MAX_CHANNEL];
+  MemoryId remoteBuffersViaPortChannel[MAX_CHANNEL];
+};
+
+union BufferRef {
+  uint8_t id;
+  BufferType type;
+};
+
 struct Operation {
   OperationType type;
   ChannelType channelType;
-  BufferType srcBufferType;
-  BufferType dstBufferType;
-  uint8_t nInputs;
-  uint8_t nOutputs;
   union {
-    // For ops which require reading from multiple remote sources
-    uint8_t inputChannelIndexes[MAX_CHANNEL_PER_OPERATION];
-    // For ops which require reading from multiple local sources
-    BufferType inputBufferType;
+    BufferRef inputBufferRefs[MAX_CHANNEL_PER_OPERATION];
     uint8_t nvlsInputIndex;
   };
   union {
-    // For ops which require writing to multiple remote destinations
-    uint8_t outputChannelIndexes[MAX_CHANNEL_PER_OPERATION];
-    // For ops which require writing to multiple local destinations
-    BufferType outputBufferType;
+    BufferRef outputBufferRefs[MAX_CHANNEL_PER_OPERATION];
     uint8_t nvlsOutputIndex;
   };
+
   union {
-    // For Barrier operation
-    struct {
-      uint32_t deviceSyncerIndex;
-      uint32_t nThreadBlocks;
-    };
     struct {
+      uint8_t channelIndexes[MAX_CHANNEL_PER_OPERATION];
       uint32_t inputOffsets[MAX_CHANNEL_PER_OPERATION];
       uint32_t outputOffsets[MAX_CHANNEL_PER_OPERATION];
-      uint32_t srcOffset;
-      uint32_t dstOffset;
-      uint32_t size;
+      uint32_t inputBufferSizes[MAX_CHANNEL_PER_OPERATION];
+      uint32_t outputBufferSizes[MAX_CHANNEL_PER_OPERATION];
+
+      uint8_t nChannels;
+      uint8_t nInputs;
+      uint8_t nOutputs;
+    };
+    struct {
+      uint32_t unitSize;
+      uint32_t maxBufferSize;
+      uint8_t nIterations;
+      uint8_t nOperations;
+    };
+    struct {
+      uint32_t deviceSyncerIndex;
+      uint32_t nThreadBlocks;
     };
   };
 };
@@ -102,6 +119,7 @@ struct __attribute__((aligned(16))) DeviceExecutionPlan {
   uint8_t nPortChannels;                // 1 bytes
   uint16_t nOperations;                 // 2 bytes
   Channels channels;                    // 2304 bytes
+  RemoteBuffers remoteBuffers;          // 192 bytes
   Operation operations[MAX_OPERATION];  // 64 * 100 = 6400 bytes
 };