From e031a4a2f4c8e16269f97ab34dff43df87c1a0c9 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Thu, 3 Jul 2025 02:16:18 -0700
Subject: [PATCH] msccl: use special send for LL on gfx950 (#1788)

[ROCm/rccl commit: ae9642d4bc03fe2609fae5743643f8f48d6291a9]
---
 projects/rccl/src/device/msccl_kernel_impl.h |  2 +-
 projects/rccl/src/device/prims_ll.h          | 47 ++++++++++++++++++++
 projects/rccl/src/device/prims_ll128.h       |  3 ++
 projects/rccl/src/device/prims_simple.h      |  3 ++
 4 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/projects/rccl/src/device/msccl_kernel_impl.h b/projects/rccl/src/device/msccl_kernel_impl.h
index 1daf10e37d..5e0c574e27 100644
--- a/projects/rccl/src/device/msccl_kernel_impl.h
+++ b/projects/rccl/src/device/msccl_kernel_impl.h
@@ -278,7 +278,7 @@ __device__ __forceinline__ void mscclRunInterpreter(
             }
 
 #endif
-          prims.send(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end.
+          prims.mscclSend(srcOffset, thisNelem); // LL.send is the only situation where there is no barrier at the end.
 
 #if defined(ENABLE_NPKIT) && defined(ENABLE_NPKIT_EVENT_MSCCL_SEND_EXIT)
             if (tid == 0) {
diff --git a/projects/rccl/src/device/prims_ll.h b/projects/rccl/src/device/prims_ll.h
index 6f6018bb86..296a458f27 100644
--- a/projects/rccl/src/device/prims_ll.h
+++ b/projects/rccl/src/device/prims_ll.h
@@ -828,4 +828,51 @@ public:
   __device__ void localCopy(T* srcs, T* dsts, int eltN) {
     return mscclGenericOp<0,1,0,0>(&srcs, 1, &dsts, 1, eltN);
   }
+
+  __device__ void mscclStoreLL(union ncclLLFifoLine* dst, uint64_t val, uint32_t flag) {
+    union ncclLLFifoLine i4;
+    i4.data1 = val & 0xffffffff;
+    i4.flag1 = flag;
+    i4.data2 = (val >> 32);
+    i4.flag2 = flag;
+    __builtin_nontemporal_store(i4.v[0], dst->v);
+    __builtin_nontemporal_store(i4.v[1], dst->v+1);
+  }
+
+  __device__ void mscclSend(intptr_t srcIx, int nelem) {
+#if defined(__gfx950__)
+    T *srcElts = userBufs[0] + srcIx;
+
+    // Always waitSend in case of cleanup
+    nelem = nelem < 0 ? 0 : nelem;
+    waitSend(divUp(nelem, EltPerLine)*sizeof(ncclLLFifoLine));
+
+    nelem -= tid*EltPerLine;
+    srcElts += tid*EltPerLine;
+    int offset = tid;
+    int eltPerTrip = nthreads*EltPerLine;
+    while (nelem > 0) {
+      int eltInLine = EltPerLine < nelem ? EltPerLine : nelem;
+
+      DataLoader dl;
+      ncclLLFifoLine line[MaxRecv];
+      uint64_t data, peerData;
+      dl.loadBegin(srcElts, eltInLine);
+      srcElts += eltPerTrip;
+      data = dl.loadFinish();
+
+      for (int i=1; i < MaxSend && i < fan.nsend(); i++)
+        mscclStoreLL(sendPtr(i)+offset, data, sendFlag(i));
+      mscclStoreLL(sendPtr(0)+offset, data, sendFlag(0));
+      nelem -= eltPerTrip;
+      offset += nthreads;
+    }
+
+    for (int i=1; i < MaxSend && i < fan.nsend(); i++)
+      incSend(i, offset);
+    incSend(0, offset);
+#else
+    LLGenericOp<0, 1, Input, -1>(srcIx, -1, nelem, false);
+#endif
+  }
 };
diff --git a/projects/rccl/src/device/prims_ll128.h b/projects/rccl/src/device/prims_ll128.h
index 5a12e7382c..400ff1fc78 100644
--- a/projects/rccl/src/device/prims_ll128.h
+++ b/projects/rccl/src/device/prims_ll128.h
@@ -587,4 +587,7 @@ public:
   __device__ void localCopy(T* srcs, T* dsts, int eltN) {
     return mscclGenericOp<0,1,0,0>(&srcs, 1, &dsts, 1, eltN);
   }
+  __device__ void mscclSend(intptr_t inpIx, int eltN) {
+    return GenericOp<0, 1, Input, -1>(inpIx, -1, eltN, false);
+  }
 };
diff --git a/projects/rccl/src/device/prims_simple.h b/projects/rccl/src/device/prims_simple.h
index 82b0d93960..f1632f8a06 100644
--- a/projects/rccl/src/device/prims_simple.h
+++ b/projects/rccl/src/device/prims_simple.h
@@ -1310,4 +1310,7 @@ public:
   __device__ __forceinline__ void localCopy(T* srcs, T* dsts, int eltN) {
     return mscclGenericOp<0,1,0,0>(&srcs, 1, &dsts, 1, eltN);
   }
+  __device__ __forceinline__ void mscclSend(intptr_t inpIx, int eltN) {
+    genericOp<0, 0, 0, 1, Input, -1>(inpIx, -1, eltN, false);
+  }
 };