WIP

Author: Binyang2014

apps/nccl/src/nccl.cu

@@ -258,28 +258,6 @@ static void registerCustomizedAlgo(ncclComm* commPtr) {
   collectionBuilder->addDefaultNativeAlgorithmBuilder("default_allreduce_packet",
                                                       reinterpret_cast<uintptr_t>(commPtr->scratchBuffer_.get()),
                                                       commPtr->scratchBufferSize_);
-
-  // std::shared_ptr<AllgatherAlgo6> allgatherAlgo6 = std::make_shared<AllgatherAlgo6>();
-  // std::shared_ptr<AllgatherAlgo8> allgatherAlgo8 =
-  //     std::make_shared<AllgatherAlgo8>(commPtr->scratchBuffer_, commPtr->scratchBufferSize_);
-  // collectionBuilder->addAlgorithmBuilder(allgatherAlgo6);
-  // // TODO(binyli): remove allgather8 algo, use nccl by default
-  // collectionBuilder->addAlgorithmBuilder(allgatherAlgo8);
-
-  // std::shared_ptr<AllreducePacket> allreduceAllpairAlgo =
-  //     std::make_shared<AllreducePacket>(commPtr->scratchBuffer_, commPtr->scratchBufferSize_);
-  // std::shared_ptr<AllreduceNvls> allreduceNvlsAlgo = std::make_shared<AllreduceNvls>();
-  // std::shared_ptr<AllreduceNvlsWithCopy> allreduceNvlsWithCopyAlgo =
-  //     std::make_shared<AllreduceNvlsWithCopy>(commPtr->scratchBuffer_, commPtr->scratchBufferSize_);
-  // std::shared_ptr<Allreduce8> allreduceAllreduce8Algo =
-  //     std::make_shared<Allreduce8>(commPtr->scratchBuffer_, commPtr->scratchBufferSize_);
-  // std::shared_ptr<AllreduceNvlsPacket> allreduceNvlsPacketAlgo =
-  //     std::make_shared<AllreduceNvlsPacket>(commPtr->scratchBuffer_, commPtr->scratchBufferSize_);
-  // collectionBuilder->addAlgorithmBuilder(allreduceAllpairAlgo);
-  // collectionBuilder->addAlgorithmBuilder(allreduceNvlsAlgo);
-  // collectionBuilder->addAlgorithmBuilder(allreduceNvlsWithCopyAlgo);
-  // collectionBuilder->addAlgorithmBuilder(allreduceAllreduce8Algo);
-  // collectionBuilder->addAlgorithmBuilder(allreduceNvlsPacketAlgo);
 }
 
 static std::pair<int, int> getDeviceComputeCapability() {
@@ -352,6 +330,9 @@ static std::shared_ptr<mscclpp::Algorithm> algoSelector(
       useNvls = false;
     }
 #endif
+    if (messageSize <= (1 << 15)) {
+      return algoMapByCollective.at(collective).at("default_allreduce_allpair_packet");
+    }
     if (messageSize <= (1 << 15) && useNvls) {
       return algoMapByCollective.at(collective).at("default_allreduce_nvls_packet");
     }

src/algorithms/algorithm.cc

@@ -3,6 +3,7 @@
 
 #include <mscclpp/algorithm.hpp>
 
+#include "algorithms/utils.hpp"
 #include "logger.hpp"
 
 namespace mscclpp {
@@ -111,9 +112,11 @@ void AlgorithmCollectionBuilder::addAlgorithmBuilder(std::shared_ptr<AlgorithmBu
   this->algoBuilders_.push_back(builder);
 }
 
-// TODO (binyli) implement this
 void AlgorithmCollectionBuilder::addDefaultNativeAlgorithmBuilder(std::string algorithmName, uintptr_t scratchBuffer,
-                                                                  size_t scratchBufferSize) {}
+                                                                  size_t scratchBufferSize) {
+  auto builder = algorithm::getDefaultNativeAlgorithmBuilder(algorithmName, scratchBuffer, scratchBufferSize);
+  this->algoBuilders_.push_back(builder);
+}
 
 void AlgorithmCollectionBuilder::setAlgorithmSelector(AlgoSelectFunc selector) { algoSelector_ = selector; }
 

src/algorithms/allreduce/allreduce_allpair_packet.cu

@@ -59,7 +59,9 @@ __global__ void allreduceAllPairs(T* buff, T* scratch, T* resultBuff, DeviceHand
   if (blockIdx.x == 0 && threadIdx.x < gridDim.x) {
     flags[threadIdx.x].read(flag, -1);
   }
-  __syncthreads();
+  if (blockIdx.x == 0) {
+    __syncthreads();
+  }
   if (threadIdx.x == 0 && blockIdx.x == 0) {
     deviceFlag++;
   }

src/algorithms/allreduce/allreduce_packet.cu

@@ -129,7 +129,9 @@ __global__ void __launch_bounds__(1024, 1)
   if (blockIdx.x == 0 && threadIdx.x < gridDim.x) {
     flags[threadIdx.x].read(flag, -1);
   }
-  __syncthreads();
+  if (blockIdx.x == 0) {
+    __syncthreads();
+  }
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_KERNEL_ALLREDUCE_ENTRY) && \
     defined(ENABLE_NPKIT_EVENT_KERNEL_ALLREDUCE_EXIT)
   NpKit::CollectGpuEventShm(NPKIT_EVENT_KERNEL_ALLREDUCE_ENTRY, 0, 0, npkit_timestamp_entry, event_buffer,

src/algorithms/utils.cc

@@ -151,20 +151,15 @@ std::shared_ptr<mscclpp::DeviceHandle<mscclpp::BaseMemoryChannel>> setupBaseMemo
   return ptr;
 }
 
-std::vector<std::shared_ptr<AlgorithmBuilder>> loadNativeAlgorithmBuilders(uintptr_t scratchBuffer,
-                                                                           size_t scratchBufferSize) {
-  std::vector<std::shared_ptr<AlgorithmBuilder>> builders;
-  builders.push_back(std::make_shared<AllreducePacket>(scratchBuffer, scratchBufferSize));
-  builders.push_back(std::make_shared<AllreduceAllpairPacket>(scratchBuffer, scratchBufferSize));
-  return builders;
-}
-
 std::shared_ptr<AlgorithmBuilder> getDefaultNativeAlgorithmBuilder(std::string algorithmName, uintptr_t scratchBuffer,
                                                                    size_t scratchBufferSize) {
   if (algorithmName == "default_allreduce_allpair_packet") {
     return std::make_shared<AllreduceAllpairPacket>(scratchBuffer, scratchBufferSize);
   }
-  return nullptr;
+  if (algorithmName == "default_allreduce_packet") {
+    return std::make_shared<AllreducePacket>(scratchBuffer, scratchBufferSize);
+  }
+  throw std::runtime_error("Unsupported default native algorithm: " + algorithmName);
 }
 }  // namespace algorithm
 

