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

src/device/all_gather.h

@@ -0,0 +1,263 @@
+/*************************************************************************
+ * 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;
+    const int *ringRanks = ring->userRanks;
+    const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLGATHER_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*int(chunkSize);
+    const ssize_t size = args->count;
+
+#if defined(ENABLE_NPKIT)
+    int npKitCtxIdx = bid;
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (tid == 0) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    T *inputBuf = (T*)args->sendbuff;
+    T *outputBuf = (T*)args->recvbuff;
+    Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
+      (tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, args->redOpArg, 0, args->connIndex, args->connIndex);
+
+#if defined(ENABLE_NPKIT)
+    if (tid == 0) {
+      prims.npKitCtxIdx = npKitCtxIdx;
+    }
+#endif
+
+    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(chunkSize, divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128);
+      realChunkSize = int(realChunkSize);
+
+      ssize_t chunkOffset = gridOffset + int(bid*realChunkSize);
+
+      /////////////// begin AllGather steps ///////////////
+      ssize_t offset;
+      int nelem = min(realChunkSize, size-chunkOffset);
+      int rankDest;
+
+      // step 0: push data to next GPU
+      rankDest = ringRanks[0];
+      offset = chunkOffset + rankDest * size;
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_SEND_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_SEND_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      if (inputBuf + chunkOffset == outputBuf + offset) { // In place
+        prims.directSend(chunkOffset, offset, nelem);
+      } else {
+        prims.directCopySend(chunkOffset, offset, nelem);
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_SEND_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_SEND_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_ENTRY)
+      if (tid == 0 && nranks > 2) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_ENTRY, nelem*(nranks-2)*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      // k-2 steps: copy to next GPU
+      for (int j=1; j<nranks-1; ++j) {
+        rankDest = ringRanks[nranks-j];
+        offset = chunkOffset + rankDest * size;
+
+        prims.directRecvCopySend(offset, nelem);
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_EXIT)
+      if (tid == 0 && nranks > 2) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_RECV_COPY_SEND_EXIT, nelem*(nranks-2)*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+      // Make final copy from buffer to dest.
+      rankDest = ringRanks[1];
+      offset = chunkOffset + rankDest * size;
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+      // Final wait/copy.
+      prims.directRecv(offset, nelem);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_DIRECT_RECV_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+
+
+    }
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_GATHER_RING_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_GATHER_RING_EXIT, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+}
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    using Proto = ProtoSimple<ALLGATHER_CHUNKSTEPS/ALLGATHER_SLICESTEPS, ALLGATHER_SLICESTEPS>;
+    runRing<T, RedOp, Proto>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllGather, 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<ncclFuncAllGather, 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<ncclFuncAllGather, 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 ssize_t rank = ncclShmem.comm.rank;
+
+    const int nThreadsBcast = args->regUsed ? (NCCL_MAX_NTHREADS - WARP_SIZE) : 4 * WARP_SIZE;
+    const int nThreadsGather = args->regUsed ? WARP_SIZE : NCCL_MAX_NTHREADS - nThreadsBcast;
+    const int tidEndGather = nThreadsGather;
+    const int tidEndBcast = tidEndGather + nThreadsBcast;
+
+    if (!args->regUsed) {
+      if (tid < tidEndGather) {
+        // Gather
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
+        Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_NVLS_ARITY, 0>, /*Direct=*/0, Proto, 0>
+          prims(tid, nThreadsGather, nvls->up, NULL, NULL, args->recvbuff,
+            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.gather(offset, nvls->nHeads * size, nelem, size, -1, 0);
+        }
+      } else if (tid < tidEndBcast) {
+        // Bcast through NVLS
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
+        Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
+          prims(tid - tidEndGather, nThreadsBcast, NULL, &nvls->down, args->sendbuff, NULL,
+            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.send(offset, nelem);
+        }
+      }
+    } else {
+      /* direct allgather */
+      if (tid < tidEndGather) {
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
+        Primitives<T, RedOp, FanSymmetric<NCCL_MAX_NVLS_ARITY>, /*Direct=*/0, Proto, 0>
+          prims(tid, nThreadsGather, nvls->up, nvls->up, NULL, NULL,
+            args->redOpArg, 0 * Proto::MaxGroupWidth, 1, 1);
+
+        /* used as sync */
+        prims.scatter(0, 0, 0, 0, -1, 0);
+
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          prims.gather(0, 0, 0, 0, -1, 0);
+        }
+      } else if (tid < tidEndBcast) {
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
+        Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
+          prims(tid - tidEndGather, nThreadsBcast, &nvls->down, &nvls->down, args->sendbuff, NULL,
+            args->redOpArg, 1 * Proto::MaxGroupWidth, 0, 0, args);
+        /* used as sync */
+        prims.recv(0, 0);
+
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t inpOffset = gridOffset + bid * chunkSize;
+          ssize_t outOffset = inpOffset + rank * size;
+          int nelem = min(chunkSize, size - inpOffset);
+          prims.directSend(inpOffset, outOffset, nelem);
+        }
+      }
+    }
+  }
+};

src/device/all_reduce.h

@@ -0,0 +1,995 @@
+/*************************************************************************
+ * 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"
+
+#if defined(ENABLE_NPKIT)
+#include "npkit/npkit.h"
+#endif
+
+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 ringIx = ring->index;
+    const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLREDUCE_CHUNKSTEPS : 1));
+    const int nranks = ncclShmem.comm.nRanks;
+    const ssize_t loopSize = nChannels*nranks*chunkSize;
+    const ssize_t size = args->count;
+
+#if defined(ENABLE_NPKIT)
+    int npKitCtxIdx = bid;
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (tid == 0) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    int minChunkSize;
+    if (Proto::Id == NCCL_PROTO_LL)
+      minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T));
+    if (Proto::Id == NCCL_PROTO_LL128) {
+      // We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
+      minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2;
+    }
+
+    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);
+
+#if defined(ENABLE_NPKIT)
+    if (tid == 0) {
+      prims.npKitCtxIdx = npKitCtxIdx;
+    }
+#endif
+
+    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*nranks));
+        realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
+      }
+      else
+        realChunkSize = min(chunkSize, divUp(size-gridOffset, nChannels*nranks*minChunkSize)*minChunkSize);
+      realChunkSize = int(realChunkSize);
+
+      auto calcOffset = [&]__device__(int chunk)->ssize_t {
+        if (Proto::Id == NCCL_PROTO_SIMPLE)
+          return gridOffset + bid*nranks*realChunkSize + chunk*realChunkSize;
+        else
+          return gridOffset + (chunk*nChannels + bid)*realChunkSize;
+      };
+      auto modRanks = [&]__device__(int r)->int {
+        return r - (r >= nranks ? nranks : 0);
+      };
+
+      ssize_t offset;
+      int nelem;
+      int chunk;
+
+      // step 0: push data to next GPU
+      chunk = modRanks(ringIx + nranks-1);
+      offset = calcOffset(chunk);
+      nelem = min(realChunkSize, size-offset);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_SEND_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_SEND_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      prims.send(offset, nelem);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_SEND_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_SEND_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+      // k-2 steps: reduce and copy to next GPU
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_RECV_REDUCE_SEND_ENTRY)
+      if (tid == 0 && nranks > 2) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_RECV_REDUCE_SEND_ENTRY, nelem*(nranks-2)*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      for (int j=2; j<nranks; ++j) {
+        chunk = modRanks(ringIx + nranks-j);
+        offset = calcOffset(chunk);
+        nelem = min(realChunkSize, size-offset);
+        prims.recvReduceSend(offset, nelem);
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_RECV_REDUCE_SEND_EXIT)
+      if (tid == 0 && nranks > 2) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_RECV_REDUCE_SEND_EXIT, nelem*(nranks-2)*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+      // step k-1: reduce this buffer and data, which will produce the final
+      // result that we store in this data and push to the next GPU
+      chunk = ringIx + 0;
+      offset = calcOffset(chunk);
+      nelem = min(realChunkSize, size-offset);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_REDUCE_COPY_SEND_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_REDUCE_COPY_SEND_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      prims.directRecvReduceCopySend(offset, offset, nelem, /*postOp=*/true);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_REDUCE_COPY_SEND_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_REDUCE_COPY_SEND_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_ENTRY)
+      if (tid == 0 && nranks > 2) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_ENTRY, nelem*(nranks-2)*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      // k-2 steps: copy to next GPU
+      for (int j=1; j<nranks-1; ++j) {
+        chunk = modRanks(ringIx + nranks-j);
+        offset = calcOffset(chunk);
+        nelem = min(realChunkSize, size-offset);
+        prims.directRecvCopySend(offset, nelem);
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_EXIT)
+      if (tid == 0 && nranks > 2) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_COPY_SEND_EXIT, nelem*(nranks-2)*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      // Make final copy from buffer to dest.
+      chunk = modRanks(ringIx + 1);
+      offset = calcOffset(chunk);
+      nelem = min(realChunkSize, size-offset);
+      prims.directRecv(offset, nelem);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_DIRECT_RECV_EXIT, nelem*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_RING_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_RING_EXIT, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+  }
+
+  template<typename T, typename RedOp, typename Proto>
+#if defined(USE_INDIRECT_FUNCTION_CALL) && !defined(__gfx940__) && !defined(__gfx941__) && !defined(__gfx942__)
+  __device__ void runTreeUpDown(ncclWorkElem *args) {
+#else
+  __device__ __attribute__((noinline)) void runTreeUpDown(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;
+    ncclTree *tree = &ncclShmem.channel.tree;
+    ssize_t chunkSize = int(
+      Proto::Id == NCCL_PROTO_SIMPLE ? args->lastChunkSize
+                   /* LL & LL128 */  : Proto::calcBytePerStep()/sizeof(T));
+    const ssize_t minChunkSize = int(
+      Proto::Id == NCCL_PROTO_SIMPLE ? nthreads*8*(sizeof(uint64_t)/sizeof(T))
+                   /* LL & LL128 */  : nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
+    const ssize_t loopSize = int(nChannels*chunkSize);
+    const ssize_t size = args->count;
+
+#if defined(ENABLE_NPKIT)
+    int npKitCtxIdx = bid;
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (tid == 0) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    if (loopSize > size)
+      chunkSize = divUp((int)size, int(nChannels*minChunkSize))*int(minChunkSize);
+
+    { // Reduce : max number of recv is 3, max number of send is 1 (binary tree + local)
+      Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto, 0> prims
+        (tid, nthreads, tree->down, &tree->up, args->sendbuff, args->recvbuff, args->redOpArg);
+
+#if defined(ENABLE_NPKIT)
+      if (tid == 0) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_REDUCE_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_REDUCE_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      if (tree->up == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.recvReduceCopy(offset, offset, nelem, /*postOp=*/true);
+        }
+      }
+      else if (tree->down[0] == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.send(offset, nelem);
+        }
+      }
+      else {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.recvReduceSend(offset, nelem);
+        }
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_REDUCE_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_REDUCE_EXIT, size*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+
+    { // Broadcast : max number of recv is 1, max number of send is 3 (binary tree + local)
+      Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto, 0> prims
+        (tid, nthreads, &tree->up, tree->down, args->sendbuff, args->recvbuff, args->redOpArg);
+
+#if defined(ENABLE_NPKIT)
+      if (tid == 0) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_BROADCAST_ENTRY)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_BROADCAST_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      if (tree->up == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.directSendFromOutput(offset, nelem);
+        }
+      }
+      else if (tree->down[0] == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.directRecv(offset, nelem);
+        }
+      }
+      else {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.directRecvCopySend(offset, nelem);
+        }
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_BROADCAST_EXIT)
+      if (tid == 0) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_BROADCAST_EXIT, size*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_UPDOWN_EXIT, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+  }
+
+  template<typename T, typename RedOp, typename Proto>
+#if defined(USE_INDIRECT_FUNCTION_CALL) && !defined(__gfx940__) && !defined(__gfx941__) && !defined(__gfx942__)
+  __device__ void runTreeSplit(ncclWorkElem *args) {
+#else
+  __device__ __attribute__((noinline)) void runTreeSplit(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;
+    ncclTree *tree = &ncclShmem.channel.tree;
+    ssize_t chunkSize = int(
+      Proto::Id != NCCL_PROTO_LL ? args->lastChunkSize
+                                 : Proto::calcBytePerStep()/sizeof(T));
+    const ssize_t minChunkSize = int(
+      Proto::Id == NCCL_PROTO_SIMPLE ? nthreads*8*(sizeof(uint64_t)/sizeof(T)) :
+      Proto::Id == NCCL_PROTO_LL     ? nthreads*(Proto::calcBytePerGrain()/sizeof(T))
+                   /* LL128 */       : nthreads*(Proto::calcBytePerGrain()/sizeof(T))/8);
+    const ssize_t loopSize = int(nChannels*chunkSize);
+    const ssize_t size = args->count;
+    int nthreadsSplit;
+    if (Proto::Id == NCCL_PROTO_SIMPLE) {
+      nthreadsSplit = nthreads/2;
+      if (nthreadsSplit >= 256) nthreadsSplit += 64;
+    } else { // LL & LL128
+      // Receiving from up to 3 sources is more compute intensive than sending
+      // to 3 dests. Use 70% for reduce and 30% for bcast.
+      nthreadsSplit = (nthreads*7/(10*WARP_SIZE))*WARP_SIZE;
+    }
+
+#if defined(ENABLE_NPKIT)
+    bool isNpKitThread = false;
+    int npKitCtxIdx = 0;
+    if (threadIdx.x == 0) {
+      isNpKitThread = true;
+      npKitCtxIdx = bid * 2;
+    } else if (tree->up != -1 && threadIdx.x == nthreadsSplit) {
+      isNpKitThread = true;
+      npKitCtxIdx = bid * 2 + 1;
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (isNpKitThread) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (isNpKitThread) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_ENTRY)
+    if (isNpKitThread) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    if (loopSize > size)
+      chunkSize = divUp((int)size, nChannels*int(minChunkSize))*int(minChunkSize);
+
+    if (tree->up == -1) {
+      // Reduce and broadcast. Max number of recv is 2, max number of send is 2
+      Primitives<T, RedOp, FanSymmetric<NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto, 0>
+        prims(tid, nthreads, tree->down, tree->down, args->sendbuff, args->recvbuff, args->redOpArg);
+
+#if defined(ENABLE_NPKIT)
+      if (isNpKitThread) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_BROADCAST_ENTRY)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_BROADCAST_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + bid*int(chunkSize);
+        int nelem = min(chunkSize, size-offset);
+        prims.directRecvReduceCopySend(offset, offset, nelem, /*doPost=*/true);
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_BROADCAST_EXIT)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_BROADCAST_EXIT, size*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+    else if (tid < nthreadsSplit) {
+      /* Reduce up. Max number of recv is 3, max number of send is 1 (binary tree + local).
+       * Why Direct=1????
+       * Answer: Because despite not performing any direct operations, the ctor
+       * must assume Direct so that it can exchange direct pointers with remote ctors
+       * that are Direct, otherwise it hangs. A cleaner solution would be to seperate
+       * into DirectRecv and DirectSend capabilities, this ctor would have both=0,
+       * but the ctor above for tree roots would be DirectRecv=0 DirectSend=1.
+       */
+      Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/0, Proto, 0>
+        prims(tid, nthreadsSplit, tree->down, &tree->up, args->sendbuff, args->recvbuff, args->redOpArg, 0*Proto::MaxGroupWidth);
+
+#if defined(ENABLE_NPKIT)
+      if (isNpKitThread) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_ENTRY)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      if (tree->down[0] == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.send(offset, nelem);
+        }
+      }
+      else {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.recvReduceSend(offset, nelem);
+        }
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_EXIT)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_REDUCE_EXIT, size*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+    else {
+      // Broadcast down. Max number of recv is 1, max number of send is 3 (binary tree + local)
+      Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/0, Proto, 0>
+        prims(tid-nthreadsSplit, nthreads-nthreadsSplit, &tree->up, tree->down, args->sendbuff, args->recvbuff,
+            args->redOpArg, 1*Proto::MaxGroupWidth);
+
+#if defined(ENABLE_NPKIT)
+      if (isNpKitThread) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_BROADCAST_ENTRY)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_BROADCAST_ENTRY, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      if (tree->down[0] == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.directRecv(offset, nelem);
+        }
+      }
+      else {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.directRecvCopySend(offset, nelem);
+        }
+      }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_BROADCAST_EXIT)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_BROADCAST_EXIT, size*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_EXIT)
+    if (isNpKitThread) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_ALL_REDUCE_TREE_SPLIT_EXIT, size*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+  }
+}
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    using Proto = ProtoSimple<ALLREDUCE_CHUNKSTEPS/ALLREDUCE_SLICESTEPS, ALLREDUCE_SLICESTEPS>;
+    runRing<T, RedOp, Proto>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runTreeUpDown<T, RedOp, ProtoSimple<1, 1>>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_DIRECT, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    static constexpr int COLLNET_COPY_THREADS = 64;
+    const int tid = threadIdx.x;
+    const int bid = args->bid;
+    const int nChannels = args->nChannels;
+    struct ncclDirect* direct = &ncclShmem.channel.collnetDirect;
+    const ssize_t chunkSize = int(args->lastChunkSize);
+    const ssize_t size = args->count;
+    const ssize_t loopSize = nChannels*direct->nHeads*chunkSize;
+
+    const int hasUp = (direct->up[0] >= 0) ? 1 : 0;
+    const int hasDn = (direct->down[0] >= 0) ? 1 : 0;
+    const int nThreadsScatter = WARP_SIZE + ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 2*COLLNET_COPY_THREADS : 0);
+    const int nThreadsGather  =             ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 1*COLLNET_COPY_THREADS : 0);
+    const int nThreadsBcast   = WARP_SIZE + ((hasUp && hasDn) ? COLLNET_COPY_THREADS : hasUp ? 0 : 1*COLLNET_COPY_THREADS);
+    const int nThreadsReduce = args->nWarps*WARP_SIZE - nThreadsScatter - nThreadsGather - nThreadsBcast;
+    const int tidStartBcast = nThreadsGather;
+    const int tidStartScatter = tidStartBcast + nThreadsBcast;
+    const int tidStartReduce = tidStartScatter + nThreadsScatter;
+
+    using Proto = ProtoSimple<1, 1>;
+
+    if (tid >= tidStartScatter && tid < tidStartReduce && hasUp) {
+      // Scatter
+      Primitives<T, RedOp, FanAsymmetric<0, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>
+        prims(tid-tidStartScatter, nThreadsScatter, NULL, direct->up, args->sendbuff, args->recvbuff,
+           args->redOpArg, 2*Proto::MaxGroupWidth, 1, 1, args);
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + bid*direct->nHeads*chunkSize;
+        int nelem = min(direct->nHeads*chunkSize, size-offset);
+        if (args->regUsed) {
+          prims.directScatter(offset, nelem, chunkSize, chunkSize, direct->headRank, direct->shift);
+        } else {
+          prims.scatter(offset, nelem, chunkSize, chunkSize, direct->headRank, direct->shift);
+        }
+      }
+    } else if (tid >= tidStartReduce && direct->out != -1) {
+      if (hasDn) {
+        // Reduce, send to network
+        Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 1>, /*Direct=*/0, Proto, 0>
+          prims(tid-tidStartReduce, nThreadsReduce, direct->down, &direct->out, args->sendbuff, args->recvbuff,
+             args->redOpArg, 3*Proto::MaxGroupWidth, 1, 1, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
+          int nelem = min(chunkSize, size-offset);
+          if (args->regUsed) {
+            prims.directRecvReduceSend(offset, nelem);
+          } else {
+            prims.recvReduceSend(offset, nelem);
+          }
+        }
+      } else {
+        // Directly send to network
+        Primitives<T, RedOp, FanAsymmetric<0, 1>, /*Direct=*/0, Proto, 0>
+          prims(tid-tidStartReduce, nThreadsReduce, nullptr, &direct->out, args->sendbuff, args->recvbuff,
+             args->redOpArg, 3*Proto::MaxGroupWidth, 1, 1);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
+          int nelem = min(chunkSize, size-offset);
+          prims.send(offset, nelem);
+        }
+      }
+    } else if (tid < tidStartBcast && hasUp) {
+      // Gather
+      Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DIRECT_ARITY, 0>, /*Direct=*/0, Proto, 0>
+        prims(tid, nThreadsGather, direct->up, NULL, args->sendbuff, args->recvbuff,
+           args->redOpArg, 0*Proto::MaxGroupWidth, 0, 0, args);
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + bid*direct->nHeads*chunkSize;
+        int nelem = min(direct->nHeads*chunkSize, size-offset);
+        prims.directGather(offset, nelem, chunkSize, chunkSize, direct->headRank, direct->shift);
+      }
+    } else if (tid >= tidStartBcast && tid < tidStartScatter && direct->out != -1) {
+      if (hasDn) {
+        // Recv from network, broadcast
+        Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DIRECT_ARITY>, /*Direct=*/0, Proto, 0>
+          prims(tid-tidStartBcast, nThreadsBcast, &direct->out, direct->down, args->sendbuff, args->recvbuff,
+             args->redOpArg, 1*Proto::MaxGroupWidth, 0, 0, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
+          int nelem = min(chunkSize, size-offset);
+          prims.recvCopyDirectSend(offset, nelem, /*postOp=*/true);
+        }
+      } else {
+        // Recv from network (no post thread needed)
+        Primitives<T, RedOp, FanAsymmetric<1, 0>, /*Direct=*/0, Proto, 0>
+          prims(tid-tidStartBcast, nThreadsBcast, &direct->out, nullptr, args->sendbuff, args->recvbuff,
+             args->redOpArg, 1*Proto::MaxGroupWidth, 0, 0);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid*direct->nHeads+direct->headRank)*chunkSize;
+          int nelem = min(chunkSize, size-offset);
+          prims.recv(offset, nelem, /*postOp=*/true);
+        }
+      }
+    }
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, 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*nvls->nHeads*chunkSize;
+    const int nranks = ncclShmem.comm.nRanks;
+    const bool hasOut = nvls->out != -1;
+    const int totalWarps = NCCL_MAX_NTHREADS/WARP_SIZE;
+    const int bcastWarps = hasOut ? (args->regUsed ? ((totalWarps - 2) >> 1) - 1 : 2) : 0;
+    const int reduceWarps = args->regUsed ? (totalWarps - bcastWarps - 2) : (hasOut ? 3 : nranks <= 6 ? 7 : 5);
+    const int scatterWarps = args->regUsed ? 1 : (totalWarps - reduceWarps - bcastWarps + 1) >> 1;
+    const int gatherWarps = args->regUsed ? 1 : (totalWarps - reduceWarps - bcastWarps) >> 1;
+
+    const int nThreadsScatter = scatterWarps*WARP_SIZE;
+    const int nThreadsGather  = gatherWarps*WARP_SIZE;
+    const int nThreadsReduce = reduceWarps*WARP_SIZE;
+    const int nThreadsBcast  = (bcastWarps)*WARP_SIZE;
+    const int tidEndScatter = nThreadsScatter;
+    const int tidEndGather = tidEndScatter + nThreadsGather;
+    const int tidEndReduce = tidEndGather + nThreadsReduce;
+    const int tidEndBcast = tidEndReduce + nThreadsBcast;
+
+    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 * nvls->nHeads * chunkSize;
+        int nelem = args->regUsed ? 0 : min(nvls->nHeads * chunkSize, size - offset);
+        prims.scatter(offset, nelem, chunkSize, chunkSize, -1, 0);
+      }
+    } else if (tid < tidEndGather) {
+      // Gather
+      using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
+      Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_NVLS_ARITY, 0>, /*Direct=*/0, Proto, 0>
+        prims(tid - tidEndScatter, nThreadsGather, nvls->up, NULL, NULL, args->recvbuff,
+          args->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1);
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + bid * nvls->nHeads * chunkSize;
+        int nelem = args->regUsed ? 0 :min(nvls->nHeads * chunkSize, size - offset);
+        prims.gather(offset, nelem, chunkSize, chunkSize, -1, 0);
+      }
+    } else if (tid < tidEndReduce && nvls->headRank != -1) {
+      if (!hasOut) {
+        // Reduce, broadcast through NVLS
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 1>;
+        Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
+          prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->down, NULL, NULL,
+            args->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
+          int nelem = min(chunkSize, size - offset);
+          prims.directRecvDirectSend(offset, offset, nelem);
+        }
+      } else {
+        // Reduce, send to network
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
+        Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
+          prims(tid - tidEndGather, nThreadsReduce, &nvls->down, &nvls->out, NULL, NULL,
+            args->redOpArg, 2 * Proto::MaxGroupWidth, 0, 1, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
+          int nelem = min(chunkSize, size - offset);
+          prims.directRecvDirectSend(offset, offset, nelem);
+        }
+      }
+    } else if (tid < tidEndBcast && nvls->headRank != -1) {
+      // Recv from network, broadcast
+      using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
+      Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
+        prims(tid - tidEndReduce, nThreadsBcast, &nvls->out, &nvls->down, NULL, NULL,
+          args->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0, args);
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
+        int nelem = min(chunkSize, size - offset);
+        prims.directRecvDirectSend(offset, offset, nelem);
+      }
+    }
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_NVLS_TREE, 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 int treeUp = nvls->treeUp;
+    const int* treeDown = nvls->treeDown;
+    const ssize_t chunkSize = int(args->lastChunkSize);
+    const ssize_t size = args->count;
+    const ssize_t loopSize = nChannels*nvls->nHeads*chunkSize;
+    const int nranks = ncclShmem.comm.nRanks;
+    const bool hasUp = treeUp != -1;
+    const int totalWarps = NCCL_MAX_NTHREADS/WARP_SIZE;
+    const int bcastWarps = hasUp ? (args->regUsed ? ((totalWarps - 2) >> 1) - 1 : 4) : 0;
+    const int reduceWarps = args->regUsed ? (totalWarps - bcastWarps - 2) : (hasUp ? 5 : nranks <= 6 ? 7 : 5);
+    const int scatterWarps = args->regUsed ? 1 : (totalWarps - reduceWarps - bcastWarps + 1) >> 1;
+    const int gatherWarps = args->regUsed ? 1 : (totalWarps - reduceWarps - bcastWarps) >> 1;
+
+    const int nThreadsScatter = scatterWarps*WARP_SIZE;
+    const int nThreadsGather  = gatherWarps*WARP_SIZE;
+    const int nThreadsReduce = reduceWarps*WARP_SIZE;
+    const int nThreadsBcast  = (bcastWarps)*WARP_SIZE;
+    const int tidEndScatter = nThreadsScatter;
+    const int tidEndGather = tidEndScatter + nThreadsGather;
+    const int tidEndReduce = tidEndGather + nThreadsReduce;
+    const int tidEndBcast = tidEndReduce + nThreadsBcast;
+
+    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 * nvls->nHeads * chunkSize;
+        int nelem = args->regUsed ? 0 : min(nvls->nHeads * chunkSize, size - offset);
+        prims.scatter(offset, nelem, chunkSize, chunkSize, -1, 0);
+      }
+    } else if (tid < tidEndGather) {
+      // Gather
+      using Proto = ProtoSimple<1, 1, COLL_UNROLL>;
+      Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_NVLS_ARITY, 0>, /*Direct=*/0, Proto, 0>
+        prims(tid - tidEndScatter, nThreadsGather, nvls->up, NULL, NULL, args->recvbuff,
+          args->redOpArg, 1 * Proto::MaxGroupWidth, 1, 1);
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + bid * nvls->nHeads * chunkSize;
+        int nelem = args->regUsed ? 0 : min(nvls->nHeads * chunkSize, size - offset);
+        prims.gather(offset, nelem, chunkSize, chunkSize, -1, 0);
+      }
+    } else if (tid < tidEndReduce && nvls->headRank != -1) {
+      if (!hasUp) {
+        // Reduce and Broadcast
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 1>;
+        Primitives<T, RedOp, FanSymmetric<3>, /*Direct=*/1, Proto, 0>
+          prims(tid - tidEndGather, nThreadsReduce, treeDown, treeDown, NULL, NULL,
+            args->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
+          int nelem = min(chunkSize, size - offset);
+          prims.directRecvDirectSend(offset, offset, nelem);
+        }
+      } else {
+        // Reduce, send to network
+        using Proto = ProtoSimple<1, 1, COLL_UNROLL, 1, 0>;
+        Primitives<T, RedOp, FanAsymmetric<3, 1>, /*Direct=*/1, Proto, 0>
+          prims(tid - tidEndGather, nThreadsReduce, treeDown, &treeUp, NULL, NULL,
+            args->redOpArg, 2 * Proto::MaxGroupWidth, 0, 0, args);
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
+          int nelem = min(chunkSize, size - offset);
+          prims.directRecvDirectSend(offset, offset, nelem);
+        }
+      }
+    } else if (tid < tidEndBcast && nvls->headRank != -1) {
+      // Recv from network, broadcast
+      using Proto = ProtoSimple<1, 1, COLL_UNROLL, 0, 1>;
+      Primitives<T, RedOp, FanAsymmetric<1, 3>, /*Direct=*/1, Proto, 0>
+        prims(tid - tidEndReduce, nThreadsBcast, &treeUp, treeDown, NULL, NULL,
+          args->redOpArg, 3 * Proto::MaxGroupWidth, 0, 0, args);
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        ssize_t offset = gridOffset + (bid * nvls->nHeads + nvls->headRank) * chunkSize;
+        int nelem = min(chunkSize, size - offset);
+        prims.directRecvDirectSend(offset, offset, nelem);
+      }
+    }
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_COLLNET_CHAIN, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    const int tid = threadIdx.x;
+    const int nthreads = args->nWarps*WARP_SIZE;
+    const int bid = args->bid;
+    const int nChannels = args->nChannels;
+    ncclTree *tree = &ncclShmem.channel.collnetChain;
+    ssize_t chunkSize = int(args->lastChunkSize);
+    const ssize_t loopSize = int(nChannels*chunkSize);
+    const int nranks = ncclShmem.comm.nRanks;
+    const ssize_t size = args->count;
+
+    int nthreadsSplit = nthreads/2;
+    if (nthreadsSplit >= 256) nthreadsSplit += 64;
+
+    int group, connIndex, send, recv, groupTid, groupNthreads;
+    using Proto = ProtoSimple<1, 1>;
+    if (tid < nthreadsSplit) {
+      // Reduce up the chain
+      group = 0;
+      connIndex = 1;
+      recv = tree->down[0];
+      send = tree->up;
+      groupTid = tid;
+      groupNthreads = nthreadsSplit;
+    } else {
+      // Broadcast down the chain
+      group = 1;
+      connIndex = 0;
+      recv = tree->up;
+      send = tree->down[0];
+      groupTid = tid - nthreadsSplit;
+      groupNthreads = nthreads-nthreadsSplit;
+    }
+
+    Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0>
+      prims(groupTid, groupNthreads, &recv, &send, args->sendbuff, args->recvbuff,
+          args->redOpArg, group*Proto::MaxGroupWidth, connIndex, connIndex);
+
+    if (tid < nthreadsSplit) {
+      if (recv == -1) {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.send(offset, nelem);
+        }
+      } else {
+        for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+          ssize_t offset = gridOffset + bid*int(chunkSize);
+          int nelem = min(chunkSize, size-offset);
+          prims.recvReduceSend(offset, nelem);
+        }
+      }
+    }
+    else {
+      if (recv == nranks) {
+        // I'm the first in the broadcast chain, I need to perform the division (postOp)
+        if (send == -1) {
+          for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+            ssize_t offset = gridOffset + bid*int(chunkSize);
+            int nelem = min(chunkSize, size-offset);
+            prims.recv(offset, nelem, /*postOp*/true);
+          }
+        } else {
+          for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+            ssize_t offset = gridOffset + bid*int(chunkSize);
+            int nelem = min(chunkSize, size-offset);
+            prims.recvCopyDirectSend(offset, nelem, /*postOp*/true);
+          }
+        }
+      } else {
+        if (send == -1) {
+          for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+            ssize_t offset = gridOffset + bid*int(chunkSize);
+            int nelem = min(chunkSize, size-offset);
+            prims.directRecv(offset, nelem);
+          }
+        } else {
+          for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+            ssize_t offset = gridOffset + bid*int(chunkSize);
+            int nelem = min(chunkSize, size-offset);
+            prims.directRecvCopySend(offset, nelem);
+          }
+        }
+      }
+    }
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, 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<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runTreeSplit<T, RedOp, ProtoLL>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runRing<T, RedOp, ProtoLL128>(args);
+    //LAUNCH_CLIQUE_KERNEL(AllReduceCliqueSplitKernel, RedOp, T, args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL128> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runTreeSplit<T, RedOp, ProtoLL128>(args);
+    //LAUNCH_CLIQUE_KERNEL(AllReduceCliqueSplitKernel, RedOp, T, args);
+  }
+};

src/device/alltoall_pivot.h

