apps/nccl: fix a bug in allreduce kernels for graph mode

allreduce7 and allreduce8 were updating the LL protocol flag
on the host side. So, it was not properly captured in graph mode.
This PR fixes the issue by updating the flag in the kernels.

Author: nusislam

apps/nccl/src/allreduce.hpp

@@ -441,11 +441,18 @@ template <typename T>
 __global__ void __launch_bounds__(32, 1)
     allreduceAllToAll(T* buff, T* scratch, T* resultBuff, mscclpp::DeviceHandle<mscclpp::MemoryChannel>* memoryChannels,
                       size_t channelDataOffset, size_t channelScratchOffset, int rank, int nRanksPerNode, int worldSize,
-                      Op op, size_t nelems, uint32_t flag) {
+                      Op op, size_t nelems, uint64_t* deviceFlag, mscclpp::DeviceSyncer* deviceSyncer) {
   // This version of allreduce only works for single nodes
   if (worldSize != nRanksPerNode) return;
   if (sizeof(T) == 2) nelems = (nelems * sizeof(T) + sizeof(T)) / sizeof(int);
   const int nPeers = nRanksPerNode - 1;
+
+  uint64_t commFlag = *deviceFlag;
+  uint32_t flag = (uint32_t) commFlag;
+
+  size_t scratchBaseOffset = (flag % 2) ? SCRATCH_SIZE/2 : 0;
+  channelScratchOffset = scratchBaseOffset;
+
   const int nBlocksPerPeer = gridDim.x / nPeers;
   const int localBlockIdx = blockIdx.x % nBlocksPerPeer;
   const int tid = threadIdx.x + localBlockIdx * blockDim.x;
@@ -478,13 +485,20 @@ __global__ void __launch_bounds__(32, 1)
     }
     dst[idx] = data;
   }
+  __syncthreads();
+
+  deviceSyncer->sync(gridDim.x);
+
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+         *deviceFlag = *deviceFlag + 1;
+  }
 }
 
 template <typename T>
 __global__ void __launch_bounds__(1024, 1)
     allreduce7(T* buff, T* scratch, T* resultBuff, mscclpp::DeviceHandle<mscclpp::MemoryChannel>* memoryChannels,
                size_t channelDataOffset, size_t channelScratchOffset, int rank, int nRanksPerNode, int worldSize, Op op,
-               size_t nelems, uint32_t flag
+               size_t nelems, uint64_t* deviceFlag, mscclpp::DeviceSyncer* deviceSyncer
 #if defined(ENABLE_NPKIT)
                ,
                NpKitEventCollectContext* npKitEventCollectContexts, uint64_t* cpuTimestamp) {
@@ -527,6 +541,13 @@ __global__ void __launch_bounds__(1024, 1)
   const int nPeers = nRanksPerNode - 1;
   const size_t nPkts = nelems / 2;
 
+  uint64_t commFlag = *deviceFlag;
+  uint32_t flag = (uint32_t) commFlag;
+
+  size_t scratchBaseOffset = (flag % 2) ? SCRATCH_SIZE/2 : 0;
+  channelScratchOffset = scratchBaseOffset;
+
+
   int nelemsPerRank = nelems / worldSize;
   if ((nelemsPerRank % 2)) nelemsPerRank = (nelemsPerRank * sizeof(T) + sizeof(T)) / sizeof(T);
 
@@ -580,6 +601,7 @@ __global__ void __launch_bounds__(1024, 1)
       channels[index].write(offset, packet);
     }
   }
+  __syncthreads();
   // step 3: get data result from scratch buffer
   mscclpp::LLPacket* dstPkt = (mscclpp::LLPacket*)((char*)scratch + scratchResultOffset);
   const int dstOffset = remoteRank * nPktsPerRank;
@@ -589,6 +611,7 @@ __global__ void __launch_bounds__(1024, 1)
     result[idx].x = data.x;
     result[idx].y = data.y;
   }
+  __syncthreads();
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_KERNEL_ALLREDUCE_ENTRY) && \
     defined(ENABLE_NPKIT_EVENT_KERNEL_ALLREDUCE_EXIT)
   NpKit::CollectGpuEventShm(NPKIT_EVENT_KERNEL_ALLREDUCE_ENTRY, 0, 0, npkit_timestamp_entry, event_buffer,
@@ -599,6 +622,11 @@ __global__ void __launch_bounds__(1024, 1)
 #if defined(ENABLE_NPKIT)
   NpKit::StoreGpuEventShm(npKitEventCollectContexts, event_buffer, event_buffer_head);
 #endif
+  deviceSyncer->sync(gridDim.x);
+
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+         *deviceFlag = *deviceFlag + 1;
+  }
 }
 
 template <typename T>
