Merge remote-tracking branch 'nccl/v2.19' into develop

2024-01-24 15:25:33 -08:00
parent 18a57bac10 8c6c595185
commit 81ddf9de89
111 changed files with 4608 additions and 2075 deletions
@@ -0,0 +1,172 @@
+/*************************************************************************
+ * Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "device.h"
+#include "collectives.h"
+#include "primitives.h"
+
+namespace {
+  template<typename T, typename RedOp, typename Proto>
+#if defined(USE_INDIRECT_FUNCTION_CALL) && !defined(__gfx940__) && !defined(__gfx941__) && !defined(__gfx942__)
+  __device__ void runRing(ncclWorkElem *args) {
+#else
+  __device__ __attribute__((noinline)) void runRing(ncclWorkElem *args) {
+#endif
+    const int tid = threadIdx.x;
+    const int nthreads = args->nWarps*WARP_SIZE;
+    const int bid = args->bid;
+    const int nChannels = args->nChannels;
+    ncclRing *ring = &ncclShmem.channel.ring;
+    int const *ringRanks = ring->userRanks;
+    const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? REDUCESCATTER_CHUNKSTEPS : 1));
+    // We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
+    const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2);
+    const int nranks = ncclShmem.comm.nRanks;
+    const ssize_t loopSize = nChannels*chunkSize;
+    const ssize_t size = args->count;
+
+    Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
+      prims(tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, args->redOpArg, 0, args->connIndex, args->connIndex);
+
+    for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+      ssize_t realChunkSize;
+      if (Proto::Id == NCCL_PROTO_SIMPLE) {
+        realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels));
+        realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
+      }
+      else if (Proto::Id == NCCL_PROTO_LL)
+        realChunkSize = size-gridOffset < loopSize ? args->lastChunkSize : chunkSize;
+      else if (Proto::Id == NCCL_PROTO_LL128)
+        realChunkSize = min(divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128, chunkSize);
+      realChunkSize = int(realChunkSize);
+
+      ssize_t chunkOffset = gridOffset + bid*int(realChunkSize);
+
+      /////////////// begin ReduceScatter steps ///////////////
+      ssize_t offset;
+      int nelem = min(realChunkSize, size-chunkOffset);
+      int rankDest;
+
+      // step 0: push data to next GPU
+      rankDest = ringRanks[nranks-1];
+      offset = chunkOffset + rankDest * size;
+      prims.send(offset, nelem);
+
+      // k-2 steps: reduce and copy to next GPU
+      for (int j=2; j<nranks; ++j) {
+        rankDest = ringRanks[nranks-j];
+        offset = chunkOffset + rankDest * size;
+        prims.recvReduceSend(offset, nelem);
+      }
+
+      // step k-1: reduce this buffer and data, which will produce the final result
+      rankDest = ringRanks[0];
+      offset = chunkOffset + rankDest * size;
+      prims.recvReduceCopy(offset, chunkOffset, nelem, /*postOp=*/true);
+    }
+  }
+}
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    using Proto = ProtoSimple<REDUCESCATTER_CHUNKSTEPS/REDUCESCATTER_SLICESTEPS, REDUCESCATTER_SLICESTEPS>;
+    runRing<T, RedOp, Proto>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runRing<T, RedOp, ProtoLL>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runRing<T, RedOp, ProtoLL128>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncReduceScatter, T, RedOp, NCCL_ALGO_NVLS, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    const int tid = threadIdx.x;
+    const int bid = args->bid;
+    const int nChannels = args->nChannels;
+    struct ncclNvls* nvls = &ncclShmem.channel.nvls;
+    const ssize_t chunkSize = int(args->lastChunkSize);
+    const ssize_t size = args->count;
+    const ssize_t loopSize = nChannels*chunkSize;
+    const int rank = ncclShmem.comm.rank;
+    const int nranks = ncclShmem.comm.nRanks;
+
+    /* if we are direct NVLS, we only need to allocate 1 warp to scatter for sync; 
+     * if not, based on #ranks, we allocate 7 or 5 warps to reduce to saturate bandwidth
+     * and the rest are allocated to scatter. */
+    const int nThreadsReduce = args->regUsed ? (NCCL_MAX_NTHREADS - WARP_SIZE) : (nranks <= 6 ? 7 * WARP_SIZE : 5 * WARP_SIZE);
+    const int nThreadsScatter = args->regUsed ? WARP_SIZE : (NCCL_MAX_NTHREADS - nThreadsReduce);
+    const int tidEndScatter = nThreadsScatter;
+    const int tidEndReduce = tidEndScatter + nThreadsReduce;
+
+    if (!args->regUsed) {
+      if (tid < tidEndScatter) {
+        // Scatter
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
+        Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
+          prims(tid, nThreadsScatter, NULL, nvls->up, args->sendbuff, NULL,
+            args->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid * chunkSize;
+          int nelem = min(chunkSize, size - offset);
+          prims.scatter(offset, nvls->nHeads * size, nelem, size, -1, 0);
+        }
+      } else if (tid < tidEndReduce) {
+        // Reduce through NVLS
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
+        Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
+          prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, NULL, NULL, args->recvbuff,
+            args->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid * chunkSize;
+          int nelem = min(chunkSize, size - offset);
+          prims.recv(offset, nelem);
+        }
+      }
+    } else {
+      if (tid < tidEndScatter) {
+        // Scatter
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
+        Primitives<T, RedOp, FanSymmetric<NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
+          prims(tid, nThreadsScatter, nvls->up, nvls->up, NULL, NULL,
+            args->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          prims.scatter(0, 0, 0, 0, -1, 0);
+        }
+
+        /* gather used as sync */
+        prims.gather(0, 0, 0, 0, -1, 0);
+      } else if (tid < tidEndReduce) {
+        // Reduce through NVLS
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
+        Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
+          prims(tid - tidEndScatter, nThreadsReduce, &nvls->down, &nvls->down, NULL, args->recvbuff,
+            args->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t outOffset = gridOffset + bid * chunkSize;
+          ssize_t inpOffset = outOffset + rank * size;
+          int nelem = min(chunkSize, size - outOffset);
+          prims.directRecvCopy(inpOffset, outOffset, nelem);
+        }
+
+        /* send for sync */
+        prims.send(0, 0);
+      }
+    }
+  }
+};