@@ -0,0 +1,82 @@
+/*************************************************************************
+ * Copyright (c) 2015-2021, NVIDIA CORPORATION. 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 nranks = ncclShmem.comm.nRanks;
+    const ncclRing *ring = &ncclShmem.channel.ring;
+    const int num_bi_rings = args->pivotA2ANumBiRings;
+    const int num_uni_rings = num_bi_rings * 2;
+    const int num_chunks = args->nChannels / 2;
+    const int chunk_id = (bid % num_bi_rings) + (bid / num_uni_rings * num_bi_rings);
+    const int elem_size = min(256, args->count & (~(args->count) + 1));
+    const ssize_t num_elems = args->count / elem_size;
+    const int num_padding_chunks = num_elems % num_chunks;
+    const ssize_t chunk_offset = elem_size * (num_elems / num_chunks * chunk_id + (chunk_id < num_padding_chunks ? chunk_id : num_padding_chunks));
+    const ssize_t chunk_size = elem_size * (num_elems / num_chunks + (chunk_id < num_padding_chunks ? 1 : 0));
+    const int pivot_direction = (bid % num_uni_rings) / num_bi_rings;
+    const ssize_t prims_size = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? ALLTOALL_PIVOT_CHUNKSTEPS : 1));
+
+    Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0> prims
+      (tid, nthreads, &ring->prev, &ring->next, args->sendbuff, args->recvbuff, /*redOpArg(ignored)=*/0);
+
+    for (int num_hops = 0; num_hops <= nranks / 2; num_hops++) {
+      const int src_rank = ring->userRanks[(nranks - num_hops) % nranks];
+      const int dst_rank = ring->userRanks[num_hops];
+      const ssize_t send_offset =
+          dst_rank * num_elems * elem_size + chunk_offset +
+          (src_rank == dst_rank ? pivot_direction * chunk_size / 2 : 0);
+      const ssize_t recv_offset =
+          src_rank * num_elems * elem_size + chunk_offset +
+          (src_rank == dst_rank ? pivot_direction * chunk_size / 2 : 0);
+      const ssize_t send_recv_size =
+          src_rank == dst_rank ?
+          (pivot_direction == 0 ? chunk_size / 2 : chunk_size - chunk_size / 2) : chunk_size;
+
+      if (num_hops == 0 && args->sendbuff != args->recvbuff) {
+        const T* sendbuff = (const T*)args->sendbuff + send_offset;
+        T* recvbuff = (T *)args->recvbuff + recv_offset;
+        reduceCopy<COLL_UNROLL, RedOp, T, 0,1, 1, 0, 1, 1, 0>(
+            tid, nthreads, 0, nullptr, false, 1, (void **)&sendbuff, 1, (void **)&recvbuff, send_recv_size);
+      } else {
+        for (ssize_t prims_offset = 0; prims_offset < send_recv_size; prims_offset += prims_size) {
+          const int prims_nelem = min(prims_size, send_recv_size - prims_offset);
+
+          // step 0: send
+          prims.send(send_offset + prims_offset, prims_nelem);
+
+          // num_hops - 1 steps: recv and copy to next gpu
+          for (int i = 0; i < num_hops - 1; i++) {
+            prims.recvSend(prims_nelem);
+          }
+
+          // final step: recv
+          prims.directRecv(recv_offset + prims_offset, prims_nelem);
+        }
+      }
+    }
+  }
+}
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncAllToAllPivot, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    using Proto = ProtoSimple<ALLTOALL_PIVOT_CHUNKSTEPS/ALLTOALL_PIVOT_SLICESTEPS, ALLTOALL_PIVOT_SLICESTEPS>;
+    runRing<T, RedOp, Proto>(args);
+  }
+};

src/device/broadcast.h

@@ -0,0 +1,111 @@
+/*************************************************************************
+ * Copyright (c) 2015-2022, NVIDIA CORPORATION. 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;
+    const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? BROADCAST_CHUNKSTEPS : 1));
+    const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
+    const ssize_t loopSize = nChannels*chunkSize;
+    const ssize_t size = args->count;
+    const int rank = ring->userRanks[0];
+    const int nextRank = ring->userRanks[1];
+    const int root = args->root;
+
+#if defined(ENABLE_NPKIT)
+    int npKitCtxIdx = bid;
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (tid == 0) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    T *inputBuf = (T*)args->sendbuff;
+    T *outputBuf = (T*)args->recvbuff;
+    Primitives<T, RedOp, FanSymmetric<1>, 0, Proto, 0>
+      prims(tid, nthreads, &ring->prev, &ring->next, inputBuf, outputBuf, args->redOpArg, 0, args->connIndex, args->connIndex);
+
+#if defined(ENABLE_NPKIT)
+    if (tid == 0) {
+      prims.npKitCtxIdx = npKitCtxIdx;
+    }
+#endif
+
+    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(chunkSize, divUp(size-gridOffset, nChannels*minChunkSizeLL128)*minChunkSizeLL128);
+      realChunkSize = int(realChunkSize);
+
+      ssize_t offset = gridOffset + int(bid*realChunkSize);
+      int nelem = min(realChunkSize, size-offset);
+
+      if (rank == root) {
+        if (inputBuf == outputBuf) {
+          prims.send(offset, nelem);
+        } else {
+          prims.copySend(offset, offset, nelem);
+        }
+      } else if (nextRank == root) {
+        prims.recv(offset, nelem);
+      } else {
+        prims.recvCopySend(offset, nelem);
+      }
+    }
+  }
+}
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    using Proto = ProtoSimple<BROADCAST_CHUNKSTEPS/BROADCAST_SLICESTEPS, BROADCAST_SLICESTEPS>;
+    runRing<T, RedOp, Proto>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncBroadcast, 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<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runRing<T, RedOp, ProtoLL128>(args);
+  }
+};

src/device/common.cu

@@ -0,0 +1,29 @@
+/*************************************************************************
+ * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "device.h"
+#include "collectives.h"
+#include "common.h"
+
+__shared__ ncclShmemData ncclShmem;
+#if __CUDA_ARCH__ < 700
+  __shared__ ulong2 ncclShmemPerWarp[ncclShmemScratchWarpSize()*(NCCL_MAX_NTHREADS/WARP_SIZE)/sizeof(ulong2)];
+#endif
+
+struct RunWorkNop {
+  __device__ void run(ncclWork *w) {}
+};
+
+__global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
+  ncclKernelMain<-1, RunWorkNop, false>(comm, channelMask, workHead);
+}
+#ifdef ENABLE_COLLTRACE
+__global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
+  ncclKernelMain<-1, RunWorkNop, true>(comm, channelMask, workHead);
+}
+#endif
+
+__device__ void ncclDevFunc_Nop() {}

src/device/common.h

@@ -0,0 +1,402 @@
+/*************************************************************************
+ * Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2023 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_DEVICE_COMMON_H_
+#define NCCL_DEVICE_COMMON_H_
+
+#include "collectives.h"
+#include "device.h"
+#include "op128.h"
+#include "device_table.h"
+#include "network/unpack/unpack_defs.h"
+
+#if defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+#define COLL_UNROLL 2
+#else
+#define COLL_UNROLL 4
+#endif
+
+#define NCCL_MAX_DEV_ARITY (NCCL_MAX_TREE_ARITY-1)  // Using balanced tree instead of split tree
+
+#define __syncwarp()
+
+#define __synclds() \
+  asm volatile("s_waitcnt lgkmcnt(0) \n s_barrier");
+
+#ifdef __GFX9__
+#define STORE(DST, SRC) \
+  { __atomic_store_n((DST), (SRC), __ATOMIC_RELAXED); }
+#else
+#define STORE(DST, SRC) \
+  { __atomic_store_n((DST), (SRC), __ATOMIC_SEQ_CST); }
+#endif
+
+#if defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__)
+#define __trace_hwreg()
+#else
+#define __trace_hwreg() \
+  asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_HW_ID)" : "=s" (collTrace->data_0));
+#endif
+#ifdef ENABLE_COLLTRACE
+  #define INC_COLL_TRACE \
+    uint32_t pos = atomicAdd(&ncclShmem.collTraceTail->tail, 1)%COLLTRACE_NUM_ITEMS; \
+    struct ncclCollTrace* collTrace = ncclShmem.collTrace+pos; \
+    collTrace->timeStamp = wall_clock64(); \
+    collTrace->bid = blockIdx.x;
+    // TODO: switch to atomicInc after llvm crash is fixed
+    // uint32_t pos = atomicInc(&ncclShmem.collTraceTail->tail, COLLTRACE_NUM_ITEMS)
+
+  #define traceKernelLaunch(launch_type) { \
+    INC_COLL_TRACE \
+    collTrace->funcIndex = ncclShmem.work.header.funcIndex; \
+    __trace_hwreg()\
+    if (ncclShmem.work.header.type == ncclWorkTypeP2p) { \
+      struct ncclWorkElemP2p *p2pElems = ncclShmem.work.p2pElems; \
+      collTrace->p2p[0].connIndex = 0; \
+      collTrace->p2pOpCount[0] = p2pElems[0].opCount; \
+      collTrace->p2p[0].ngroups = p2pElems[0].ngroups; \
+      collTrace->p2p[0].nWarps = p2pElems[0].nWarps; \
+      collTrace->p2p[0].warpStart = p2pElems[0].warpStart; \
+      collTrace->p2p[0].peer = p2pElems[0].p2pType == ncclWorkP2pTypeRecv ? (uint16_t)(p2pElems[0].peer) : -1; \
+      collTrace->p2p[1].connIndex = 0; \
+      collTrace->p2pOpCount[1] = p2pElems[1].opCount; \
+      collTrace->p2p[1].ngroups = p2pElems[1].ngroups; \
+      collTrace->p2p[1].nWarps = p2pElems[1].nWarps; \
+      collTrace->p2p[1].warpStart = p2pElems[1].warpStart; \
+      collTrace->p2p[1].peer = p2pElems[1].p2pType == ncclWorkP2pTypeSend ? (uint16_t)(p2pElems[1].peer) : -1; \
+      collTrace->type = (launch_type) | ncclCollTraceP2pElemType; \
+    } else if (ncclShmem.work.header.type == ncclWorkTypeColl) { \
+      struct ncclWorkElem *elems = ncclShmem.work.elems; \
+      collTrace->opCount = elems[0].opCount; \
+      collTrace->coll.nWarps = elems[0].nWarps; \
+      collTrace->coll.bid = elems[0].bid; \
+      collTrace->coll.nChannels = elems[0].nChannels; \
+      collTrace->type = (launch_type) | ncclCollTraceCollElemType; \
+    } \
+  }
+  #define traceKernelEnd(end_type)  { \
+    INC_COLL_TRACE \
+    if (ncclShmem.work.header.type == ncclWorkTypeP2p) { \
+      struct ncclWorkElemP2p *p2pElems = ncclShmem.work.p2pElems; \
+      collTrace->p2pOpCount[0] = p2pElems[0].opCount; \
+      collTrace->p2pOpCount[1] = p2pElems[1].opCount; \
+    } else if (ncclShmem.work.header.type == ncclWorkTypeColl) { \
+      struct ncclWorkElem *elems = ncclShmem.work.elems; \
+      collTrace->opCount = elems[0].opCount; \
+    } \
+    collTrace->type = end_type; \
+  }
+  #define traceData(data2, data4, data8_0, data8_1) { \
+    INC_COLL_TRACE \
+    collTrace->funcIndex = data2; \
+    collTrace->data_0 = data4; \
+    collTrace->opCount = data8_0; \
+    collTrace->data_1 = data8_1; \
+    collTrace->type = ncclCollTraceDataType; \
+  }
+#else
+#define traceKernelLaunch(launch_type)
+#define traceKernelEnd(end_type)
+#define traceData(data2, data4, data8_0, data8_1)
+#endif
+
+struct ncclShmemGroup {
+  ncclConnInfo *recvConns[NCCL_MAX_NVLS_ARITY];
+  ncclConnInfo *sendConns[NCCL_MAX_NVLS_ARITY];
+  void* srcs[NCCL_MAX_NVLS_ARITY+1];
+  void* dsts[NCCL_MAX_NVLS_ARITY+1];
+  uint64_t barrier;
+  uint64_t barrier_next[NCCL_MAX_GROUPS];
+  union {
+    unpackGroupShmem unpack;
+  } devicePlugin;
+};
+
+#define LDS_NUM_EVENTS 64
+
+struct ncclShmemData {
+  struct ncclShmemGroup groups[NCCL_MAX_GROUPS];
+  uint64_t redOpArgs[NCCL_MAX_NVLS_ARITY+1];
+  int channelId;
+  int aborted;
+  alignas(16) struct ncclDevComm comm;
+  alignas(16) struct ncclDevChannel channel;
+  alignas(16) struct ncclWork work;
+  alignas(16) union {
+    unpackShmem unpack;
+  } devicePlugin;
+#ifdef ENABLE_COLLTRACE
+  struct ncclCollTrace* collTrace;
+  union ncclCollTraceTail* collTraceTail;
+#endif
+#ifdef ENABLE_PROFILING
+  struct ncclProf prof;
+#endif
+#if defined(ENABLE_NPKIT)
+  NpKitEvent event_buffer[LDS_NUM_EVENTS];
+  uint64_t event_buffer_head;
+#endif
+};
+static_assert(offsetof(struct ncclShmemData, work)%16 == 0, "ncclShmem.work needs to be 16B aligned");
+
+extern __shared__ ncclShmemData ncclShmem;
+#if __CUDA_ARCH__ >= 700
+  extern __shared__ ulong2 ncclShmemPerWarp[/*ncclShmemDynamicSize()/sizeof(ulong2)*/];
+#else
+  extern __shared__ ulong2 ncclShmemPerWarp[ncclShmemScratchWarpSize()*(NCCL_MAX_NTHREADS/WARP_SIZE)/sizeof(ulong2)];
+#endif
+
+__device__ inline void* ncclScratchForWarp(int warp) {
+  return (char*)ncclShmemPerWarp + warp*ncclShmemScratchWarpSize();
+}
+
+#ifdef ENABLE_PROFILING
+#define __insert_timestamp(line_num) do { \
+      if (ncclShmem.prof.count < PROFILE_NUM_ITEMS) { \
+        ncclShmem.prof.elem[ncclShmem.prof.count].line = line_num; \
+        ncclShmem.prof.elem[ncclShmem.prof.count].timeStamp = wall_clock64(); \
+        ncclShmem.prof.count++; \
+      } \
+    } while(0);
+#else
+#define __insert_timestamp(line_num)
+#endif
+
+// Copy 16-byte aligned data. You must call with at least `(bytes+15)/16` threads.
+inline __device__ void copyToShmem16(int tid, void* dst, void const* src, int bytes) {
+  int offset = 16*tid;
+  if (offset < bytes) {
+    ulong2 *src2, *dst2;
+    src2 = (ulong2*)((char const*)src + offset);
+    dst2 = (ulong2*)((char*)dst + offset);
+    dst2->x = src2->x;
+    dst2->y = src2->y;
+  }
+}
+
+template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
+struct RunWorkElement {
+  __device__ void run(ncclWorkElem*) {
+    // Put NOT IMPLEMENTED behavior here.
+  }
+};
+
+template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto>
+struct RunWork {
+  // This __forceinline__ is necessary. The compiler was inserting a function call
+  // here from the LL ncclKernel.
+  __device__ __forceinline__ void run(ncclWork *w) {
+    int wid = threadIdx.x / WARP_SIZE;
+    ncclWorkElem* we = w->header.type == ncclWorkTypeRegColl ? &w->regElems[0].elem : &w->elems[0];
+    int stride = w->header.type == ncclWorkTypeRegColl ? sizeof(ncclWorkElemReg) : sizeof(ncclWorkElem);
+    #pragma unroll 1
+    while ((char*)we + stride <= (char*)(w+1) && we->isUsed) {
+      if (wid < we->nWarps) {
+        RunWorkElement<Fn, T, RedOp, Algo, Proto>().run(we);
+      }
+      we = (ncclWorkElem*)((char*)we + stride);
+    }
+  }
+};
+
+static __forceinline__ __device__ void ncclRedopPtrDeref(struct ncclWorkElem* we) {
+  if (we->isUsed && we->redOpArgIsPtr) {
+    /* redOpArg is a pointer to the scalar value, so we'll dereference it
+     * here so that redOpArg holds the bits of the scalar going forward.
+     * The tricky thing is we don't know its type T since that's encoded in
+     * the funcIndex. Because it would be difficult to get sizeof(T) from
+     * funcIndex, we'll cheat and just dereference the largest possible size
+     * given the alignment of the pointer. We might be reading in more bytes
+     * than we need but that's harmless.
+     */
+    if (we->redOpArg%2 != 0)
+      we->redOpArg = *reinterpret_cast<uint8_t*>(we->redOpArg);
+    else if (we->redOpArg%4 != 0)
+      we->redOpArg = *reinterpret_cast<uint16_t*>(we->redOpArg);
+    else if (we->redOpArg%8 != 0)
+      we->redOpArg = *reinterpret_cast<uint32_t*>(we->redOpArg);
+    else
+      we->redOpArg = *reinterpret_cast<uint64_t*>(we->redOpArg);
+  }
+}
+
+template<int SpecializedFnId, typename SpecializedRunWork, bool COLLTRACE>
+__device__ void ncclKernelMain(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) {
+  const int tid = threadIdx.x;
+  int x = tid;
+  switch (tid/WARP_SIZE) {
+  case 0:
+    if (channelMask & (1ull<<x)) {
+      int y = __popcll(channelMask & ((1ull<<x)-1));
+      if (blockIdx.x == y) ncclShmem.channelId = x;
+    }
+    if (WARP_SIZE < MAXCHANNELS) {
+      x = WARP_SIZE + tid;
+      if (channelMask & (1ull<<x)) {
+        int y = __popcll(channelMask & ((1ull<<x)-1));
+        if (blockIdx.x == y) ncclShmem.channelId = x;
+      }
+    }
+    break;
+  case 1:
+    if (tid < WARP_SIZE + NCCL_MAX_GROUPS)
+      ncclShmem.groups[tid-WARP_SIZE].barrier = 0;
+    break;
+  case 2:
+    if (tid < 2*WARP_SIZE + NCCL_MAX_GROUPS*NCCL_MAX_GROUPS)
+      ncclShmem.groups[(tid-2*WARP_SIZE)/NCCL_MAX_GROUPS].barrier_next[(tid-2*WARP_SIZE)%NCCL_MAX_GROUPS] = 0;
+    break;
+  case 3:
+    /* set abort flag to 0 */
+    if (tid == 3*WARP_SIZE) ncclShmem.aborted = 0;
+    break;
+  default:
+    break;
+  }
+  __synclds(); // publish ncclShmem.channelId
+  // To map blockId to channelId, we need the n'th set bit of channelMask which
+  // is the inverse of counting the number of set bits among the the first n.
+  int channelId = ncclShmem.channelId;
+
+  if (true) {
+    void *dst, *src;
+    int bytes;
+    // Use first 3 warps to load comm, channel, and work into shmem
+    switch (tid/WARP_SIZE) {
+    case 0:
+      dst = &ncclShmem.comm;
+      src = comm;
+      bytes = sizeof(ncclDevComm);
+      static_assert(sizeof(ncclDevComm) <= 16*WARP_SIZE, "ncclDevComm cannot be loaded by a single warp in one insn.");
+      break;
+    case 1:
+      // Get address of channel without incurring indirect load from ncclDevComm::channels
+      dst = &ncclShmem.channel;
+      src = &((ncclDevCommAndChannels*)comm)->channels[channelId];
+      bytes = sizeof(ncclDevChannel);
+      static_assert(sizeof(ncclDevChannel) <= 16*WARP_SIZE, "ncclDevChannel cannot be loaded by a single warp in one insn.");
+      break;
+    case 2:
+      dst = &ncclShmem.work;
+      src = workHead + blockIdx.x;
+      bytes = sizeof(ncclWork);
+      static_assert(sizeof(ncclWork) <= 16*WARP_SIZE, "ncclWork cannot be loaded by a single warp in one insn.");
+      break;
+    default:
+      bytes = 0;
+      break;
+    }
+    if (bytes) copyToShmem16(tid%WARP_SIZE, dst, src, bytes);
+  }
+#ifdef ENABLE_COLLTRACE
+  if (tid == 0) {
+    ncclShmem.collTrace = comm->collTrace + COLLTRACE_NUM_ITEMS*ncclShmem.channelId;
+    ncclShmem.collTraceTail = comm->collTraceTail + ncclShmem.channelId;
+  }
+#endif
+  __synclds(); // publish shmem
+#ifdef ENABLE_PROFILING
+  if (tid == 0) {
+    ncclShmem.prof.count = 0;
+    ncclShmem.prof.seq = ncclShmem.comm.devProf[blockIdx.x].seq;
+  }
+#endif
+  if (tid == 0) __insert_timestamp(__LINE__);
+  if (COLLTRACE && tid == 0) traceKernelLaunch(ncclCollTraceKernelLaunchType);
+
+  while (true) {
+    // Notify host that all fifo reads are complete.
+    if (tid == 0 && ncclShmem.work.header.isLast && ncclShmem.work.header.inFifo) {
+      *ncclShmem.channel.workFifoDone = ncclShmem.work.header.doneAcks;
+    }
+
+    __syncwarp();
+    if (ncclShmem.work.header.type == ncclWorkTypeColl) {
+      if (tid < NCCL_MAX_WORK_ELEMENTS) ncclRedopPtrDeref(&ncclShmem.work.elems[tid]);
+    } else if (ncclShmem.work.header.type == ncclWorkTypeRegColl) {
+      if (tid < NCCL_MAX_WORK_ELEMENTS_REG) ncclRedopPtrDeref(&ncclShmem.work.regElems[tid].elem);
+    }
+    __synclds();
+
+    if (tid == 0) __insert_timestamp(__LINE__);
+
+    if (0 <= SpecializedFnId && ncclShmem.work.header.funcIndex == (unsigned)SpecializedFnId) {
+      SpecializedRunWork().run(&ncclShmem.work);
+    } else {
+#ifdef USE_INDIRECT_FUNCTION_CALL
+      ncclDevFuncTable[ncclShmem.work.header.funcIndex]();
+#else
+      NCCL_CALL_FUNCTIONS(ncclShmem.work.header.funcIndex);
+#endif
+    }
+
+    int workIxNext = ncclShmem.work.header.workNext;
+    __synclds();
+    if (ncclShmem.work.header.isLast) break;
+
+    copyToShmem16(tid, &ncclShmem.work, workHead + workIxNext, sizeof(ncclWork));
+
+    { // Check whether the last operation was aborted and make sure all threads exit
+      int aborted = tid == 0 ? *comm->abortFlag : 0;
+      if (__any(aborted)) { // publish ncclShmem.work
+        traceKernelEnd(ncclCollTraceAbortType);
+        break;
+      }
+    }
+    if (COLLTRACE && tid == 0) traceKernelLaunch(ncclCollTraceCollLaunchType);
+  }
+  if (COLLTRACE && tid == 0) traceKernelEnd(ncclCollTraceKernelEndType);
+
+#ifdef ENABLE_PROFILING
+  if (ncclShmem.comm.devProf->seq < PROFILE_NUM_LAUNCHES) {
+    __synclds();
+    copyToShmem16(tid, ncclShmem.comm.devProf+MAXCHANNELS*ncclShmem.prof.seq+blockIdx.x, &ncclShmem.prof, sizeof(struct ncclProf));
+    if (tid == 0) ncclShmem.comm.devProf[blockIdx.x].seq++;
+  }
+#endif
+}
+
+__global__ void ncclDevKernel_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);
+#ifdef ENABLE_COLLTRACE
+__global__ void ncclDevKernelDebug_Generic(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead);
+#endif
+
+#ifdef USE_INDIRECT_FUNCTION_CALL
+__device__ void ncclDevFunc_Nop();
+#else
+__device__ __attribute__((noinline)) void ncclDevFunc_Nop();
+#endif
+
+#ifdef ENABLE_COLLTRACE
+#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
+  __global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \
+    ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, false>(comm, channelMask, workHead); \
+  } \
+  \
+  __global__ void ncclDevKernelDebug_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \
+    ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, true>(comm, channelMask, workHead); \
+  }
+#else
+#define DEFINE_ncclDevKernel(suffix, coll, redop, ty, algo, proto, specializedFnId) \
+  __global__ void ncclDevKernel_##suffix(struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead) { \
+    ncclKernelMain<specializedFnId, RunWork<coll, ty, redop<ty>, algo, proto>, false>(comm, channelMask, workHead); \
+  }
+#endif
+
+#ifdef USE_INDIRECT_FUNCTION_CALL
+#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
+  __device__ void ncclDevFunc_##suffix() { \
+    RunWork<coll, ty, redop<ty>, algo, proto>().run(&ncclShmem.work); \
+  }
+#else
+#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto) \
+  __device__ __attribute__((noinline)) void ncclDevFunc_##suffix() { \
+    RunWork<coll, ty, redop<ty>, algo, proto>().run(&ncclShmem.work); \
+  }
+#endif
+
+#endif

src/device/common_kernel.h