src/include/algorithms/utils.hpp

@@ -27,8 +27,8 @@ constexpr int MAX_NRANKS_PER_NODE = 8;
 constexpr int SCRATCH_SIZE = 2 * 1024 * 1024 * 70;  // double buffer * 35 thread-blocks * 8 ranks * 256KB = 70MB
 static bool mscclppDisableChannelCache = env()->disableChannelCache;
 
-__device__ DeviceSyncer deviceSyncer;
-__constant__ DeviceSemaphore deviceSemaphore[NUM_SEMAPHORES];
+static __device__ DeviceSyncer deviceSyncer;
+static __constant__ DeviceSemaphore deviceSemaphore[NUM_SEMAPHORES];
 
 std::vector<RegisteredMemory> setupRemoteMemories(std::shared_ptr<Communicator> comm, int rank,
                                                   RegisteredMemory localMemory);
@@ -61,9 +61,6 @@ std::vector<BaseMemoryChannel> setupBaseMemoryChannels(
 std::shared_ptr<DeviceHandle<BaseMemoryChannel>> setupBaseMemoryChannelDeviceHandles(
     const std::vector<BaseMemoryChannel>& baseMemoryChannels);
 
-std::vector<std::shared_ptr<AlgorithmBuilder>> loadNativeAlgorithmBuilders(uintptr_t scratchBuffer,
-                                                                           size_t scratchBufferSize);
-
 std::shared_ptr<AlgorithmBuilder> getDefaultNativeAlgorithmBuilder(std::string algorithmName, uintptr_t scratchBuffer,
                                                                    size_t scratchBufferSize);
 }  // namespace algorithm