Merge remote-tracking branch 'nccl/master' into develop

[ROCm/rccl commit: 36343be84f]
2025-01-23 11:48:18 -06:00
parent e3a29f5eab 5ca1b6c160
commit 1b000665df
122 changed files with 9009 additions and 4702 deletions
@@ -13,12 +13,13 @@
 #include "rccl_float8.h"
 #include <hip/hip_bfloat16.h>
 #include "nccl_common.h"
-#include "align.h"
-#include "collectives.h"
+#include "bitops.h"
 #if defined(ENABLE_NPKIT)
 #include "npkit/npkit_struct.h"
 #endif
+#include <algorithm>
 #include <stdint.h>
+#include <sys/types.h>

 extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+2];

@@ -31,6 +32,12 @@ extern const char* funcNames[FUNC_INDEX_TOTAL];
 #define NCCL_MAX_OPS 2048
 #define NCCL_STEPS 8

+#ifdef __CUDA_ARCH__
+  #define NCCL_CUDA_ARCH __CUDA_ARCH__
+#else
+  #define NCCL_CUDA_ARCH 0
+#endif
+
 #include "net_device.h"

 enum ncclDevRedOp_t {
@@ -61,12 +68,13 @@ union ncclLLFifoLine {
 };

 #define WARP_SIZE warpSize
-
 #define MAXCHANNELS 128
 #define CHANNEL_LIMIT 16
-
+#define NCCL_MAX_LOCAL_RANKS 64
 #define NCCL_MAX_NTHREADS 256
+#define NCCL_MIN_NTHREADS (4*WARP_SIZE)
 #define NCCL_SIMPLE_MAX_NTHREADS NCCL_MAX_NTHREADS
+#define NCCL_SIMPLE_EXTRA_GROUP_IF_NTHREADS_GE (3*WARP_SIZE)
 #define NCCL_LL_MAX_NTHREADS NCCL_MAX_NTHREADS
 #define NCCL_LL_LINES_PER_THREAD 8
 #ifdef TEST_LL_CLEANUP
@@ -97,6 +105,9 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
 #define NCCL_IPC_READ     0x10
 #define NCCL_NVLS_MIN_POLL 0x20

+// Number of named barriers supported by CUDA
+#define NCCL_MAX_GROUPS (NCCL_MAX_NTHREADS/WARP_SIZE)
+
 #define NCCL_MAX_COLLNET_SIZE (1L << 29)

 enum ncclRegBufferType {
@@ -218,124 +229,153 @@ struct ncclDevComm;

 #pragma pack(push)  /* push current alignment to stack */
 #pragma pack(8)     /* set alignment to 8 bytes boundary */
-/* ncclWork is to be a power of two, currently 8x64 bytes, */
-/* to make sure reads to host from the CUDA kernel are aligned. */
-/* Make sure to adjust padding at the end of ncclWorkElem. */
-#define NCCL_WORK_SIZE 256

-enum ncclWorkType : uint8_t {
-   ncclWorkTypeUnused=0,
-   ncclWorkTypeColl=1,
-   ncclWorkTypeP2p=2,
-   ncclWorkTypeRegColl=3
-};
-enum ncclWorkP2PType : uint8_t {
-  ncclWorkP2pTypeUnused=0,
-  ncclWorkP2pTypeSend,
-  ncclWorkP2pTypeRecv
+struct alignas(16) ncclDevWorkP2p {
+  void *sendAddr, *recvAddr;
+  size_t sendBytes, recvBytes;
+  int sendRank, recvRank;
+  // From the part index, nP2pChannels, and channelBase the device code can
+  // calculate which part of the transfer a channel is responsible for.
+  uint8_t nP2pChannels; // Always equal to comm->p2pnChannels
+  uint8_t channelBase; // Channel owning first part.
+  // Zero channels indicates no work in that direction.
+  uint8_t nSendChannels, nRecvChannels;
+  // Chunk size stored in 8 bits via u32fp8Encode/Decode.
+  uint8_t sendChunkSize_u32fp8, recvChunkSize_u32fp8;
+
+  uint8_t sendProtoLL:1, recvProtoLL:1;
+  uint8_t sendRegistered:1, recvRegistered:1;
+
+  uint8_t connIndex:2;
 };

-struct ncclWorkHeader {
-  union {
-    int32_t workNext;  // when isLast=0: Offset from kernel argument workHead
-    uint32_t doneAcks; // when isLast=1: Monotonic (mod 1<<32) ack value to send back.
-  };
-  uint16_t funcIndex;
-  uint8_t isLast:1; // last work for this kernel
-  uint8_t inFifo:1; // is this work in the fifo
-  enum ncclWorkType type;
-};
+// Compute the subset of the data transfer corresponding to the given part index.
+inline __host__ __device__ void ncclP2pPartBounds(int nParts, int part, size_t bytes, size_t* partBeg, size_t* partEnd) {
+  size_t partBytes = alignUp(divUp(bytes, nParts), 4<<10);
+  #if __CUDA_ARCH__
+    *partBeg = min((part+0)*partBytes, bytes);
+    *partEnd = min((part+1)*partBytes, bytes);
+  #else
+    *partBeg = std::min<size_t>((part+0)*partBytes, bytes);
+    *partEnd = std::min<size_t>((part+1)*partBytes, bytes);
+  #endif
+}

-struct ncclWorkElem {
-  union {
-    uint8_t flagBits;
-    struct {
-      uint8_t isUsed:1, redOpArgIsPtr:1, oneNode:1;
-    };
-  };
-  uint8_t regUsed;
-  uint8_t nWarps;
-  uint8_t direct;
+// implemented in channel.h
+inline __host__ uint8_t ncclP2pChannelBaseForRound(struct ncclComm* comm, int p2pRound);

+// ncclP2pChannelToPart and ncclP2pChannelForPart are inverses. The device code
+// uses ncclP2pChannelToPart to determine which part "this" channel is responsible for.
+inline __host__ int ncclP2pChannelForPart(int nP2pChannels, int base, int part, int nParts) {
+  return (base * nParts + part) & (nP2pChannels-1);
+}
+inline __device__ int ncclP2pChannelToPart(int nP2pChannels, int base, int channel, int nParts) {
+  return (channel - base * nParts) & (nParts-1);
+}
+
+struct alignas(16) ncclDevWorkColl {
+  // Running on channels [channelLo..channelHi], hi is inclusive.
+  //   nChannels == (channelHi - channelLo) + 1
+  uint32_t channelLo:8, channelHi:8;
+  uint32_t nWarps:8;
+  uint32_t redOpArgIsPtr:1, regUsed:2, oneNode:1, direct:4;
  uint32_t root:30, connIndex:2;
-  const void *sendbuff;
-  void *recvbuff;
-
-  uint64_t count:39, opCount:25;
+  uint16_t pivotA2ANumBiRings;
+  void* recvbuff;
+  void* sendbuff;
+  union {
+    // Continuous-byte-distribution scheduling. The lo and hi channels are of
+    // different size than the channels in the middle.
+    struct {
+      size_t countLo, countMid, countHi;
+      // Chunk counts where units are ncclProtoGrainSize(protocol) bytes
+      uint64_t chunkGrainsLo:21, chunkGrainsMid:21, chunkGrainsHi:21;
+    } cbd;
+    // Collnet scheduling. All channels divide work evenly.
+    struct {
+      size_t count; // Total size, not divided per channel.
+      uint32_t chunkCount;
+    } collnet;
+  };
  uint64_t redOpArg;
-  uint64_t chunkCount:25, workCount:39;
-  union {
-    struct {
-      uint64_t lastChunkCount:25;
-      uint64_t workOffset:39;
-    };
-    struct {
-      uint32_t nChannels;
-      uint16_t bid;
-      // Pivot A2A kernel computes chunk size itself.
-      // Instead, it needs the number of bidirectional rings.
-      uint16_t pivotA2ANumBiRings;
-    };
-  };
 };

-static_assert((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWorkElem)))/sizeof(ncclWorkElem) == 4, "Sanity check: NCCL_MAX_WORK_ELEMENTS == 4");
-#define NCCL_MAX_WORK_ELEMENTS 1

-struct ncclWorkElemP2p {
-  struct {
-    int32_t peer:28;
-    uint32_t connIndex:2;
-    int32_t proto:2;
-  };
-  union {
-    uint16_t flagBit;
-    struct {
-      enum ncclWorkP2PType p2pType:4;
-      uint8_t nWarps:4;
-      uint8_t warpStart:4;
-      uint8_t ngroups:4;
-    };
-  };
-  uint8_t reg:1;
-  uint16_t opCount:15;
-  // Important not to use any fields with greater than 4-byte alignment since
-  // we need sizeof(ncclWorkElemP2p)==28, but that would be padded up to 32 if
-  // there were 8-byte fields.
-  //void* buff;
-  uint32_t buffHi32, buffLo32; // buff = buffHi32<<32 | buffLo32;
-  //size_t count;
-  uint32_t countHi32, countLo32; // count = countHi32<<32 | countLo32;
-  int chunkSize;
-};
+__host__ __device__ constexpr int ncclProtoGrainSize(int proto) {
+  return proto == NCCL_PROTO_LL ? 16 :
+         proto == NCCL_PROTO_LL128 ? WARP_SIZE*NCCL_LL128_SHMEM_ELEMS_PER_THREAD/NCCL_LL128_LINEELEMS*NCCL_LL128_DATAELEMS*sizeof(uint64_t) :
+         proto == NCCL_PROTO_SIMPLE ? 512 :
+         -1;
+}

-static_assert(((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWorkElemP2p)))/sizeof(ncclWorkElemP2p)) == 8, "Sanity check: NCCL_MAX_WORK_ELEMENTS_P2P == 8");
-#define NCCL_MAX_WORK_ELEMENTS_P2P 2
+template<typename Int>
+__host__ __device__ inline void ncclCollCbdPart(
+    struct ncclDevWorkColl* work, uint32_t channelId, int proto, int eltSize,
+    Int* count, Int* partOffset, Int* partCount, Int* chunkCount
+  ) {
+  int eltPerGrain = ncclProtoGrainSize(proto)/eltSize;
+  int nMidChannels = work->channelHi - work->channelLo - 1;
+  // We can assum that nMidChannels<0 implies countMid==0, which let's us assume
+  // that countMid*nMidChannels == 0.
+  if (count != nullptr) {
+    *count = work->cbd.countLo + work->cbd.countMid*nMidChannels + work->cbd.countHi;
+  }
+  if (channelId == work->channelLo) {
+    *partOffset = 0;
+    *partCount = work->cbd.countLo;
+    *chunkCount = work->cbd.chunkGrainsLo*eltPerGrain;
+  } else if (channelId == work->channelHi) {
+    *partOffset = work->cbd.countLo + nMidChannels*work->cbd.countMid;
+    *partCount = work->cbd.countHi;
+    *chunkCount = work->cbd.chunkGrainsHi*eltPerGrain;
+  } else {
+    int mid = channelId - work->channelLo - 1;
+    *partOffset = work->cbd.countLo + mid*work->cbd.countMid;
+    *partCount = work->cbd.countMid;
+    *chunkCount = work->cbd.chunkGrainsMid*eltPerGrain;
+  }
+}