@@ -0,0 +1,264 @@
+/*************************************************************************
+ * 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
+ ************************************************************************/
+
+#ifndef NCCL_COMMON_KERNEL_H_
+#define NCCL_COMMON_KERNEL_H_
+
+#include "device.h"
+#include "op128.h"
+#include "reduce_kernel.h"
+#include <cstdio>
+#include <cstdint>
+
+#include <hip/hip_runtime.h>
+
+#define __syncwarp()
+
+// Define min for ssize_t
+inline __device__ int min(int a, ssize_t b) { return (a < b) ? a : b; }
+
+inline __device__ int loadInt(int* ptr) {
+  int v;
+  v = atomicAdd((unsigned long long *)ptr, 0);
+  return v;
+}
+
+template<typename RedFn, typename T, int Unroll, int BytePerPack,
+         int MultimemSrcs, int MinSrcs, int MaxSrcs,
+         int MultimemDsts, int MinDsts, int MaxDsts, int PreOpSrcs,
+         typename IntBytes>
+__device__ __forceinline__ void reduceCopyPacks(
+    int nThreads, int &thread,
+    uint64_t redArg, uint64_t *preOpArgs, bool postOp,
+    int nSrcs, void **srcPtrs, int nDsts, void **dstPtrs,
+    IntBytes &nBytesBehind, IntBytes &nBytesAhead
+  ) {
+  static_assert(std::is_signed<IntBytes>::value, "IntBytes must be a signed integral type.");
+  //if (BytePerPack == 0) __trap();
+
+  // A hunk is the amount of contiguous data a warp consumes per loop iteration
+  // assuming all threads partake.
+  constexpr int BytePerHunk = Unroll*WARP_SIZE*BytePerPack;
+  int nWarps = nThreads/WARP_SIZE;
+  int warp = thread/WARP_SIZE;
+  int lane = thread%WARP_SIZE;
+
+  // This thread's initial position.
+  IntBytes threadBytesBehind = nBytesBehind + (warp*BytePerHunk + lane*BytePerPack);
+  IntBytes threadBytesAhead = nBytesAhead - (warp*BytePerHunk + lane*BytePerPack);
+  // Number of hunks to be consumed over all warps.
+  IntBytes nHunksAhead = nBytesAhead/(BytePerHunk + !BytePerHunk);
+  // Advance collective position.
+  nBytesBehind += nHunksAhead*BytePerHunk;
+  nBytesAhead -= nHunksAhead*BytePerHunk;
+  if (Unroll==1 && BytePerPack <= nBytesAhead) {
+    // Only Unroll=1 can do partial hunks (where not all threads partake).
+    nHunksAhead += 1;
+    nBytesBehind += nBytesAhead - (nBytesAhead%(BytePerPack + !BytePerPack));
+    nBytesAhead = nBytesAhead%(BytePerPack + !BytePerPack);
+  }
+  nHunksAhead -= warp;
+
+  RedFn redFn(redArg);
+  uintptr_t minSrcs[MinSrcs + !MinSrcs];
+  uintptr_t minDsts[MinDsts + !MinDsts];
+  #pragma unroll
+  for (int s=0; s < MinSrcs; s++)
+    minSrcs[s] = cvta_to_global(srcPtrs[s]) + threadBytesBehind;
+  #pragma unroll
+  for (int d=0; d < MinDsts; d++)
+    minDsts[d] = cvta_to_global(dstPtrs[d]) + threadBytesBehind;
+
+  // We dictate loop termination condition according to whether partial hunks
+  // can be handled or not.
+  while (Unroll==1 ? (BytePerPack <= threadBytesAhead) : (0 < nHunksAhead)) {
+    BytePack<BytePerPack> acc[Unroll];
+
+    { RedFn preFn(0 < PreOpSrcs ? preOpArgs[0] : 0);
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        if (0 < MultimemSrcs) {
+          // applyLoadMultimem uses relaxed semantics for same reason we use volatile below.
+          acc[u] = applyLoadMultimem<RedFn, BytePerPack>(redFn, minSrcs[0]);
+        } else {
+          // Use volatile loads in case credits are polled for with volatile (instead of acquire).
+          acc[u] = ld_volatile_global<BytePerPack>(minSrcs[0]);
+          if (0 < PreOpSrcs) acc[u] = applyPreOp(preFn, acc[u]);
+        }
+        minSrcs[0] += WARP_SIZE*BytePerPack;
+      }
+    }
+
+    #pragma unroll Unroll
+    for (int s=1; s < MinSrcs; s++) {
+      BytePack<BytePerPack> tmp[Unroll];
+      RedFn preFn(s < PreOpSrcs ? preOpArgs[s] : 0);
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        if (s < MultimemSrcs) {
+          // applyLoadMultimem uses relaxed semantics for same reason we use volatile below.
+          acc[u] = applyLoadMultimem<RedFn, BytePerPack>(redFn, minSrcs[s]);
+        } else {
+          // Use volatile loads in case credits are polled for with volatile (instead of acquire).
+          tmp[u] = ld_volatile_global<BytePerPack>(minSrcs[s]);
+        }
+        minSrcs[s] += WARP_SIZE*BytePerPack;
+      }
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        if (s < PreOpSrcs) tmp[u] = applyPreOp(preFn, tmp[u]);
+        acc[u] = applyReduce(redFn, acc[u], tmp[u]);
+      }
+    }
+
+    for (int s=MinSrcs; (MinSrcs < MaxSrcs) && (s < MaxSrcs) && (s < nSrcs); s++) {
+      uintptr_t src = cvta_to_global(srcPtrs[s]) + threadBytesBehind;
+      BytePack<BytePerPack> tmp[Unroll];
+      RedFn preFn(s < PreOpSrcs ? preOpArgs[s] : 0);
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        // Use volatile loads in case credits are polled for with volatile (instead of acquire).
+        tmp[u] = ld_volatile_global<BytePerPack>(src);
+        src += WARP_SIZE*BytePerPack;
+      }
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        if (s < PreOpSrcs) tmp[u] = applyPreOp(preFn, tmp[u]);
+        acc[u] = applyReduce(redFn, acc[u], tmp[u]);
+      }
+    }
+
+    if (postOp) {
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++)
+        acc[u] = applyPostOp(redFn, acc[u]);
+    }
+
+    #pragma unroll Unroll
+    for (int d=0; d < MinDsts; d++) {
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        if (d < MultimemDsts) {
+          multimem_st_global(minDsts[d], acc[u]);
+        } else {
+          st_global<BytePerPack>(minDsts[d], acc[u]);
+        }
+        minDsts[d] += WARP_SIZE*BytePerPack;
+      }
+    }
+    for (int d=MinDsts; (MinDsts < MaxDsts) && (d < MaxDsts) && (d < nDsts); d++) {
+      uintptr_t dst = cvta_to_global(dstPtrs[d]) + threadBytesBehind;
+      #pragma unroll Unroll
+      for (int u=0; u < Unroll; u++) {
+        st_global<BytePerPack>(dst, acc[u]);
+        dst += WARP_SIZE*BytePerPack;
+      }
+    }
+
+    nWarps = nThreads/WARP_SIZE;
+    #pragma unroll
+    for (int s=0; s < MinSrcs; s++) minSrcs[s] += (nWarps-1)*BytePerHunk;
+    #pragma unroll
+    for (int d=0; d < MinDsts; d++) minDsts[d] += (nWarps-1)*BytePerHunk;
+    threadBytesBehind += nWarps*BytePerHunk;
+    threadBytesAhead -= nWarps*BytePerHunk;
+    nHunksAhead -= nWarps;
+  }
+
+  nWarps = nThreads/WARP_SIZE;
+  warp = thread/WARP_SIZE;
+  lane = thread%WARP_SIZE;
+  // The last loop iteration could have been partial, i.e. not taken by all
+  // threads. The threads that weren't included need an extra subtraction to
+  // make the value warp uniform.
+  if (Unroll==1 && nHunksAhead > 0) nHunksAhead -= nWarps;
+  // Rotate warps so the warp which got the least work here will be warp 0.
+  // This effectively assigns: warp = (warp-nHunks+nWarps)%nWarps;
+  warp = -nHunksAhead;
+  thread = warp*WARP_SIZE + lane;
+}
+
+template<int Unroll, typename RedFn, typename T,
+         int MultimemSrcs, int MinSrcs, int MaxSrcs,
+         int MultimemDsts, int MinDsts, int MaxDsts, int PreOpSrcs,
+         typename IntBytes>
+__device__ __forceinline__ void reduceCopy(
+    int thread, int nThreads,
+    uint64_t redArg, uint64_t *preOpArgs, bool postOp,
+    int nSrcs, void **srcPtrs, int nDsts, void **dstPtrs,
+    IntBytes nElts
+  ) {
+  static_assert(MultimemSrcs <= MinSrcs && MultimemDsts <= MinDsts, "Multimem pointers cannot exceed respective Min values.");
+  //int nWarps = nThreads/WARP_SIZE;
+  //int warp = thread/WARP_SIZE;
+  int lane = thread%WARP_SIZE;
+  // If a multimem src is present then our biggest pack size is limited to what
+  // is supported for this redfn/type.
+  constexpr int BigPackSize = (MultimemSrcs == 0) ? 16 : LoadMultimem_BigPackSize<RedFn>::BigPackSize;
+
+  IntBytes nBytesBehind = 0;
+  IntBytes nBytesAhead = nElts*sizeof(T);
+
+  #if __cpp_if_constexpr
+  if constexpr (BigPackSize > sizeof(T)) {
+  #else
+  if (BigPackSize > sizeof(T)) {
+  #endif
+    // Check that all pointers are BigPackSize aligned.
+    bool aligned = true;
+    if (lane < nSrcs) aligned &= 0 == cvta_to_global(srcPtrs[lane]) % (BigPackSize + !BigPackSize);
+    if (lane < nDsts) aligned &= 0 == cvta_to_global(dstPtrs[lane]) % (BigPackSize + !BigPackSize);
+    aligned = !(__any(!aligned));
+    if (aligned) {
+#if defined(__gfx90a__)
+      reduceCopyPacks<RedFn, T, ((MinSrcs > 1) ? 2 : Unroll), BigPackSize,
+        MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+        (nThreads, thread, redArg, preOpArgs, postOp,
+         nSrcs, srcPtrs, nDsts, dstPtrs, nBytesBehind, nBytesAhead);
+#else
+      reduceCopyPacks<RedFn, T, Unroll*((MinSrcs == 1 && MinDsts == 1) ? 2 : 1), BigPackSize,
+        MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+        (nThreads, /*&*/thread, redArg, preOpArgs, postOp,
+         nSrcs, srcPtrs, nDsts, dstPtrs, /*&*/nBytesBehind, /*&*/nBytesAhead);
+#endif
+      if (nBytesAhead == 0) return;
+
+      reduceCopyPacks<RedFn, T, /*Unroll=*/1, BigPackSize,
+        MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+        (nThreads, /*&*/thread, redArg, preOpArgs, postOp,
+         nSrcs, srcPtrs, nDsts, dstPtrs, /*&*/nBytesBehind, /*&*/nBytesAhead);
+      if (nBytesAhead == 0) return;
+    }
+  }
+
+#if defined(__gfx90a__)
+  if (MinSrcs > 1) {
+    reduceCopyPacks<RedFn, T, Unroll/2*(16/sizeof(T))/2, sizeof(T),
+    MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+    (nThreads, thread, redArg, preOpArgs, postOp,
+     nSrcs, srcPtrs, nDsts, dstPtrs, nBytesBehind, nBytesAhead);
+  } else {
+    reduceCopyPacks<RedFn, T, Unroll*(16/sizeof(T))/2, /*BytePerPack=*/sizeof(T),
+    MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+    (nThreads, /*&*/thread, redArg, preOpArgs, postOp,
+     nSrcs, srcPtrs, nDsts, dstPtrs, /*&*/nBytesBehind, /*&*/nBytesAhead);
+  }
+#else
+  reduceCopyPacks<RedFn, T, Unroll*(16/sizeof(T))/2, /*BytePerPack=*/sizeof(T),
+    MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+    (nThreads, /*&*/thread, redArg, preOpArgs, postOp,
+     nSrcs, srcPtrs, nDsts, dstPtrs, /*&*/nBytesBehind, /*&*/nBytesAhead);
+#endif
+  if (nBytesAhead == 0) return;
+
+  reduceCopyPacks<RedFn, T, /*Unroll=*/1, /*BytePerPack=*/sizeof(T),
+    MultimemSrcs, MinSrcs, MaxSrcs, MultimemDsts, MinDsts, MaxDsts, PreOpSrcs>
+    (nThreads, /*&*/thread, redArg, preOpArgs, postOp,
+     nSrcs, srcPtrs, nDsts, dstPtrs, /*&*/nBytesBehind, /*&*/nBytesAhead);
+}
+
+#endif // COMMON_KERNEL_H_

src/device/msccl_kernel_impl.h

@@ -0,0 +1,434 @@
+/*************************************************************************
+ * Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+#ifndef MSSCLKERNELIMPL_H
+#define MSSCLKERNELIMPL_H
+
+#include "device.h"
+#include "primitives.h"
+#include "collectives.h"
+
+#include "msccl/msccl_struct.h"
+#include "msccl/msccl_kernel.h"
+
+extern __shared__ struct mscclShmemData mscclShmem;
+
+#define MSCCL_MAX_ITER 65536
+
+// flags are a 3-tuple of (workindex, gridoffset_iter, step) and it follows a lexicographical order. a threadblock is ahead of another iff its flag is ahead
+#define COMPUTE_FLAG(__WORKINDEX__,__GRIDOFFSET_ITER__,__STEP__) \
+  MSCCL_MAX_ITER*MSCCL_MAX_NUM_STEPS*(uint64_t)__WORKINDEX__ + ((uint64_t)__GRIDOFFSET_ITER__ * MSCCL_MAX_NUM_STEPS + (uint64_t)__STEP__)
+
+#define GET_WORKINDEX_FROM_FLAG(__FLAG__) \
+  (__FLAG__) / (MSCCL_MAX_ITER*MSCCL_MAX_NUM_STEPS)
+
+#ifdef ENABLE_COLLTRACE
+  #define INC_COLL_TRACE \
+    uint32_t pos = atomicAdd(&ncclShmem.collTraceTail->tail, 1)%COLLTRACE_NUM_ITEMS; \
+    struct ncclCollTrace* collTrace = ncclShmem.collTrace+pos; \
+    collTrace->timeStamp = wall_clock64(); \
+    collTrace->bid = blockIdx.x;
+    // TODO: switch to atomicInc after llvm crash is fixed
+    // uint32_t pos = atomicInc(&ncclShmem.collTraceTail->tail, COLLTRACE_NUM_ITEMS)
+
+  #define traceData(data2, data4, data8_0, data8_1) { \
+    INC_COLL_TRACE \
+    collTrace->funcIndex = data2; \
+    collTrace->data_0 = data4; \
+    collTrace->opCount = data8_0; \
+    collTrace->data_1 = data8_1; \
+    collTrace->type = ncclCollTraceDataType; \
+  }
+#else
+#define traceData(data2, data4, data8_0, data8_1)
+#endif
+
+inline __device__ static void barrier(int nthreads) {
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+  assert(nthreads == NCCL_MAX_NTHREADS);
+  __asm__ __volatile__("s_waitcnt vmcnt(0) lgkmcnt(0)\ns_barrier");
+#else
+  asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(15));
+#endif
+}
+
+// Copy 8-byte aligned data. You must call with at least `(bytes+7)/8` threads.
+inline __device__ static void copyToShmem8(int tid, void* dst, void const* src, int bytes) {
+  int offset = sizeof(uint32_t) * tid;
+  if (offset < bytes) {
+    uint32_t *src2 = (uint32_t*)((char const*)src + offset);
+    uint32_t *dst2 = (uint32_t*)((char*)dst + offset);
+    *dst2 = *src2;
+    offset += WARP_SIZE*sizeof(uint32_t);
+  }
+}
+
+__device__ __forceinline__ static void threadBlockCopy(
+  uint32_t *dst, uint32_t const *src, uint64_t size, int tid, int nthreads) {
+  for (int i = tid; i < size; i += nthreads) {
+    dst[i] = src[i];
+  }
+}
+
+#define MSCCL_REDUCE_UNROLL_LOOP_A(numloops, BytePerPack) \
+  for (int r = 0; r < numloops; r++) { \
+    srcOffset = srcBaseOffset + (ssize_t)mscclShmem.mscclTB.reductionSrcOffsets[t->reductionPointer+r] * sizePerMscclChunk; \
+    reduceInput = ld_volatile_global<BytePerPack>((uintptr_t)(srcPointer + srcOffset)); \
+    o = applyReduce(redFn, reduceInput, o); \
+  }
+
+template<typename T, typename RedOp, int BytePerPack>
+__device__ __forceinline__ static void mscclReduce(int c, int numReductions, int currIdx, ssize_t sizePerMscclChunk, RedOp redFn,
+  struct mscclTransmission* t, ssize_t gridOffset, ssize_t &srcOffset, ssize_t dstOffset, T *srcPointer, T *dstPointer) {
+  const int elemsPerPack = BytePerPack/sizeof(T);
+  T* dstIndex = dstPointer + dstOffset + currIdx*elemsPerPack;
+  BytePack<BytePerPack> reduceInput;
+  BytePack<BytePerPack> o = ld_volatile_global<BytePerPack>((uintptr_t)dstIndex);
+  ssize_t srcBaseOffset = gridOffset + (ssize_t)c * sizePerMscclChunk + currIdx*elemsPerPack;
+  switch (numReductions) {
+    case 7:
+      #pragma unroll
+      MSCCL_REDUCE_UNROLL_LOOP_A(7, BytePerPack);
+      break;
+#if defined(__gfx90a__)
+    case 15:
+      #pragma unroll
+      MSCCL_REDUCE_UNROLL_LOOP_A(15, BytePerPack);
+      break;
+#endif
+    default:
+      MSCCL_REDUCE_UNROLL_LOOP_A(numReductions, BytePerPack);
+      break;
+  }
+  st_global<BytePerPack>((uintptr_t)dstIndex, o);
+}
+
+
+template<typename T, typename RedOp, typename Proto, bool fullOps>
+__device__ __forceinline__ void mscclRunInterpreter(
+  struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work) {
+  const int tid = threadIdx.x;
+  const int bid = blockIdx.x;
+  const int nthreads = NCCL_MAX_NTHREADS;
+
+#if defined(ENABLE_NPKIT)
+  uint64_t timestamp_entry = 0;
+  if (tid == 0) {
+     timestamp_entry = NPKIT_GET_GPU_TIMESTAMP();
+  }
+#endif
+  // initialize mscclShmem.mscclTB
+  threadBlockCopy(
+    (uint32_t *)&mscclShmem.mscclTB, (uint32_t *)(algo->mscclTBs + bid),
+    sizeof(struct mscclThreadBlock) / sizeof(uint32_t), tid, nthreads);
+  __synclds(); // publish mscclShmem.mscclTB.channelId
+
+  // initialize ncclShmem and mscclShmem.work
+  int channelId = mscclShmem.mscclTB.channelId;
+  {
+    void *dst, *src;
+    int bytes = 0;
+    // Use first 3 warps to load comm, channel, and work into shmem
+    switch (tid/WARP_SIZE) {
+    case 0:
+      dst = &ncclShmem.comm;
+      src = comm;
+      bytes = sizeof(ncclDevComm);
+      break;
+    case 1:
+      // Get address of channel without incurring indirect load from ncclDevComm::channels
+      dst = &ncclShmem.channel;
+      src = &((ncclDevCommAndChannels*)comm)->channels[channelId];
+      bytes = sizeof(ncclDevChannel);
+      break;
+    case 2:
+      dst = &mscclShmem.work;
+      src = work + blockIdx.x;
+      bytes = sizeof(mscclWork);
+      break;
+    case 3:
+      /* set abort flag to 0 */
+      if (tid%WARP_SIZE == 0) ncclShmem.aborted = 0;
+#ifdef ENABLE_COLLTRACE
+      else if (tid%WARP_SIZE == 1) ncclShmem.collTrace = comm->collTrace + COLLTRACE_NUM_ITEMS*channelId;
+      else if (tid%WARP_SIZE == 2) ncclShmem.collTraceTail = comm->collTraceTail + channelId;
+#endif
+      break;
+    default:
+      break;
+    }
+    copyToShmem8(tid%WARP_SIZE, dst, src, bytes);
+  }
+
+#if defined(ENABLE_NPKIT)
+  int npKitCtxIdx = bid;
+  int xcc_id = 0;
+  if (tid == 0) {
+    ncclShmem.event_buffer_head = 0;
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+    asm volatile ("s_getreg_b32 %0, hwreg(HW_REG_XCC_ID)" : "=s" (xcc_id));
+#endif
+  }
+#endif
+  __synclds(); // publish shmem
+
+  if (fullOps && tid == 0) {
+    traceData(__LINE__, mscclShmem.work.fnIndex, (uint64_t)mscclShmem.work.sendBuff, 0);
+  }
+
+  if (tid == 0)
+    *mscclShmem.work.workFifoDone = mscclShmem.work.workFifoDoneAck;
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+  if (tid == 0) {
+    uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+    NpKit::CollectGpuEventLDS(NPKIT_EVENT_TIME_SYNC_CPU, 0, xcc_id, *cpuTimestamp);
+  }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+  if (tid == 0) {
+    NpKit::CollectGpuEventLDS(NPKIT_EVENT_TIME_SYNC_GPU, 0, xcc_id, NPKIT_GET_GPU_TIMESTAMP());
+  }
+#endif
+
+  // User pointers for primitives
+  T* thisInput = (T*)mscclShmem.work.sendBuff;
+  T* thisOutput = (T*)mscclShmem.work.recvBuff;
+  T* thisScratch = (T*)mscclShmem.work.scratchBuffer;
+  int recvPeer = mscclShmem.mscclTB.recvPeer;
+  int sendPeer = mscclShmem.mscclTB.sendPeer;
+
+  const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? MSCCL_CHUNKSTEPS : 1));
+  int minChunkSize;
+  if (Proto::Id == NCCL_PROTO_LL)
+    minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T));
+  if (Proto::Id == NCCL_PROTO_LL128) {
+    // We should not need the final /2 but it makes performance much, much smoother. Might be a bug somewhere.
+    minChunkSize = nthreads*(Proto::calcBytePerGrain()/sizeof(T))/2;
+  }
+
+  RedOp redFn(mscclShmem.work.redOpArg);
+  Primitives<T, RedOp, FanAsymmetric<1,1>, 1, Proto, 0> prims
+    (tid, nthreads, &recvPeer, &sendPeer, thisInput, thisOutput, mscclShmem.work.redOpArg);
+
+#if defined(ENABLE_NPKIT)
+  if (tid == 0) {
+    prims.npKitCtxIdx = npKitCtxIdx;
+  }
+#endif
+
+  const ssize_t sizePerMscclChunk = mscclShmem.work.sizePerMscclChunk;
+  uint32_t maxAllowedCount = mscclShmem.work.maxAllowedCount;
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RUN_ENTRY)
+  if (tid == 0) {
+    NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RUN_ENTRY, mscclShmem.work.sizePerMscclChunk*mscclShmem.work.nChunksPerLoop, xcc_id, timestamp_entry);
+  }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_INIT_ENTRY)
+  if (tid == 0) {
+    NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_INIT_ENTRY, 0, xcc_id, timestamp_entry);
+  }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_INIT_EXIT)
+  if (tid == 0) {
+    NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_INIT_EXIT, 0, xcc_id, NPKIT_GET_GPU_TIMESTAMP());
+  }
+#endif
+
+  // msccl flags all start out with 0. this is used as a part of the flag to make sure different work items deal with different synchronization flags
+  // this still needs more work. when we make a way around the queue, the flag might have been set to undesired values. will be fixed in subsequent versions.
+  const int64_t workIndex = mscclShmem.work.workIndex;
+  volatile struct mscclFlag* mscclFlags = mscclShmem.work.syncFlags;
+  for (ssize_t gridOffset = 0, iter = 0; gridOffset < sizePerMscclChunk; gridOffset += chunkSize, iter++) {
+    ssize_t realChunkSize;
+    if (Proto::Id == NCCL_PROTO_SIMPLE) {
+      realChunkSize = min(chunkSize, sizePerMscclChunk-gridOffset);
+      realChunkSize = roundUp(realChunkSize, nthreads*sizeof(uint64_t)/sizeof(T));
+    }
+    else
+      realChunkSize = min(chunkSize, divUp(sizePerMscclChunk-gridOffset, minChunkSize)*minChunkSize);
+    realChunkSize = int(realChunkSize);
+    int nelem = min(realChunkSize, sizePerMscclChunk-gridOffset);
+
+    ssize_t srcOffset, dstOffset;
+    T *srcPointer, *dstPointer;
+    int step = 0;
+    for (int i = 0; i < mscclShmem.mscclTB.nSteps; i++){
+      struct mscclTransmission* t = &mscclShmem.mscclTB.transmissions[i];
+      // first wait if there is a dependence
+      int16_t numDependencies = t->numDependencies;
+      if (numDependencies > 0){
+        if (tid < numDependencies) {
+          int16_t dependentPointer = t->dependencePointer;
+          int8_t dependentBid = mscclShmem.mscclTB.dependentBid[dependentPointer+tid];
+          int16_t dependentStep = mscclShmem.mscclTB.dependentStep[dependentPointer+tid];
+          uint64_t goalFlag = COMPUTE_FLAG(workIndex, iter, dependentStep);
+          while (true){
+            uint64_t curFlag = __atomic_load_n(&(mscclFlags + dependentBid)->flag, __ATOMIC_RELAXED);
+            if (curFlag >= goalFlag && GET_WORKINDEX_FROM_FLAG(curFlag) == workIndex) break;
+          }
+        }
+        step += numDependencies-1;
+        barrier(nthreads);
+      }
+
+      srcPointer = (t->srcBuffer == MSCCL_INPUT_BUFFER) ? thisInput : ((t->srcBuffer == MSCCL_OUTPUT_BUFFER) ? thisOutput : thisScratch);
+      dstPointer = (t->dstBuffer == MSCCL_INPUT_BUFFER) ? thisInput : ((t->dstBuffer == MSCCL_OUTPUT_BUFFER) ? thisOutput : thisScratch);
+      prims.setDataPtrs(srcPointer, dstPointer);
+
+      int count = t->count;
+      for (int c = 0; c < count; c += maxAllowedCount) {
+        srcOffset = gridOffset + (ssize_t) (t->srcOffset+c) * sizePerMscclChunk;
+        dstOffset = gridOffset + (ssize_t) (t->dstOffset+c) * sizePerMscclChunk;
+        int thisCount = min(maxAllowedCount, count - c);
+        int thisNelem = nelem * thisCount;
+        if (t->type == MSCCL_SEND) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_SEND_ENTRY)
+            if (tid == 0) {
+              NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_SEND_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+            }
+
+#endif	
+          prims.send(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end.
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_SEND_EXIT)
+            if (tid == 0) {
+              NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_SEND_EXIT, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+            }
+#endif
+        }
+        else if (t->type == MSCCL_RECV) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RECV_ENTRY)
+            if (tid == 0) {
+              NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+            }
+#endif	
+          prims.recv(dstOffset, thisNelem);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RECV_EXIT)
+            if (tid == 0) {
+              NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_EXIT, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+            }
+#endif
+        }
+        else if (t->type == MSCCL_REDUCE) {
+          int numReductions = t->numReductions;
+          int currIdx = tid;
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_REDUCE_ENTRY)
+          if (tid == 0) {
+            NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_REDUCE_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+          }
+#endif
+          dstOffset = gridOffset + (ssize_t) (t->dstOffset+c) * sizePerMscclChunk;
+          // process 16-byte packed elements
+          const int elemsPerPack = 16/sizeof(T);
+          while (currIdx < thisNelem/elemsPerPack) {
+            mscclReduce<T, RedOp, 16>(c, numReductions, currIdx, sizePerMscclChunk, redFn, t, gridOffset, srcOffset, dstOffset, srcPointer, dstPointer);
+            currIdx += nthreads;
+          }
+          // process remaining elements
+          currIdx = tid + (thisNelem/elemsPerPack)*elemsPerPack;
+          if (currIdx < thisNelem) {
+            mscclReduce<T, RedOp, sizeof(T)>(c, numReductions, currIdx, sizePerMscclChunk, redFn, t, gridOffset, srcOffset, dstOffset, srcPointer, dstPointer);
+          }
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_REDUCE_EXIT)
+          if (tid == 0) {
+            NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_REDUCE_EXIT, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+          }
+#endif
+          barrier(nthreads);
+          if (c == 0) step += (numReductions-1); // only advance step once!
+        }
+        else if (fullOps && t->type == MSCCL_RECV_COPY_SEND)
+          prims.recvCopySend(dstOffset, thisNelem);
+        else if (fullOps && t->type == MSCCL_RECV_REDUCE_SEND)
+          prims.recvReduceSend(srcOffset, thisNelem);
+        else if (fullOps && t->type == MSCCL_RECV_REDUCE_COPY_SEND)
+          prims.recvReduceCopySend(srcOffset, dstOffset, thisNelem);
+        else if (fullOps && t->type == MSCCL_RECV_REDUCE_COPY) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RECV_REDUCE_COPY_ENTRY)
+          if (tid == 0) {
+            NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_REDUCE_COPY_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+          }
+#endif
+          prims.recvReduceCopy(srcOffset, dstOffset, thisNelem);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RECV_REDUCE_COPY_EXIT)
+          if (tid == 0) {
+            NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RECV_REDUCE_COPY_EXIT, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
+          }
+#endif
+        }
+        else if (t->type == MSCCL_LOCAL_COPY)
+          prims.localCopy(srcPointer+srcOffset, dstPointer+dstOffset, thisNelem);
+        else
+          return;
+      }
+      if (t->hasDependence && tid == nthreads-1)
+        __atomic_store_n(&mscclFlags[bid].flag, (uint64_t) COMPUTE_FLAG(workIndex, iter, step), t->type == MSCCL_REDUCE ? __ATOMIC_RELEASE : __ATOMIC_RELAXED);
+      step++;
+    }
+  }
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_RUN_EXIT)
+  if (tid == 0) {
+    NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_RUN_EXIT, mscclShmem.work.sizePerMscclChunk*mscclShmem.work.nChunksPerLoop, xcc_id, NPKIT_GET_GPU_TIMESTAMP());
+  }
+#endif
+#if defined(ENABLE_NPKIT)
+  __synclds();
+  NpKitEventCollectContext* ctx = ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx;
+  copyToShmem16(tid, ctx->event_buffer+ctx->event_buffer_head, ncclShmem.event_buffer, sizeof(NpKitEvent)*ncclShmem.event_buffer_head);
+  if (tid == 0) ctx->event_buffer_head += ncclShmem.event_buffer_head;
+#endif
+
+  if (fullOps && tid == 0) {
+    traceData(__LINE__, mscclShmem.work.fnIndex, (uint64_t)mscclShmem.work.sendBuff, 0);
+  }
+}
+
+#define MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, type, fullOps) \
+__global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, LL, fullOps)(struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work) { \
+  mscclRunInterpreter<type, Func##devredop<type>, ProtoLL, fullOps>(comm, algo, work); \
+} \
+__global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, LL128, fullOps)(struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work) { \
+  mscclRunInterpreter<type, Func##devredop<type>, ProtoLL128, fullOps>(comm, algo, work); \
+} \
+__global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple, fullOps)(struct ncclDevComm* comm, struct mscclAlgo* algo, struct mscclWork* work) { \
+  mscclRunInterpreter<type, Func##devredop<type>, ProtoSimple<MSCCL_CHUNKSTEPS/MSCCL_SLICESTEPS, MSCCL_SLICESTEPS>, fullOps>(comm, algo, work); \
+}
+
+#define MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(devredop, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int8_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint8_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int32_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint32_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int64_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint64_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, half, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, float, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, double, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps)
+
+#define MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(devredop, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int8_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint8_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int32_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint32_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int64_t, fullOps) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, uint64_t, fullOps)
+
+#define MSCCL_IMPL_KERNEL_ENTRY_FUNC() \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Sum, false) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(Prod, false) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(MinMax, false) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP(PreMulSum, false) \
+  MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(SumPostDiv, false)
+
+#endif

src/device/network/unpack/unpack.h

@@ -0,0 +1,282 @@
+/*************************************************************************
+ * Copyright (c) 2023, Google LLC.  All rights reserved.
+ * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+#ifndef NET_DEVICE_UNPACK_H
+#define NET_DEVICE_UNPACK_H
+
+#include "unpack_defs.h"
+
+#include "op128.h"
+#include "align.h"
+#include "device.h"
+#include "common.h"
+
+// #define ALIGNED_LOAD
+
+inline __device__ void load64gpu(const uint64_t* ptr, uint64_t &v) {
+  #if __CUDA_ARCH__ >= 700
+      asm volatile("ld.relaxed.gpu.u64 {%0}, [%1];"
+      : "=l"(v) : "l"(ptr));
+  #else
+#ifdef NEED_CHECKING
+      asm volatile("ld.volatile.global.u64 {%0}, [%1];"
+      : "=l"(v) : "l"(ptr));
+#endif
+  #endif
+}
+
+#define PAGE_META_SIZE 16
+#define META_LOAD_SIZE 16
+#define DATA_LOAD_SIZE 16
+
+// Map internal association of handle with group and peer index (called once at init time)
+inline __device__ void ncclNetDeviceUnpackSetup(void* ohandle, const int group, const int index) {
+  struct unpackNetDeviceHandle* handle = (struct unpackNetDeviceHandle*) ohandle;
+  ncclShmem.groups[group].devicePlugin.unpack.g_meta[index] = handle->meta;
+  ncclShmem.devicePlugin.unpack.bounce_buf = handle->bounce_buf;
+  ncclShmem.groups[group].devicePlugin.unpack.head = handle->head;
+}
+
+inline __device__ void ncclNetDeviceIncrementHead(const int group) {
+  ncclShmem.groups[group].devicePlugin.unpack.head++;
+}
+
+inline __device__ void ncclNetDeviceSaveHead(void* ohandle, const int group) {
+  struct unpackNetDeviceHandle* handle = (struct unpackNetDeviceHandle*) ohandle;
+  handle->head = ncclShmem.groups[group].devicePlugin.unpack.head;
+}
+
+template <uint8_t sz>
+inline __device__ void bulkLoad(const int t, const uint32_t len, char* cpy_src, char* cpy_dst, BytePack<sz> *reg, const int w, loadMeta* g_meta, loadMeta* s_meta, uint32_t src_off, uint64_t dst_off){
+  bulkLoad<1>(t, len, cpy_src, cpy_dst, reg, w, g_meta, s_meta, src_off, dst_off);
+}
+
+template <>
+inline __device__ void bulkLoad<1>(const int t, const uint32_t len, char* cpy_src, char* cpy_dst, BytePack<1> reg[16], const int w, loadMeta* g_meta, loadMeta* s_meta, uint32_t src_off, uint64_t dst_off){
+  uint64_t data_s;
+  for (data_s = t * DATA_LOAD_SIZE; data_s + DATA_LOAD_SIZE - 1 < len; data_s += WARP_SIZE * DATA_LOAD_SIZE) {
+
+#ifdef ALIGNED_LOAD
+    load128 ((uint64_t*)(cpy_src + data_s), reg.u64[0], reg.u64[1]);
+#else
+#pragma unroll
+    for (int i=0; i<16; i++) {
+      reg[i] = ld_volatile_global<1>((uintptr_t)((uint8_t*)(cpy_src + data_s) + i));
+    }
+#endif
+
+#pragma unroll
+    for (int i=0; i<16; i++) {
+      st_global<1>((uintptr_t)((uint8_t*)(cpy_dst + data_s) + i), reg[i]);
+    }
+  }
+}
+
+template <>
+inline __device__ void bulkLoad<2>(const int t, const uint32_t len, char* cpy_src, char* cpy_dst, BytePack<2> reg[8], const int w, loadMeta* g_meta, loadMeta* s_meta, uint32_t src_off, uint64_t dst_off){
+  uint64_t data_s;
+  for (data_s = t * DATA_LOAD_SIZE; data_s + DATA_LOAD_SIZE - 1 < len; data_s += WARP_SIZE * DATA_LOAD_SIZE) {
+#ifdef ALIGNED_LOAD
+    load128 ((uint64_t*)(cpy_src + data_s), reg.u64[0], reg.u64[1]);
+#else
+#pragma unroll
+    for (int i=0; i<8; i++) {
+      reg[i] = ld_volatile_global<2>((uintptr_t)((uint16_t*)(cpy_src + data_s) + i));
+    }
+#endif
+
+
+#pragma unroll
+    for (int i=0; i<8; i++) {
+      st_global<2>((uintptr_t)((uint16_t*)(cpy_dst + data_s) + i), reg[i]);
+    }
+  }
+}
+
+template <>
+inline __device__ void bulkLoad<4>(const int t, const uint32_t len, char* cpy_src, char* cpy_dst, BytePack<4> reg[4], const int w, loadMeta* g_meta, loadMeta* s_meta, uint32_t src_off, uint64_t dst_off){
+  uint64_t data_s;
+  for (data_s = t * DATA_LOAD_SIZE; data_s + DATA_LOAD_SIZE - 1 < len; data_s += WARP_SIZE * DATA_LOAD_SIZE) {
+#ifdef ALIGNED_LOAD
+    load128 ((uint64_t*)(cpy_src + data_s), reg.u64[0], reg.u64[1]);
+#else
+#pragma unroll
+    for (int i=0; i<4; i++) {
+      reg[i] = ld_volatile_global<4>((uintptr_t)((uint32_t *)(cpy_src + data_s) + i));
+    }
+#endif
+
+#pragma unroll
+    for (int i=0; i<4; i++) {
+      st_global<4>((uintptr_t)((uint32_t*)(cpy_dst + data_s) + i), reg[i]);
+    }
+  }
+}
+
+template <>
+inline __device__ void bulkLoad<8>(const int t, const uint32_t len, char* cpy_src, char* cpy_dst, BytePack<8> reg[2], const int w, loadMeta* g_meta, loadMeta* s_meta, uint32_t src_off, uint64_t dst_off){
+  uint64_t data_s;
+  for (data_s = t * DATA_LOAD_SIZE; data_s + DATA_LOAD_SIZE - 1 < len; data_s += WARP_SIZE * DATA_LOAD_SIZE) {
+#ifdef ALIGNED_LOAD
+    load128 ((uint64_t*)(cpy_src + data_s), reg.u64[0], reg.u64[1]);
+#else
+#pragma unroll
+    for (int i=0; i<2; i++) {
+      reg[i] = ld_volatile_global<8>((uintptr_t)((uint64_t*)(cpy_src + data_s) + i));
+    }
+#endif
+
+#pragma unroll
+    for (int i=0; i<2; i++) {
+      st_global<8>((uintptr_t)((uint64_t*)(cpy_dst + data_s) + i), reg[i]);
+    }
+  }
+}
+
+template <>
+inline __device__ void bulkLoad<16>(const int t, const uint32_t len, char* cpy_src, char* cpy_dst, BytePack<16> reg[1], const int w, loadMeta* g_meta, loadMeta* s_meta, uint32_t src_off, uint64_t dst_off){
+  uint64_t data_s;
+  for (data_s = t * DATA_LOAD_SIZE; data_s + DATA_LOAD_SIZE - 1 < len; data_s += WARP_SIZE * DATA_LOAD_SIZE) {
+    reg[0] = ld_volatile_global<16>((uintptr_t)(cpy_src + data_s));
+    st_global<16>((uintptr_t)(cpy_dst + data_s), reg[0]);
+  }
+}
+
+#ifndef PAGE_SIZE
+#define PAGE_SIZE 4096
+#endif
+inline __device__ int ppw(const int nbytes, int nw) {
+  int v = DIVUP(nbytes, SLICE_PAGE_SIZE);
+  v = DIVUP(v, nw);
+  while (v > WARP_SHM_PAGE_CNT) {
+    v = DIVUP(v, 2);
+  }
+  return v;
+}
+
+// This function is called by all threads
+// Pack data from the internal iovec to the supplied flat buffer using all the
+// threads
+template <int Recv>
+inline __device__ void ncclNetDeviceUnpack(
+    const int tid, const int tidInBlock, const int nworkers, const int group, int mask, int Src, int workSize);
+
+template <>
+inline __device__ void ncclNetDeviceUnpack</*Recv=*/0>(
+    const int tid, const int tidInBlock, const int nworkers, const int group, int mask, int Src, int workSize) {
+  // send unpack empty
+}
+
+inline __device__ void ncclNetDeviceUnpackInner(
+    const int tid, const int tidInBlock, const int nworkers, const int group, const int index,
+    void *src, const int nbytes, const uint64_t step);
+
+template <>
+inline __device__ void ncclNetDeviceUnpack</*Recv=*/1>(
+    const int tid, const int tidInBlock, const int nworkers, const int group, int mask, int Src, int workSize) {
+
+  while (mask != 0) {
+    int ix = __ffs(mask)-1; // Get the first set bit of the mask (this should correlate to a peer index)
+    mask &= mask-1; // Drop the first set bit of the mask
+
+    // Pack data from the internal iovec to the supplied flat srcs buffer using all the threads
+    // + Src is necessary in the case of accessing the user buffer directly
+    ncclNetDeviceUnpackInner(tid, tidInBlock, nworkers, group /* in case they need to use split warps shared memory partitioning*/,
+        ix, ncclShmem.groups[group].srcs[ix + Src], workSize, ncclShmem.groups[group].devicePlugin.unpack.head);
+  }
+}
+
+inline __device__ void ncclNetDeviceUnpackInner(
+    const int tid, const int tidInBlock, const int nworkers, const int group, const int index,
+    void *src, const int nbytes, const uint64_t step) {
+  // from src/collectives/device/common_kernel.h
+  const int w = tid / WARP_SIZE;        // Warp number
+  const int nw = nworkers / WARP_SIZE;  // Number of warps
+  const int t = tid % WARP_SIZE;        // Thread (inside the warp)
+
+  BytePack<16> reg;
+  loadMeta meta;
+
+  uint64_t head;
+  struct netUnpackMeta* g_meta_struct;
+  void* bounce_buf;
+
+  loadMeta* g_meta;
+  loadMeta* s_meta;
+  uint64_t meta_cnt;
+
+  // hack head use per-warp
+  head          = step;
+  g_meta_struct = ncclShmem.groups[group].devicePlugin.unpack.g_meta[index];
+  bounce_buf    = ncclShmem.devicePlugin.unpack.bounce_buf;
+
+  __syncwarp();
+
+  head %= NCCL_NET_DEVICE_UNPACK_MAX_QUEUE_DEPTH;
+
+  g_meta = g_meta_struct->mem[head];
+
+  // Currently, even/odd groups perform send/recv separately. We don't really need space for send side.
+  // Total size is N page per warp * 16 B per page * 20 WARPS max = 320 * N bytes, N == WARP_SHM_PAGE_CNT
+  static_assert(ncclShmemScratchWarpSize() >= WARP_SHM_SIZE, "Each warp must have enough scratch space");
+  s_meta = (loadMeta*) ncclScratchForWarp(tidInBlock / WARP_SIZE); // (loadMeta*) (ncclShmem.devicePlugin.unpack.meta + shm_off);
+
+  load64gpu(g_meta_struct->cnt + head, meta_cnt);
+
+  int PPW = ppw(nbytes, nw);
+
+  for (uint64_t meta_s = w * PPW; meta_s < meta_cnt; meta_s += nw * PPW) {
+
+    uint64_t iter_meta_cnt = meta_cnt - meta_s;
+    iter_meta_cnt = iter_meta_cnt < PPW ? iter_meta_cnt : PPW;
+
+    // TODO: this load size needs to work if not aligned, but since the two are both 16...
+    if (t < PPW * PAGE_META_SIZE / META_LOAD_SIZE && t < iter_meta_cnt) {  // avoid last iter load garbage data
+      load128((const uint64_t*) (g_meta + (meta_s + t)), reg.u64[0], reg.u64[1]);
+
+      storeShmem128(shmemCvtPtr((uint64_t *)(s_meta + (w * PPW + t))), reg.u64[0], reg.u64[1]);
+    }
+
+    __syncwarp();
+
+    for (int x = 0; x < iter_meta_cnt; x++) {
+      int meta_idx = x + w * PPW;
+      
+      // load page offs
+      loadShmem128(shmemCvtPtr((uint64_t*) (s_meta + meta_idx)), meta.r64[0], meta.r64[1]);
+
+      if (meta.len >= DATA_LOAD_SIZE) {
+        // fast path, but need to adapt to alignment issue
+
+        // bulk copy data
+        uint8_t align_off = (meta.src_off | meta.dst_off) % DATA_LOAD_SIZE;
+        align_off = align_off & -align_off;  // keep the lowest bit
+        if (align_off == 0) {  // 0x16
+          bulkLoad<16>(t, meta.len, (char*) bounce_buf + meta.src_off, (char*) src + meta.dst_off, &reg, w, g_meta, s_meta, meta.src_off, meta.dst_off);
+        } else if (align_off & 0x8) {
+          bulkLoad<8>(t, meta.len, (char*) bounce_buf + meta.src_off, (char*) src + meta.dst_off, (BytePack<8>*) &reg, w, g_meta, s_meta, meta.src_off, meta.dst_off);
+        } else if (align_off & 0x4) {
+          bulkLoad<4>(t, meta.len, (char*) bounce_buf + meta.src_off, (char*) src + meta.dst_off, (BytePack<4>*) &reg, w, g_meta, s_meta, meta.src_off, meta.dst_off);
+        } else if (align_off & 0x2) {
+          bulkLoad<2>(t, meta.len, (char*) bounce_buf + meta.src_off, (char*) src + meta.dst_off, (BytePack<2>*) &reg, w, g_meta, s_meta, meta.src_off, meta.dst_off);
+        } else { // if (align_off & 0x1)
+          bulkLoad<1>(t, meta.len, (char*) bounce_buf + meta.src_off, (char*) src + meta.dst_off, (BytePack<1>*) &reg, w, g_meta, s_meta, meta.src_off, meta.dst_off);
+        }
+      }
+
+      // must be less than 16 bytes
+      if (t < meta.len % DATA_LOAD_SIZE) {
+        volatile char* cpy_src = (char*) bounce_buf + meta.src_off + (meta.len / DATA_LOAD_SIZE) * DATA_LOAD_SIZE + t;
+        volatile char* cpy_dst = (char*) src        + meta.dst_off + (meta.len / DATA_LOAD_SIZE) * DATA_LOAD_SIZE + t;
+        *cpy_dst = *cpy_src;
+      }
+    }
+
+    __syncwarp();
+  }
+}
+
+#endif  // NET_DEVICE_UNPACK_DEFS_H_

