[RCCL] [JIRA ID : ROCM-25479] LL128: generate separate kernels for UBR

projects/rccl/cmake/scripts/add_unroll.sh

@@ -37,6 +37,34 @@ if [[ "$HIP_FILE" =~ .*/src/device/.*\.h ]]; then
   perl -pi -e 's/(runRing<T, RedOp, Proto, USE_ACC, COLL_UNROLL.*?)>/\1, Pipeline>/' "$HIP_FILE"
   perl -pi -e 's/(runTreeSplit<T, RedOp, Proto, USE_ACC, COLL_UNROLL.*?)>/\1, Pipeline>/' "$HIP_FILE"
   perl -pi -e 's/(runTreeUpDown<T, RedOp, Proto, USE_ACC, COLL_UNROLL.*?)>/\1, Pipeline>/' "$HIP_FILE"
-  sed -i "s/\\(struct RunWorkBatch<ncclFunc[^>]*\\)>*/\\1, USE_ACC, COLL_UNROLL, Pipeline>/" "$HIP_FILE"
-  sed -i "s/\\(RunWorkColl<[^,]*,[^,]*,[^,]*,[^,]*,[^>]*\\)>/\\1, USE_ACC, COLL_UNROLL, Pipeline>/" "$HIP_FILE"
+  sed -i "s/\\(struct RunWorkBatch<ncclFunc[^>]*\\)>*/\\1, USE_ACC, COLL_UNROLL, Pipeline, UserRegMode>/" "$HIP_FILE"
+  sed -i "s/\\(RunWorkColl<[^,]*,[^,]*,[^,]*,[^,]*,[^>]*\\)>/\\1, USE_ACC, COLL_UNROLL, Pipeline, UserRegMode>/" "$HIP_FILE"
+
+  # Declare the UserRegMode template parameter on RunWorkColl / RunWorkBatch
+  # partial specialization headers (those immediately followed by a
+  # `struct RunWork...`). Partial specializations may not carry default template
+  # arguments, so this is added without a default (the primary templates in
+  # common.h provide the default). Must run after the USE_ACC/COLL_UNROLL/Pipeline
+  # header expansion above. The optional params between RedOp and USE_ACC cover
+  # specializations that carry an extra `int Proto` (e.g. AllGatherV).
+  perl -0777 -pi -e 's/(template\s*<typename T, typename RedOp[^>]*?, int USE_ACC, int COLL_UNROLL, int Pipeline)>(\s*\nstruct RunWork)/\1, int UserRegMode>\2/g' "$HIP_FILE"
+
+  # Thread a compile-time UserRegMode parameter onto runRing so LL/LL128 can select
+  # their user-buffer access path (0=runtime, 1=registered, 2=non-registered)
+  # without a runtime dual code path. Append it (defaulted) to every (already
+  # expanded) runRing/runTree header, right after the Pipeline / isNetOffload tail.
+  perl -pi -e 's/(template\s*<typename T, typename RedOp, typename Proto[^>]*int Pipeline[^>]*)>/\1, int UserRegMode = 0>/g' "$HIP_FILE"
+
+  # LL128 reg/noreg split: LL (ProtoLL) keeps the baseline single runtime path and
+  # LL128 is generated as two separate kernels (see generate.py). The RunWorkColl
+  # specialization is instantiated per UserRegMode, so these dispatch helpers just
+  # forward the compile-time UserRegMode template parameter into runRing/runTreeSplit.
+  #   all_reduce ring : <T, RedOp, Proto, RCCLMetadata, USE_ACC, COLL_UNROLL, Pipeline, UserRegMode>
+  #   all_reduce tree : <T, RedOp, Proto, USE_ACC, COLL_UNROLL, Pipeline, UserRegMode>
+  #   all_gather ring : <T, RedOp, Proto, USE_ACC, COLL_UNROLL, Pipeline, isNetOffload, UserRegMode>
+  #   broadcast  ring : <T, RedOp, Proto, USE_ACC, COLL_UNROLL, Pipeline, UserRegMode>
+  perl -pi -e 's/runARRingLL128<T, RedOp>\(/runRing<T, RedOp, ProtoLL128, RCCL_METADATA_EMPTY, USE_ACC, COLL_UNROLL, 0, UserRegMode>(/g' "$HIP_FILE"
+  perl -pi -e 's/runARTreeLL128<T, RedOp>\(/runTreeSplit<T, RedOp, ProtoLL128, USE_ACC, COLL_UNROLL, 0, UserRegMode>(/g' "$HIP_FILE"
+  perl -pi -e 's/runAGRingLL128<T, RedOp>\(/runRing<T, RedOp, ProtoLL128, USE_ACC, COLL_UNROLL, 0, false, UserRegMode>(/g' "$HIP_FILE"
+  perl -pi -e 's/runBcastRingLL128<T, RedOp>\(/runRing<T, RedOp, ProtoLL128, USE_ACC, COLL_UNROLL, 0, UserRegMode>(/g' "$HIP_FILE"
 fi