@@ -741,15 +769,15 @@ template <typename T>
 cudaError_t allreduce(T* buff, T* scratch, T* resultBuff, mscclpp::DeviceHandle<mscclpp::MemoryChannel>* memoryChannels,
                       mscclpp::DeviceHandle<mscclpp::MemoryChannel>* memoryOutChannels, size_t channelInOffset,
                       size_t channelOutOffset, size_t channelScratchOffset, int rank, int nRanksPerNode, int worldSize,
-                      Op op, size_t nelems, cudaStream_t stream) {
+                      Op op, size_t nelems, cudaStream_t stream, uint64_t* deviceFlag, mscclpp::DeviceSyncer* syncer) {
   static uint32_t flag = 1;
 
   if (sizeof(T) * nelems < worldSize * sizeof(int)) {
     int nBlocks = 7;
     int nThreadsPerBlock = 32;
     allreduceAllToAll<<<nBlocks, nThreadsPerBlock, 0, stream>>>(buff, scratch, resultBuff, memoryChannels,
                                                                 channelInOffset, channelScratchOffset, rank,
-                                                                nRanksPerNode, worldSize, op, nelems, flag++);
+                                                                nRanksPerNode, worldSize, op, nelems, deviceFlag, syncer);
   } else if (sizeof(T) * nelems <= (1 << 20)) {
     int nBlocks = 28;
     int nThreadsPerBlock = 1024;
@@ -765,7 +793,7 @@ cudaError_t allreduce(T* buff, T* scratch, T* resultBuff, mscclpp::DeviceHandle<
 #else
     allreduce7<<<nBlocks, nThreadsPerBlock, 0, stream>>>(buff, scratch, resultBuff, memoryChannels, channelInOffset,
                                                          channelScratchOffset, rank, nRanksPerNode, worldSize, op,
-                                                         nelems, flag++);
+                                                         nelems, deviceFlag, syncer);
 #endif
   } else {
     int nBlocks = 35;

apps/nccl/src/nccl.cu

@@ -194,6 +194,9 @@ struct ncclComm {
   uint32_t numScratchBuff;
   uint32_t buffFlag;
 
+  uint64_t* deviceFlag;
+  mscclpp::DeviceSyncer *syncer;
+
   void* mscclppNcclComm;
 };
 
@@ -384,23 +387,25 @@ static ncclResult_t ncclAllReduceFallback(const void* sendbuff, void* recvbuff,
     case ncclFloat16:
       CUDACHECK(allreduce((half*)sendbuff, (half*)comm->scratchBuff.get(), (half*)recvbuff, memoryChannels,
                           memoryOutChannels, offsetIn, offsetOut, offsetScratch, rank, NRANKS_PER_NODE,
-                          comm->comm->bootstrap()->getNranks(), reduceOp, count, stream));
+                          comm->comm->bootstrap()->getNranks(), reduceOp, count, stream, comm->deviceFlag, comm->syncer));
       break;
     case ncclFloat32:
       CUDACHECK(allreduce((float*)sendbuff, (float*)comm->scratchBuff.get(), (float*)recvbuff, memoryChannels,
                           memoryOutChannels, offsetIn, offsetOut, offsetScratch, comm->comm->bootstrap()->getRank(),
-                          NRANKS_PER_NODE, comm->comm->bootstrap()->getNranks(), reduceOp, count, stream));
+                          NRANKS_PER_NODE, comm->comm->bootstrap()->getNranks(), reduceOp, count, stream, 
+			  comm->deviceFlag, comm->syncer));
       break;
     case ncclBfloat16:
       CUDACHECK(allreduce((__bfloat16*)sendbuff, (__bfloat16*)comm->scratchBuff.get(), (__bfloat16*)recvbuff,
                           memoryChannels, memoryOutChannels, offsetIn, offsetOut, offsetScratch, rank, NRANKS_PER_NODE,
-                          comm->comm->bootstrap()->getNranks(), reduceOp, count, stream));
+                          comm->comm->bootstrap()->getNranks(), reduceOp, count, stream, comm->deviceFlag, comm->syncer));
       break;
     case ncclInt32:
     case ncclUint32:
       CUDACHECK(allreduce((int*)sendbuff, (int*)comm->scratchBuff.get(), (int*)recvbuff, memoryChannels,
                           memoryOutChannels, offsetIn, offsetOut, offsetScratch, comm->comm->bootstrap()->getRank(),
-                          NRANKS_PER_NODE, comm->comm->bootstrap()->getNranks(), reduceOp, count, stream));
+                          NRANKS_PER_NODE, comm->comm->bootstrap()->getNranks(), reduceOp, count, stream, 
+			  comm->deviceFlag, comm->syncer));
       break;
     default:
       WARN("datatype is invalid, datatype: %d", datatype);
@@ -524,6 +529,13 @@ static void ncclCommInitRankFallbackSingleNode(ncclComm* commPtr, std::shared_pt
   commPtr->scratchBuff = mscclpp::GpuBuffer<char>(SCRATCH_SIZE).memory();
   commPtr->remoteScratchRegMemories =
       setupRemoteMemories(commPtr->comm, rank, commPtr->scratchBuff.get(), SCRATCH_SIZE, mscclpp::Transport::CudaIpc);
+
+  hipMalloc((void**)&(commPtr->syncer), sizeof(mscclpp::DeviceSyncer));
+  hipMemset((void*)(commPtr->syncer), 0, sizeof(mscclpp::DeviceSyncer));
+
+  uint64_t initFlag = 1;
+  hipMalloc((void**)&(commPtr->deviceFlag), sizeof(uint64_t));
+  hipMemcpy((void*)(commPtr->deviceFlag), &initFlag, sizeof(uint64_t), hipMemcpyHostToDevice);
 }
 
 NCCL_API ncclResult_t ncclGetVersion(int* version) {