src/device/network/unpack/unpack_defs.h

@@ -0,0 +1,61 @@
+/*************************************************************************
+ * Copyright (c) 2023, Google LLC.  All rights reserved.
+ * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+#ifndef NET_DEVICE_UNPACK_DEFS_H
+#define NET_DEVICE_UNPACK_DEFS_H
+
+#include <stdint.h>
+
+#include "device.h"
+
+#define NCCL_NET_DEVICE_UNPACK_MAX_QUEUE_DEPTH 16
+
+union alignas(16) loadMeta {
+  uint64_t r64[2];
+  struct {
+    uint32_t src_off;
+    uint32_t len;
+    uint64_t dst_off;
+  };
+};
+static_assert(sizeof(union loadMeta) == 16, "Must be 16-byte aligned");
+
+/****** global memory ******/
+
+#define NET_UNPACK_MAX_QUEUE_DEPTH 16  // MAX_REQUESTS
+#define NET_UNPACK_MAX_SLICE_SIZE 4194304  // 4MB per Irecv call
+#define SLICE_PAGE_SIZE 4096
+#define NET_UNPACK_MAX_SLICE_PAGES \
+  (NET_UNPACK_MAX_SLICE_SIZE / SLICE_PAGE_SIZE * 2)  // * 2 for slack, wasteful..
+
+struct netUnpackMeta {
+  loadMeta mem[NCCL_NET_DEVICE_UNPACK_MAX_QUEUE_DEPTH][NET_UNPACK_MAX_SLICE_PAGES];
+  uint64_t cnt[NCCL_NET_DEVICE_UNPACK_MAX_QUEUE_DEPTH];
+};
+
+struct unpackNetDeviceHandle {
+  struct netUnpackMeta *meta;  // mapped
+  void* bounce_buf;
+  uint64_t head;
+};
+
+/****** shared memory ******/
+
+#define NET_UNPACK_MAX_GROUPS 16 // Forked from NCCL_MAX_GROUPS in devcomm.h
+#define NET_UNPACK_MAX_NPEERS 2  // The most you should have is 2 network peers per-group (indexed by index)
+#define WARP_SHM_PAGE_CNT 4
+#define WARP_SHM_SIZE (WARP_SHM_PAGE_CNT * sizeof(union loadMeta))
+struct unpackShmem {
+  void* bounce_buf;
+};
+
+struct unpackGroupShmem {
+  int unpackNetDeviceIndexMask; // We store a single unpackNetDeviceIndex because only one peer can be network recv
+  uint64_t head;
+  struct netUnpackMeta* g_meta[NET_UNPACK_MAX_NPEERS]; // head of handle to index into meta for meta copy
+};
+
+#endif // NET_DEVICE_UNPACK_DEFS_H_

src/device/onerank.cu

@@ -0,0 +1,79 @@
+/*************************************************************************
+ * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include "alloc.h"
+#include "collectives.h"
+#include "common_kernel.h"
+#include "common.h"
+#include <cuda_runtime.h>
+
+namespace {
+  template<typename RedOp>
+  __global__ __launch_bounds__(512, 1)
+  void oneRankReduce(void* dst, void* src, size_t nElts, uint64_t redOpArg, bool redOpArgIsPtr) {
+    using T = typename RedOp::EltType;
+    int tid = threadIdx.x;
+    int tn = blockDim.x;
+    int bid = blockIdx.x;
+    int bn = gridDim.x;
+
+    // each block/channel gets a roughly equal segment of 16 byte packs
+    constexpr int EltPerPack = 16/sizeof(T);
+    intptr_t i0 = (bid+0)*alignUp(nElts/bn, EltPerPack);
+    intptr_t i1 = (bid+1)*alignUp(nElts/bn, EltPerPack);
+    i0 = min(i0, nElts);
+    i1 = min(i1, nElts);
+    src = (T*)src + i0;
+    dst = (T*)dst + i0;
+
+    if (redOpArgIsPtr) {
+      if (redOpArg%2 != 0) {
+        redOpArg = *reinterpret_cast<uint8_t*>(redOpArg);
+      } else if (redOpArg%4 != 0) {
+        redOpArg = *reinterpret_cast<uint16_t*>(redOpArg);
+      } else if (redOpArg%8 != 0) {
+        redOpArg = *reinterpret_cast<uint32_t*>(redOpArg);
+      } else {
+        redOpArg = *reinterpret_cast<uint64_t*>(redOpArg);
+      }
+    }
+    reduceCopy<COLL_UNROLL, RedOp, T, 0,1,1, 0,1,1, /*PreOpSrcs=*/1>
+      (tid, tn, redOpArg, &redOpArg, true, 1, &src, 1, &dst, i1-i0);
+  }
+}
+
+ncclResult_t ncclLaunchOneRank(void* dst, void const* src, size_t nElts, struct ncclDevRedOpFull redOp, ncclDataType_t eltType, cudaStream_t stream) {
+  size_t eltSize = ncclTypeSize(eltType);
+  if (redOp.op != ncclDevPreMulSum) {
+    if (dst != src) {
+      NCCLCHECK(ncclCudaMemcpyAsync((char*)dst, (char*)src, nElts*eltSize, stream));
+    }
+    return ncclSuccess;
+  }
+
+  void const* kernel;
+  switch (eltType) {
+  case ncclInt8:     kernel = (void const*)&oneRankReduce<FuncPreMulSum<int8_t>>; break;
+  case ncclUint8:    kernel = (void const*)&oneRankReduce<FuncPreMulSum<uint8_t>>; break;
+  case ncclInt32:    kernel = (void const*)&oneRankReduce<FuncPreMulSum<int32_t>>; break;
+  case ncclUint32:   kernel = (void const*)&oneRankReduce<FuncPreMulSum<uint32_t>>; break;
+  case ncclInt64:    kernel = (void const*)&oneRankReduce<FuncPreMulSum<int64_t>>; break;
+  case ncclUint64:   kernel = (void const*)&oneRankReduce<FuncPreMulSum<uint64_t>>; break;
+  case ncclFloat16:  kernel = (void const*)&oneRankReduce<FuncPreMulSum<half>>; break;
+  #if defined(RCCL_BFLOAT16)
+  case ncclBfloat16: kernel = (void const*)&oneRankReduce<FuncPreMulSum<rccl_bfloat16>>; break;
+  #endif
+  case ncclFloat32:  kernel = (void const*)&oneRankReduce<FuncPreMulSum<float>>; break;
+  case ncclFloat64:  kernel = (void const*)&oneRankReduce<FuncPreMulSum<double>>; break;
+  default: return ncclInvalidArgument;
+  }
+  dim3 grid = {0, 1, 1};
+  grid.x = std::min(32, (int)divUp(nElts*eltSize, 16<<10));
+  dim3 block = {512, 1, 1};
+  void* args[5] = {&dst, &src, &nElts, &redOp.scalarArg, &redOp.scalarArgIsPtr};
+  CUDACHECK(cudaLaunchKernel(kernel, grid, block, args, 0, stream));
+  return ncclSuccess;
+}

src/device/op128.h

@@ -0,0 +1,452 @@
+/*************************************************************************
+ * Copyright (c) 2016-2019, NVIDIA CORPORATION. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef OP128_H_
+#define OP128_H_
+
+#include <type_traits>
+
+inline __device__ void load128(const uint64_t* ptr, uint64_t &v0, uint64_t &v1) {
+  v0 = __builtin_nontemporal_load(ptr);
+  v1 = __builtin_nontemporal_load(ptr+1);
+}
+
+inline __device__ void store128(uint64_t* ptr, uint64_t v0, uint64_t v1) {
+  __builtin_nontemporal_store(v0, ptr);
+  __builtin_nontemporal_store(v1, ptr+1);
+}
+
+inline __device__ uint64_t* shmemCvtPtr(volatile uint64_t* shmemGenericPtr) {
+  return (uint64_t*)shmemGenericPtr;
+}
+
+inline __device__ void loadShmem128(uint64_t* shmemAsmPtr, uint64_t &v0, uint64_t &v1) {
+  v0 = *(shmemAsmPtr);
+  v1 = *(shmemAsmPtr+1);
+}
+
+inline __device__ void storeShmem128(uint64_t* shmemAsmPtr, uint64_t v0, uint64_t v1) {
+  *(shmemAsmPtr) = v0;
+  *(shmemAsmPtr+1) = v1;
+}
+
+template<typename T>
+inline __device__ void loadShmemMisaligned128(T *ptr, uint64_t &v0, uint64_t &v1) {
+  union {
+    uint32_t tmp4[4];
+    uint64_t tmp8[2];
+  };
+  if(sizeof(T) < 4) {
+    uint32_t *ptr4 = reinterpret_cast<uint32_t*>(reinterpret_cast<uintptr_t>(ptr) & -uintptr_t(4));
+    #pragma unroll
+    for(int e=0; e < 4; e++) {
+      // Produce 4 bytes of sub-register type by reading 2 4-byte
+      // aligned values and shifting.
+      uint32_t lo, hi;
+      lo = __builtin_nontemporal_load(ptr4+e+0);
+      hi = __builtin_nontemporal_load(ptr4+e+1);
+      tmp4[e] = __funnelshift_r(lo, hi, 8*(int(reinterpret_cast<uintptr_t>(ptr))%4));
+    }
+  }
+  else if(sizeof(T) == 4) {
+    #pragma unroll
+    for(int e=0; e < 4; e++)
+      tmp4[e] = __builtin_nontemporal_load(ptr+e);
+  }
+  else /*sizeof(T)==8*/ {
+    #pragma unroll
+    for(int e=0; e < 2; e++)
+      tmp8[e] = __builtin_nontemporal_load(ptr+e);
+  }
+  v0 = tmp8[0];
+  v1 = tmp8[1];
+}
+
+
+template<typename T>
+__device__ __forceinline__ uint32_t cvta_to_shared(T* ptr) {
+  return (uint32_t)(uint64_t)(ptr);
+}
+template<typename T>
+__device__ __forceinline__ uintptr_t cvta_to_global(T* ptr) {
+  return (uintptr_t)(ptr);
+}
+
+template<typename T>
+__device__ __forceinline__ T* cvta_from_shared(uint32_t shptr) {
+  return (T*)shptr;
+}
+template<typename T>
+__device__ __forceinline__ T* cvta_from_global(uintptr_t gptr) {
+  return (T*)gptr;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// BytePack<Size>: struct of bytes.
+
+template<int Size>
+union BytePack;
+template<>
+union BytePack<0> {};
+template<>
+union BytePack<1> {
+  uint8_t u8, native;
+};
+template<>
+union BytePack<2> {
+  BytePack<1> half[2];
+  uint8_t u8[2];
+  uint16_t u16, native;
+};
+template<>
+union BytePack<4> {
+  BytePack<2> half[2];
+  uint8_t u8[4];
+  uint16_t u16[2];
+  uint32_t u32, native;
+};
+template<>
+union BytePack<8> {
+  BytePack<4> half[2];
+  uint8_t u8[8];
+  uint16_t u16[4];
+  uint32_t u32[2];
+  uint64_t u64, native;
+};
+template<>
+union alignas(16) BytePack<16> {
+  BytePack<8> half[2];
+  uint8_t u8[16];
+  uint16_t u16[8];
+  uint32_t u32[4];
+  uint64_t u64[2];
+  ulong2 ul2, native;
+#if !defined(USE_INDIRECT_FUNCTION_CALL) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+  inline __device__ BytePack<16>& operator=(BytePack<16> other) {
+    u64[0] = other.u64[0];
+    u64[1] = other.u64[1];
+    return *this;
+  }
+#endif
+};
+
+template<typename T>
+struct BytePackOf {
+  static constexpr int Size = sizeof(T);
+  using Pack = BytePack<Size>;
+};
+template<>
+struct BytePackOf<BytePack<0>> {
+  static constexpr int Size = 0;
+  using Pack = BytePack<0>;
+};
+
+template<typename T>
+__device__ __forceinline__ typename BytePackOf<T>::Pack toPack(T value)  {
+  union { typename BytePackOf<T>::Pack p; T v; };
+  v = value;
+  return p;
+}
+
+template<typename T>
+__device__ __forceinline__ T fromPack(typename BytePackOf<T>::Pack pack)  {
+  union { typename BytePackOf<T>::Pack p; T v; };
+  p = pack;
+  return v;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Load/store of BytePack<?> using integral addresses.
+
+template<int Size> __device__ BytePack<Size> ld_global(uintptr_t addr);
+template<int Size> __device__ BytePack<Size> ld_volatile_global(uintptr_t addr);
+//template<int Size> __device__ BytePack<Size> ld_shared(uint32_t addr);
+//template<int Size> __device__ BytePack<Size> ld_volatile_shared(uint32_t addr);
+template<int Size> __device__ void st_global(uintptr_t addr, BytePack<Size> value);
+//template<int Size> __device__ void st_shared(uint32_t addr, BytePack<Size> value);
+
+template<> __device__ __forceinline__ BytePack<0> ld_global<0>(uintptr_t addr) { return {}; }
+template<> __device__ __forceinline__ BytePack<0> ld_volatile_global<0>(uintptr_t addr) { return {}; }
+//template<> __device__ __forceinline__ BytePack<0> ld_shared<0>(uint32_t addr) { return {}; }
+//template<> __device__ __forceinline__ BytePack<0> ld_volatile_shared<0>(uint32_t addr) { return {}; }
+template<> __device__ __forceinline__ void st_global<0>(uintptr_t addr, BytePack<0> value) {}
+//template<> __device__ __forceinline__ void st_shared<0>(uint32_t addr, BytePack<0> value) {}
+
+#ifdef NEED_CHECKING
+template<int Size> __device__ BytePack<Size> ld_global(uintptr_t addr);
+template<int Size> __device__ BytePack<Size> ld_shared(uint32_t addr);
+template<int Size> __device__ BytePack<Size> ld_volatile_global(uintptr_t addr);
+template<int Size> __device__ BytePack<Size> ld_volatile_shared(uint32_t addr);
+template<int Size> __device__ BytePack<Size> ld_relaxed_gpu_global(uintptr_t addr);
+template<int Size> __device__ void st_global(uintptr_t addr, BytePack<Size> value);
+template<int Size> __device__ void st_shared(uint32_t addr, BytePack<Size> value);
+template<int Size> __device__ void st_relaxed_gpu_global(uintptr_t addr, BytePack<Size> value);
+
+template<> __device__ __forceinline__ BytePack<0> ld_global<0>(uintptr_t addr) { return {}; }
+template<> __device__ __forceinline__ BytePack<0> ld_shared<0>(uint32_t addr) { return {}; }
+template<> __device__ __forceinline__ BytePack<0> ld_volatile_global<0>(uintptr_t addr) { return {}; }
+template<> __device__ __forceinline__ BytePack<0> ld_volatile_shared<0>(uint32_t addr) { return {}; }
+template<> __device__ __forceinline__ BytePack<0> ld_relaxed_gpu_global<0>(uintptr_t addr) { return {}; }
+template<> __device__ __forceinline__ void st_global<0>(uintptr_t addr, BytePack<0> value) {}
+template<> __device__ __forceinline__ void st_shared<0>(uint32_t addr, BytePack<0> value) {}
+template<> __device__ __forceinline__ void st_relaxed_gpu_global<0>(uintptr_t addr, BytePack<0> value) {}
+#endif
+
+// Used to define implementations for above prototypes.
+#define DEFINE_ld_st__size_space(bytes, data_cxx_ty, data_ptx_ty, data_reg_ty, space, addr_cxx_ty, addr_reg_ty) \
+  template<> \
+  __device__ __forceinline__ BytePack<bytes> ld_##space<bytes>(addr_cxx_ty addr) { \
+    data_cxx_ty tmp; \
+    tmp = *((data_cxx_ty *)addr); \
+    BytePack<bytes> ans; \
+    ans.native = tmp; \
+    return ans; \
+  } \
+  template<> \
+  __device__ __forceinline__ BytePack<bytes> ld_volatile_##space<bytes>(addr_cxx_ty addr) { \
+    data_cxx_ty tmp; \
+    tmp =  __builtin_nontemporal_load((data_cxx_ty *)addr); \
+    BytePack<bytes> ans; \
+    ans.native = tmp; \
+    return ans; \
+  } \
+  template<> \
+  __device__ __forceinline__ void st_##space<bytes>(addr_cxx_ty addr, BytePack<bytes> value) { \
+    data_cxx_ty tmp = value.native; \
+    *((data_cxx_ty *)addr) = tmp; \
+  }
+
+#ifdef NEED_CHECKING
+#if __CUDA_ARCH__ >= 700
+  #define PTX_relaxed_gpu "relaxed.gpu"
+#else
+  #define PTX_relaxed_gpu "volatile"
+#endif
+
+#define DEFINE_ld_st_gpu_relaxed__size(bytes, data_cxx_ty, data_ptx_ty, data_reg_ty) \
+  template<> \
+  __device__ __forceinline__ BytePack<bytes> ld_relaxed_gpu_global<bytes>(uintptr_t addr) { \
+    data_cxx_ty tmp; \
+    asm("ld." PTX_relaxed_gpu ".global." #data_ptx_ty " %0, [%1];" : "="#data_reg_ty(tmp) : "l"(addr)); \
+    BytePack<bytes> ans; \
+    ans.native = tmp; \
+    return ans; \
+  } \
+  template<> \
+  __device__ __forceinline__ void st_relaxed_gpu_global<bytes>(uintptr_t addr, BytePack<bytes> value) { \
+    data_cxx_ty tmp = value.native; \
+    asm volatile("st." PTX_relaxed_gpu ".global." #data_ptx_ty " [%0], %1;" :: "l"(addr), #data_reg_ty(tmp) : "memory"); \
+  }
+#endif
+
+#define DEFINE_ld_st__size(bytes, data_cxx_ty, data_ptx_ty, data_reg_ty) \
+  DEFINE_ld_st__size_space(bytes, data_cxx_ty, data_ptx_ty, data_reg_ty, global, uintptr_t, l)
+  // DEFINE_ld_st__size_space(bytes, data_cxx_ty, data_ptx_ty, data_reg_ty, shared, uint32_t, r)
+  // DEFINE_ld_st_gpu_relaxed__size(bytes, data_cxx_ty, data_ptx_ty, data_reg_ty)
+
+// Single-byte types use 4-byte registers since there is no 1-byte register
+// character for asm blocks. See https://docs.nvidia.com/cuda/inline-ptx-assembly/index.html#constraints
+DEFINE_ld_st__size(1, uint32_t, b8, r)
+DEFINE_ld_st__size(2, uint16_t, b16, h)
+DEFINE_ld_st__size(4, uint32_t, b32, r)
+DEFINE_ld_st__size(8, uint64_t, b64, l)
+
+#undef DEFINE_ld_st__size_space
+#undef DEFINE_ld_st__size
+
+#define DEFINE_ld_st_16__space(space, addr_cxx_ty, addr_reg_ty) \
+  template<> \
+  __device__ __forceinline__ BytePack<16> ld_##space<16>(addr_cxx_ty addr) { \
+    BytePack<16> ans; \
+    ans.u64[0] = *((uint64_t*)addr); \
+    ans.u64[1] = *((uint64_t*)addr+1); \
+    return ans; \
+  } \
+  template<> \
+  __device__ __forceinline__ BytePack<16> ld_volatile_##space<16>(addr_cxx_ty addr) { \
+    BytePack<16> ans; \
+    ans.u64[0] = __builtin_nontemporal_load((uint64_t*)addr); \
+    ans.u64[1] = __builtin_nontemporal_load((uint64_t*)addr+1); \
+    return ans; \
+  } \
+  template<> \
+  __device__ __forceinline__ void st_##space<16>(addr_cxx_ty addr, BytePack<16> value) { \
+    __builtin_nontemporal_store(value.u64[0], (uint64_t*)addr); \
+    __builtin_nontemporal_store(value.u64[1], (uint64_t*)addr+1); \
+  }
+
+DEFINE_ld_st_16__space(global, uintptr_t, l)
+// DEFINE_ld_st_16__space(shared, uint32_t, r)
+#undef DEFINE_ld_st_16
+
+#ifdef NEED_CHECKING
+template<>
+__device__ __forceinline__ BytePack<16> ld_relaxed_gpu_global<16>(uintptr_t addr) {
+  BytePack<16> ans;
+  asm("ld." PTX_relaxed_gpu ".global.v2.b64 {%0,%1}, [%2];" : "=l"(ans.u64[0]), "=l"(ans.u64[1]) : "l"(addr));
+  return ans;
+}
+template<>
+__device__ __forceinline__ void st_relaxed_gpu_global<16>(uintptr_t addr, BytePack<16> value) {
+  asm volatile("st." PTX_relaxed_gpu ".global.v2.b64 [%0], {%1,%2};" :: "l"(addr), "l"(value.u64[0]), "l"(value.u64[1]) : "memory");
+}
+
+#undef PTX_relaxed_gpu
+#endif
+
+////////////////////////////////////////////////////////////////////////////////
+// Atomic load/store using c++ pointers.
+
+__device__ __forceinline__ uint64_t ld_volatile_global(uint64_t *ptr) {
+  uint64_t ans;
+  ans = __builtin_nontemporal_load(ptr);
+  return ans;
+}
+__device__ __forceinline__ uint64_t ld_relaxed_sys_global(uint64_t *ptr) {
+  uint64_t ans;
+  ans = __builtin_nontemporal_load(ptr);
+  return ans;
+}
+#ifdef NEED_CHECKING
+__device__ __forceinline__ uint64_t ld_relaxed_gpu_global(uint64_t *ptr) {
+  uint64_t ans;
+  #if __CUDA_ARCH__ >= 700
+    asm("ld.relaxed.gpu.global.u64 %0, [%1];" : "=l"(ans) : "l"(cvta_to_global(ptr)));
+  #else
+    asm("ld.volatile.global.u64 %0, [%1];" : "=l"(ans) : "l"(cvta_to_global(ptr)));
+  #endif
+  return ans;
+}
+#endif
+__device__ __forceinline__ uint64_t ld_acquire_sys_global(uint64_t *ptr) {
+  uint64_t ans;
+  ans = __atomic_load_n(ptr ,__ATOMIC_SEQ_CST);
+  return ans;
+}
+
+__device__ __forceinline__ void st_volatile_global(uint64_t *ptr, uint64_t val) {
+  __builtin_nontemporal_store(val, ptr);
+}
+__device__ __forceinline__ void st_relaxed_sys_global(uint64_t *ptr, uint64_t val) {
+  __builtin_nontemporal_store(val, ptr);
+}
+__device__ __forceinline__ void st_release_sys_global(uint64_t *ptr, uint64_t val) {
+  __atomic_store_n(ptr, val, __ATOMIC_SEQ_CST);
+}
+
+__device__ __forceinline__ void fence_acq_rel_sys() {
+    //asm volatile("membar.sys;" ::: "memory");
+}
+__device__ __forceinline__ void fence_acq_rel_gpu() {
+    //asm volatile("membar.gl;" ::: "memory");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Multimem stores of BytePack<?>.
+
+template<int Size>
+__device__ __forceinline__ void multimem_st_global(uintptr_t addr, BytePack<Size> val);
+
+#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010
+template<>
+__device__ __forceinline__ void multimem_st_global<0>(uintptr_t addr, BytePack<0> val) {
+  // nop
+}
+template<>
+__device__ __forceinline__ void multimem_st_global<1>(uintptr_t addr, BytePack<1> val) {
+  asm volatile("st.global.b8 [%0], %1;" :: "l"(addr), "r"((uint32_t)val.u8) : "memory");
+}
+template<>
+__device__ __forceinline__ void multimem_st_global<2>(uintptr_t addr, BytePack<2> val) {
+  asm volatile("st.global.b16 [%0], %1;" :: "l"(addr), "h"(val.u16) : "memory");
+}
+template<>
+__device__ __forceinline__ void multimem_st_global<4>(uintptr_t addr, BytePack<4> val) {
+  asm volatile("multimem.st.global.b32 [%0], %1;" :: "l"(addr), "r"(val.u32) : "memory");
+}
+template<>
+__device__ __forceinline__ void multimem_st_global<8>(uintptr_t addr, BytePack<8> val) {
+  asm volatile("multimem.st.global.b64 [%0], %1;" :: "l"(addr), "l"(val.u64) : "memory");
+}
+template<>
+__device__ __forceinline__ void multimem_st_global<16>(uintptr_t addr, BytePack<16> val) {
+  asm volatile("multimem.st.global.v4.f32 [%0], {%1,%2,%3,%4};"
+    :: "l"(addr), "r"(val.u32[0]), "r"(val.u32[1]), "r"(val.u32[2]), "r"(val.u32[3])
+    : "memory");
+}
+#else
+template<int Size>
+__device__ __forceinline__ void multimem_st_global(uintptr_t addr, BytePack<Size> val) {
+  // nop
+}
+#endif
+
+#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010
+// Warp-uniform memory copy from shared address (not generic) to global memory.
+// The number of bytes copied is `min(MaxBytes, nBytesAhead)`, a negative value
+// is interpeted as zero. EltSize is the guaranteed alignment of the addresses and sizes.
+template<int EltSize, int MaxBytes, bool Multimem, typename IntBytes>
+__device__ __forceinline__ void copyGlobalShared_WarpUnrolled(
+    int lane, uintptr_t dstAddr, uint32_t srcAddr, IntBytes nBytesAhead
+  ) {
+  static_assert(std::is_signed<IntBytes>::value, "`IntBytes` must be a signed integral type.");
+  int nBytes = min(nBytesAhead, (IntBytes)MaxBytes);
+  int nFrontBytes = min(nBytes, (16 - int(dstAddr%16))%16);
+  int nMiddleBytes = (nBytes-nFrontBytes) & -16;
+  int nBackBytes = (nBytes-nFrontBytes) % 16;
+
+  { int backLane = WARP_SIZE-1 - lane;
+    bool hasFront = lane*EltSize < nFrontBytes;
+    bool hasBack = backLane*EltSize < nBackBytes;
+    int offset = hasFront ? lane*EltSize : (nBytes - (backLane+1)*EltSize);
+    if (hasFront | hasBack) {
+      BytePack<EltSize> tmp = ld_shared<EltSize>(srcAddr+offset);
+      // Can't use multimem_st since it doesn't support EltSize==2
+      st_global<EltSize>(dstAddr+offset, tmp);
+    }
+  }
+
+  srcAddr += nFrontBytes;
+  int srcMisalign = EltSize < 4 ? (srcAddr%4) : 0;
+  srcAddr += -srcMisalign + lane*16;
+  dstAddr += nFrontBytes + lane*16;
+  nMiddleBytes -= lane*16;
+  #pragma unroll
+  for (int u=0; u < divUp(MaxBytes, WARP_SIZE*16); u++) {
+    if (nMiddleBytes <= 0) break;
+    union {
+      BytePack<4> b4[4];
+      BytePack<16> b16;
+    };
+    b4[0] = ld_shared<4>(srcAddr + 0*4);
+    b4[1] = ld_shared<4>(srcAddr + 1*4);
+    b4[2] = ld_shared<4>(srcAddr + 2*4);
+    b4[3] = ld_shared<4>(srcAddr + 3*4);
+    if (srcMisalign != 0) {
+      BytePack<4> b4_4 = ld_shared<4>(srcAddr + 4*4);
+      b4[0].u32 = __funnelshift_r(b4[0].u32, b4[1].u32, srcMisalign*8);
+      b4[1].u32 = __funnelshift_r(b4[1].u32, b4[2].u32, srcMisalign*8);
+      b4[2].u32 = __funnelshift_r(b4[2].u32, b4[3].u32, srcMisalign*8);
+      b4[3].u32 = __funnelshift_r(b4[3].u32, b4_4.u32, srcMisalign*8);
+    }
+    if (Multimem) multimem_st_global<16>(dstAddr, b16);
+    else          st_global<16>(dstAddr, b16);
+
+    srcAddr += WARP_SIZE*16;
+    dstAddr += WARP_SIZE*16;
+    nMiddleBytes -= WARP_SIZE*16;
+  }
+}
+#else
+template<int EltSize, int MaxBytes, bool Multimem, typename IntBytes>
+__device__ __forceinline__ void copyGlobalShared_WarpUnrolled(
+    int lane, uintptr_t dstAddr, uint32_t srcAddr, IntBytes nBytesAhead
+  ) {
+  // nop
+}
+#endif
+
+#endif

src/device/primitives.h

@@ -0,0 +1,170 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef NCCL_PRIMITIVES_H_
+#define NCCL_PRIMITIVES_H_
+
+#include <type_traits>
+#include "reduce_kernel.h" // for reduction funcs
+#include "common_kernel.h"
+#include "common.h"
+
+#define NCCL_SPINS_BEFORE_CHECK_ABORT 1000000
+
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+#define barrier_by_group() do { \
+  if (nthreads == NCCL_MAX_NTHREADS) { \
+    __builtin_amdgcn_s_barrier(); \
+  } else { \
+    const int w = threadIdx.x/WARP_SIZE; \
+    const int wid = threadIdx.x%WARP_SIZE; \
+    if (wid == 0) { \
+      barrier_next[w] += nthreads/WARP_SIZE; \
+      atomicAdd((unsigned long long *)barriers, 1); \
+      while (atomicAdd((unsigned long long *)barriers, 0) < barrier_next[w]) __builtin_amdgcn_s_sleep(1); \
+      __asm__ __volatile__("s_wakeup"); \
+    } \
+  } \
+} while (0)
+#else
+#define barrier_by_group() do { \
+  if (nthreads == NCCL_MAX_NTHREADS) { \
+    __threadfence(); __builtin_amdgcn_s_barrier(); \
+  } else { \
+    const int w = threadIdx.x/WARP_SIZE; \
+    const int wid = threadIdx.x%WARP_SIZE; \
+    __threadfence(); \
+    if (wid == 0) { \
+      barrier_next[w] += nthreads/WARP_SIZE; \
+      atomicAdd((unsigned long long *)barriers, 1); \
+      while (atomicAdd((unsigned long long *)barriers, 0) < barrier_next[w]) __builtin_amdgcn_s_sleep(1); \
+      __asm__ __volatile__("s_wakeup"); \
+    } \
+  } \
+} while (0)
+#endif
+
+/* Protocol classes: ProtoSimple, ProtoLL, ProtoLL128
+ * We use these as template args to the Primtiives class instead of integral
+ * enums (e.g. NCCL_PROTO_LL) because for SIMPLE we need to carry a few extra
+ * numbers. Also these types hold methods which let us compute numbers important
+ * to how that protocol operates with a consistent interface so that our
+ * algorithm code can operate protocol parametrically.
+ */
+template<int SlicePerChunk_1, int StepPerSlice_1, int Unroll_1 = COLL_UNROLL, int MultimemSrcs_1 = 0, int MultimemDsts_1 = 0>
+struct ProtoSimple {
+  static constexpr int Id = NCCL_PROTO_SIMPLE;
+  static constexpr int SlicePerChunk = SlicePerChunk_1;
+  static constexpr int StepPerSlice = StepPerSlice_1;
+  static constexpr int Unroll = Unroll_1;
+  static constexpr int MultimemSrcs = MultimemSrcs_1;
+  static constexpr int MultimemDsts = MultimemDsts_1;
+
+  // Data bytes (no flags etc) in one step of the fifo queue.
+  __device__ static int calcBytePerStep() {
+    return ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS;
+  }
+  // Granularity of data bytes transferred per thread.
+  __device__ static int calcBytePerGrain() {
+    return sizeof(uint64_t); // Bogus value? Nobody queries this metric for simple.
+  }
+  // Group width is how many consecutive group values a subchannel occupies.
+  static constexpr int MaxGroupWidth = 1;
+};
+
+struct ProtoLL {
+  static constexpr int Id = NCCL_PROTO_LL;
+
+  // Data bytes (no flags etc) in one step of the fifo queue.
+  __device__ static int calcBytePerStep() {
+    return ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/2; // Half is data
+  }
+  // Granularity of data bytes transferred per thread.
+  __device__ static int calcBytePerGrain() {
+    return sizeof(uint64_t); // One 16-byte line has 8-bytes of data
+  }
+  // Group width is how many consecutive group values a subchannel occupies.
+  static constexpr int MaxGroupWidth = 1;
+};
+
+struct ProtoLL128 {
+  static constexpr int Id = NCCL_PROTO_LL128;
+
+  // Data bytes (no flags etc) in one step of the fifo queue.
+  __device__ static int calcBytePerStep() {
+    return (ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS)*NCCL_LL128_DATAELEMS/NCCL_LL128_LINEELEMS;
+  }
+  // Granularity of data bytes transferred per thread.
+  __device__ static int calcBytePerGrain() {
+    return NCCL_LL128_SHMEM_ELEMS_PER_THREAD*NCCL_LL128_DATAELEMS*sizeof(uint64_t)/NCCL_LL128_LINEELEMS;
+  }
+  // Group width is how many consecutive group values a subchannel occupies.
+  static constexpr int MaxGroupWidth = 1;
+};
+
+/* Fan (as in fan-in & fan-out) classes hold recv and send counts. The template
+ * arguments are static bounds on the maximum values. Asymmetric counts are
+ * independent. Symmetric is a static guarantee that nrecv==nsend, so it only
+ * stores one value at runtime. This optimization save 32-bit register, but more
+ * importantly uses fewer predicate registers when unrolling loops.
+ */
+template<int MaxRecv_, int MaxSend_>
+struct FanAsymmetric {
+  static constexpr int MaxRecv = MaxRecv_, MaxSend = MaxSend_;
+  int nr, ns;
+  FanAsymmetric() = default;
+  __device__ FanAsymmetric(int nrecv, int nsend): nr(nrecv), ns(nsend) {
+    // assert(nrecv <= MaxRecv && nsend <= MaxSend);
+  }
+  __device__ int nrecv() const { return MaxRecv ? nr : 0; }
+  __device__ int nsend() const { return MaxSend ? ns : 0; }
+};
+
+template<int MaxArity>
+struct FanSymmetric {
+  static constexpr int MaxRecv = MaxArity, MaxSend = MaxArity;
+  int n;
+  FanSymmetric() = default;
+  __device__ FanSymmetric(int nrecv, int nsend): n(nrecv) {
+    // assert(nrecv == nsend && nrecv <= MaxArity);
+  }
+  __device__ int nrecv() const { return n; }
+  __device__ int nsend() const { return n; }
+};
+
+// The primitives class. Specialized per protocol in the other headers.
+template<typename T, typename RedOp, typename Fan, int Direct, typename Proto, int P2p>
+class Primitives;
+
+// Used by LL & LL128 to implement direct members in the naive way.
+template<typename RealPrimitives>
+struct PrimitivesWithoutDirect {
+  __device__ void directSend(intptr_t inpIx, intptr_t outIx, int eltN) {
+    static_cast<RealPrimitives*>(this)->send(inpIx, eltN);
+  }
+  __device__ void directSendFromOutput(intptr_t outIx, int eltN) {
+    static_cast<RealPrimitives*>(this)->sendFromOutput(outIx, eltN);
+  }
+  __device__ void directRecv(intptr_t outIx, int eltN) {
+    static_cast<RealPrimitives*>(this)->recv(outIx, eltN, /*postOp=*/false);
+  }
+  __device__ void directCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    static_cast<RealPrimitives*>(this)->copySend(inpIx, outIx, eltN, postOp);
+  }
+  __device__ void directRecvCopySend(intptr_t outIx, int eltN) {
+    static_cast<RealPrimitives*>(this)->recvCopySend(outIx, eltN, /*postOp=*/false);
+  }
+  __device__ void directRecvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    // Direct is only for the send part
+    static_cast<RealPrimitives*>(this)->recvReduceCopySend(inpIx, outIx, eltN, postOp);
+  }
+};
+
+#include "prims_simple.h"
+#include "prims_ll.h"
+#include "prims_ll128.h"
+#endif