projects/rccl/src/device/all_gather.h

@@ -54,9 +54,9 @@ __device__ __attribute__((noinline)) void runRing(int tid, int nthreads, struct
       // Coverity reports that the callee treats &ring->next as an array.  However, due to the use of
       // FanSymmetric<1>, only the first element is ever accessed, so it's fine.
       // coverity[callee_ptr_arith:FALSE]
-    Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0, isNetOffload> prims(
-      tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex,
-      work->connIndex, work, nullptr, isNetOffload ? NCCL_MAX_NET_SIZE : 0);
+    Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0, isNetOffload, RCCL_METADATA_EMPTY, 0, 0, UserRegMode>
+      prims(tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex,
+            work->connIndex, work, nullptr, isNetOffload ? NCCL_MAX_NET_SIZE : 0);
 
     for (size_t elemOffset = 0; elemOffset < partCount; elemOffset += chunkCount) {
         /////////////// begin AllGather steps ///////////////
@@ -153,7 +153,9 @@ struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
 template <typename T, typename RedOp>
 struct RunWorkColl<ncclFuncAllGather, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
   __device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
-    runRing<T, RedOp, ProtoLL128>(tid, nthreads, work);
+    // LL128 is generated as separate registered/non-registered kernels; the
+    // compile-time UserRegMode is forwarded by cmake/scripts/add_unroll.sh.
+    runAGRingLL128<T, RedOp>(tid, nthreads, work);
   }
 };
 

projects/rccl/src/device/all_reduce.h