-struct ncclWorkElemReg {
-  struct ncclWorkElem elem;
+struct alignas(16) ncclDevWorkCollReg {
+  struct ncclDevWorkColl coll;
  void* dnInputs[NCCL_MAX_DIRECT_ARITY+1];
  void* dnOutputs[NCCL_MAX_DIRECT_ARITY+1];
  void* upOutputs[NCCL_MAX_DIRECT_ARITY+1];
 };

-#define NCCL_MAX_WORK_ELEMENTS_REG ((NCCL_WORK_SIZE - alignUp(sizeof(ncclWorkHeader), alignof(ncclWorkElemReg)))/sizeof(ncclWorkElemReg))
-static_assert(NCCL_MAX_WORK_ELEMENTS_REG == 1, "Sanity check: NCCL_MAX_WORK_ELEMENTS_REG == 1");
-
-// Number of named barriers supported by CUDA
-#define NCCL_MAX_GROUPS (NCCL_MAX_NTHREADS/WARP_SIZE)
-
-struct ncclWork {
-  struct ncclWorkHeader header;
-  union {
-    char pad[NCCL_WORK_SIZE - sizeof(struct ncclWorkHeader)];
-    struct ncclWorkElem elems[NCCL_MAX_WORK_ELEMENTS];
-    struct ncclWorkElemP2p p2pElems[NCCL_MAX_WORK_ELEMENTS_P2P];
-    struct ncclWorkElemReg regElems[NCCL_MAX_WORK_ELEMENTS_REG];
-  };
+enum ncclDevWorkType: uint8_t {
+  ncclDevWorkTypeP2p,
+  ncclDevWorkTypeColl,
+  ncclDevWorkTypeCollReg
+};
+
+constexpr size_t ncclDevWorkSize(enum ncclDevWorkType type) {
+  return type == ncclDevWorkTypeP2p ? sizeof(ncclDevWorkP2p) :
+         type == ncclDevWorkTypeColl ? sizeof(ncclDevWorkColl) : sizeof(ncclDevWorkCollReg);
+}
+
+#define NCCL_MAX_DEV_WORK_BATCH_BYTES 128
+#define NCCL_MAX_DEV_WORK_BATCH_COLLS (NCCL_MAX_DEV_WORK_BATCH_BYTES/sizeof(ncclDevWorkColl))
+#define NCCL_MAX_DEV_WORK_P2P_PER_BATCH 1
+struct alignas(16) ncclDevWorkBatch {
+  union {
+    struct {
+      // nextExtends: should next one be merged into this one.
+      // nextJump=0: end of this channel's batch list
+      // nextJump>0: batches[thisIndex+nextJump] is next batch in this list
+      uint32_t nextJump:14, nextExtends:1;
+      uint32_t workType:2, funcId:15;
+    };
+    // Unioning bitfields with underlying type hints compiler to emit the best
+    // SASS LD/ST accesses.
+    uint32_t flags;
+  };
+  // Rolling offset in fifo where this batch's work structs begin
+  uint32_t offsetBase;
+  // Set of relative offsets from offsetBase for this channel's subset of the batch:
+  // For each bit index i in offsetMask, find work at fifo offset: offsetBase + i*sizeof(WorkStructType)
+  uint64_t offsetBitset;
 };
-static_assert(sizeof(struct ncclWork) == NCCL_WORK_SIZE, "Sanity check: sizeof(struct ncclWork) == NCCL_WORK_SIZE");
-static_assert(sizeof(struct ncclWork)%16 == 0, "Sanity check: sizeof(struct ncclWork)%16 == 0");

 struct ncclDevChannelPeer {
  // Stripped version of ncclChannelPeer where we only keep the ncclConnInfo
@@ -378,26 +418,27 @@ struct ncclCollTrace {
  uint8_t bid;
  int16_t funcIndex;
  uint32_t data_0;
-  uint8_t channelId;
-  uint64_t timeStamp:56;
-  union {
-    uint64_t opCount;
-    uint32_t p2pOpCount[2];
-  };
+  uint64_t timeStamp;
+  uint64_t opCount;
  union {
    uint64_t data_1;
    struct {
      uint8_t nWarps;
-      uint8_t bid;
      uint8_t nChannels;
+      uint8_t bid;
+      uint8_t root;
    } coll;
    struct {
-      int16_t peer;
-      uint8_t ngroups:4;
-      uint8_t connIndex:4;
-      uint8_t warpStart:4;
-      uint8_t nWarps:4;
-    } p2p[2];
+      uint8_t sendRank;
+      uint8_t recvRank;
+      uint8_t nP2pChannels;
+      uint8_t nSendChannels;
+      uint8_t nRecvChannels;
+      uint8_t channelBase;
+      uint8_t connIndex;
+      uint8_t sendProtoLL:1;
+      uint8_t recvProtoLL:1;
+    } p2p;
  };
 };
 static_assert(sizeof(struct ncclCollTrace) == 8*sizeof(int), "ncclCollTrace must have a pow2 size");
@@ -428,10 +469,10 @@ struct ncclDevComm {
  int nNodes;
  int buffSizes[NCCL_NUM_PROTOCOLS];
  int p2pChunkSize;
+  int p2pnChannelsPerPeer;

-  // Operation list for aggregation
-  int workFifoDepth;
-  struct ncclWork* workFifoHeap; // may be cudaHost or GDR memory
+  // Work fifo return credits
+  uint32_t* workConsumed/*[MAXCHANNELS]*/;

  int* collNetDenseToUserRank;

@@ -462,11 +503,41 @@ struct alignas(16) ncclDevCommAndChannels {
  struct ncclDevChannel channels[MAXCHANNELS];
 };

-#ifdef __CUDA_ARCH__
-  #define NCCL_CUDA_ARCH __CUDA_ARCH__
-#else
-  #define NCCL_CUDA_ARCH 0
-#endif
+enum ncclDevWorkStorageType: uint8_t {
+  ncclDevWorkStorageTypeArgs=0,
+  ncclDevWorkStorageTypeFifo=1,
+  ncclDevWorkStorageTypePersistent=2
+};
+
+struct channelMasks {
+  uint64_t masks[MAXCHANNELS/64];
+};
+
+struct alignas(16) ncclDevKernelArgs {
+  struct ncclDevComm* comm;
+  struct channelMasks channelMask;
+  enum ncclDevWorkStorageType workStorageType;
+  uint32_t workMask;
+  void* workBuf;
+  // A channel's first batch is at `blockIdx.x`. Use `nextJump` to follow rest of list.
+  // struct ncclDevWorkBatch batches[];
+};
+
+__host__ __device__ constexpr int ncclMaxKernelArgsSize(/*int cudaDriver, */int cudaArch=NCCL_CUDA_ARCH) {
+  //return (cudaArch < 700 || cudaDriver < 12010) ? 4<<10 : (32<<10)-4;
+  return 4<<10;
+}
+
+template<size_t capacity>
+struct alignas(16) ncclDevKernelArgsStorage {
+  union {
+    struct ncclDevKernelArgs args;
+    ulong2 storage[capacity/sizeof(ulong2)];
+  };
+};
+
+typedef ncclDevKernelArgsStorage<(4<<10)> ncclDevKernelArgs4K;
+//typedef ncclDevKernelArgsStorage<(32<<10)-4> ncclDevKernelArgs31K;

 template<typename T>
 __host__ __device__ constexpr T min_constexpr(T a) { return a; }
@@ -482,6 +553,10 @@ __host__ __device__ constexpr T max_constexpr(T a, T b, Ts ...c) {
  return max_constexpr<T>((a > b ? a : b), c...);
 }

+constexpr int ncclDevMaxChannelsForArgsBytes(size_t argsBytes) {
+  return min_constexpr<size_t>(MAXCHANNELS, (argsBytes - sizeof(struct ncclDevKernelArgs))/sizeof(struct ncclDevWorkBatch));
+}
+
 // Calculate the unroll factor given:
 // * bytePerPack: number of bytes accessed per instruction
 // * insns: max permissible unroll value
@@ -523,6 +598,16 @@ __host__ __device__ constexpr int ncclShmemDynamicSize(int cudaArch = NCCL_CUDA_
  return cudaArch < 700 ? 0 : ncclShmemScratchWarpSize(cudaArch)*(NCCL_MAX_NTHREADS/WARP_SIZE);
 }

+// Host-side table of kernel function pointers.
+extern int const ncclDevKernelCount;
+extern void* const ncclDevKernelList[/*ncclDevKernelCount*/];
+
+// Table of most specialized kernel function to run given func index.
+extern int const ncclDevFuncIdCount;
+extern int const ncclDevFuncRowToId[];
+extern void* const ncclDevKernelForFunc[/*funcIndex*/];
+extern bool const ncclDevKernelForFuncIsSpecialized[/*funcIndex*/];
+
 // Launch a one-rank reduction on stream.
 ncclResult_t ncclLaunchOneRank(void* dst, void const* src, size_t nElts, struct ncclDevRedOpFull redOp, ncclDataType_t type, cudaStream_t stream);