src/device/prims_ll.h

@@ -0,0 +1,775 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#if defined(ENABLE_NPKIT)
+#include "npkit/npkit.h"
+#endif
+
+#ifdef __GFX11__
+#define LL_STORE(SRC, DST) \
+  __atomic_store_n((DST), (SRC), __ATOMIC_RELAXED)
+#define LL_LOAD(SRC) \
+  __atomic_load_n(SRC, __ATOMIC_RELAXED)
+#else
+#define LL_STORE(SRC, DST) \
+  __builtin_nontemporal_store((SRC), (DST))
+#define LL_LOAD(SRC) \
+  __builtin_nontemporal_load(SRC)
+#endif
+
+template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
+class Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>:
+  public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL, P2p>> {
+
+  // In the case of Fan::MaxRecv == 0, we need to force MaxRecv to 1 for this to compile
+  // This is because of a recv buffer which is allocated to MaxRecv length in send-only cases
+  static constexpr int MaxRecv = Fan::MaxRecv > 1 ? Fan::MaxRecv : 1;
+  static constexpr int MaxSend = Fan::MaxSend;
+  static constexpr int Input=0, Output=1;
+  RedOp redOp;
+  const int tid;
+  const int nthreads;
+  const int wid;
+  const int group;
+  const int stepLines;
+  Fan fan;
+  T *userBufs[2];
+  struct ncclConnInfo* recvConn = NULL;
+  volatile uint64_t* recvConnHeadPtr = NULL;
+  uint64_t recvConnHead;
+
+  struct ncclConnInfo* sendConn = NULL;
+  volatile int* sendConnFifoPtr = NULL;
+  volatile uint64_t* sendConnHeadPtr = NULL;
+  uint64_t sendConnHead;
+  uint64_t sendConnHeadCache; // Cache last seen value
+
+  uint64_t recvStep[MaxRecv];
+  uint64_t sendStep[MaxSend];
+  union ncclLLFifoLine* recvBuff[MaxRecv];
+  union ncclLLFifoLine* sendBuff[MaxSend];
+
+#if defined(ENABLE_NPKIT)
+public:
+  int npKitCtxIdx = 0;
+  uint64_t npKitDataProcessEntryTime = 0;
+  uint64_t npKitDataProcessExitTime = 0;
+  uint64_t npKitDataProcessTotalTime = 0;
+private:
+#endif
+
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+  uint64_t npKitWaitRecvDataProcessSize = 0;
+  uint64_t npKitWaitRecvEntryTime = 0;
+  uint64_t npKitWaitRecvExitTime = 0;
+  uint64_t npKitWaitRecvTotalTime = 0;
+#endif
+
+  inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*stepLines; }
+  inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*stepLines; }
+  inline __device__ union ncclLLFifoLine* recvPtr(int i) { return recvBuff[i]+recvOffset(i); }
+  inline __device__ union ncclLLFifoLine* sendPtr(int i) { return sendBuff[i]+sendOffset(i); }
+  inline __device__ uint32_t recvFlag(int i) { return NCCL_LL_FLAG(recvStep[i]+1); }
+  inline __device__ uint32_t sendFlag(int i) { return NCCL_LL_FLAG(sendStep[i]+1); }
+
+  uint64_t* barriers;
+  uint64_t* barrier_next;
+
+  inline __device__ void barrier() {
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+    if (nthreads != WARP_SIZE)
+      barrier_by_group();
+#else
+    asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(15-group));
+#endif
+  }
+
+  uint32_t abort = 0;
+
+  inline __device__ int checkAbort(int &spins, int send) {
+    spins++;
+    if (abort == 0 && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
+      abort = __atomic_load_n((ncclShmem.comm.abortFlag), __ATOMIC_SEQ_CST);
+      spins = 0;
+    }
+    return abort;
+  }
+
+  inline __device__ void waitSend(int nbytes) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_WAIT_SEND_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_LL_WAIT_SEND_ENTRY, nbytes, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    if (sendConnHeadPtr) {
+      int spins = 0;
+      while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
+        __builtin_amdgcn_s_sleep(1);
+        sendConnHeadCache = atomicAdd((unsigned long long *)sendConnHeadPtr, 0);
+        if (checkAbort(spins, 1)) break;
+      }
+      __asm__ __volatile__("s_wakeup");
+      if (sendConnFifoPtr) {
+        int size = ((sendConnHead & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) ? stepLines*sizeof(union ncclLLFifoLine) : nbytes;
+        __atomic_store_n(sendConnFifoPtr+sendConnHead%NCCL_STEPS, (size), __ATOMIC_RELAXED);
+      }
+      sendConnHead += 1;
+    }
+    barrier();
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_WAIT_SEND_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_LL_WAIT_SEND_EXIT, nbytes, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+
+  inline __device__ void incRecv(int i) {
+    recvStep[i] += 1;
+  }
+  inline __device__ void postRecv() {
+    barrier();
+    if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
+  }
+
+  inline __device__ void incSend(int i, int offset) {
+    // LL Cleanup : write all flags in the slice to make sure we don't have
+    // data corruption when flag loops over.
+    if ((sendStep[i] & NCCL_LL_CLEAN_MASK) == NCCL_LL_CLEAN_MASK) {
+      for (int o = offset; o<stepLines; o+=nthreads) storeLL(sendPtr(i)+o, 0, sendFlag(i));
+    }
+    sendStep[i]++;
+  }
+
+  __device__ uint64_t readLL(int offset, int i) {
+    union ncclLLFifoLine* src = recvPtr(i) + offset;
+    uint32_t flag = recvFlag(i);
+    uint32_t data1, flag1, data2, flag2;
+    int spins = 0;
+
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+    int npkitWaitRecvSpins = 0;
+    if (tid == 0) {
+      npKitWaitRecvEntryTime = NPKIT_GET_GPU_TIMESTAMP();
+    }
+#endif
+
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+    union ncclLLFifoLine i4;
+    do {
+      i4.v[0] = LL_LOAD(src->v);
+      i4.v[1] = LL_LOAD(src->v+1);
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+      npkitWaitRecvSpins++;
+#endif
+      if (checkAbort(spins, 0)) break;
+    } while ((i4.flag1 != flag) || (i4.flag2 != flag));
+    uint64_t val64 = (uint64_t)(i4.data1) + (((uint64_t)i4.data2) << 32);
+#else
+    do {
+      asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(data1), "=r"(flag1), "=r"(data2), "=r"(flag2) : "l"(&src->i4));
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+      npkitWaitRecvSpins++;
+#endif
+      if (checkAbort(spins, 0)) break;
+    } while ((flag1 != flag) || (flag2 != flag));
+    uint64_t val64 = data1 + (((uint64_t)data2) << 32);
+#endif
+
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+    if (tid == 0) {
+      npKitWaitRecvExitTime = NPKIT_GET_GPU_TIMESTAMP();
+      npKitWaitRecvTotalTime += (npKitWaitRecvExitTime - npKitWaitRecvEntryTime) * (npkitWaitRecvSpins - 1) / npkitWaitRecvSpins;
+    }
+#endif
+
+    return val64;
+  }
+
+  template<int BeginIx>
+  __device__ void readLLBeginAll(int offset, ncclLLFifoLine(&line)[MaxRecv]) {
+    #pragma unroll
+    for (int i=BeginIx; i < MaxRecv; i++) {
+      if (i < fan.nrecv()) {
+        union ncclLLFifoLine* src = recvPtr(i) + offset;
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+        line[i].v[0] = LL_LOAD(src->v);
+        line[i].v[1] = LL_LOAD(src->v+1);
+#else
+        asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(line[i].data1), "=r"(line[i].flag1), "=r"(line[i].data2), "=r"(line[i].flag2) : "l"(&src->i4));
+#endif
+      }
+    }
+  }
+  __device__ uint64_t readLLFinish(int offset, ncclLLFifoLine(&line)[MaxRecv], int i) {
+    union ncclLLFifoLine* src = recvPtr(i) + offset;
+    uint32_t flag = recvFlag(i);
+    int spins = 0;
+
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+    int npkitWaitRecvSpins = 0;
+    if (tid == 0) {
+      npKitWaitRecvEntryTime = NPKIT_GET_GPU_TIMESTAMP();
+    }
+#endif
+
+    do {
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+      line[i].v[0] = LL_LOAD(src->v);
+      line[i].v[1] = LL_LOAD(src->v+1);
+#else
+      asm("ld.volatile.global.v4.u32 {%0,%1,%2,%3}, [%4];" : "=r"(line[i].data1), "=r"(line[i].flag1), "=r"(line[i].data2), "=r"(line[i].flag2) : "l"(&src->i4));
+#endif
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+      npkitWaitRecvSpins++;
+#endif
+      if (checkAbort(spins, 0)) break;
+    } while(line[i].flag1 != flag || line[i].flag2 != flag);
+    uint64_t val64 = line[i].data1 + (((uint64_t)line[i].data2) << 32);
+
+#if defined(ENABLE_NPKIT) && (defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT) || defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME))
+    if (tid == 0) {
+      npKitWaitRecvExitTime = NPKIT_GET_GPU_TIMESTAMP();
+      npKitWaitRecvTotalTime += (npKitWaitRecvExitTime - npKitWaitRecvEntryTime) * (npkitWaitRecvSpins - 1) / npkitWaitRecvSpins;
+    }
+#endif
+
+    return val64;
+  }
+
+  __device__ void storeLL(union ncclLLFifoLine* dst, uint64_t val, uint32_t flag) {
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+    union ncclLLFifoLine i4;
+    i4.data1 = val & 0xffffffff;
+    i4.flag1 = flag;
+    i4.data2 = (val >> 32);
+    i4.flag2 = flag;
+    LL_STORE(i4.v[0], dst->v);
+    LL_STORE(i4.v[1], dst->v+1);
+#else
+    asm volatile("st.volatile.global.v4.u32 [%0], {%1,%2,%3,%4};" :: "l"(&dst->i4), "r"((uint32_t)val), "r"(flag), "r"((uint32_t)(val >> 32)), "r"(flag));
+#endif
+  }
+
+  static constexpr int EltPerLine = sizeof(uint64_t)/sizeof(T);
+
+  template<typename U>
+  __device__ static U load(U *src) {
+    union {
+      U elt;
+      uint8_t u1;
+      uint16_t u2;
+      uint32_t u4;
+      uint64_t u8;
+    };
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+    if(sizeof(U) == 1)
+      u1 = LL_LOAD((uint8_t*)src);
+    else if(sizeof(U) == 2)
+      u2 = LL_LOAD((uint16_t*)src);
+    else if(sizeof(U) == 4)
+      u4 = LL_LOAD((uint32_t*)src);
+    else
+      u8 = LL_LOAD((uint64_t*)src);
+#else
+    if(sizeof(U) == 1)
+      asm("ld.volatile.global.b8 %0,[%1];" : "=r"(u4) : "l"(src));
+    else if(sizeof(U) == 2)
+      asm("ld.volatile.global.b16 %0,[%1];" : "=h"(u2) : "l"(src));
+    else if(sizeof(U) == 4)
+      asm("ld.volatile.global.b32 %0,[%1];" : "=r"(u4) : "l"(src));
+    else
+      asm("ld.volatile.global.b64 %0,[%1];" : "=l"(u8) : "l"(src));
+#endif
+    return elt;
+  }
+
+  template<typename U>
+  __device__ static void store(U *dst, U val) {
+    union {
+      U elt;
+      uint8_t u1;
+      uint16_t u2;
+      uint32_t u4;
+      uint64_t u8;
+    };
+    elt = val;
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+    if(sizeof(U) == 1)
+      LL_STORE(u1, (uint8_t*)dst);
+    else if(sizeof(U) == 2)
+      LL_STORE(u2, (uint16_t*)dst);
+    else if(sizeof(U) == 4)
+      LL_STORE(u4, (uint32_t*)dst);
+    else
+      LL_STORE(u8, (uint64_t*)dst);
+#else
+    if(sizeof(U) == 1)
+      asm("st.volatile.global.b8 [%0],%1;" :: "l"(dst), "r"(u4));
+    else if(sizeof(U) == 2)
+      asm("st.volatile.global.b16 [%0],%1;" :: "l"(dst), "h"(u2));
+    else if(sizeof(U) == 4)
+      asm("st.volatile.global.b32 [%0],%1;" :: "l"(dst), "r"(u4));
+    else
+      asm("st.volatile.global.b64 [%0],%1;" :: "l"(dst), "l"(u8));
+#endif
+  }
+
+  struct DataLoader {
+    int misalign;
+    union {
+      uint32_t u4[sizeof(T) <= 2 ? 3 : 2];
+      uint64_t u8;
+      T elt[EltPerLine];
+    };
+
+    __device__ void loadBegin(T *src, int eltN) {
+      if (sizeof(T) <= 2) {
+        misalign = reinterpret_cast<uintptr_t>(src)%4;
+        uint32_t *p = reinterpret_cast<uint32_t*>(reinterpret_cast<uintptr_t>(src) & -uintptr_t(4));
+        u4[0] = load(p+0);
+        u4[1] = misalign + eltN*sizeof(T) > 4 ? load(p+1) : 0;
+        // u4[2] would be simpler, but that throws warnings on some compilers
+        u4[sizeof(T) <= 2 ? 2 : 0] = misalign + eltN*sizeof(T) > 8 ? load(p+2) : 0;
+      }
+      else {
+        #pragma unroll
+        for(int i=0; i < EltPerLine; i++) {
+          if(i==0 || i < eltN)
+            elt[i] = load(src + i);
+        }
+      }
+    }
+
+    __device__ uint64_t loadFinish() {
+      if (sizeof(T) <= 2) {
+        u4[0] = __funnelshift_r(u4[0], u4[1], 8*misalign);
+        // u4[2] would be simpler, but that throws warnings on some compilers
+        u4[1] = __funnelshift_r(u4[1], u4[sizeof(T) <= 2 ? 2 : 0], 8*misalign);
+      }
+      return u8;
+    }
+  };
+
+  __device__ void storeData(T *dst, uint64_t val, int eltN) {
+    union {
+      uint64_t u8;
+      T elt[EltPerLine];
+    };
+    u8 = val;
+    #pragma unroll
+    for(int i=0; i < EltPerLine; i++) {
+      if (i==0 || i < eltN)
+        //store(dst+i, elt[i]);
+        dst[i] = elt[i];
+    }
+  }
+
+  __device__ void mscclStoreData(T *dst, uint64_t val, int eltN) {
+    union {
+      uint64_t u8;
+      T elt[EltPerLine];
+    };
+    u8 = val;
+    #pragma unroll
+    for(int i=0; i < EltPerLine; i++) {
+      if (i==0 || i < eltN)
+        store(dst+i, elt[i]);
+        // dst[i] = elt[i];
+    }
+  }
+
+  template <int RECV, int SEND, int SrcBuf, int DstBuf>
+  __device__ void LLGenericOp(intptr_t srcIx, intptr_t dstIx, int nelem, bool postOp) {
+    constexpr int SRC = SrcBuf != -1 ? 1 : 0;
+    constexpr int DST = DstBuf != -1 ? 1 : 0;
+    T *srcElts = SrcBuf == -1 ? nullptr : userBufs[SrcBuf] + srcIx;
+    T *dstElts = DstBuf == -1 ? nullptr : userBufs[DstBuf] + dstIx;
+
+    // Always waitSend in case of cleanup
+    nelem = nelem < 0 ? 0 : nelem;
+    if (SEND) waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine));
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT)
+    if (tid == 0) {
+      npKitWaitRecvTotalTime = 0;
+      npKitWaitRecvDataProcessSize = nelem*sizeof(T);
+      NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY,
+          npKitWaitRecvDataProcessSize, 0, NPKIT_GET_GPU_TIMESTAMP(), ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+    if (tid == 0) {
+      npKitWaitRecvTotalTime = 0;
+      npKitDataProcessEntryTime = NPKIT_GET_GPU_TIMESTAMP();
+    }
+#endif
+
+    nelem -= tid*EltPerLine;
+    srcElts += tid*EltPerLine;
+    dstElts += tid*EltPerLine;
+    int offset = tid;
+    int eltPerTrip = nthreads*EltPerLine;
+    while (nelem > 0) {
+      int eltInLine = EltPerLine < nelem ? EltPerLine : nelem;
+
+      DataLoader dl;
+      ncclLLFifoLine line[MaxRecv];
+      uint64_t data, peerData;
+      if (SRC) {
+        dl.loadBegin(srcElts, eltInLine);
+        srcElts += eltPerTrip;
+      }
+      if (RECV) {
+        readLLBeginAll<1>(offset, line);
+        peerData = readLL(offset, 0);
+      }
+      if (SRC) {
+        data = dl.loadFinish();
+        if (SrcBuf == Input) data = applyPreOp(redOp, data);
+      }
+      if (RECV) {
+        data = !SRC ? peerData : applyReduce(redOp, peerData, data);
+        #pragma unroll MaxRecv
+        for (int i=1; i < MaxRecv && i < fan.nrecv(); i++) {
+          peerData = readLLFinish(offset, line, i);
+          data = applyReduce(redOp, peerData, data);
+        }
+      }
+
+      if (postOp) data = applyPostOp(redOp, data);
+
+      // Send : inter-node, then intra-node, then local
+      if (SEND) {
+        for (int i=1; i < MaxSend && i < fan.nsend(); i++)
+          storeLL(sendPtr(i)+offset, data, sendFlag(i));
+        storeLL(sendPtr(0)+offset, data, sendFlag(0));
+      }
+      if (DST) {
+        storeData(dstElts, data, eltInLine);
+        dstElts += eltPerTrip;
+      }
+      nelem -= eltPerTrip;
+      offset += nthreads;
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+    if (tid == 0) {
+      npKitDataProcessExitTime = NPKIT_GET_GPU_TIMESTAMP();
+      npKitDataProcessTotalTime += npKitDataProcessExitTime - npKitDataProcessEntryTime - npKitWaitRecvTotalTime;
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_ENTRY) && defined(ENABLE_NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_LL_DATA_PROCESS_EXIT,
+          npKitWaitRecvDataProcessSize, npKitWaitRecvTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    if (RECV) {
+      for (int i=0; i < MaxRecv; i++) incRecv(i);
+      postRecv();
+    }
+    if (SEND) {
+      for (int i=1; i < MaxSend && i < fan.nsend(); i++)
+        incSend(i, offset);
+      incSend(0, offset);
+    }
+  }
+
+  template <int REDUCE, int COPY, int MULTISRCS, int MULTIDSTS>
+  __device__ __forceinline__ void mscclGenericOp(T** srcs, int nsrcs, T** dsts, int ndsts, int nelem) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_GENERIC_OP_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_MSCCL_GENERIC_OP_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    nelem = nelem < 0 ? 0 : nelem;
+    T *srcElts = srcs[0];
+    T *dstElts = dsts[0];
+    nelem -= tid*EltPerLine;
+    srcElts += tid*EltPerLine;
+    dstElts += tid*EltPerLine;
+    if (MULTISRCS){
+      for (int i = 1; i < nsrcs; i++){
+        srcs[i] += tid*EltPerLine;
+      }
+    }
+    if (MULTIDSTS){
+      for (int i = 1; i < ndsts; i++){
+        dsts[i] += tid*EltPerLine;
+      }
+    }
+    int offset = tid;
+    int eltPerTrip = nthreads*EltPerLine;
+    while (nelem > 0) {
+      int eltInLine = EltPerLine < nelem ? EltPerLine : nelem;
+
+      DataLoader dl;
+      uint64_t data;
+      dl.loadBegin(srcElts, eltInLine);
+      srcElts += eltPerTrip;
+      data = dl.loadFinish();
+      if (REDUCE) {
+        uint64_t dataD;
+        dl.loadBegin(dstElts, eltInLine);
+        dataD = dl.loadFinish();
+        dataD = applyReduce(redOp, dataD, data);
+        if (MULTISRCS){
+          for (int i = 1; i < nsrcs; i++){
+            dl.loadBegin(srcs[i], eltInLine);
+            srcs[i] += eltPerTrip;
+            data = dl.loadFinish();
+            dataD = applyReduce(redOp, dataD, data);
+          }
+        }
+        mscclStoreData(dstElts, dataD, eltInLine);
+        dstElts += eltPerTrip;
+      }
+      if (COPY){
+        mscclStoreData(dstElts, data, eltInLine);
+        dstElts += eltPerTrip;
+        if (MULTIDSTS){
+          for (int i = 1; i < ndsts; i++){
+            dl.loadBegin(srcs[i], eltInLine);
+            srcs[i] += eltPerTrip;
+            data = dl.loadFinish();
+            mscclStoreData(dsts[i], data, eltInLine);
+            dsts[i] += eltPerTrip;
+          }
+        }
+      }
+      nelem -= eltPerTrip;
+      offset += nthreads;
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_GENERIC_OP_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_MSCCL_GENERIC_OP_EXIT, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    barrier();
+  }
+
+  __device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
+    recvBuff[i] = (union ncclLLFifoLine*)conn->buffs[NCCL_PROTO_LL];
+    recvStep[i] = conn->step;
+    if (wid == i) recvConn = conn;
+  }
+  __device__ __forceinline__ void loadRecvSync() {
+    if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv()) {
+      recvConnHeadPtr = recvConn->head;
+      recvConnHead = recvConn->step;
+    }
+  }
+
+  __device__ __forceinline__ void loadSendConn(struct ncclConnInfo* conn, int i) {
+    sendBuff[i] = (union ncclLLFifoLine*)conn->buffs[NCCL_PROTO_LL];
+    sendStep[i] = conn->step;
+    if (wid == i) sendConn = conn;
+  }
+  __device__ __forceinline__ void loadSendSync() {
+    if (tid < fan.nsend()) {
+      sendConnHeadPtr = sendConn->head;
+      sendConnHeadCache = *sendConnHeadPtr;
+      sendConnHead = sendConn->step;
+      sendConnFifoPtr = sendConn->sizesFifo;
+    }
+  }
+
+ public:
+  __device__  Primitives(
+      const int tid, const int nthreads, int const *recvPeers, int const *sendPeers,
+      void const *inputBuf, void *outputBuf, uint64_t redOpArg, uint8_t group=0,
+      uint8_t connIndexRecv=0, uint8_t connIndexSend=0, struct ncclWorkElem* e = nullptr, int stepSize_=0
+    ):
+    redOp(redOpArg),
+    tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), group(group),
+    stepLines(ncclShmem.comm.buffSizes[NCCL_PROTO_LL]/NCCL_STEPS/sizeof(ncclLLFifoLine)) {
+    auto *channel = &ncclShmem.channel;
+    barriers = &ncclShmem.groups[group].barrier;
+    barrier_next = ncclShmem.groups[group].barrier_next;
+    // If we are going to support oneshot collNet + LL, then we would need to add connector index here
+    int nrecv=0, nsend=0;
+    // We compare with Fan::MaxRecv here because this->MaxRecv is always at least 1
+    while (nrecv < Fan::MaxRecv && recvPeers[nrecv] >= 0) {
+      loadRecvConn(&channel->peers[recvPeers[nrecv]]->recv[connIndexRecv], nrecv);
+      nrecv++;
+    }
+    while (nsend < MaxSend && sendPeers[nsend] >= 0) {
+      loadSendConn(&channel->peers[sendPeers[nsend]]->send[connIndexSend], nsend);
+      nsend++;
+    }
+    this->fan = Fan(nrecv, nsend);
+    loadRecvSync();
+    loadSendSync();
+    setDataPtrs(inputBuf, outputBuf);
+  }
+
+  __device__ ~Primitives() {
+    // Save steps for the next operation
+    if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv())
+      recvConn->step = recvConnHead;
+    if (tid < fan.nsend())
+      sendConn->step = sendConnHead;
+    // Ensure all steps written back
+    barrier();
+  }
+
+  __device__ void setDataPtrs(void const *inputBuf, void *outputBuf) {
+    userBufs[Input] = (T*)inputBuf;
+    userBufs[Output] = (T*)outputBuf;
+  }
+
+  __device__ void moveDataPtrs(intptr_t delta) {
+    userBufs[Input] += delta;
+    userBufs[Output] += delta;
+  }
+
+  __device__ void send(intptr_t inpIx, int eltN) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void sendFromOutput(intptr_t outIx, int eltN) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_FROM_OUTPUT_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_FROM_OUTPUT_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<0, 1, Output, -1>(outIx, -1, eltN, false);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_FROM_OUTPUT_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_FROM_OUTPUT_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void recvReduceSend(intptr_t inpIx, int eltN) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_REDUCE_SEND_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_REDUCE_SEND_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<1, 1, Input, -1>(inpIx, -1, eltN, false);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_REDUCE_SEND_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_REDUCE_SEND_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void recvReduceCopy(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_REDUCE_COPY_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_REDUCE_COPY_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<1, 0, Input, Output>(inpIx, outIx, eltN, postOp);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_REDUCE_COPY_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_REDUCE_COPY_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void copySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_COPY_SEND_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_COPY_SEND_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<0, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_COPY_SEND_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_COPY_SEND_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_COPY_SEND_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_COPY_SEND_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<1, 1, -1, Output>(-1, outIx, eltN, postOp);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_COPY_SEND_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_COPY_SEND_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void recvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_REDUCE_COPY_SEND_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_REDUCE_COPY_SEND_ENTRY, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+    LLGenericOp<1, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_RECV_REDUCE_COPY_SEND_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_RECV_REDUCE_COPY_SEND_EXIT, eltN*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+  }
+  __device__ void recvSend(int eltN) {
+    return LLGenericOp<1, 1, -1, -1>(-1, -1, eltN, false);
+  }
+
+  // MSCCL primitives
+  __device__ void sendWithBarrier(intptr_t inpIx, int eltN) {
+    send(inpIx, eltN);
+    // This is the only primitive.instruction where there is no barrier at the end, add it
+    barrier();
+  }
+  __device__ void localCopy(T* srcs, T* dsts, int eltN) {
+    return mscclGenericOp<0,1,0,0>(&srcs, 1, &dsts, 1, eltN);
+  }
+};

src/device/prims_ll128.h

