From 1557a1f25885191ca87a92a6141c535fe39e2c73 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Fri, 3 Nov 2023 10:38:02 -0700
Subject: [PATCH] msccl: use 32-bit LDS access and add RCCL_MSCCL_FORCE_FULLOPS
 (#953)

[ROCm/rccl commit: bb843459432b31b22478c21bc7b1340175dbf463]
---
 .../collectives/device/msccl_kernel_impl.h    | 19 ++++++++-----------
 projects/rccl/src/misc/msccl/msccl_setup.cc   |  5 ++++-
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/projects/rccl/src/collectives/device/msccl_kernel_impl.h b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
index 9699b1d029..1c875b7011 100644
--- a/projects/rccl/src/collectives/device/msccl_kernel_impl.h
+++ b/projects/rccl/src/collectives/device/msccl_kernel_impl.h
@@ -100,16 +100,17 @@ inline __device__ static void barrier(int nthreads) {
 
 // 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 = 8 * tid;
+  int offset = sizeof(uint32_t) * tid;
   if (offset < bytes) {
-    uint64_t *src2 = (uint64_t*)((char const*)src + offset);
-    uint64_t *dst2 = (uint64_t*)((char*)dst + offset);
+    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(
-  uint64_t *dst, uint64_t const *src, uint64_t size, int tid, int nthreads) {
+  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];
   }
@@ -143,8 +144,8 @@ __device__ __forceinline__ void mscclRunInterpreter(
 #endif
   // initialize mscclShmem.mscclTB
   threadBlockCopy(
-    (uint64_t *)&mscclShmem.mscclTB, (uint64_t *)(algo->mscclTBs + bid),
-    sizeof(struct mscclThreadBlock) / sizeof(uint64_t), tid, nthreads);
+    (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
@@ -158,20 +159,17 @@ __device__ __forceinline__ void mscclRunInterpreter(
       dst = &ncclShmem.comm;
       src = comm;
       bytes = sizeof(ncclDevComm);
-      static_assert(sizeof(ncclDevComm) <= sizeof(uint64_t) * 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) <= sizeof(uint64_t) * WARP_SIZE, "ncclDevChannel cannot be loaded by a single warp in one insn.");
       break;
     case 2:
       dst = &mscclShmem.work;
       src = work + blockIdx.x;
       bytes = sizeof(mscclWork);
-      static_assert(sizeof(mscclWork) <= sizeof(uint64_t) * WARP_SIZE, "mscclWork cannot be loaded by a single warp in one insn.");
       break;
     case 3:
       /* set abort flag to 0 */
@@ -194,6 +192,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
   }
 #endif
   __synclds(); // publish shmem
+
   if (tid == 0)
     *mscclShmem.work.workFifoDone = mscclShmem.work.workFifoDoneAck;
 
@@ -210,8 +209,6 @@ __device__ __forceinline__ void mscclRunInterpreter(
   }
 #endif
 
-  __synclds(); // publish shmem
-  
   // User pointers for primitives
   T* thisInput = (T*)mscclShmem.work.sendBuff;
   T* thisOutput = (T*)mscclShmem.work.recvBuff;
diff --git a/projects/rccl/src/misc/msccl/msccl_setup.cc b/projects/rccl/src/misc/msccl/msccl_setup.cc
index 5222cf39c6..5c8dff5106 100644
--- a/projects/rccl/src/misc/msccl/msccl_setup.cc
+++ b/projects/rccl/src/misc/msccl/msccl_setup.cc
@@ -367,6 +367,8 @@ static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) {
   }
 }
 
+RCCL_PARAM(MscclForceFullOps, "MSCCL_FORCE_FULLOPS", 0);
+
 ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count,
     ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo,
     ncclComm_t comm, hipStream_t stream) {
@@ -433,7 +435,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
                         (1<<MSCCL_RECV_REDUCE_COPY) |
                         (1<<MSCCL_LOCAL_COPY);
   //check if need full ops msccl kernel
-  if (hostAlgo->typeMask & fullOpMask) fnIndex += sizeof(mscclKernelEntries)/sizeof(void *)/2;
+  if ((hostAlgo->typeMask & fullOpMask) || rcclParamMscclForceFullOps())
+    fnIndex += sizeof(mscclKernelEntries)/sizeof(void *)/2;
   void *func = mscclKernelEntries[fnIndex];
   if (enableDoneEvent) {
     CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, comm->doneEvent, 0));