From df128879a6b903e2422a7561b6c1923d79378f3b Mon Sep 17 00:00:00 2001
From: Ziyue Yang <ziyyang@microsoft.com>
Date: Thu, 16 Nov 2023 00:18:44 +0800
Subject: [PATCH] Optimize MSCCL reduce primitive switching for gfx942 (#962)

* Optimize reduce primitive switching for gfx942

* address comment
---
 src/collectives/device/msccl_kernel_impl.h | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/collectives/device/msccl_kernel_impl.h b/src/collectives/device/msccl_kernel_impl.h
index a66a083362..99778488ed 100644
--- a/src/collectives/device/msccl_kernel_impl.h
+++ b/src/collectives/device/msccl_kernel_impl.h
@@ -366,19 +366,28 @@ __device__ __forceinline__ void mscclRunInterpreter(
         }
         else if (t->type == MSCCL_REDUCE) {
           int numReductions = t->numReductions;
-          if (thisNelem < nthreads){
+          int currIdx = tid;
+#if defined(__gfx942__)
+          if (Proto::Id == NCCL_PROTO_LL) {
+#else
+          if (thisNelem < nthreads) {
+#endif
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_REDUCE_ENTRY)
             if (tid == 0) {
               NpKit::CollectGpuEventLDS(NPKIT_EVENT_MSCCL_REDUCE_ENTRY, thisNelem*sizeof(T), 0, NPKIT_GET_GPU_TIMESTAMP());
             }
 #endif
 
-            if (tid < thisNelem){
+#if defined(__gfx942__)
+            while (currIdx < thisNelem) {
+#else
+            if (currIdx < thisNelem) {
+#endif
               dstOffset = gridOffset + (ssize_t) (t->dstOffset+c) * sizePerMscclChunk;
-              T* dstIndex = dstPointer + dstOffset + tid;
+              T* dstIndex = dstPointer + dstOffset + currIdx;
               T reduceInput;
               T o = load(dstIndex);
-              ssize_t srcBaseOffset = gridOffset + (ssize_t)c * sizePerMscclChunk + tid;
+              ssize_t srcBaseOffset = gridOffset + (ssize_t)c * sizePerMscclChunk + currIdx;
               switch (numReductions) {
                 case 7:
                   #pragma unroll
@@ -395,6 +404,9 @@ __device__ __forceinline__ void mscclRunInterpreter(
                   break;
               }
               store(dstIndex, o);
+#if defined(__gfx942__)
+              currIdx += nthreads;
+#endif
             }
 
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_REDUCE_EXIT)