@@ -0,0 +1,584 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#if defined(ENABLE_NPKIT)
+#include "npkit/npkit.h"
+#endif
+
+#define NCCL_LL128_FLAGTHREAD (NCCL_LL128_LINEELEMS-1)
+
+#ifndef RCCL_USE_WBINVL1_VOL
+#if defined(__GFX8__) || defined(__GFX9__)
+#define RCCL_USE_WBINVL1_VOL 1
+#else
+#define RCCL_USE_WBINVL1_VOL 0
+#endif
+#endif
+
+template<typename T, typename RedOp, typename Fan, int Direct, int P2p>
+class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>:
+  public PrimitivesWithoutDirect<Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p>> {
+
+  static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
+  static constexpr int Input=0, Output=1;
+  RedOp redOp;
+  const int tid;
+  const int nthreads;
+  const int wid;
+  const int stepSize;
+  const int warp;
+  const int warpInBlock; // warp index in thread block
+  const bool flagThread;
+  const int group;
+  Fan fan;
+  T *userBufs[2];
+  struct ncclConnInfo* recvConn = NULL;
+  volatile uint64_t* recvConnHeadPtr = NULL;
+  uint64_t recvConnHead;
+
+  struct ncclConnInfo* sendConn = NULL;
+  volatile int* sendConnFifoPtr = NULL;
+  volatile uint64_t* sendConnTailPtr = NULL;
+  uint64_t sendConnTail;
+  volatile uint64_t* sendConnHeadPtr = NULL;
+  uint64_t sendConnHead;
+  uint64_t sendConnHeadCache; // Cache last seen value
+
+  uint64_t recvStep[MaxRecv];
+  uint64_t sendStep[MaxSend];
+  uint64_t* recvBuff[MaxRecv];
+  uint64_t* sendBuff[MaxSend];
+
+  inline __device__ int recvOffset(int i) { return (recvStep[i]%NCCL_STEPS)*stepSize; }
+  inline __device__ int sendOffset(int i) { return (sendStep[i]%NCCL_STEPS)*stepSize; }
+  inline __device__ uint64_t* recvPtr(int i) { return recvBuff[i]+recvOffset(i); }
+  inline __device__ uint64_t* sendPtr(int i) { return sendBuff[i]+sendOffset(i); }
+  inline __device__ uint64_t recvFlag(int i) { return recvStep[i]+1; }
+  inline __device__ uint64_t sendFlag(int i) { return sendStep[i]+1; }
+
+  uint64_t* barriers;
+  uint64_t* barrier_next;
+
+#if defined(ENABLE_NPKIT)
+public:
+  int npKitCtxIdx = 0;
+  uint64_t npKitDataProcessEntryTime = 0;
+  uint64_t npKitDataProcessExitTime = 0;
+  uint64_t npKitDataProcessTotalTime = 0;
+private:
+#endif
+
+  inline __device__ void barrier() {
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
+  if (nthreads != WARP_SIZE)
+    barrier_by_group();
+#else
+   asm volatile ("bar.sync %1, %0;" :: "r"(nthreads), "r"(15-group));
+#endif
+  }
+
+  uint32_t abort = 0;
+  uint32_t* sync;
+
+  inline __device__ int checkAbort(int &spins, int i, int send) {
+    spins++;
+    if (abort == 0 && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
+      abort = __atomic_load_n(ncclShmem.comm.abortFlag, __ATOMIC_SEQ_CST);
+      spins = 0;
+    }
+    return abort;
+  }
+
+  inline __device__ void waitSend(int nbytes) {
+    if (sendConnHeadPtr) {
+      int spins = 0;
+      while (sendConnHeadCache + NCCL_STEPS < sendConnHead + 1) {
+        __builtin_amdgcn_s_sleep(1);
+        sendConnHeadCache = __atomic_load_n(sendConnHeadPtr, __ATOMIC_RELAXED);
+        if (checkAbort(spins, wid, 1)) break;
+      }
+      __asm__ __volatile__("s_wakeup");
+      if (sendConnFifoPtr) {
+        __atomic_store_n(sendConnFifoPtr+sendStep[wid]%NCCL_STEPS, nbytes, __ATOMIC_RELAXED);
+      }
+      sendConnHead += 1;
+    }
+  }
+
+  inline __device__ void postRecv() {
+    if (recvConnHeadPtr) STORE(recvConnHeadPtr, recvConnHead += 1);
+  }
+  inline __device__ void postSend() {
+    if (sendConnTailPtr) { STORE((unsigned long long *)sendConnTailPtr, sendConnTail += 1); }
+  }
+
+  template<int WordPerThread>
+  __device__ __forceinline__ void loadRegsBegin(uint64_t(&regs)[WordPerThread], T const *src, int eltN) {
+    constexpr int EltPer16B = 16/sizeof(T);
+    /* We are aligned to 16 bytes, so load directly to registers no shmem.
+     * Flag threads load half as much data which gets shuffled to the even
+     * registers during Finish. The point of splitting into two phases is to
+     * defer that shuffle, which incurs a dependency stall, until after other
+     * memops are launched by the caller.
+     */
+    #pragma unroll
+    for(int g=0; g < WordPerThread/2; g++) {
+      int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4);
+      if(!flagThread || g%2==0) {
+        if(ix*EltPer16B < eltN) {
+          if(reinterpret_cast<uintptr_t>(src)%4 == 0) {
+            regs[2*g+0] = __builtin_nontemporal_load((uint64_t*)(src + ix*EltPer16B));
+            regs[2*g+1] = __builtin_nontemporal_load((uint64_t*)(src + ix*EltPer16B)+1);
+          } else {
+            union {
+              uint64_t regs64[WordPerThread];
+              uint32_t regs32[WordPerThread*2];
+              uint16_t regs16[WordPerThread*4];
+              uint8_t regs8[WordPerThread*8];
+            };
+            if (sizeof(T) == 8) {
+              uint64_t *src64 = (uint64_t*)(src+ix*EltPer16B);
+              for (int i=0; i < 2; i++)
+                regs64[2*g+i] = __builtin_nontemporal_load(src64+i);
+            } else if (sizeof(T) == 4) {
+              uint32_t *src32 = (uint32_t*)(src+ix*EltPer16B);
+              for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++)
+                regs32[2*g+i] = __builtin_nontemporal_load(src32+i);
+            } else if (sizeof(T) == 2) {
+              uint16_t *src16 = (uint16_t*)(src+ix*EltPer16B);
+              for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++)
+                regs16[2*g+i] = __builtin_nontemporal_load(src16+i);
+            } else if (sizeof(T) == 1) {
+              uint8_t *src8 = (uint8_t*)(src+ix*EltPer16B);
+              for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T); i++)
+                regs8[2*g+i] = __builtin_nontemporal_load(src8+i);
+            }
+            regs[2*g+0] = regs64[2*g+0];
+            regs[2*g+1] = regs64[2*g+1];
+          }
+        }
+      }
+    }
+  }
+
+  template<int WordPerThread>
+  __device__ __forceinline__ void loadRegsFinish(uint64_t(&regs)[WordPerThread]) {
+    // Move data out of flag registers into the vacant registers.
+    #pragma unroll
+    for (int g=1; g < WordPerThread/2; g+=2) {
+      if (flagThread) regs[2*g] = regs[2*g-1];
+    }
+  }
+
+  template<int WordPerThread>
+  __device__ __forceinline__ void storeRegs(T *dst, uint64_t(&regs)[WordPerThread], int eltN) {
+    constexpr int EltPer16B = 16/sizeof(T);
+    // Reverse Finish() register permuatation.
+    #pragma unroll
+    for (int g=1; g < WordPerThread/2; g+=2) {
+      if (flagThread) regs[2*g-1] = regs[2*g];
+    }
+    // Write to dst if 4-byte aligned, shmem otherwise.
+    int misalignment = reinterpret_cast<uintptr_t>(dst)%4;
+    #pragma unroll
+    for(int g=0; g < WordPerThread/2; g++) {
+      int ix = g*WARP_SIZE - 16*(g/2) + wid - (g%2)*(wid/4);
+      if (!flagThread || g%2==0) {
+        if(misalignment == 0 && (ix+1)*EltPer16B <= eltN) {
+          __builtin_nontemporal_store(regs[2*g+0], (uint64_t*)(dst + ix*EltPer16B));
+          __builtin_nontemporal_store(regs[2*g+1], (uint64_t*)(dst + ix*EltPer16B)+1);
+        } else {
+          union {
+            uint64_t regs64[WordPerThread];
+            uint32_t regs32[WordPerThread*2];
+            uint16_t regs16[WordPerThread*4];
+            uint8_t regs8[WordPerThread*8];
+          };
+          regs64[2*g+0] = regs[2*g+0];
+          regs64[2*g+1] = regs[2*g+1];
+          int remaining = eltN - ix*EltPer16B;
+          if (sizeof(T) == 8) {
+            uint64_t *dst64 = (uint64_t*)(dst+ix*EltPer16B);
+            for (int i=0; i < 2 && i < remaining; i++)
+              __builtin_nontemporal_store(regs64[2*g+i], dst64+i);
+          } else if (sizeof(T) == 4) {
+            uint32_t *dst32 = (uint32_t*)(dst+ix*EltPer16B);
+            for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++)
+              __builtin_nontemporal_store(regs32[2*g+i], dst32+i);
+          } else if (sizeof(T) == 2) {
+            uint16_t *dst16 = (uint16_t*)(dst+ix*EltPer16B);
+            for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++)
+              __builtin_nontemporal_store(regs16[2*g+i], dst16+i);
+          } else if (sizeof(T) == 1) {
+            uint8_t *dst8 = (uint8_t*)(dst+ix*EltPer16B);
+            for (int i=0; i < 2*sizeof(uint64_t)/sizeof(T) && i < remaining; i++)
+              __builtin_nontemporal_store(regs8[2*g+i], dst8+i);
+          }
+        }
+      }
+    }
+  }
+
+  #define WARP_MASK 0xffffffff
+
+  template <int ELEMS_PER_THREAD, int RECV, int SEND, int SrcBuf, int DstBuf>
+  __device__ __forceinline__ void recvReduceSendCopy(uint64_t(&v)[ELEMS_PER_THREAD], int ll128Offset, bool postOp) {
+    constexpr int SRC = SrcBuf != -1 ? 1 : 0;
+    uint64_t vr[ELEMS_PER_THREAD];
+
+    __syncwarp();
+    /************************ Wait first recv ********************/
+    if (RECV) {
+      uint64_t* ptr = recvPtr(0)+ll128Offset;
+      uint64_t flag = recvFlag(0);
+      bool needReload;
+      int spins = 0;
+      do {
+        needReload = false;
+        #pragma unroll
+        for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+          vr[u] = __builtin_nontemporal_load(ptr+u*WARP_SIZE);
+          vr[u+1] = __builtin_nontemporal_load(ptr+u*WARP_SIZE+1);
+          needReload |= flagThread && (vr[u+1] != flag);
+        }
+        needReload &= (0 == checkAbort(spins, 0, 0));
+      } while (__any(needReload));
+    }
+
+    /************* Finish register load **************/
+    if (SRC) {
+      // By deferring register shuffle here we've overlapped spinning on first
+      // peer's data with memory loads of src data.
+      loadRegsFinish(v);
+      if (SrcBuf == Input) {
+        #pragma unroll
+        for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+          v[u] = applyPreOp(redOp, v[u]);
+          if (!flagThread)
+            v[u+1] = applyPreOp(redOp, v[u+1]);
+        }
+      }
+    }
+
+    /************************ Recv rest *********************/
+    if (RECV) {
+      { // Consume data from first recv
+        uint64_t* ptr = recvPtr(0)+ll128Offset;
+        #pragma unroll
+        for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+          v[u]   = SRC ? applyReduce(redOp, vr[u], v[u]) : vr[u];
+          v[u+1] = SRC ? applyReduce(redOp, vr[u+1], v[u+1]) : vr[u+1];
+        }
+      }
+
+      for (int i=1; i<MaxRecv && i<fan.nrecv(); i++) {
+        uint64_t flag = recvFlag(i);
+        uint64_t* ptr = recvPtr(i)+ll128Offset;
+        bool needReload;
+        int spins = 0;
+        do {
+          needReload = false;
+          #pragma unroll
+          for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+            vr[u] = __builtin_nontemporal_load(ptr+u*WARP_SIZE);
+            vr[u+1] = __builtin_nontemporal_load(ptr+u*WARP_SIZE+1);
+            needReload |= flagThread && (vr[u+1] != flag);
+          }
+          needReload &= (0 == checkAbort(spins, i, 0));
+        } while (__any(needReload));
+
+        #pragma unroll
+        for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+          v[u]   = applyReduce(redOp, vr[u], v[u]);
+          v[u+1] = applyReduce(redOp, vr[u+1], v[u+1]);
+        }
+      }
+    }
+    /********************** End Recv ************************/
+
+    if (postOp) {
+      #pragma unroll
+      for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+        v[u]   = applyPostOp(redOp, v[u]);
+        v[u+1] = applyPostOp(redOp, v[u+1]);
+      }
+    }
+
+#if RCCL_USE_WBINVL1_VOL
+    if (tid == 0) __builtin_amdgcn_buffer_wbinvl1();
+#endif
+    /************************ Send **************************/
+    if (SEND) {
+      for (int i=1; i<MaxSend && i<fan.nsend(); i++) {
+        uint64_t flag = sendFlag(i);
+        uint64_t* ptr = sendPtr(i)+ll128Offset;
+        #pragma unroll
+        for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+          __builtin_nontemporal_store(v[u], ptr+u*WARP_SIZE);
+          __builtin_nontemporal_store(flagThread ? flag : v[u+1], ptr+u*WARP_SIZE+1);
+        }
+      }
+      uint64_t flag = sendFlag(0);
+      uint64_t* ptr = sendPtr(0)+ll128Offset;
+      #pragma unroll
+      for (int u=0; u<ELEMS_PER_THREAD; u+=2) {
+        __builtin_nontemporal_store(v[u], ptr+u*WARP_SIZE);
+        __builtin_nontemporal_store(flagThread ? flag : v[u+1], ptr+u*WARP_SIZE+1);
+      }
+    }
+    /********************** End Send ************************/
+  }
+
+  static constexpr int WireWordPerSlice = WARP_SIZE*NCCL_LL128_SHMEM_ELEMS_PER_THREAD;
+  static constexpr int DataEltPerSlice = (WireWordPerSlice - WireWordPerSlice/NCCL_LL128_LINEELEMS)*(sizeof(uint64_t)/sizeof(T));
+
+  template <int RECV, int SEND, int SrcBuf, int DstBuf>
+  __device__ __forceinline__ void GenericOp(intptr_t srcIx, intptr_t dstIx, int nelem, bool postOp) {
+    constexpr int SRC = SrcBuf != -1 ? 1 : 0;
+    constexpr int DST = DstBuf != -1 ? 1 : 0;
+    T const *srcPtr = SrcBuf == -1 ? nullptr : userBufs[SrcBuf] + srcIx;
+    T       *dstPtr = DstBuf == -1 ? nullptr : userBufs[DstBuf] + dstIx;
+    int wireOffset = WireWordPerSlice*warp + 2*wid;
+    const int nwarps = nthreads/WARP_SIZE;
+    nelem = nelem < 0 ? 0 : nelem;
+
+    if (SEND) waitSend(divUp(nelem, DataEltPerSlice)*WireWordPerSlice*sizeof(uint64_t));
+    barrier();
+    nelem -= DataEltPerSlice*warp;
+    srcPtr += DataEltPerSlice*warp;
+    dstPtr += DataEltPerSlice*warp;
+    while (nelem > 0) {
+      const int eltInSlice = min(nelem, DataEltPerSlice);
+      uint64_t regs[NCCL_LL128_SHMEM_ELEMS_PER_THREAD];
+      if (SRC) loadRegsBegin(regs, srcPtr, eltInSlice);
+      recvReduceSendCopy<NCCL_LL128_SHMEM_ELEMS_PER_THREAD, RECV, SEND, SrcBuf, DstBuf>(regs, wireOffset, postOp);
+      if (DST) storeRegs(dstPtr, regs, eltInSlice);
+
+      wireOffset += WireWordPerSlice*nwarps;
+      srcPtr += DataEltPerSlice*nwarps;
+      dstPtr += DataEltPerSlice*nwarps;
+      nelem -= DataEltPerSlice*nwarps;
+    }
+
+    barrier();
+    if (SEND) for (int i=0; i < MaxSend; i++) sendStep[i] += 1;
+    if (SEND) postSend();
+    if (RECV) for (int i=0; i < MaxRecv; i++) recvStep[i] += 1;
+    if (RECV) postRecv();
+  }
+
+  template <int REDUCE, int COPY, int MULTISRCS, int MULTIDSTS>
+  __device__ __forceinline__ void mscclGenericOp(T** srcs, int nsrcs, T** dsts, int ndsts, int nelem) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_GENERIC_OP_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_MSCCL_GENERIC_OP_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    T const *srcPtr = srcs[0];
+    T       *dstPtr = dsts[0];
+    int wireOffset = WireWordPerSlice*warp + 2*wid;
+    const int nwarps = nthreads/WARP_SIZE;
+    nelem = nelem < 0 ? 0 : nelem;
+
+    nelem -= DataEltPerSlice*warp;
+    srcPtr += DataEltPerSlice*warp;
+    dstPtr += DataEltPerSlice*warp;
+    if (MULTISRCS){
+      for (int i = 1; i < nsrcs; i++){
+        srcs[i] += DataEltPerSlice*warp;
+      }
+    }
+    if (MULTIDSTS){
+      for (int i = 1; i < ndsts; i++){
+        dsts[i] += DataEltPerSlice*warp;
+      }
+    }
+    while (nelem > 0) {
+      const int eltInSlice = min(nelem, DataEltPerSlice);
+      uint64_t regs[NCCL_LL128_SHMEM_ELEMS_PER_THREAD];
+      loadRegsBegin(regs, srcPtr, eltInSlice);
+      loadRegsFinish(regs);
+      if (REDUCE){
+        uint64_t regsD[NCCL_LL128_SHMEM_ELEMS_PER_THREAD];
+        loadRegsBegin(regsD, dstPtr, eltInSlice);
+        loadRegsFinish(regsD);
+        #pragma unroll
+        for (int u=0; u<NCCL_LL128_SHMEM_ELEMS_PER_THREAD; u+=2) {
+          regsD[u] = applyReduce(redOp, regs[u], regsD[u]);
+          if (!flagThread)
+            regsD[u+1] = applyReduce(redOp, regs[u+1], regsD[u+1]);
+        }
+        if (MULTISRCS){
+          for (int i = 1; i < nsrcs; i++){
+            loadRegsBegin(regs, srcs[i], eltInSlice);
+            loadRegsFinish(regs);
+            for (int u=0; u<NCCL_LL128_SHMEM_ELEMS_PER_THREAD; u+=2) {
+              regsD[u] = applyReduce(redOp, regs[u], regsD[u]);
+              if (!flagThread)
+                regsD[u+1] = applyReduce(redOp, regs[u+1], regsD[u+1]);
+            }
+          }
+        }
+        storeRegs(dstPtr, regsD, eltInSlice);
+      }
+      if (COPY){
+        storeRegs(dstPtr, regs, eltInSlice);
+        if (MULTIDSTS){
+          for (int i = 1; i < nsrcs; i++){
+            loadRegsBegin(regs, srcs[i], eltInSlice);
+            loadRegsFinish(regs);
+            storeRegs(dsts[i], regs, eltInSlice);
+          }
+        }
+      }
+
+      wireOffset += WireWordPerSlice*nwarps;
+      srcPtr += DataEltPerSlice*nwarps;
+      dstPtr += DataEltPerSlice*nwarps;
+      if (MULTISRCS){
+        for (int i = 1; i < nsrcs; i++){
+          srcs[i] += DataEltPerSlice*nwarps;
+        }
+      }
+      if (MULTIDSTS){
+        for (int i = 1; i < ndsts; i++){
+          dsts[i] += DataEltPerSlice*nwarps;
+        }
+      }
+      nelem -= DataEltPerSlice*nwarps;
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_GENERIC_OP_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_MSCCL_GENERIC_OP_EXIT, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    barrier();
+  }
+
+  __device__ __forceinline__ void loadRecvConn(struct ncclConnInfo* conn, int i) {
+    recvBuff[i] = (uint64_t*)conn->buffs[NCCL_PROTO_LL128];
+    recvStep[i] = conn->step;
+    if (wid == i) recvConn = conn;
+  }
+  __device__ __forceinline__ void loadRecvSync() {
+    if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv()) {
+      recvConnHeadPtr = recvConn->head;
+      recvConnHead = recvConn->step;
+    }
+  }
+
+  __device__ __forceinline__ void loadSendConn(struct ncclConnInfo* conn, int i) {
+    sendBuff[i] = (uint64_t*)conn->buffs[NCCL_PROTO_LL128];
+    sendStep[i] = conn->step;
+    if (wid == i) sendConn = conn;
+  }
+  __device__ __forceinline__ void loadSendSync() {
+    if (tid < fan.nsend()) {
+      sendConnHeadPtr = sendConn->head;
+      sendConnHeadCache = *sendConnHeadPtr;
+      sendConnHead = sendConn->step;
+      sendConnFifoPtr = sendConn->sizesFifo;
+    }
+    if (tid >= nthreads-WARP_SIZE && wid<fan.nsend()) {
+      if (sendConn->sizesFifo) {
+        sendConnTailPtr = sendConn->tail;
+        sendConnTail = sendConn->step;
+      }
+    }
+  }
+
+public:
+  __device__ Primitives(
+      const int tid, const int nthreads, int const *recvPeers, int const *sendPeers,
+      void const *inputBuf, void *outputBuf, uint64_t redOpArg, uint8_t group=0,
+      uint8_t connIndexRecv=0, uint8_t connIndexSend=0, struct ncclWorkElem* e = nullptr, int stepSize_=0
+    ):
+    redOp(redOpArg),
+    tid(tid), nthreads(nthreads), wid(tid%WARP_SIZE), warp(tid/WARP_SIZE),
+    warpInBlock(threadIdx.x/WARP_SIZE),
+    flagThread((tid%4)==3), group(group),
+    stepSize(ncclShmem.comm.buffSizes[NCCL_PROTO_LL128]/NCCL_STEPS/sizeof(uint64_t)) {
+    auto *channel = &ncclShmem.channel;
+    barriers = &ncclShmem.groups[group].barrier;
+    barrier_next = ncclShmem.groups[group].barrier_next;
+    int nrecv=0, nsend=0;
+    while (nrecv < MaxRecv && recvPeers[nrecv] >= 0) {
+      loadRecvConn(&channel->peers[recvPeers[nrecv]]->recv[connIndexRecv], nrecv);
+      nrecv++;
+    }
+    while (nsend < MaxSend && sendPeers[nsend] >= 0) {
+      loadSendConn(&channel->peers[sendPeers[nsend]]->send[connIndexSend], nsend);
+      nsend++;
+    }
+    this->fan = Fan(nrecv, nsend);
+    loadRecvSync();
+    loadSendSync();
+    setDataPtrs(inputBuf, outputBuf);
+  }
+
+  __device__ ~Primitives() {
+    // Save steps for the next operation
+    if (tid >= nthreads-WARP_SIZE && wid < fan.nrecv())
+      recvConn->step = recvConnHead;
+    if (tid < fan.nsend())
+      sendConn->step = sendConnHead;
+    // Ensure all steps written back
+    barrier();
+  }
+
+  __device__ void setDataPtrs(void const *inputBuf, void *outputBuf) {
+    userBufs[Input] = (T*)inputBuf;
+    userBufs[Output] = (T*)outputBuf;
+  }
+
+  __device__ void moveDataPtrs(intptr_t delta) {
+    userBufs[Input] += delta;
+    userBufs[Output] += delta;
+  }
+
+  __device__ void send(intptr_t inpIx, int eltN) {
+    return GenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false);
+  }
+  __device__ void sendFromOutput(intptr_t outIx, int eltN) {
+    return GenericOp<0, 1, Output, -1>(outIx, -1, eltN, false);
+  }
+  __device__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
+    return GenericOp<1, 0, -1, Output>(-1, outIx, eltN, postOp);
+  }
+  __device__ void recvReduceSend(intptr_t inpIx, int eltN) {
+    return GenericOp<1, 1, Input, -1>(inpIx, -1, eltN, false);
+  }
+  __device__ void recvReduceCopy(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    return GenericOp<1, 0, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+  __device__ void copySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    return GenericOp<0, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+  __device__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
+    return GenericOp<1, 1, -1, Output>(-1, outIx, eltN, postOp);
+  }
+  __device__ void recvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    return GenericOp<1, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+  __device__ void recvSend(int eltN) {
+    return GenericOp<1, 1, -1, -1>(-1, -1, eltN, false);
+  }
+
+  // MSCCL primitives
+  __device__ void sendWithBarrier(intptr_t inpIx, int eltN) {
+    send(inpIx, eltN);
+  }
+  __device__ void localCopy(T* srcs, T* dsts, int eltN) {
+    return mscclGenericOp<0,1,0,0>(&srcs, 1, &dsts, 1, eltN);
+  }
+};

src/device/prims_simple.h

