From 7cde92deff3db149f738893ed3e2018c60095ac8 Mon Sep 17 00:00:00 2001
From: "Wen-Heng (Jack) Chung" <whchung@gmail.com>
Date: Thu, 22 Sep 2022 13:02:56 -0500
Subject: [PATCH] Tweak unroll factors.

[ROCm/rccl commit: 84054c3b30db9f7b68b9d3a11aa2197540df2482]
---
 projects/rccl/src/collectives/device/prims_ll.h     | 4 ++--
 projects/rccl/src/collectives/device/prims_simple.h | 5 -----
 projects/rccl/tools/rccl-prim-test/copy_kernel.h    | 4 +---
 3 files changed, 3 insertions(+), 10 deletions(-)

diff --git a/projects/rccl/src/collectives/device/prims_ll.h b/projects/rccl/src/collectives/device/prims_ll.h
index 0c96ca9022..388cec161d 100644
--- a/projects/rccl/src/collectives/device/prims_ll.h
+++ b/projects/rccl/src/collectives/device/prims_ll.h
@@ -183,7 +183,7 @@ private:
 
   template<int BeginIx>
   __device__ void readLLBeginAll(int offset, ncclLLFifoLine(&line)[MaxRecv]) {
-    #pragma unroll 1
+    #pragma unroll
     for (int i=BeginIx; i < MaxRecv; i++) {
       if (i < fan.nrecv()) {
         union ncclLLFifoLine* src = recvPtr(i) + offset;
@@ -412,7 +412,7 @@ private:
       }
       if (RECV) {
         data = !SRC ? peerData : MULTI<RedOp,T>()(redOp, peerData, data);
-        #pragma unroll 1
+        #pragma unroll
         for (int i=1; i < MaxRecv && i < fan.nrecv(); i++) {
           peerData = readLLFinish(offset, line, i);
           data = MULTI<RedOp,T>()(redOp, peerData, data);
diff --git a/projects/rccl/src/collectives/device/prims_simple.h b/projects/rccl/src/collectives/device/prims_simple.h
index aafb0ec9b5..b4c11fa9a7 100644
--- a/projects/rccl/src/collectives/device/prims_simple.h
+++ b/projects/rccl/src/collectives/device/prims_simple.h
@@ -200,7 +200,6 @@ private:
       //     barrier();
       //     post();
       //   } // Since we no longer unroll, new branch added here
-      #pragma unroll 1
       do {
         sliceSize = sliceSize < nelem-offset ? sliceSize : nelem-offset;
         if (Src && (flags & (SrcBuf==Input ? RoleInput : RoleOutput)))
@@ -339,7 +338,6 @@ private:
     // slices are all empty. Since empty slices are the uncommon case, and
     // worker perf is the limiter, perf-wise this loop is effectively unentered,
     // hence just a single branch insn.
-    #pragma unroll 1
     while (slice < SlicePerChunk) {
       sliceSize = sliceSize < nelem-offset ? sliceSize : nelem-offset;
       { // Only workers could have Wait roles so we know the slice must be empty
@@ -371,7 +369,6 @@ private:
     int sliceSize = stepSize*StepPerSlice;
     int dataSize = max(DIVUP(peerElem, 16*SlicePerChunk)*16, sliceSize/32);  // per-peer slice size
 
-    #pragma unroll 1
     for (int slice=0; slice<SlicePerChunk; ++slice) {
       int realSize = max(0, min(dataSize, peerElem-offset));
       if (tid < nworkers) {
@@ -383,7 +380,6 @@ private:
           // realSize is not accurate here; but intra-node does not rely on sizes FIFO
           waitPeer<0, DirectSend, 0, 1, 1, 0>(0, inpIx, offset, realSize);
           subBarrier();
-          #pragma unroll 1
           // Loop over peers
           for (int j=0; j<fan.nsend(); j++) {
             int i = (j+shift)%fan.nsend();
@@ -409,7 +405,6 @@ private:
             // Since waitPeer sets srcs[0] to output buffer + offset, we are doing a direct-write based recv
             // Do nothing
           } else {
-            #pragma unroll 1
             for (int j=0; j<fan.nrecv(); j++) {
               int i = (j+shift)%fan.nrecv();
               peerOffset = i*peerElem;
diff --git a/projects/rccl/tools/rccl-prim-test/copy_kernel.h b/projects/rccl/tools/rccl-prim-test/copy_kernel.h
index f14035e8be..7ea328c3eb 100644
--- a/projects/rccl/tools/rccl-prim-test/copy_kernel.h
+++ b/projects/rccl/tools/rccl-prim-test/copy_kernel.h
@@ -172,7 +172,6 @@ __device__ void ReduceCopy128bMulti(const int w, const int nw, const int t,
       for (int u = 0; u < UNROLL; ++u) Fetch128(vals2[u], srcs[i]+u*WARP_SIZE);
       for (int u = 0; u < UNROLL; ++u) MULTI128<FUNC, T>()(vals[u], vals2[u]);
     }
-    #pragma unroll 1
     for (int i=MINSRCS; i<MAXSRCS && i<nsrcs; i++) {
       Pack128 vals2[UNROLL];
       for (int u = 0; u < UNROLL; ++u) Fetch128(vals2[u], srcs[i]+u*WARP_SIZE);
@@ -183,7 +182,6 @@ __device__ void ReduceCopy128bMulti(const int w, const int nw, const int t,
     for (int i = 0; i < MINDSTS; i++) {
       for (int u = 0; u < UNROLL; ++u) Store128(dsts[i]+u*WARP_SIZE, vals[u]);
     }
-    #pragma unroll 1
     for (int i=MINDSTS; i<MAXDSTS; i++) {
       if (i<ndsts) {
         for (int u = 0; u < UNROLL; ++u) Store128(dsts[i]+u*WARP_SIZE, vals[u]);
@@ -264,4 +262,4 @@ __device__ __forceinline__ void ReduceOrCopyMulti(const int tid, const int nthre
   ReduceCopyMulti<FUNC, T, 1, MINSRCS, MAXSRCS, MINDSTS, MAXDSTS>(w, nw, t, nsrcs, srcs, ndsts, dsts, offset, Nrem);
 }
 
-#endif
\ No newline at end of file
+#endif