@@ -10,6 +10,10 @@
 #include "primitives.h"
 
 namespace {
+// NOTE: the build-time transformer cmake/scripts/add_unroll.sh appends
+// ", int USE_ACC, int COLL_UNROLL, int Pipeline, int UserRegMode = 0" to this
+// template header. UserRegMode selects the LL128 user-buffer access path at
+// compile time (see prims_ll128.h): 0=runtime, 1=registered, 2=non-registered.
 template <typename T, typename RedOp, typename Proto, int RCCLMetadata>
 #if defined(USE_INDIRECT_FUNCTION_CALL) && !defined(__gfx942__) && !defined(__gfx950__) && !defined(__gfx1250__)
 __device__ void runRing(int tid, int nthreads, struct ncclDevWorkColl* work) {
@@ -43,9 +47,9 @@ __device__ __attribute__((noinline)) void runRing(int tid, int nthreads, struct
     // Coverity reports that the callee treats &ring->next as an array.  However, due to the use of
     // FanSymmetric<1>, only the first element is ever accessed, so it's fine.
     // coverity[callee_ptr_arith:FALSE]
-  Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0, false, RCCLMetadata, Pipeline, USE_ACC> prims(
-    tid, nthreads, &ring->prev, &ring->next, work->sendbuff, work->recvbuff, work->redOpArg, 0, work->connIndex,
-    work->connIndex, work);
+  Primitives<T, RedOp, FanSymmetric<1>, /*Direct=*/1, Proto, 0, false, RCCLMetadata, Pipeline, USE_ACC, UserRegMode>
+    prims(tid, nthreads, &ring->prev, &ring->next, work->sendbuff, work->recvbuff, work->redOpArg, 0, work->connIndex,
+          work->connIndex, work);
 
   for (ssize_t elemOffset = 0; elemOffset < channelCount; elemOffset += loopCount) {
     ssize_t remCount = channelCount - elemOffset;
@@ -193,8 +197,9 @@ __device__ __attribute__((noinline)) void runTreeSplit(int tid, int nthreads, st
 
   if (tree->up == -1) {
       // Reduce and broadcast. Max number of recv is 2, max number of send is 2
-    Primitives<T, RedOp, FanSymmetric<NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto, 0, false, 0, Pipeline, USE_ACC> prims(
-      tid, nthreads, tree->down, tree->down, work->sendbuff, work->recvbuff, work->redOpArg, 0, 0, 0, work);
+    Primitives<T, RedOp, FanSymmetric<NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto, 0, false, 0, Pipeline, USE_ACC,
+               UserRegMode>
+      prims(tid, nthreads, tree->down, tree->down, work->sendbuff, work->recvbuff, work->redOpArg, 0, 0, 0, work);
 
     for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
       offset = gridOffset + elemOffset;
@@ -214,7 +219,8 @@ __device__ __attribute__((noinline)) void runTreeSplit(int tid, int nthreads, st
       // Coverity reports that the callee treats &tree->up as an array.  However, due to the use of
       // FanAsymmetric<n, 1>, only the first element is ever accessed, so it's fine.
       // coverity[callee_ptr_arith:FALSE]
-    Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/1, Proto, 0, false, 0, Pipeline, USE_ACC>
+    Primitives<T, RedOp, FanAsymmetric<NCCL_MAX_DEV_ARITY, 1>, /*Direct=*/1, Proto, 0, false, 0, Pipeline, USE_ACC,
+               UserRegMode>
       prims(tid, nthreadsSplit, tree->down, &tree->up, work->sendbuff, work->recvbuff, work->redOpArg,
             0 * Proto::MaxGroupWidth, 0, 0, work);
 
@@ -237,7 +243,8 @@ __device__ __attribute__((noinline)) void runTreeSplit(int tid, int nthreads, st
       // Coverity reports that the callee treats &tree->up as an array.  However, due to the use of
       // FanAsymmetric<1, n>, only the first element is ever accessed, so it's fine.
       // coverity[callee_ptr_arith:FALSE]
-    Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto, 0, false, 0, Pipeline, USE_ACC>
+    Primitives<T, RedOp, FanAsymmetric<1, NCCL_MAX_DEV_ARITY>, /*Direct=*/1, Proto, 0, false, 0, Pipeline, USE_ACC,
+               UserRegMode>
       prims(tid - nthreadsSplit, nthreads - nthreadsSplit, &tree->up, tree->down, work->sendbuff, work->recvbuff,
             work->redOpArg, 1 * Proto::MaxGroupWidth, 0, 0, work);
 
@@ -841,13 +848,17 @@ struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL> {
 template <typename T, typename RedOp>
 struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
   __device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
-    runRing<T, RedOp, ProtoLL128, RCCL_METADATA_EMPTY>(tid, nthreads, work);
+    // LL128 is generated as separate registered/non-registered kernels; the
+    // compile-time UserRegMode is forwarded by cmake/scripts/add_unroll.sh into
+    // runRing<..., USE_ACC, COLL_UNROLL, /*Pipeline=*/0, UserRegMode>.
+    runARRingLL128<T, RedOp>(tid, nthreads, work);
   }
 };
 
 template <typename T, typename RedOp>
 struct RunWorkColl<ncclFuncAllReduce, T, RedOp, NCCL_ALGO_TREE, NCCL_PROTO_LL128> {
   __device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
-    runTreeSplit<T, RedOp, ProtoLL128>(tid, nthreads, work);
+    // See the LL128 ring note above; add_unroll.sh forwards UserRegMode.
+    runARTreeLL128<T, RedOp>(tid, nthreads, work);
   }
 };

projects/rccl/src/device/broadcast.h

@@ -48,9 +48,9 @@ __device__ __attribute__((noinline)) void runRing(int tid, int nthreads, struct
       // Coverity reports that the callee treats &ring->next as an array.  However, due to the use of
       // FanSymmetric<1>, only the first element is ever accessed, so it's fine.
       // coverity[callee_ptr_arith:FALSE]
-    Primitives<T, RedOp, FanSymmetric<1>, 1, Proto, 0> prims(tid, workNthreads, &ring->prev, &ring->next, inputBuf,
-                                                             outputBuf, work->redOpArg, 0, work->connIndex,
-                                                             work->connIndex, work);
+    Primitives<T, RedOp, FanSymmetric<1>, 1, Proto, 0, false, RCCL_METADATA_EMPTY, 0, 0, UserRegMode>
+      prims(tid, workNthreads, &ring->prev, &ring->next, inputBuf, outputBuf, work->redOpArg, 0, work->connIndex,
+            work->connIndex, work);
 
     for (size_t elemOffset = 0; elemOffset < channelCount; elemOffset += chunkCount) {
       offset = gridOffset + elemOffset;
@@ -99,6 +99,8 @@ struct RunWorkColl<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL> {
 template <typename T, typename RedOp>
 struct RunWorkColl<ncclFuncBroadcast, T, RedOp, NCCL_ALGO_RING, NCCL_PROTO_LL128> {
   __device__ __forceinline__ void run(int tid, int nthreads, struct ncclDevWorkColl* work) {
-    runRing<T, RedOp, ProtoLL128>(tid, nthreads, work);
+    // LL128 is generated as separate registered/non-registered kernels; the
+    // compile-time UserRegMode is forwarded by cmake/scripts/add_unroll.sh.
+    runBcastRingLL128<T, RedOp>(tid, nthreads, work);
   }
 };

projects/rccl/src/device/common.h

@@ -333,14 +333,16 @@ __device__ __forceinline__ unsigned long long int globaltimer() {
 #endif
 }
 
-template <ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline>
+template <ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline,
+          int UserRegMode = 0>
 struct RunWorkColl {
   __device__ void run(int tid, int tn, struct ncclDevWorkColl* work) {
     // Put NOT IMPLEMENTED behavior here.
   }
 };
 
-template <ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline>
+template <ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline,
+          int UserRegMode = 0>
 struct RunWorkBatch;
 
 // Specialized for P2p in sendrecv.h
@@ -351,7 +353,8 @@ template <typename T, typename RedOp, int Proto>
 struct RunWorkBatch<ncclFuncAllGatherV, T, RedOp, NCCL_ALGO_RING, Proto>;
 
 // Specialized here for non-P2p (Coll and CollReg)
-template <ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline>
+template <ncclFunc_t Fn, typename T, typename RedOp, int Algo, int Proto, int USE_ACC, int COLL_UNROLL, int Pipeline,
+          int UserRegMode>
 struct RunWorkBatch {
   // This __forceinline__ is necessary. The compiler was inserting a function call
   // here from the LL ncclKernel.
@@ -650,14 +653,14 @@ __global__ void ncclDevKernel_Generic_32(ncclDevKernelArgsDefaultStorage NCCL_GR
 
 // noinline iff RCCL_DEVICE_LINKER (each devfunc is a standalone shard).
 #ifdef RCCL_DEVICE_LINKER
-#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, acc, pipeline, unroll) \
+#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, acc, pipeline, unroll, userreg) \
   __device__ __attribute__((noinline)) void ncclDevFunc_##suffix() { \
-    RunWorkBatch<coll, ty, redop<ty>, algo, proto, acc, unroll, pipeline>().run(); \
+    RunWorkBatch<coll, ty, redop<ty>, algo, proto, acc, unroll, pipeline, userreg>().run(); \
   }
 #else
-#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, acc, pipeline, unroll) \
+#define DEFINE_ncclDevFunc(suffix, coll, redop, ty, algo, proto, acc, pipeline, unroll, userreg) \
   __device__ void ncclDevFunc_##suffix() { \
-    RunWorkBatch<coll, ty, redop<ty>, algo, proto, acc, unroll, pipeline>().run(); \
+    RunWorkBatch<coll, ty, redop<ty>, algo, proto, acc, unroll, pipeline, userreg>().run(); \
   }
 #endif