@@ -0,0 +1,863 @@
+/*************************************************************************
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2023 Advanced Micro Devices, Inc. All rights reserved.
+ * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#if defined(ENABLE_NPKIT)
+#include "npkit/npkit.h"
+#endif
+
+#include "msccl/msccl_struct.h"
+#include "network/unpack/unpack.h"
+
+template<typename T, typename RedOp, typename Fan, int Direct,
+         int SlicePerChunk, int StepPerSlice, int Unroll, int P2p, int MultimemSrcs, int MultimemDsts>
+class Primitives<
+    T, RedOp, Fan, Direct, ProtoSimple<SlicePerChunk, StepPerSlice, Unroll, MultimemSrcs, MultimemDsts>, P2p
+  > {
+  static constexpr int MaxRecv = Fan::MaxRecv, MaxSend = Fan::MaxSend;
+  static constexpr int Input=0, Output=1;
+  static constexpr int RoleInput = 0x01,
+                       RoleOutput = 0x02,
+                       RoleWaitRecv = 0x04,
+                       RoleWaitSend = 0x08,
+                       RolePostSend = 0x10,
+                       RolePostRecv = 0x20,
+                       Aborted = 0x40,
+                       OffsFifoEnabled = 0x80,
+                       SizesFifoEnabled = 0x100,
+                       DirectWrite = 0x200,
+                       DirectRead = 0x400,
+                       ThreadsSynced = 0x800,
+                       NvlsMinPolling = 0x1000,
+                       NetDeviceUnpack = 0x2000,
+                       AnyNetDeviceUnpack = 0x4000,
+                       NvlsDirectRead = 0x8000,
+                       NvlsDirectWrite = 0x10000;
+  const int tid, tidInBlock;
+  const int nthreads;
+  int nworkers;
+  const int stepSize;
+  Fan fan;
+  int index; // Peer index I'm responsible for
+  int flags;
+  int group;
+  uint64_t step;
+  int *connOffsFifoPtr;   // (flags & OffsFifoEnabled)
+  union {
+    T *userBuff;            // (flags & (RoleInput|RoleOutput))
+    T *connEltsFifo;        // !(flags & (RoleInput|RoleOutput))
+  };
+  union {
+    int volatile *connSizesFifoPtr; //  (flags & SizesFifoEnabled)
+    T *directBuff;                  // !(flags & SizesFifoEnabled)
+  };
+  uint64_t *connStepPtr;
+  uint64_t connStepCache; // Cache last seen value of (*connStepPtr)
+  uint64_t* barriers;
+  uint64_t* barrier_next;
+  uint32_t* next_hdp_reg;
+  void*    mhandle;
+  void*    netDeviceHandle;
+
+#if defined(ENABLE_NPKIT)
+public:
+  int npKitCtxIdx = 0;
+  uint64_t npKitDataProcessEntryTime = 0;
+  uint64_t npKitDataProcessExitTime = 0;
+  uint64_t npKitDataProcessTotalTime = 0;
+private:
+#endif
+
+  // Don't use barrier 0 as it's used by the final sync
+  inline __device__ void barrier() {
+    flags |= ThreadsSynced;
+    if (nthreads == WARP_SIZE)
+      __syncwarp();
+    else
+      barrier_by_group();
+  }
+
+  inline __device__ void subBarrier() {
+    barrier();
+  }
+
+  inline __device__ bool checkAbort(int &spins) {
+    spins++;
+    if (!(flags & Aborted) && spins == NCCL_SPINS_BEFORE_CHECK_ABORT) {
+      if (__atomic_load_n(ncclShmem.comm.abortFlag, __ATOMIC_SEQ_CST)) {
+        flags |= Aborted;
+        ncclShmem.aborted = 1;
+      }
+      spins = 0;
+    }
+    return flags & Aborted;
+  }
+
+  inline __device__ uint64_t loadStepValue(uint64_t* ptr) {
+    #if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010
+    if (flags & NvlsMinPolling) {
+      uint64_t ans;
+      asm("multimem.ld_reduce.acquire.sys.global.min.u64 %0, [%1];" : "=l"(ans) : "l"(cvta_to_global(ptr)));
+      return ans;
+    }
+    #endif
+    // volatile is faster than acquire but not as correct. Make sure reduceCopy
+    // loads data using volatile so it doesn't see stale data in L1.
+    return __atomic_load_n(ptr, __ATOMIC_RELAXED);
+  }
+
+  template <int DirectRecv, int DirectSend, int Recv, int Send, int Src, int Dst>
+  __device__ __forceinline__ void waitPeer(intptr_t srcIx, intptr_t dstIx, int offset, int nelts) {
+    const bool isSendNotRecv = (Send && Recv) ? (flags & RoleWaitSend) : Send;
+    const bool noRecvWait = DirectRecv && Src && (flags & DirectRead);        // no wait when directly reading from remote input
+    const bool noSendWait = DirectSend && (flags & (DirectRead|DirectWrite)); // no wait in empty send (e.g. directScatter) or direct remote write
+    if (((flags & (Recv*RoleWaitRecv)) && !noRecvWait) ||
+        ((flags & (Send*RoleWaitSend)) && !noSendWait)) {
+      int spins = 0;
+      while (connStepCache + (isSendNotRecv ? NCCL_STEPS : 0) < step + StepPerSlice) {
+        __builtin_amdgcn_s_sleep(1);
+        connStepCache = loadStepValue(connStepPtr);
+        if (checkAbort(spins)) break;
+        //if (spins == 0) printf("r=%d b=%d t=%d SPUN OUT got=%d want=%d\n", ncclShmem.comm.rank, blockIdx.x, threadIdx.x, int(connStepCache + (isSendNotRecv ? NCCL_STEPS : 0)), int(step+StepPerSlice));
+        if (spins == 0) traceData(__LINE__, threadIdx.x, int(connStepCache + (isSendNotRecv ? NCCL_STEPS : 0)), int(step+StepPerSlice));
+      }
+      __asm__ __volatile__("s_wakeup");
+    }
+
+    if (flags & (Recv*RoleWaitRecv | Send*RoleWaitSend)) {
+      if (isSendNotRecv && (flags & SizesFifoEnabled))
+        __atomic_store_n(connSizesFifoPtr+step%NCCL_STEPS, nelts*sizeof(T), __ATOMIC_RELAXED);
+
+      void **ptrs = isSendNotRecv ? (ncclShmem.groups[group].dsts + Dst)
+                                  : (ncclShmem.groups[group].srcs + Src);
+      if (flags & OffsFifoEnabled)
+        ptrs[index] = connEltsFifo + loadInt(connOffsFifoPtr + (step%NCCL_STEPS))/sizeof(T);
+      else if (isSendNotRecv && DirectSend) {
+        if (flags & (DirectWrite | NvlsDirectWrite)) {
+          ptrs[index] = directBuff + dstIx + offset;
+        } else if (flags & DirectRead) {  // empty send
+          ptrs[index] = nullptr;
+        } else {
+          ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
+        }
+      } else if (!isSendNotRecv && DirectRecv) {
+        if (flags & (DirectRead | NvlsDirectRead)) {
+          ptrs[index] = directBuff + srcIx + offset;
+        } else if (flags & DirectWrite) {
+          ptrs[index] = directBuff + dstIx + offset;  // send to next from my output buffer
+        } else {
+          ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
+        }
+      }
+      else {
+        ptrs[index] = connEltsFifo + (step%NCCL_STEPS)*stepSize;
+      }
+      if ((flags & (AnyNetDeviceUnpack)) && (flags & (Recv*RoleWaitRecv))) {
+        ncclNetDeviceIncrementHead(group);
+      }
+      step += StepPerSlice;
+    }
+  }
+
+  template<int Recv, int Send>
+  inline __device__ void postPeer(bool dataStored) {
+    if (Send && (flags & RolePostSend) && dataStored)
+#ifdef __GFX9__
+    __threadfence();
+#else
+    __threadfence_system();
+#endif
+
+    if ((flags & Send*RolePostSend) && next_hdp_reg)
+      STORE((unsigned int *)next_hdp_reg, 0x1);
+
+    if (flags & (Recv*RolePostRecv | Send*RolePostSend)) {
+      step += StepPerSlice;
+      STORE(connStepPtr, step);
+    }
+  }
+
+  template <int DirectRecv1, int DirectSend1, int Recv, int Send, int SrcBuf, int DstBuf>
+  __device__ __forceinline__ void genericOp(
+      intptr_t srcIx, intptr_t dstIx, int nelem, bool postOp
+    ) {
+    constexpr int DirectRecv = /*1 &&*/ Direct && DirectRecv1;
+    constexpr int DirectSend = /*1 &&*/ Direct && DirectSend1;
+    constexpr int Src = SrcBuf != -1;
+    constexpr int Dst = DstBuf != -1;
+
+    nelem = nelem < 0 ? 0 : nelem;
+    int sliceSize = stepSize*StepPerSlice;
+    sliceSize = max(divUp(nelem, 16*SlicePerChunk)*16, sliceSize/32);
+    int slice = 0;
+    int offset = 0;
+
+    if (tid < nworkers && offset < nelem) {
+      // Worker-only loop for non-empty slices. Non-workers and empty slices are
+      // processed in the loop following this if block. The benefit of splitting
+      // the loop like this is we pull two branches out of the critical path.
+      // Using "number of branch insns (taken or not) encountered dynamically"
+      // as the performance metric, then:
+      //   perf_orig = 2*numslices
+      //   perf_new = 2+numslices
+      // So the new code and old code behave the same for numslices=2, and for
+      // numslices>2 the new code is superior. And note that in the case
+      // numslices=1, the loop is trivially unrollable (single iteration) so we
+      // don't incur that that tail branch and we still have perf_new=2.
+      //
+      // ORIGINAL CODE:
+      //   unrolled for(slices) {
+      //     if(worker) { // This branch removed
+      //       wait();
+      //       subBarrier();
+      //       if(slice not empty) // This branch removed
+      //         ReduceCopyMulti();
+      //     }
+      //     barrier();
+      //     post();
+      //   } // Since we no longer unroll, new branch added here
+      #pragma unroll 1
+      do {
+        sliceSize = sliceSize < nelem-offset ? sliceSize : nelem-offset;
+        if (Src && (flags & (SrcBuf==Input ? RoleInput : RoleOutput)))
+          ncclShmem.groups[group].srcs[0] = userBuff + srcIx + offset;
+        if (Dst && (flags & (DstBuf==Input ? RoleInput : RoleOutput)))
+          ncclShmem.groups[group].dsts[0] = userBuff + dstIx + offset;
+        waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(srcIx, dstIx, offset, sliceSize);
+        subBarrier();
+        /* if user abort the kernel, we don't need to actually perform copy/reduce; just set size
+         * to 0 to avoid unnecessary workload. */
+        int workSize = ncclShmem.aborted ? 0 : sliceSize;
+        if (flags & AnyNetDeviceUnpack) {
+          ncclNetDeviceUnpack<Recv>(tid, tidInBlock, nworkers, group, ncclShmem.groups[group].devicePlugin.unpack.unpackNetDeviceIndexMask, Src, workSize);
+          // Sync here to make sure all workers are reading from the updated srcs)
+          subBarrier();
+        }
+
+        if (DirectRecv && ncclShmem.groups[group].srcs[0] == ncclShmem.groups[group].dsts[0]
+            /* NVLS can have srcs[0] == dsts[0], but we cannot enter this "if branch",
+             * so we need to check whether MultimemSrcs and MultimemDsts are 0. */
+            && MultimemSrcs == 0 && MultimemDsts == 0) {
+          // We can only have one direct receive. Since srcs[0] == dstPtr+offset, skip one copy
+          if (Send) {
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY)
+            if (tid == 0) {
+              NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY, sliceSize*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+                  ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+            }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+            if (tid == 0) {
+              npKitDataProcessEntryTime = NPKIT_GET_GPU_TIMESTAMP();
+            }
+#endif
+
+            reduceCopy<Unroll, RedOp, T, 0, 1, 1, 0, 1, MaxSend, /*PreOpSrcs*/0>
+              (tid, nworkers, /*redArg*/0, /*preOpArgs*/nullptr, /*postOp*/false,
+               1, ncclShmem.groups[group].srcs,
+               fan.nsend(), ncclShmem.groups[group].dsts+1,
+               workSize);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+            if (tid == 0) {
+              npKitDataProcessExitTime = NPKIT_GET_GPU_TIMESTAMP();
+              npKitDataProcessTotalTime += npKitDataProcessExitTime - npKitDataProcessEntryTime;
+            }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT)
+            if (tid == 0) {
+              NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT, sliceSize*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+                  ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+            }
+#endif
+
+          }
+        } else if (DirectSend && !DirectRecv && SrcBuf != Input && ncclShmem.groups[group].dsts[Dst] == nullptr) {
+          // For broadcast in CollNet to do empty send
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY)
+          if (tid == 0) {
+            NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY, sliceSize*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+                ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+          }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+          if (tid == 0) {
+            npKitDataProcessEntryTime = NPKIT_GET_GPU_TIMESTAMP();
+          }
+#endif
+
+          reduceCopy<Unroll, RedOp, T, 0, 1, 1, 0, 1, 1, /*PreOpSrcs*/0>
+            (tid, nworkers, ncclShmem.redOpArgs[0],  nullptr, postOp,
+             Recv, ncclShmem.groups[group].srcs,
+             Dst, ncclShmem.groups[group].dsts,
+             workSize);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+          if (tid == 0) {
+            npKitDataProcessExitTime = NPKIT_GET_GPU_TIMESTAMP();
+            npKitDataProcessTotalTime += npKitDataProcessExitTime - npKitDataProcessEntryTime;
+          }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT)
+          if (tid == 0) {
+            NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT, sliceSize*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+                ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+          }
+#endif
+
+        } else {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY)
+          if (tid == 0) {
+            NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY, sliceSize*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+                ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+          }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+          if (tid == 0) {
+            npKitDataProcessEntryTime = NPKIT_GET_GPU_TIMESTAMP();
+          }
+#endif
+
+          constexpr int PreOpSrcs = SrcBuf != Input ? 0 :
+                                    DirectRecv*MaxRecv == NCCL_MAX_DIRECT_ARITY ? (1+NCCL_MAX_DIRECT_ARITY) : 1;
+          reduceCopy<Unroll, RedOp, T,
+            MultimemSrcs, Recv+Src, Recv*MaxRecv+Src,
+            MultimemDsts, Send+Dst, Send*MaxSend+Dst, PreOpSrcs>
+            (tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, postOp,
+             Recv*fan.nrecv()+Src, ncclShmem.groups[group].srcs,
+             Send*fan.nsend()+Dst, ncclShmem.groups[group].dsts,
+             workSize);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_PRIM_COLLECT_DATA_PROCESS_TIME)
+          if (tid == 0) {
+            npKitDataProcessExitTime = NPKIT_GET_GPU_TIMESTAMP();
+            npKitDataProcessTotalTime += npKitDataProcessExitTime - npKitDataProcessEntryTime;
+          }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT)
+          if (tid == 0) {
+            NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT, sliceSize*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+                ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+          }
+#endif
+
+        }
+        barrier(); // This barrier has a counterpart in following loop
+        postPeer<Recv, Send>(0 < sliceSize);
+        offset += sliceSize;
+        slice += 1;
+      } while (slice < SlicePerChunk && offset < nelem);
+    }
+
+    // Non-workers come straight here. Workers too but only once the remaining
+    // slices are all empty. Since empty slices are the uncommon case, and
+    // worker perf is the limiter, perf-wise this loop is effectively unentered,
+    // hence just a single branch insn.
+    #pragma unroll 1
+    while (slice < SlicePerChunk) {
+      sliceSize = sliceSize < nelem-offset ? sliceSize : nelem-offset;
+      { // Only workers could have Wait roles so we know the slice must be empty
+        // since we've exited the loop above.
+        waitPeer<DirectRecv, DirectSend, Recv, Send, Src, Dst>(0, 0, 0, 0);
+      }
+      barrier(); // Has couterpart in preceding worker-only loop.
+      postPeer<Recv, Send>(0 < sliceSize);
+      offset += sliceSize;
+      slice += 1;
+    }
+  }
+
+  template <int REDUCE, int COPY, int MULTISRCS, int MULTIDSTS>
+  __device__ __forceinline__ void mscclGenericOp(T** srcs, int nsrcs, T** dsts, int ndsts, int nelem) {
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_GENERIC_OP_ENTRY)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_MSCCL_GENERIC_OP_ENTRY, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    nelem = nelem < 0 ? 0 : nelem;
+    if (tid < nworkers) {
+      if (REDUCE){
+        srcs[nsrcs] = dsts[0];
+        nsrcs++;
+        if (MULTISRCS){
+          reduceCopy<Unroll, RedOp, T, 0, 3, MSCCL_MAX_REDUCE_FUSION, 0, 1, 1, 0>
+            (tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, false, nsrcs, (void **)srcs, 1, (void **)dsts, nelem);
+        } else {
+          reduceCopy<Unroll, RedOp, T, 0, 2, 2, 0, 1, 1, 0>
+            (tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, false, 2, (void **)srcs, 1, (void **)dsts, nelem);
+        }
+      }
+      if (COPY){
+        reduceCopy<Unroll, RedOp, T, 0, 1, 1, 0, 1, 1, 0>
+          (tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, false, 1, (void **)srcs, 1, (void **)dsts, nelem);
+        if (MULTISRCS) {
+          for (int i = 1; i < nsrcs; i++){
+            reduceCopy<Unroll, RedOp, T, 0, 1, 1, 0, 1, 1, 0>
+              (tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, false, 1, (void **)&srcs[i], 1, (void **)&dsts[i], nelem);
+          }
+        }
+      }
+    }
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_GENERIC_OP_EXIT)
+    if (tid == 0) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_MSCCL_GENERIC_OP_EXIT, nelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    barrier();
+  }
+
+  // Scatter/Gather generic op
+  // skip: my own rank order in the buffer chunks
+  // shift: peer offset to avoid all ranks sending to or receiving from same peer
+  template <int DirectRecv1, int DirectSend1, int Recv, int Send>
+  __device__ __forceinline__ void
+  ScatterGatherOp(intptr_t inpIx, intptr_t outIx, ssize_t totalElem, int peerElem, ssize_t peerOffset, int skip, int shift, bool postOp) {
+    constexpr int DirectRecv = /*1 &&*/ Direct && DirectRecv1;
+    constexpr int DirectSend = /*1 &&*/ Direct && DirectSend1;
+    int offset = 0; // slice offset
+    int sliceSize = stepSize*StepPerSlice;
+    int dataSize = max(DIVUP(peerElem, 16*SlicePerChunk)*16, sliceSize/32);  // per-peer slice size
+
+    #pragma unroll 1
+    for (int slice=0; slice<SlicePerChunk; ++slice) {
+      ssize_t realSize = max(0, min(dataSize, peerElem-offset));
+      bool fenceNeeded = false;
+      if (tid < nworkers) {
+        if (Send) {
+          // Scatter pre-scales data of input buffer only in non-Direct case
+          constexpr int PreOpSrcs = DirectSend ? 0 : 1;
+          if (flags & RoleInput) ncclShmem.groups[group].srcs[0] = userBuff + inpIx + offset;
+          // realSize is not accurate here; but intra-node does not rely on sizes FIFO
+          waitPeer<0, DirectSend, 0, 1, 1, 0>(0, inpIx, offset, realSize);
+          subBarrier();
+          #pragma unroll 1
+          // Loop over peers
+          for (int j=0; j<fan.nsend(); j++) {
+            int i = (j+shift)%fan.nsend();
+            ssize_t pOffset = i*peerOffset;
+            // Skip the data I am responsible of reducing myself
+            if (skip >= 0 && i >= skip) pOffset += peerElem;
+            void* src0 = (T*)ncclShmem.groups[group].srcs[0] + pOffset;
+            ssize_t realPeerSize = min(realSize, totalElem-pOffset);
+            if (realPeerSize > 0 && ncclShmem.groups[group].dsts[i] != nullptr) {
+              reduceCopy<Unroll, RedOp, T, 0, 1, 1, 0, 1, 1, PreOpSrcs>(tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, false, 1, &src0, 1, ncclShmem.groups[group].dsts+i, realPeerSize);
+              // Mark for threadfence at the end
+              fenceNeeded |= true;
+            }
+          }
+        } else if (Recv) {
+          if (flags & RoleOutput) ncclShmem.groups[group].dsts[0] = userBuff + outIx + offset;
+          ssize_t pOffset = index*peerOffset;
+          if (skip >= 0 && index >= skip) pOffset += peerElem;
+          // Adjust remote index with peer offset in case we are directly pulling from peer's output buffer
+          waitPeer<DirectRecv, 0, 1, 0, 0, 1>(outIx+pOffset, outIx+pOffset, offset, realSize);
+          subBarrier();
+          #pragma unroll 1
+          for (int j=0; j<fan.nrecv(); j++) {
+            int i = (j+shift)%fan.nrecv();
+            pOffset = i*peerOffset;
+            if (skip >= 0 && i >= skip) pOffset += peerElem;
+            void* dst0 = (T*)ncclShmem.groups[group].dsts[0] + pOffset;
+            ssize_t realPeerSize = min(realSize, totalElem-pOffset);
+            if (DirectRecv && ncclShmem.groups[group].srcs[i] == dst0) realPeerSize = 0;
+            if (realPeerSize > 0) reduceCopy<Unroll, RedOp, T, 0,1,1, 0,1,1, /*PreOpSrcs=*/0>(tid, nworkers, ncclShmem.redOpArgs[0], ncclShmem.redOpArgs, postOp, 1, ncclShmem.groups[group].srcs+i, 1, &dst0, realPeerSize);
+          }
+        }
+      }
+      fenceNeeded = __any(fenceNeeded);
+      postPeer<Recv, Send>(fenceNeeded);
+      offset += realSize;
+    }
+  }
+
+  __device__ __forceinline__ void loadRecvConn(ncclDevChannelPeer *peer, int connIndex, struct ncclWorkElem* e) {
+    if (flags & (RoleWaitRecv|RolePostRecv)) {
+      auto *conn = &peer->recv[connIndex];
+      if (conn->netDeviceHandle.netDeviceType == NCCL_NET_DEVICE_UNPACK) {
+        // handle must be a device ptr
+        netDeviceHandle = conn->netDeviceHandle.handle;
+        // Cache the handle
+        ncclNetDeviceUnpackSetup(netDeviceHandle, group, index);
+        flags |= NetDeviceUnpack;
+      }
+      step = conn->step;
+      step = roundUp(step, SlicePerChunk*StepPerSlice);
+      if (flags & RolePostRecv) {
+        connStepPtr = conn->head;
+        STORE(connStepPtr, step); // Return credits in case we rounded up.
+      }
+      if (flags & RoleWaitRecv) {
+        ncclShmem.groups[group].recvConns[index] = conn; // WaitRecv role saves since that's who needs it in setDataPtrs()
+        flags |= (conn->flags & NCCL_NVLS_MIN_POLL) ? NvlsMinPolling : 0;
+        connStepPtr = conn->tail;
+        connStepCache = loadStepValue(connStepPtr);
+        flags |= (conn->offsFifo != nullptr) ? OffsFifoEnabled : 0;
+        if (Direct) {
+          // User buffers have been registered
+          if ((conn->flags & (NCCL_IPC_READ|NCCL_IPC_WRITE)) && e != nullptr && e->regUsed) {
+            if (connIndex == 1 && P2p == 0) {
+              flags |= DirectRead;  // scatter-reduce use direct pull
+            } else {
+              flags |= (e->direct & NCCL_DIRECT_WRITE) ? DirectWrite :
+                       (e->direct & NCCL_DIRECT_READ)  ? DirectRead  : 0;
+            }
+          } else if (conn->flags & (NCCL_DIRECT_WRITE|NCCL_DIRECT_READ)) {
+            if (connIndex == 1 && P2p == 0) {
+              flags |= DirectRead;  // scatter-reduce use direct pull
+            } else {
+              // direct read not allowed in non-register case
+              // otherwise, in one-to-multi send, we could mix empty send and intermediate send
+              flags |= (conn->flags & NCCL_DIRECT_WRITE) ? DirectWrite : 0;
+            }
+          } else if ((conn->flags & NCCL_NVLS_MIN_POLL) && e != nullptr && e->regUsed) {
+            /* NVLS direct */
+            flags |= NvlsDirectRead;
+          }
+        }
+        if (flags & OffsFifoEnabled)
+          connOffsFifoPtr = conn->offsFifo;
+        connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
+      }
+    }
+  }
+
+  __device__ __forceinline__ void loadSendConn(ncclDevChannelPeer *peer, int connIndex, struct ncclWorkElem* e) {
+    if (flags & (RoleWaitSend|RolePostSend)) {
+      auto *conn = &peer->send[connIndex];
+      step = conn->step;
+      step = roundUp(step, SlicePerChunk*StepPerSlice);
+      if (flags & RolePostSend) {
+        connStepPtr = conn->tail;
+	      next_hdp_reg = conn->next_hdp_reg;
+        connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
+      }
+      if (flags & RoleWaitSend) {
+        ncclShmem.groups[group].sendConns[index] = conn; // WaitSend role saves since that's who needs it in setDataPtrs()
+        flags |= (conn->flags & NCCL_NVLS_MIN_POLL) ? NvlsMinPolling : 0;
+        connStepPtr = conn->head;
+        connStepCache = loadStepValue(connStepPtr);
+        flags |= (conn->offsFifo != nullptr) ? OffsFifoEnabled : 0;
+        if (flags & OffsFifoEnabled)
+          connOffsFifoPtr = conn->offsFifo;
+        connEltsFifo = (T*)conn->buffs[NCCL_PROTO_SIMPLE];
+
+        if (conn->sizesFifo != nullptr) {
+          flags |= SizesFifoEnabled;
+          connSizesFifoPtr = conn->sizesFifo;
+        } else if (Direct) {
+          // User buffers have been registered
+          if ((conn->flags & (NCCL_IPC_READ|NCCL_IPC_WRITE)) && e != nullptr && e->regUsed) {
+            if (connIndex == 1 && P2p == 0) {
+              flags |= DirectRead;  // scatter-reduce use direct pull
+            } else {
+              flags |= (e->direct & NCCL_DIRECT_WRITE) ? DirectWrite :
+                       (e->direct & NCCL_DIRECT_READ)  ? DirectRead  : 0;
+            }
+          } else if (conn->flags & (NCCL_DIRECT_WRITE|NCCL_DIRECT_READ)) {
+            if (connIndex == 1 && P2p == 0) {
+              flags |= DirectRead;  // scatter-reduce use direct pull
+            } else {
+              // direct read not allowed in non-register case
+              // otherwise, in one-to-multi send, we could mix empty send and intermediate send
+              flags |= (conn->flags & NCCL_DIRECT_WRITE) ? DirectWrite : 0;
+            }
+          } else if ((conn->flags & NCCL_NVLS_MIN_POLL) && e != nullptr && e->regUsed) {
+            /* NVLS direct */
+            flags |= NvlsDirectWrite;
+          }
+        }
+      }
+    }
+  }
+
+ public:
+  __forceinline__ __device__ Primitives(
+      int tid, int nthreads, int const *recvPeers, int const *sendPeers,
+      void const *inputBuf, void *outputBuf, uint64_t redOpArg, uint8_t group=0,
+      uint8_t connIndexRecv = 0, uint8_t connIndexSend = 0, struct ncclWorkElem* e = nullptr, int stepSize_=0
+    ):
+    tid(tid), nthreads(nthreads), tidInBlock(threadIdx.x), group(group),
+    stepSize(stepSize_ == 0 ? ncclShmem.comm.buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS/sizeof(T) : stepSize_) {
+
+    // For send operations, we need an extra warp to overlap the threadfence and the copy
+    barriers = &ncclShmem.groups[group].barrier;
+    barrier_next = ncclShmem.groups[group].barrier_next;
+    this->nworkers = nthreads;
+
+    int nrecv=0, nsend=0;
+    while (nrecv < MaxRecv && recvPeers[nrecv] != -1) nrecv++;
+    while (nsend < MaxSend && sendPeers[nsend] != -1) nsend++;
+    this->fan = Fan(nrecv, nsend);
+
+    constexpr int ThreadPerSync = 8;
+    static_assert(MaxSend <= ThreadPerSync && MaxRecv <= ThreadPerSync, "Not enough threads to cover all peers");
+
+    int g = tid / ThreadPerSync;
+    int ng = nthreads / ThreadPerSync;
+    index = tid % ThreadPerSync;
+    flags = 0;
+    if (g == 0) {
+      if (index < nrecv) flags |= RoleWaitRecv;
+      if (index == nrecv) flags |= RoleInput;
+    } else if (g == 1) {
+      if (index < nsend) flags |= RoleWaitSend;
+      if (index == nsend) flags |= RoleOutput;
+    } else if (g == ng - 2) {
+      if (index < nrecv) flags |= RolePostRecv;
+    } else if (g == ng - 1) {
+      if (index < nsend) flags |= RolePostSend;
+    }
+
+    int peer = 0;
+    if (flags & (RoleWaitRecv|RolePostRecv)) peer = recvPeers[index];
+    if (flags & (RoleWaitSend|RolePostSend)) peer = sendPeers[index];
+
+    loadRecvConn(ncclShmem.channel.peers[peer], connIndexRecv, e);
+    loadSendConn(ncclShmem.channel.peers[peer], connIndexSend, e);
+
+    // if (barrierAny(flags & NetDeviceUnpack)) {
+    //   flags |= AnyNetDeviceUnpack;
+    //   // g == 0 is the first ThreadPerSync # of threads of this warp
+    //   // g == 0 is also the RoleWaitRecv threads of this group, thus the thread ID will correlate to the peer index
+    //   if (g == 0) {
+    //     uint32_t mask = __ballot_sync((1U << ThreadPerSync) - 1, (flags & NetDeviceUnpack) ? 1 : 0);
+
+    //     // We only want to update the shared memory variable with a single thread
+    //     if (tid == 0) {
+    //       ncclShmem.groups[this->group].devicePlugin.unpack.unpackNetDeviceIndexMask = mask;
+    //     }
+    //   }
+    // }
+
+    setDataPtrs(inputBuf, outputBuf, redOpArg, (struct ncclWorkElemReg*)e);
+  }
+
+  __forceinline__ __device__ ~Primitives() {
+    // Ensure ncclShmem.groups[].send/recvConns are available
+    if (!(flags & ThreadsSynced))
+      barrier();
+    // Save steps for the next operation
+    if (flags & (RolePostSend|RolePostRecv)) {
+      auto *conns = (flags & RolePostSend) ? ncclShmem.groups[group].sendConns : ncclShmem.groups[group].recvConns;
+      conns[index]->step = step;
+    }
+    
+    if ((flags & (AnyNetDeviceUnpack)) && (flags & (RoleWaitRecv))) {
+      ncclNetDeviceSaveHead(netDeviceHandle, group);
+    }
+    barrier();
+  }
+
+  __device__ void setDataPtrs(void const *inputBuf, void *outputBuf, uint64_t redOpArg, struct ncclWorkElemReg* e) {
+    if (flags & RoleInput) {
+      userBuff = (T*)inputBuf;
+      ncclShmem.redOpArgs[0] = redOpArg;  // scaler for local input
+    }
+    if (flags & RoleOutput) userBuff = (T*)outputBuf;
+    bool recvProvider = flags == (flags|RoleWaitRecv|DirectWrite);
+    bool sendAcceptor = (flags == (flags|RoleWaitSend|DirectWrite)) || (flags == (flags|RoleWaitSend|NvlsDirectWrite));
+    bool sendProvider = flags == (flags|RoleWaitSend|DirectRead); // sender provides direct buffer (to be fetched)
+    bool recvAcceptor = flags == (flags|RoleWaitRecv|DirectRead) || (flags == (flags|RoleWaitRecv|NvlsDirectRead)); // receiver accepts direct buffer
+    int regUsed = e != nullptr ? e->elem.regUsed : 0;
+
+    if (Direct && recvProvider) {
+      int spins = 0;
+      void *volatile *slot = ncclShmem.groups[group].recvConns[index]->ptrExchange;
+      // Wait for consumer to consume previous value before trampling it.
+      if (slot) {
+        while ((void *)atomicAdd((unsigned long long *) slot,0) != nullptr && !checkAbort(spins));
+        directBuff = (T*)outputBuf;
+        // Encode pointer by XOR'ing against some address they definitely wouldn't send
+        // since we want to allow them sending us nullptr while not colliding with
+        // the empty slot value.
+        *slot = reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(directBuff) ^ reinterpret_cast<uintptr_t>(slot));
+      }
+    }
+    if (Direct && sendAcceptor) {
+      int spins = 0;
+      void *volatile *slot = ncclShmem.groups[group].sendConns[index]->ptrExchange;
+      void *ptr;
+      while (slot) {
+        ptr = (void *)atomicAdd((unsigned long long *) slot,0);
+        if (ptr != nullptr || checkAbort(spins)) break;
+      }
+
+      if (slot) {
+        directBuff = regUsed ? (T*)(e->dnOutputs[index]) :
+                   reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(ptr) ^ reinterpret_cast<uintptr_t>(slot));
+        *slot = nullptr;
+      } else {
+        /* slot is NULL, it must be regUsed == 1 */
+        directBuff = (T*)e->dnOutputs[index];
+      }
+    }
+    if (Direct && sendProvider) {
+      int spins = 0;
+      void *volatile *slot = ncclShmem.groups[group].sendConns[index]->ptrExchange;
+      volatile uint64_t* argSlot0 = ncclShmem.groups[group].sendConns[index]->redOpArgExchange;
+      volatile uint64_t* argSlot1 = ncclShmem.groups[group].sendConns[index]->redOpArgExchange+1;
+      // Wait for consumer to consume previous value before trampling it.
+      if (slot && argSlot0 && argSlot1) {
+        while (((void *)atomicAdd((unsigned long long *) slot,0) != nullptr || *argSlot0 != 0 || *argSlot1 !=0) && !checkAbort(spins));
+        // If there is no recv, then we are directly pulling from input buffer (e.g. directScatter)
+        // Otherwise, we are pulling from output buffer (e.g. recvCopyDirectSend)
+        directBuff = MaxRecv == 0 ? (T*)inputBuf : (T*)outputBuf;
+        // Exchange pre-scalers for use in direct pull
+        *argSlot0 = (uint64_t(1)<<32) | (uint32_t)redOpArg;
+        *argSlot1 = (uint64_t(1)<<32) | (uint32_t)(redOpArg>>32);
+        // Encode pointer by XOR'ing against some address they definitely wouldn't send
+        // since we want to allow them sending us nullptr while not colliding with
+        // the empty slot value.
+        *slot = reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(directBuff) ^ reinterpret_cast<uintptr_t>(slot));
+      }
+    }
+    if (Direct && recvAcceptor) {
+      int spins = 0;
+      void *volatile *slot = ncclShmem.groups[group].recvConns[index]->ptrExchange;
+      volatile uint64_t* argSlot0 = ncclShmem.groups[group].recvConns[index]->redOpArgExchange;
+      volatile uint64_t* argSlot1 = ncclShmem.groups[group].recvConns[index]->redOpArgExchange+1;
+      void *ptr;
+      while (slot) {
+        ptr = (void *)atomicAdd((unsigned long long *) slot,0);
+        if (ptr != nullptr || checkAbort(spins)) break;
+      }
+
+      if (slot && argSlot0 && argSlot1) {
+        directBuff = regUsed ? (T*)(MaxSend == 0 ? e->upOutputs[index] : e->dnInputs[index]) :
+          reinterpret_cast<T*>(reinterpret_cast<uintptr_t>(ptr) ^ reinterpret_cast<uintptr_t>(slot));
+        if (MaxSend != 0) { // reduce group rather than gather group
+          // Store scalers for remote inputs
+          uint64_t arg0, arg1;
+          while (true) {
+            arg0 = *argSlot0;
+            arg1 = *argSlot1;
+            if ((arg0 != 0 && arg1 != 0) || checkAbort(spins)) break;
+          }
+          ncclShmem.redOpArgs[1 + index] = ((arg1 & 0xffffffff) << 32) | (arg0 & 0xffffffff);
+        }
+        *argSlot0 = 0; *argSlot1 = 0;
+        *slot = nullptr;
+      } else {
+        directBuff = (T*)e->dnInputs[index];
+      }
+    }
+  }
+
+  __device__ void moveDataPtrs(intptr_t delta) {
+    if (flags & (RoleInput|RoleOutput))
+      userBuff += delta;
+  }
+
+  // Set MSCCL data pointers
+  __device__ __forceinline__ void setDataPtrs(void const *inputBuf, void *outputBuf) {
+    if (flags & RoleInput) userBuff = (T*)inputBuf;
+    if (flags & RoleOutput) userBuff = (T*)outputBuf;
+  }
+
+  __device__ __forceinline__ void send(intptr_t inpIx, int eltN) {
+    genericOp<0, 0, 0, 1, Input, -1>(inpIx, -1, eltN, false);
+  }
+  __device__ __forceinline__ void sendFromOutput(intptr_t outIx, int eltN) {
+    genericOp<0, 0, 0, 1, Output, -1>(outIx, -1, eltN, false);
+  }
+  __device__ __forceinline__ void directSend(intptr_t inpIx, intptr_t outIx, int eltN) {
+    genericOp<0, 1, 0, 1, Input, -1>(inpIx, outIx, eltN, false);
+  }
+  __device__ __forceinline__ void directSendFromOutput(intptr_t outIx, int eltN) {
+    genericOp<0, 1, 0, 1, Output, -1>(outIx, outIx, eltN, false);
+  }
+
+  __device__ __forceinline__ void recv(intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 0, 1, 0, -1, Output>(-1, outIx, eltN, postOp);
+  }
+  __device__ __forceinline__ void directRecv(intptr_t outIx, int eltN) {
+    genericOp<1, 0, 1, 0, -1, Output>(-1, outIx, eltN, /*postOp=*/false);
+  }
+  __device__ __forceinline__ void directRecvCopy(intptr_t inpIx, intptr_t outIx, int eltN) {
+    genericOp<1, 0, 1, 0, -1, Output>(inpIx, outIx, eltN, /*postOp=*/false);
+  }
+
+  __device__ __forceinline__ void copySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 0, 0, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+  __device__ __forceinline__ void directCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 1, 0, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+
+  __device__ __forceinline__ void recvSend(int eltN, bool postOp=false) {
+    genericOp<0, 0, 1, 1, -1, -1>(-1, -1, eltN, postOp);
+  }
+  __device__ __forceinline__ void recvCopySend(intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 0, 1, 1, -1, Output>(-1, outIx, eltN, postOp);
+  }
+  __device__ __forceinline__ void directRecvCopySend(intptr_t outIx, int eltN) {
+    genericOp<1, 1, 1, 1, -1, Output>(-1, outIx, eltN, false);
+  }
+  __device__ __forceinline__ void directRecvDirectSend(intptr_t inpIx, intptr_t outIx, int eltN) {
+    genericOp<1, 1, 1, 1, -1, -1>(inpIx, outIx, eltN, false);
+  }
+  __device__ __forceinline__ void recvCopyDirectSend(intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 1, 1, 1, -1, Output>(-1, outIx, eltN, postOp);
+  }
+
+  __device__ __forceinline__ void recvReduceCopy(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 0, 1, 0, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+
+  __device__ __forceinline__ void recvReduceSend(intptr_t inpIx, int eltN, bool postOp=false) {
+    genericOp<0, 0, 1, 1, Input, -1>(inpIx, -1, eltN, postOp);
+  }
+  __device__ __forceinline__ void directRecvReduceSend(intptr_t inpIx, int eltN, bool postOp=false) {
+    genericOp<1, 0, 1, 1, Input, -1>(inpIx, -1, eltN, postOp);
+  }
+
+  __device__ __forceinline__ void recvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    genericOp<0, 0, 1, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+  __device__ __forceinline__ void directRecvReduceCopySend(intptr_t inpIx, intptr_t outIx, int eltN, bool postOp=false) {
+    // Direct is only for the send part
+    genericOp<0, 1, 1, 1, Input, Output>(inpIx, outIx, eltN, postOp);
+  }
+
+  __device__ __forceinline__ void
+  scatter(intptr_t inpIx, ssize_t totalElem, int peerElem, ssize_t peerOffset, int skip, int shift) {
+    ScatterGatherOp<0, 0, 0, 1>(inpIx, -1, totalElem, peerElem, peerOffset, skip, shift, /*postOp=*/false);
+  }
+  __device__ __forceinline__ void
+  directScatter(intptr_t inpIx, ssize_t totalElem, int peerElem, ssize_t peerOffset, int skip, int shift) {
+    ScatterGatherOp<0, 1, 0, 1>(inpIx, -1, totalElem, peerElem, peerOffset, skip, shift, /*postOp=*/false);
+  }
+
+  __device__ __forceinline__ void
+  gather(intptr_t outIx, ssize_t totalElem, int peerElem, ssize_t peerOffset, int skip, int shift, bool postOp=false) {
+    ScatterGatherOp<0, 0, 1, 0>(-1, outIx, totalElem, peerElem, peerOffset, skip, shift, postOp);
+  }
+  __device__ __forceinline__ void
+  directGather(intptr_t outIx, ssize_t totalElem, int peerElem, ssize_t peerOffset, int skip, int shift) {
+    ScatterGatherOp<1, 0, 1, 0>(-1, outIx, totalElem, peerElem, peerOffset, skip, shift, /*postOp=*/false);
+  }
+
+  // MSCCL primitives
+  __device__ __forceinline__ void sendWithBarrier(intptr_t inpIx, int eltN) {
+    send(inpIx, eltN);
+  }
+  __device__ __forceinline__ void localCopy(T* srcs, T* dsts, int eltN) {
+    return mscclGenericOp<0,1,0,0>(&srcs, 1, &dsts, 1, eltN);
+  }
+};

src/device/reduce.h

@@ -0,0 +1,96 @@
+/*************************************************************************
+ * 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;
+    const ssize_t chunkSize = int(Proto::calcBytePerStep()/sizeof(T) * (Proto::Id == NCCL_PROTO_SIMPLE ? REDUCE_CHUNKSTEPS : 1));
+    const ssize_t minChunkSizeLL128 = int(nthreads*(Proto::calcBytePerGrain()/sizeof(T)));
+    const int nranks = ncclShmem.comm.nRanks;
+    const ssize_t loopSize = nChannels*chunkSize;
+    const ssize_t size = args->count;
+    const int rank = ncclShmem.comm.rank;
+    const int prevRank = ring->userRanks[nranks-1];
+    const int root = args->root;
+
+    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);
+
+    auto calcChunkSize = [&]__device__(ssize_t gridOffset)->int {
+      int 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);
+      return realChunkSize;
+    };
+
+    if (prevRank == root) {
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        int realChunkSize = calcChunkSize(gridOffset);
+        ssize_t offset = gridOffset + bid*realChunkSize;
+        int nelem = min(realChunkSize, size-offset);
+        prims.send(offset, nelem);
+      }
+    }
+    else if (rank == root) {
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        int realChunkSize = calcChunkSize(gridOffset);
+        ssize_t offset = gridOffset + bid*realChunkSize;
+        int nelem = min(realChunkSize, size-offset);
+        prims.recvReduceCopy(offset, offset, nelem, /*postOp=*/true);
+      }
+    }
+    else {
+      for (ssize_t gridOffset = 0; gridOffset < size; gridOffset += loopSize) {
+        int realChunkSize = calcChunkSize(gridOffset);
+        ssize_t offset = gridOffset + bid*realChunkSize;
+        int nelem = min(realChunkSize, size-offset);
+        prims.recvReduceSend(offset, nelem);
+      }
+    }
+  }
+}
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    using Proto = ProtoSimple<REDUCE_CHUNKSTEPS/REDUCE_SLICESTEPS, REDUCE_SLICESTEPS>;
+    runRing<T, RedOp, Proto>(args);
+  }
+};
+
+template<typename T, typename RedOp>
+struct RunWorkElement<ncclFuncReduce, 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<ncclFuncReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
+  __device__ __forceinline__ void run(ncclWorkElem *args) {
+    runRing<T, RedOp, ProtoLL128>(args);
+  }
+};

src/device/reduce_kernel.h

