Fix MSCCL work FIFO out-of-bound issue (#935)

[ROCm/rccl commit: 4c117e5335]

projects/rccl/src/include/msccl/msccl_struct.h

@@ -15,7 +15,7 @@
 
 #define MSCCL_MAX_NUM_STEPS 64
 #define MSCCL_MAX_NUM_THREAD_BLOCKS_PER_CHANNEL 32
-#define MSCCL_MAX_NUM_THREAD_BLOCKS (MSCCL_MAX_NUM_THREAD_BLOCKS_PER_CHANNEL * MAXCHANNELS)
+#define MSCCL_MAX_NUM_THREAD_BLOCKS 64
 #define MSCCL_MAX_COUNT 72 // max concurrent number of msccl chunk transmission
 #define MSCCL_MAX_REDUCE_FUSION 16
 #define MSCCL_MAX_NUM_ALGOS 1024
@@ -216,7 +216,7 @@ struct mscclStatus {
   struct mscclWork* workFifo;
   uint32_t* workFifoDone;
   uint32_t workFifoSent;
-  uint32_t workFifoSentPerChannel[MAXCHANNELS];
+  uint32_t workFifoSentPerThreadBlock[MSCCL_MAX_NUM_THREAD_BLOCKS];
   uint32_t workFifoAckdMin;
 };
 

projects/rccl/src/misc/msccl/msccl_lifecycle.cc

@@ -192,10 +192,10 @@ ncclResult_t mscclInit(ncclComm_t comm) {
     status.needsProxy = false;
     status.workFifoDepth = MSCCL_WORK_FIFO_DEPTH;
     NCCLCHECK(ncclCudaCalloc(&status.workFifo, status.workFifoDepth, nullptr, true));
-    NCCLCHECK(ncclCudaHostCalloc(&status.workFifoDone, MAXCHANNELS));
+    NCCLCHECK(ncclCudaHostCalloc(&status.workFifoDone, MSCCL_MAX_NUM_THREAD_BLOCKS));
     status.workFifoSent = 0;
-    for (int i = 0; i < MAXCHANNELS; i++) {
-      status.workFifoSentPerChannel[i] = 0;
+    for (int i = 0; i < MSCCL_MAX_NUM_THREAD_BLOCKS; i++) {
+      status.workFifoSentPerThreadBlock[i] = 0;
     }
     status.workFifoAckdMin = 0;
     mscclSchedulerTriedLoadAlgo = false;

projects/rccl/src/misc/msccl/msccl_setup.cc

@@ -330,28 +330,28 @@ static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) {
     while (1) {
       // We have to poll for notifications from device.
       uint32_t* doneLive = status.workFifoDone;
-      uint32_t ackd[MAXCHANNELS];
-      for (int c=0; c < MAXCHANNELS; c++) {
+      uint32_t ackd[MSCCL_MAX_NUM_THREAD_BLOCKS];
+      for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) {
         ackd[c] = __atomic_load_n(&doneLive[c], __ATOMIC_RELAXED);
       }
       // Compiler-only fence to prevent fusion of loops to encourage dense loads.
       __atomic_signal_fence(__ATOMIC_SEQ_CST);
 
       uint32_t ackdAll = status.workFifoSent;
-      for (int c=0; c < MAXCHANNELS; c++) {
+      for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) {
         // ackdAll is min over all non-quiesced channels
-        if (ackd[c] != status.workFifoSentPerChannel[c])
+        if (ackd[c] != status.workFifoSentPerThreadBlock[c])
           ackdAll = rollingMin32(ackdAll, ackd[c]);
       }
 
       // Compiler only fence to prevent fusion of loops to encourage dense stores.
       __atomic_signal_fence(__ATOMIC_SEQ_CST);
 
-      for (int c=0; c < MAXCHANNELS; c++) {
+      for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) {
         // Advance counter on quiesced channels so they don't lag behind
         // too far where they could get lost in 32-bit wraparound.
-        if (ackd[c] == status.workFifoSentPerChannel[c]) {
-          status.workFifoSentPerChannel[c] = ackdAll;
+        if (ackd[c] == status.workFifoSentPerThreadBlock[c]) {
+          status.workFifoSentPerThreadBlock[c] = ackdAll;
           __atomic_store_n(&doneLive[c], ackdAll, __ATOMIC_RELAXED);
         }
       }
@@ -416,7 +416,7 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
   for (int i = 0; i < numBlocks; i++) {
     work.workFifoDoneAck = workFifoSent + i;
     work.workFifoDone = status.workFifoDone + i;
-    status.workFifoSentPerChannel[i] = workFifoSent + i;
+    status.workFifoSentPerThreadBlock[i] = workFifoSent + i;
     status.workFifo[(workFifoSent + i) & workFifoIdxMask] = work;
   }
   status.workFifoSent = workFifoSent + numBlocks;