From dbcba2923b95423e6f7e8f572386cc806fd9f864 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Fri, 3 Nov 2023 16:06:49 -0700
Subject: [PATCH] Use parallel init of LDS and adjust P2P channels for gfx94x
 (#943)

* Use parallel init of LDS and adjust P2P channels for gfx94x

* Move another init to parallel

* Fix NCCL_NCHANNELS_PER_PEER setting
---
 src/collectives/device/common.h | 34 ++++++++++++++++++++-------------
 src/graph/paths.cc              | 13 ++++---------
 2 files changed, 25 insertions(+), 22 deletions(-)

diff --git a/src/collectives/device/common.h b/src/collectives/device/common.h
index 77da1ded1d..4acd171f43 100644
--- a/src/collectives/device/common.h
+++ b/src/collectives/device/common.h
@@ -468,17 +468,10 @@ template<ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int FnI
 __forceinline__ __device__ void ncclKernel(
     struct ncclDevComm* comm, uint64_t channelMask, struct ncclWork* workHead
   )  {
-  int tid = threadIdx.x;
-  if (tid == 0) {
-    for (auto i = 0; i < NCCL_MAX_GROUPS; i++) {
-      ncclShmem.groups[i].barrier = 0;
-      for (auto j = 0; j < NCCL_MAX_GROUPS; j++) ncclShmem.groups[i].barrier_next[j] = 0;
-    }
-  }
-  // 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.
-  if (tid < WARP_SIZE) {
-    int x = tid;
+  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;
@@ -490,11 +483,26 @@ __forceinline__ __device__ void ncclKernel(
         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;
-  /* set abort flag to 0 */
-  if (tid == 0) ncclShmem.aborted = 0;
 
   if (true) {
     void *dst, *src;
diff --git a/src/graph/paths.cc b/src/graph/paths.cc
index 5df512b108..7f373b2f4f 100644
--- a/src/graph/paths.cc
+++ b/src/graph/paths.cc
@@ -761,12 +761,8 @@ static ncclResult_t ncclTopoGetNchannels(struct ncclTopoSystem* system, int g /*
     // Local rank
     path = system->nodes[GPU].nodes[peer].paths[GPU]+g;
     if (path->type == PATH_NVL) {
-      if (ncclParamNChannelsPerPeer() == -2) {
-        float nvlBw = ncclTopoXGMISpeed(system->nodes[GPU].nodes[g].gpu.gcn);
-        *nChannels = 2*std::max(1, (int)(path->bw / nvlBw));
-      } else {
-        *nChannels = ncclParamNChannelsPerPeer();
-      }
+      float nvlBw = ncclTopoXGMISpeed(system->nodes[GPU].nodes[g].gpu.gcn);
+      *nChannels = (IsArchMatch(system->nodes[GPU].nodes[0].gpu.gcn, "gfx94") ? 4 : 2)*std::max(1, (int)(path->bw / nvlBw));
     } else {
       *nChannels = 2;
     }
@@ -816,10 +812,9 @@ ncclResult_t ncclTopoComputeP2pChannels(struct ncclComm* comm) {
     // Adjust P2P channels on Rome
     comm->p2pnChannelsPerPeer = 2;
     comm->p2pnChannels = 2;
-  }
-  else {
+  } else {
     // Round to next pow2 nChannelsPerPeer and nChannels
-    comm->p2pnChannelsPerPeer = nextPow2(minChannels);
+    comm->p2pnChannelsPerPeer = (ncclParamNChannelsPerPeer() == -2 ? nextPow2(minChannels) : ncclParamNChannelsPerPeer());
     comm->p2pnChannels = nextPow2(comm->p2pnChannels);
   }