@@ -0,0 +1,711 @@
+/*************************************************************************
+ * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+
+#ifndef NCCL_REDUCE_KERNEL_H_
+#define NCCL_REDUCE_KERNEL_H_
+
+#include "op128.h"
+#include <limits>
+#include <type_traits>
+
+template<typename T>
+struct IsFloatingPoint: std::false_type {};
+template<>
+struct IsFloatingPoint<half>: std::true_type {};
+#if defined(RCCL_BFLOAT16)
+template<>
+struct IsFloatingPoint<rccl_bfloat16>: std::true_type {};
+#endif
+template<>
+struct IsFloatingPoint<float>: std::true_type {};
+template<>
+struct IsFloatingPoint<double>: std::true_type {};
+
+////////////////////////////////////////////////////////////////////////////////
+// The reduction function classes. All classes must:
+//  1. Expose the `EltType` typedef.
+//  2. Have constructor taking no arguments (default constructible).
+//  3. Have constructor taking `uint64_t opArg`.
+
+template<typename T>
+struct FuncCopy { using EltType = T; __device__ FuncCopy(uint64_t opArg=0) {}; };
+template<typename T>
+struct FuncSum  { using EltType = T; __device__ FuncSum(uint64_t opArg=0) {}; };
+template<typename T>
+struct FuncProd { using EltType = T; __device__ FuncProd(uint64_t opArg=0) {}; };
+template<typename T>
+struct FuncMinMax {
+  using EltType = T;
+  BytePack<sizeof(T)> xormask; // only used by integers
+  bool isMinNotMax; // only used by floats
+  __device__ FuncMinMax(uint64_t opArg=0) {
+    xormask.native = opArg;
+    isMinNotMax = (opArg&1)==0;
+  }
+};
+template<typename T> struct FuncPreMulSum;
+template<typename T> struct FuncSumPostDiv;
+
+////////////////////////////////////////////////////////////////////////////////
+// Trait classes for reduction functions. Given a function (FuncSum, etc.)
+// and a number of elements in a pack, will reduce, preOp, or postOp a pack
+// of elements. These classes are intended to be specialized for specific
+// combinations of reduction function and pack size.
+
+template<typename Fn, int EltPerPack>
+struct Apply_Reduce /*{
+  static BytePack<EltPerPack*sizeof(T)> reduce(
+    Fn fn, BytePack<EltPerPack*sizeof(T)> a, BytePack<EltPerPack*sizeof(T)> b
+  );
+}*/;
+template<typename Fn, int EltPerPack>
+struct Apply_PreOp/*{
+  static constexpr bool IsIdentity;
+  static BytePack<EltPerPack*sizeof(T)> preOp(Fn fn, BytePack<EltPerPack*sizeof(T)> a);
+}*/;
+template<typename Fn, int EltPerPack>
+struct Apply_PostOp/*{
+  static constexpr bool IsIdentity;
+  static BytePack<EltPerPack*sizeof(T)> postOp(Fn fn, BytePack<EltPerPack*sizeof(T)> a);
+}*/;
+template<typename Fn>
+struct LoadMultimem_BigPackSize/*{
+  // If non-zero, then this and sizeof(T) are valid pack sizes for LoadMultimem,
+  // otherwise there are no valid pack sizes for LoadMultimem.
+  static constexpr int BigPackSize = 0;
+}*/;
+template<typename Fn, int BytePerPack>
+struct Apply_LoadMultimem/*{
+  static BytePack<BytePerPack> load(Fn fn, uintptr_t addr);
+}*/;
+
+////////////////////////////////////////////////////////////////////////////////
+// Public API for calling the trait classes. These take the data elements as a
+// pack of any type, which could be a BytePack<?> or any integral type (uint64_t,
+// uint32_t, etc.), and will return a new pack where each element has been
+// transformed appropriately.
+
+template<typename Fn, typename Pack>
+__device__ __forceinline__ Pack applyReduce(Fn fn, Pack a, Pack b) {
+  return fromPack<Pack>(
+    Apply_Reduce<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
+      ::reduce(fn, toPack(a), toPack(b))
+  );
+}
+
+template<typename Fn, typename Pack>
+__device__ __forceinline__ Pack applyPreOp(Fn fn, Pack a) {
+  return fromPack<Pack>(
+    Apply_PreOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
+      ::preOp(fn, toPack(a))
+  );
+}
+
+template<typename Fn, typename Pack>
+__device__ __forceinline__ Pack applyPostOp(Fn fn, Pack a) {
+  return fromPack<Pack>(
+    Apply_PostOp<Fn, BytePackOf<Pack>::Size/sizeof(typename Fn::EltType)>
+      ::postOp(fn, toPack(a))
+  );
+}
+
+template<typename Fn, int BytePerPack>
+__device__ __forceinline__ BytePack<BytePerPack> applyLoadMultimem(Fn fn, uintptr_t addr) {
+  return Apply_LoadMultimem<Fn, BytePerPack>::load(fn, addr);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Apply_Reduce
+
+// Nonsensical base case
+template<typename Fn>
+struct Apply_Reduce<Fn, /*EltPerPack=*/0> {
+  __device__ static BytePack<0> reduce(Fn fn, BytePack<0> a, BytePack<0> b) {
+    return  {};
+  }
+};
+
+// General recursive definition (EltPerPack > 1). This is how we iterate over
+// all elements in a pack of any size, by breaking it into halves. Eventually
+// we'll hit a base case (a more specific template specialization which takes
+// precedence).
+template<typename Fn, int EltPerPack>
+struct Apply_Reduce {
+  template<int Size>
+  __device__ static BytePack<Size> reduce(Fn fn, BytePack<Size> a, BytePack<Size> b) {
+    a.half[0] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[0], b.half[0]);
+    a.half[1] = Apply_Reduce<Fn, EltPerPack/2>::reduce(fn, a.half[1], b.half[1]);
+    return a;
+  }
+};
+
+// Base case definitions (EltPerPack == 1)
+template<typename T>
+struct Apply_Reduce<FuncCopy<T>, /*EltPerPack=*/1> {
+  __device__ static BytePack<sizeof(T)> reduce(FuncCopy<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
+    return a;
+  }
+};
+template<typename T>
+struct Apply_Reduce<FuncSum<T>, /*EltPerPack=*/1> {
+  __device__ static BytePack<sizeof(T)> reduce(FuncSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
+    return toPack<T>(fromPack<T>(a) + fromPack<T>(b));
+  }
+};
+template<typename T>
+struct Apply_Reduce<FuncProd<T>, /*EltPerPack=*/1> {
+  __device__ static BytePack<sizeof(T)> reduce(FuncProd<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
+    return toPack<T>(fromPack<T>(a) * fromPack<T>(b));
+  }
+};
+template<typename T>
+struct Apply_Reduce<FuncMinMax<T>, /*EltPerPack=*/1> {
+  __device__ static BytePack<sizeof(T)> reduce(FuncMinMax<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
+    return (a.native ^ fn.xormask.native) < (b.native ^ fn.xormask.native) ? a : b;
+  }
+};
+
+// Optimizations for specfic types and element count combinations:
+template<>
+struct Apply_Reduce<FuncSum<uint8_t>, /*EltPerPack=*/4> {
+  __device__ static BytePack<4> reduce(FuncSum<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
+    constexpr uint32_t even = 0x00ff00ffu;
+    uint32_t x = (a.native &  even) + (b.native &  even);
+    uint32_t y = (a.native & ~even) + (b.native & ~even);
+    //a.native = (x & even) | (y & ~even);
+    a.native = __byte_perm(x, y, 0x7250);
+    return a;
+  }
+};
+
+template<>
+struct Apply_Reduce<FuncMinMax<uint8_t>, /*EltPerPack=*/4> {
+  __device__ static BytePack<4> reduce(FuncMinMax<uint8_t> fn, BytePack<4> a, BytePack<4> b) {
+    constexpr uint32_t ones = 0x01010101u;
+    constexpr uint32_t even = 0x00ff00ffu; // even byte mask
+    // Replicate xormask to all bytes
+    uint32_t x = fn.xormask.native * ones;
+    // Transform inputs by xormask
+    uint32_t ax = a.native ^ x;
+    uint32_t bx = b.native ^ x;
+    // Use 9-bit arithmetic to compute d=a-b
+    uint32_t d0 = (ax    & even) + (~bx      & even) + ones;
+    uint32_t d1 = (ax>>8 & even) + (~(bx>>8) & even) + ones;
+    // Move sign bit of each 9-bit delta into the least bit of origin byte
+    //uint32_t s = (d0>>8 & ones & even) | (d1 & ones & ~even);
+    uint32_t s = __byte_perm(d0, d1, 0x7351) & ones;
+    // Broadcast least bit across whole byte
+    s *= 0xffu;
+    // Compose result by selecting bytes via: signbit(a-b)==1 ? a : b
+    a.native = (a.native & s) | (b.native & ~s);
+    return a;
+  }
+};
+
+#ifdef NEED_CHECKING
+template<>
+struct Apply_Reduce<FuncProd<uint8_t>, /*EltPerPack=*/4> {
+  __device__ static BytePack<4> reduce(FuncProd<uint8_t> fn, BytePack<4> apack, BytePack<4> bpack) {
+    uint32_t a = apack.native;
+    uint32_t b = bpack.native;
+    uint32_t ab0 = (a*b) & 0xffu;
+    asm("mad.lo.u32 %0, %1, %2, %0;" : "+r"(ab0) : "r"(a&0xff00u), "r"(b&0xff00u));
+    uint32_t ab1;
+    asm("mul.hi.u32 %0, %1, %2;"     : "=r"(ab1) : "r"(a&0xff0000), "r"(b&0xff0000));
+    asm("mad.hi.u32 %0, %1, %2, %0;" : "+r"(ab1) : "r"(a&0xff000000u), "r"(b&0xff000000u));
+    apack.native = __byte_perm(ab0, ab1, 0x6420);
+    return apack;
+  }
+};
+#endif
+
+#define SPECIALIZE_REDUCE(Fn, T, EltPerPack, Vec, expr_of_fn_x_y) \
+  template<> \
+  struct Apply_Reduce<Fn<T>, EltPerPack> { \
+    __device__ __forceinline__ static BytePack<sizeof(Vec)> reduce( \
+        Fn<T> fn, BytePack<sizeof(Vec)> a, BytePack<sizeof(Vec)> b \
+      ) { \
+      Vec x = fromPack<Vec>(a); \
+      Vec y = fromPack<Vec>(b); \
+      return toPack<Vec>(expr_of_fn_x_y); \
+    } \
+  };
+
+SPECIALIZE_REDUCE(FuncMinMax, float, 1, float, fn.isMinNotMax ? fminf(x, y) : fmaxf(x, y))
+SPECIALIZE_REDUCE(FuncMinMax, double, 1, double, fn.isMinNotMax ? fmin(x, y) : fmax(x, y))
+
+#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
+  SPECIALIZE_REDUCE(FuncSum, half, 1, half, __hadd(x, y))
+  SPECIALIZE_REDUCE(FuncSum, half, 2, half2, __hadd2(x, y))
+  SPECIALIZE_REDUCE(FuncProd, half, 1, half, __hmul(x, y))
+  SPECIALIZE_REDUCE(FuncProd, half, 2, half2, __hmul2(x, y))
+#else
+  SPECIALIZE_REDUCE(FuncSum, half, 1, half, __float2half(__half2float(x) + __half2float(y)))
+  SPECIALIZE_REDUCE(FuncProd, half, 1, half, __float2half(__half2float(x) * __half2float(y)))
+#endif
+
+#if __CUDA_ARCH__ >= 800
+  SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, fn.isMinNotMax ? __hmin(x, y) : __hmax(x, y))
+  SPECIALIZE_REDUCE(FuncMinMax, half, 2, half2, fn.isMinNotMax ? __hmin2(x, y) : __hmax2(x, y))
+#else
+  SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, __float2half(fn.isMinNotMax ? fminf(__half2float(x), __half2float(y)) : fmaxf(__half2float(x), __half2float(y))))
+#endif
+
+#if defined(RCCL_BFLOAT16)
+#if __CUDA_ARCH__ >= 800
+  SPECIALIZE_REDUCE(FuncSum, __nv_bfloat16, 1, __nv_bfloat16, __hadd(x, y))
+  SPECIALIZE_REDUCE(FuncSum, __nv_bfloat16, 2, __nv_bfloat162, __hadd2(x, y))
+  SPECIALIZE_REDUCE(FuncProd, __nv_bfloat16, 1, __nv_bfloat16, __hmul(x, y))
+  SPECIALIZE_REDUCE(FuncProd, __nv_bfloat16, 2, __nv_bfloat162, __hmul2(x, y))
+  SPECIALIZE_REDUCE(FuncMinMax, __nv_bfloat16, 1, __nv_bfloat16, fn.isMinNotMax ? __hmin(x, y) : __hmax(x, y))
+  SPECIALIZE_REDUCE(FuncMinMax, __nv_bfloat16, 2, __nv_bfloat162, fn.isMinNotMax ? __hmin2(x, y) : __hmax2(x, y))
+#else
+  SPECIALIZE_REDUCE(FuncSum, rccl_bfloat16, 1, rccl_bfloat16, (rccl_bfloat16)((float)(x) + (float)(y)))
+  SPECIALIZE_REDUCE(FuncProd, rccl_bfloat16, 1, rccl_bfloat16, (rccl_bfloat16)((float)(x) * (float)(y)))
+  SPECIALIZE_REDUCE(FuncMinMax, rccl_bfloat16, 1, rccl_bfloat16, (rccl_bfloat16)(fn.isMinNotMax ? fminf((float)(x), (float)(y)) : fmaxf((float)(x), (float)(y))))
+#endif
+#endif
+
+#undef SPECIALIZE_REDUCE
+
+////////////////////////////////////////////////////////////////////////////////
+// Apply_PreOp
+
+// General recursive definition (EltPerPack > 1)
+template<typename Fn, int EltPerPack>
+struct Apply_PreOp {
+  static constexpr bool IsIdentity = Apply_PreOp<Fn, EltPerPack/2>::IsIdentity;
+  template<int Size>
+  __device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
+    #if __cpp_if_constexpr
+    if constexpr(!IsIdentity) {
+    #else
+    if (!IsIdentity) {
+    #endif
+      // The `if (!IsIdentity)` condition is not strictly necessary, but it may help
+      // compiler in that it won't have to tear a register apart for no reason
+      // just to put it back together again.
+      a.half[0] = Apply_PreOp<Fn, EltPerPack/2>::preOp(fn, a.half[0]);
+      a.half[1] = Apply_PreOp<Fn, EltPerPack/2>::preOp(fn, a.half[1]);
+    }
+    return a;
+  }
+};
+// Base case definition (EltPerPack == 1), by default is identity function.
+template<typename Fn>
+struct Apply_PreOp<Fn, /*EltPerPack=*/1> {
+  static constexpr bool IsIdentity = true;
+  template<int Size>
+  __device__ static BytePack<Size> preOp(Fn fn, BytePack<Size> a) {
+    return a;
+  }
+};
+// Base case definition (EltPerPack == 0), is nonsense!
+template<typename Fn>
+struct Apply_PreOp<Fn, /*EltPerPack=*/0> {
+  static constexpr bool IsIdentity = true;
+  __device__ static BytePack<0> preOp(Fn fn, BytePack<0> a) {
+    return {};
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Apply_PostOp
+
+// General recursive definition (EltPerPack > 1)
+template<typename Fn, int EltPerPack>
+struct Apply_PostOp {
+  static constexpr bool IsIdentity = Apply_PostOp<Fn, EltPerPack/2>::IsIdentity;
+  template<int Size>
+  __device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
+    #if __cpp_if_constexpr
+    if constexpr(!IsIdentity) {
+    #else
+    if (!IsIdentity) {
+    #endif
+      // The `if (!IsIdentity)` condition is not strictly necessary, but it may help
+      // compiler in that it won't have to tear a register apart for no reason
+      // just to put it back together again.
+      a.half[0] = Apply_PostOp<Fn, EltPerPack/2>::postOp(fn, a.half[0]);
+      a.half[1] = Apply_PostOp<Fn, EltPerPack/2>::postOp(fn, a.half[1]);
+    }
+    return a;
+  }
+};
+// Base case definition (EltPerPack == 1), by default is identity function.
+template<typename Fn>
+struct Apply_PostOp<Fn, /*EltPerPack=*/1> {
+  static constexpr bool IsIdentity = true;
+  template<int Size>
+  __device__ static BytePack<Size> postOp(Fn fn, BytePack<Size> a) {
+    return a;
+  }
+};
+// Base case definition (EltPerPack == 0), is nonsense!
+template<typename Fn>
+struct Apply_PostOp<Fn, /*EltPerPack=*/0> {
+  static constexpr bool IsIdentity = true;
+  __device__ static BytePack<0> postOp(Fn fn, BytePack<0> a) {
+    return {};
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////
+// FuncPreMulSum
+
+// General definition for all integral types, float, and double.
+template<typename T>
+struct FuncPreMulSum {
+  using EltType = T;
+  T scalar;
+  __device__ FuncPreMulSum(uint64_t opArg=0) {
+    union { uint64_t u64; T val; };
+    u64 = opArg;
+    scalar = val;
+  }
+};
+
+template<>
+struct FuncPreMulSum<half> {
+  using EltType = half;
+#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
+  half2 scalar;
+  __device__ FuncPreMulSum(uint64_t opArg=0) {
+    union { uint64_t u64; half val; };
+    u64 = opArg;
+    scalar.x = val;
+    scalar.y = val;
+  }
+#else
+  float scalar;
+  __device__ FuncPreMulSum(uint64_t opArg=0) {
+    union { uint64_t u64; half val; };
+    u64 = opArg;
+    scalar = __half2float(val);
+  }
+#endif
+};
+
+#if defined(RCCL_BFLOAT16)
+  template<>
+  struct FuncPreMulSum<rccl_bfloat16> {
+    using EltType = rccl_bfloat16;
+  #if __CUDA_ARCH__ >= 800
+    __nv_bfloat162 scalar;
+    __device__ FuncPreMulSum(uint64_t opArg=0) {
+      union { uint64_t u64; __nv_bfloat16 val; };
+      u64 = opArg;
+      scalar.x = val;
+      scalar.y = val;
+    }
+  #else
+    float scalar;
+    __device__ FuncPreMulSum(uint64_t opArg=0) {
+      union { uint64_t u64; rccl_bfloat16 val; };
+      u64 = opArg;
+      scalar = (float)(val);
+    }
+  #endif
+  };
+#endif
+
+template<typename T>
+struct Apply_Reduce<FuncPreMulSum<T>, /*EltPerPack=*/1> {
+  __device__ static BytePack<sizeof(T)> reduce(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
+    // FuncPreMulSum reduce dispatches to FuncSum.
+    return Apply_Reduce<FuncSum<T>, 1>::reduce(FuncSum<T>(), a, b);
+  }
+};
+
+// PreOp of FuncPreMulSum for integral types, float, and double.
+template<typename T>
+struct Apply_PreOp<FuncPreMulSum<T>, /*EltPerPack=*/1> {
+  static constexpr bool IsIdentity = false;
+  __device__ static BytePack<sizeof(T)> preOp(FuncPreMulSum<T> fn, BytePack<sizeof(T)> a) {
+    return toPack<T>(fromPack<T>(a) * fn.scalar);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Apply_PreOp of FuncPreMulSum for float16.
+
+template<>
+struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/1> {
+  static constexpr bool IsIdentity = false;
+  __device__ static BytePack<sizeof(half)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half)> a) {
+    #if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
+      return toPack<half>(__hmul(fromPack<half>(a), fn.scalar.x));
+    #else
+      return toPack<half>(__float2half(__half2float(fromPack<half>(a)) * fn.scalar));
+    #endif
+  }
+};
+#if __CUDA_ARCH__ >= 530 && __CUDA_ARCH__ != 610
+  template<>
+  struct Apply_PreOp<FuncPreMulSum<half>, /*EltPerPack=*/2> {
+    static constexpr bool IsIdentity = false;
+    __device__ static BytePack<sizeof(half2)> preOp(FuncPreMulSum<half> fn, BytePack<sizeof(half2)> a) {
+      return toPack<half2>(__hmul2(fromPack<half2>(a), fn.scalar));
+    }
+  };
+#endif
+
+////////////////////////////////////////////////////////////////////////////////
+// Apply_PreOp of FuncPreMulSum for bfloat16.
+
+#if defined(RCCL_BFLOAT16)
+  template<>
+  struct Apply_PreOp<FuncPreMulSum<rccl_bfloat16>, /*EltPerPack=*/1> {
+    static constexpr bool IsIdentity = false;
+    __device__ static BytePack<sizeof(rccl_bfloat16)> preOp(
+        FuncPreMulSum<rccl_bfloat16> fn, BytePack<sizeof(rccl_bfloat16)> a
+      ) {
+      #if __CUDA_ARCH__ >= 800
+        return toPack<__nv_bfloat16>(__hmul(fromPack<__nv_bfloat16>(a), fn.scalar.x));
+      #else
+        return toPack<rccl_bfloat16>((rccl_bfloat16)((float)(fromPack<rccl_bfloat16>(a)) * fn.scalar));
+      #endif
+    }
+  };
+  #if __CUDA_ARCH__ >= 800
+    template<>
+    struct Apply_PreOp<FuncPreMulSum<rccl_bfloat16>, /*EltPerPack=*/2> {
+      static constexpr bool IsIdentity = false;
+      __device__ static BytePack<sizeof(__nv_bfloat162)> preOp(
+          FuncPreMulSum<__nv_bfloat16> fn, BytePack<sizeof(__nv_bfloat162)> a
+        ) {
+        return toPack<__nv_bfloat162>(__hmul2(fromPack<__nv_bfloat162>(a), fn.scalar));
+      }
+    };
+  #endif
+#endif
+
+////////////////////////////////////////////////////////////////////////////////
+// FuncSumPostDiv
+
+template<typename T, bool IsFloating=IsFloatingPoint<T>::value>
+struct FuncSumPostDiv_IntOnly;
+
+template<typename T>
+struct FuncSumPostDiv: FuncSumPostDiv_IntOnly<T> {
+  __device__ FuncSumPostDiv(uint64_t opArg=0):
+    FuncSumPostDiv_IntOnly<T>(opArg) {
+  }
+};
+
+template<typename T>
+struct FuncSumPostDiv_IntOnly<T, /*IsFloating=*/false>: FuncSum<T> {
+  using EltType = T;
+  int divisor;
+  __device__ FuncSumPostDiv_IntOnly(uint64_t opArg=0): divisor(opArg) {}
+};
+
+template<typename T>
+struct FuncSumPostDiv_IntOnly<T, /*IsFloating=*/true> {
+  static_assert(sizeof(T)!=sizeof(T), "FuncSumPostDiv is only for implementing ncclAvg on integral types.");
+};
+
+template<typename T>
+struct Apply_Reduce<FuncSumPostDiv<T>, /*EltPerPack=*/1>:
+    Apply_Reduce<FuncSum<T>, 1> {
+  __device__ static BytePack<sizeof(T)> reduce(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a, BytePack<sizeof(T)> b) {
+    // FuncSumPostDiv reduce dispatches to FuncSum.
+    return Apply_Reduce<FuncSum<T>, 1>::reduce(FuncSum<T>(), a, b);
+  }
+};
+
+template<typename T>
+struct Apply_PostOp<FuncSumPostDiv<T>, /*EltPerPack=*/1> {
+  static constexpr bool IsIdentity = false;
+  __device__ static BytePack<sizeof(T)> postOp(FuncSumPostDiv<T> fn, BytePack<sizeof(T)> a) {
+    return toPack<T>(fromPack<T>(a) / fn.divisor);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Apply_LoadMultimem
+
+#define SIZEOF_BytePack_field_u16 2
+#define PTX_REG_BytePack_field_u16 "h"
+
+#define SIZEOF_BytePack_field_u32 4
+#define PTX_REG_BytePack_field_u32 "r"
+
+#define SIZEOF_BytePack_field_u64 8
+#define PTX_REG_BytePack_field_u64 "l"
+
+#define DEFINE_Apply_LoadMultimem_sum(T, ptx_ty, pack_field) \
+  template<> \
+  struct Apply_LoadMultimem<FuncSum<T>, SIZEOF_BytePack_field_##pack_field> { \
+    static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
+    __device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
+      BytePack<PackSize> ans; \
+      asm("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
+        : "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
+        : "l"(addr)); \
+      return ans; \
+    } \
+  };
+#define DEFINE_Apply_LoadMultimem_minmax(T, ptx_ty, pack_field) \
+  template<> \
+  struct Apply_LoadMultimem<FuncMinMax<T>, SIZEOF_BytePack_field_##pack_field> { \
+    static constexpr int PackSize = SIZEOF_BytePack_field_##pack_field; \
+    __device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
+      BytePack<PackSize> ans; \
+      if (fn.isMinNotMax) { \
+        asm("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
+          : "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
+          : "l"(addr)); \
+      } else { \
+        asm("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
+          : "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field) \
+          : "l"(addr)); \
+      } \
+      return ans; \
+    } \
+  };
+
+#define DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
+  template<> \
+  struct Apply_LoadMultimem<FuncSum<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
+    static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
+    __device__ static BytePack<PackSize> load(FuncSum<T> fn, uintptr_t addr) { \
+      BytePack<PackSize> ans; \
+      asm("multimem.ld_reduce.relaxed.sys.global.add.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
+        : "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
+          "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
+          "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
+          "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
+        : "l"(addr)); \
+      return ans; \
+    } \
+  };
+#define DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
+  template<> \
+  struct Apply_LoadMultimem<FuncMinMax<T>, 4*(SIZEOF_BytePack_field_##pack_field)> { \
+    static constexpr int PackSize = 4*(SIZEOF_BytePack_field_##pack_field); \
+    __device__ static BytePack<PackSize> load(FuncMinMax<T> fn, uintptr_t addr) { \
+      BytePack<PackSize> ans; \
+      if (fn.isMinNotMax) { \
+        asm("multimem.ld_reduce.relaxed.sys.global.min.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
+          : "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
+            "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
+            "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
+            "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
+          : "l"(addr)); \
+      } else { \
+        asm("multimem.ld_reduce.relaxed.sys.global.max.v4." #ptx_ty " {%0,%1,%2,%3}, [%4];" \
+          : "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[0]), \
+            "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[1]), \
+            "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[2]), \
+            "=" PTX_REG_BytePack_field_##pack_field(ans.pack_field[3]) \
+          : "l"(addr)); \
+      } \
+      return ans; \
+    } \
+  };
+
+#define DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(T, ptx_ty, pack_field) \
+  DEFINE_Apply_LoadMultimem_sum_v4(T, ptx_ty, pack_field) \
+  template<> \
+  struct Apply_LoadMultimem<FuncSum<T>, sizeof(T)> { \
+    __device__ static BytePack<sizeof(T)> load(FuncSum<T> fn, uintptr_t addr) { \
+      BytePack<2*sizeof(T)> tmp; \
+      asm("multimem.ld_reduce.relaxed.sys.global.add." #ptx_ty " %0, [%1];" \
+        : "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
+        : "l"(addr & -uintptr_t(sizeof(T)))); \
+      return tmp.half[(addr/sizeof(T))%2]; \
+    } \
+  };
+#define DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(T, ptx_ty, pack_field) \
+  DEFINE_Apply_LoadMultimem_minmax_v4(T, ptx_ty, pack_field) \
+  template<> \
+  struct Apply_LoadMultimem<FuncMinMax<T>, sizeof(T)> { \
+    __device__ static BytePack<sizeof(T)> load(FuncMinMax<T> fn, uintptr_t addr) { \
+      BytePack<2*sizeof(T)> tmp; \
+      if (fn.isMinNotMax) { \
+        asm("multimem.ld_reduce.relaxed.sys.global.min." #ptx_ty " %0, [%1];" \
+          : "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
+          : "l"(addr & -uintptr_t(sizeof(T)))); \
+      } else { \
+        asm("multimem.ld_reduce.relaxed.sys.global.max." #ptx_ty " %0, [%1];" \
+          : "=" PTX_REG_BytePack_field_##pack_field(tmp.pack_field) \
+          : "l"(addr & -uintptr_t(sizeof(T)))); \
+      } \
+      return tmp.half[(addr/sizeof(T))%2]; \
+    } \
+  };
+
+template<typename Fn, int BytePerPack>
+struct Apply_LoadMultimem {
+  __device__ static BytePack<BytePerPack> load(Fn fn, uintptr_t addr) {
+    //__trap();
+    return {};
+  }
+};
+
+#if __CUDA_ARCH__ >= 900 && CUDART_VERSION >= 12010
+  template<typename Fn>
+  struct LoadMultimem_BigPackSize {
+    using T = typename Fn::EltType;
+    static constexpr bool IsSum = std::is_same<Fn, FuncSum<T>>::value ||
+                                  std::is_same<Fn, FuncPreMulSum<T>>::value ||
+                                  std::is_same<Fn, FuncSumPostDiv<T>>::value;
+    static constexpr bool IsMinMax = std::is_same<Fn, FuncMinMax<T>>::value;
+    static constexpr bool IsFloat = IsFloatingPoint<T>::value;
+    static constexpr int BigPackSize =
+      IsFloat && IsSum && sizeof(T) < 8 ? 16 :
+      IsFloat && IsSum ? 8 :
+      IsFloat && IsMinMax && sizeof(T)==2 ? 16 :
+      !IsFloat && (IsSum||IsMinMax) && sizeof(T)>=4 ? sizeof(T) :
+      /*multimem.ld_reduce not supported:*/ 0;
+  };
+
+  DEFINE_Apply_LoadMultimem_sum(uint32_t, u32, u32)
+  DEFINE_Apply_LoadMultimem_minmax(uint32_t, u32, u32)
+
+  DEFINE_Apply_LoadMultimem_sum(int32_t, s32, u32)
+  DEFINE_Apply_LoadMultimem_minmax(int32_t, s32, u32)
+
+  DEFINE_Apply_LoadMultimem_sum(uint64_t, u64, u64)
+  DEFINE_Apply_LoadMultimem_minmax(uint64_t, u64, u64)
+
+  DEFINE_Apply_LoadMultimem_sum(int64_t, u64, u64)
+  DEFINE_Apply_LoadMultimem_minmax(int64_t, s64, u64)
+
+  DEFINE_Apply_LoadMultimem_sum(float, f32, u32)
+  DEFINE_Apply_LoadMultimem_sum_v4(float, f32, u32)
+
+  DEFINE_Apply_LoadMultimem_sum(double, f64, u64)
+
+  DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(half, f16x2, u32)
+  DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(half, f16x2, u32)
+
+  #if defined(RCCL_BFLOAT16)
+    DEFINE_Apply_LoadMultimem_sum_v4x2_and_subhalf(rccl_bfloat16, bf16x2, u32)
+    DEFINE_Apply_LoadMultimem_minmax_v4x2_and_subhalf(rccl_bfloat16, bf16x2, u32)
+  #endif
+#else
+  template<typename Fn>
+  struct LoadMultimem_BigPackSize {
+    static constexpr int BigPackSize = 0;
+  };
+#endif
+
+#undef DEFINE_Apply_LoadMultimem
+#undef DEFINE_Apply_LoadMultimem_v4
+#undef DEFINE_Apply_LoadMultimem_v4x2_and_subhalf
+#undef SIZEOF_BytePack_field_u64
+#undef PTX_REG_BytePack_field_u64
+#undef SIZEOF_BytePack_field_u32
+#undef PTX_REG_BytePack_field_u32
+#undef SIZEOF_BytePack_field_u16
+#undef PTX_REG_BytePack_field_u16
+
+#endif // REDUCE_KERNEL_H_

src/device/reduce_scatter.h

@@ -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);
+      }
+    }
+  }
+};

src/device/sendrecv.h

@@ -0,0 +1,233 @@
+/*************************************************************************
+ * 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"
+#if defined(ENABLE_NPKIT)
+#include "npkit/npkit.h"
+#endif
+
+template<typename T, typename RedOp>
+struct RunWork<ncclFuncSendRecv, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_SIMPLE> {
+  template<typename Proto>
+  __device__ void runSend(const int tid, const int nthreads, const uint8_t group, struct ncclWorkElemP2p* args) {
+    void* buff = reinterpret_cast<void*>(uintptr_t(args->buffHi32)<<32 | args->buffLo32);
+    ssize_t count = reinterpret_cast<size_t>(size_t(args->countHi32)<<32 | args->countLo32);
+
+#if defined(ENABLE_NPKIT)
+    bool isNpKitThread = (tid == 0);
+    int npKitCtxIdx = blockIdx.x * NCCL_MAX_WORK_ELEMENTS_P2P + group;
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (isNpKitThread) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (isNpKitThread) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    if (args->peer == ncclShmem.comm.rank) {
+      struct ncclWorkElemP2p* recvArgs = args-1;
+      void* recvBuff = reinterpret_cast<void*>(uintptr_t(recvArgs->buffHi32)<<32 | recvArgs->buffLo32);
+      if (buff != recvBuff) {
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_LOCAL_COPY_ENTRY)
+        if (isNpKitThread) {
+          NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_RECV_LOCAL_COPY_ENTRY, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+              ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY)
+        if (isNpKitThread) {
+          NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_ENTRY, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+              ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        }
+#endif
+
+#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+        reduceCopy<COLL_UNROLL*2, RedOp, T, 0,1,1, 0,1,1, /*PreOpSrcs=*/0>
+          (tid, nthreads, 0, nullptr, false, 1, &buff, 1, &recvBuff, count);
+#else
+        reduceCopy<COLL_UNROLL, RedOp, T, 0,1,1, 0,1,1, /*PreOpSrcs=*/0>
+          (tid, nthreads, 0, nullptr, false, 1, &buff, 1, &recvBuff, count);
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT)
+        if (isNpKitThread) {
+          NpKit::CollectGpuEvent(NPKIT_EVENT_PRIM_SIMPLE_REDUCE_OR_COPY_MULTI_EXIT, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+              ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_LOCAL_COPY_EXIT)
+        if (isNpKitThread) {
+          NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_RECV_LOCAL_COPY_EXIT, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+              ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        }
+#endif
+      }
+    } else {
+      int chunkSize = args->chunkSize/sizeof(T);
+      if (args->proto == NCCL_PROTO_LL) chunkSize /= 2;
+      int const peer = args->peer;
+      Primitives<T, RedOp, FanAsymmetric<0, 1>, 1, Proto, 1> prims
+        (tid, nthreads, nullptr, &peer, buff, nullptr, /*redOpArg(ignored)=*/0, group, 1, 1, nullptr, ncclShmem.comm.p2pChunkSize/sizeof(T));
+
+#if defined(ENABLE_NPKIT)
+      if (isNpKitThread) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_SEND_ENTRY)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_RECV_SEND_ENTRY, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      size_t offset = 0;
+      do {
+        int nelem = min(size_t(chunkSize), count-offset);
+        prims.directSend(offset, offset, nelem);
+        offset += nelem;
+      } while(offset < count);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_SEND_EXIT)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_RECV_SEND_EXIT, count*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+  }
+
+  template<typename Proto>
+  __device__ void runRecv(const int tid, const int nthreads, const uint8_t group, struct ncclWorkElemP2p* args) {
+#if defined(ENABLE_NPKIT)
+    bool isNpKitThread = (tid == 0);
+    int npKitCtxIdx = blockIdx.x * NCCL_MAX_WORK_ELEMENTS_P2P + group;
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_CPU)
+    if (isNpKitThread) {
+      uint64_t* cpuTimestamp = ncclShmem.comm.cpuTimestamp;
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_CPU, 0, 0, *cpuTimestamp,
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_TIME_SYNC_GPU)
+    if (isNpKitThread) {
+      NpKit::CollectGpuEvent(NPKIT_EVENT_TIME_SYNC_GPU, 0, 0, NPKIT_GET_GPU_TIMESTAMP(),
+          ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+    }
+#endif
+
+    if (args->peer != ncclShmem.comm.rank) {
+      void* buff = reinterpret_cast<void*>(uintptr_t(args->buffHi32)<<32 | args->buffLo32);
+      ssize_t count = reinterpret_cast<size_t>(size_t(args->countHi32)<<32 | args->countLo32);
+      int chunkSize = args->chunkSize/sizeof(T);
+      if (args->proto == NCCL_PROTO_LL) chunkSize /= 2; // This is to account for chunkEffectiveSize
+      int const peer = args->peer;
+      Primitives<T, RedOp, FanAsymmetric<1, 0>, 1, Proto, 1> prims
+        (tid, nthreads, &peer, nullptr, nullptr, buff, /*redOpArg(ignored)=*/0, group, 1, 1, nullptr, ncclShmem.comm.p2pChunkSize/sizeof(T));
+
+#if defined(ENABLE_NPKIT)
+      if (isNpKitThread) {
+        prims.npKitCtxIdx = npKitCtxIdx;
+      }
+#endif
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_RECV_ENTRY)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_RECV_RECV_ENTRY, count*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+        prims.npKitDataProcessTotalTime = 0;
+      }
+#endif
+
+      size_t offset = 0;
+      do {
+        int nelem = min(size_t(chunkSize), count-offset);
+        prims.directRecv(offset, nelem);
+        offset += nelem;
+      } while(offset < count);
+
+#if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_SEND_RECV_RECV_EXIT)
+      if (isNpKitThread) {
+        NpKit::CollectGpuEvent(NPKIT_EVENT_SEND_RECV_RECV_EXIT, count*sizeof(T), prims.npKitDataProcessTotalTime, NPKIT_GET_GPU_TIMESTAMP(),
+            ncclShmem.comm.npKitEventCollectContexts + npKitCtxIdx);
+      }
+#endif
+
+    }
+  }
+
+#if defined(USE_INDIRECT_FUNCTION_CALL) && !defined(__gfx940__) && !defined(__gfx941__) && !defined(__gfx942__)
+  __device__  void run(ncclWork *work) {
+#else
+  __device__  __attribute__((noinline)) void run(ncclWork *work) {
+#endif
+    struct ncclWorkElemP2p* args = work->p2pElems;
+    int ngroups = args->ngroups;
+    int tid = threadIdx.x;
+    int wid = tid / WARP_SIZE;
+    // This has to work even for groups of 2.5 warps (which is 8 groups, and means 3
+    // warps for send, 2 warps for recv).
+    // warpStarts were rounded thanks to int division, but for group number we need to round the other way around
+    // So we mirror wid then mirror again the group.
+    #define NWARPS (NCCL_MAX_NTHREADS/WARP_SIZE)
+    uint8_t group = ngroups-1- (NWARPS-1-wid) * ngroups / NWARPS;
+    args += group;
+    tid -= args->warpStart * WARP_SIZE;
+    int nthreads = args->nWarps * WARP_SIZE;
+
+    if (args->p2pType == ncclWorkP2pTypeUnused) return;
+    if (tid >= nthreads || args->peer == -1) return;
+
+    // Select Proto here
+    // This is to allow the same kernel to run multiple primitives on different warps (thread groups)
+    if ((group%2) == 0) {
+      if (args->proto == NCCL_PROTO_LL) {
+        runRecv<ProtoLL>(tid, nthreads, group, args);
+      } else {
+#if defined(__gfx90a__)
+        runRecv<ProtoSimple<1,1,8>>(tid, nthreads, group, args);
+#elif defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+        runRecv<ProtoSimple<1,1,4>>(tid, nthreads, group, args);
+#else
+        runRecv<ProtoSimple<1,1>>(tid, nthreads, group, args);
+#endif
+      }
+    } else {
+      if (args->proto == NCCL_PROTO_LL) {
+        runSend<ProtoLL>(tid, nthreads, group, args);
+      } else {
+#if defined(__gfx90a__)
+        runSend<ProtoSimple<1,1,8>>(tid, nthreads, group, args);
+#elif defined(__gfx908__) || defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__)
+        runSend<ProtoSimple<1,1,4>>(tid, nthreads, group, args);
+#else
+        runSend<ProtoSimple<1,1>>(tid, nthreads, group, args);
+#endif
+      }
+    }
+  }
+};