diff --git a/projects/rccl/src/include/msccl/msccl_struct.h b/projects/rccl/src/include/msccl/msccl_struct.h index 19a4c375ca..5daaaa883e 100644 --- a/projects/rccl/src/include/msccl/msccl_struct.h +++ b/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; }; diff --git a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc index 4b32e98837..e1d8e6e3ef 100644 --- a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc +++ b/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; diff --git a/projects/rccl/src/misc/msccl/msccl_setup.cc b/projects/rccl/src/misc/msccl/msccl_setup.cc index b6f259b8a6..b683805a98 100644 --- a/projects/rccl/src/misc/msccl/msccl_setup.cc +++ b/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;