Merge pull request #137 from gilbertlee-amd/GenericOpFix

Fix for GenericOp device primitive bug

[ROCm/rccl commit: 60279867b3]

projects/rccl/src/collectives/collectives.h

@@ -57,18 +57,12 @@
 DECL_ALL_COLLS
 
 // CHUNKSIZE must be a multiple of SLICESIZE
-//#define ALLREDUCE_SLICESTEPS (NCCL_STEPS/4)
-//#define ALLREDUCE_CHUNKSTEPS (NCCL_STEPS/2)
-//#define ALLGATHER_SLICESTEPS (NCCL_STEPS/4)
-//#define ALLGATHER_CHUNKSTEPS (NCCL_STEPS/2)
-//#define REDUCESCATTER_SLICESTEPS (NCCL_STEPS/4)
-//#define REDUCESCATTER_CHUNKSTEPS (NCCL_STEPS/2)
-#define ALLREDUCE_SLICESTEPS 4
-#define ALLREDUCE_CHUNKSTEPS 4
-#define ALLGATHER_SLICESTEPS 4
-#define ALLGATHER_CHUNKSTEPS 4
-#define REDUCESCATTER_SLICESTEPS 4
-#define REDUCESCATTER_CHUNKSTEPS 4
+#define ALLREDUCE_SLICESTEPS (NCCL_STEPS/4)
+#define ALLREDUCE_CHUNKSTEPS (NCCL_STEPS/2)
+#define ALLGATHER_SLICESTEPS (NCCL_STEPS/4)
+#define ALLGATHER_CHUNKSTEPS (NCCL_STEPS/2)
+#define REDUCESCATTER_SLICESTEPS (NCCL_STEPS/4)
+#define REDUCESCATTER_CHUNKSTEPS (NCCL_STEPS/2)
 #define BROADCAST_SLICESTEPS 1
 #define BROADCAST_CHUNKSTEPS 1
 #define REDUCE_SLICESTEPS 1

projects/rccl/src/collectives/device/primitives.h

@@ -189,14 +189,18 @@ class ncclPrimitives {
         }
       }
       exitIfAbortBarrier(abort, abortCount);
-      if (tid == 0) FOR_SEND(postSendSize, realSize*sizeof(T));
-      if (SEND) __threadfence_system();
-      if (tid == 0) FOR_SEND(postSend);
-      if (tid == 0) FOR_RECV(postRecv);
+      if (tid == 0)
+      {
+        FOR_SEND(postSendSize, realSize*sizeof(T));
+        __threadfence_system();
+        FOR_SEND(postSend);
+        FOR_RECV(postRecv);
+      }
+
+      for (int i=0; i<RECV*NRECV+SRC; i++) srcs[i] += sliceSize;
+      for (int i=0; i<SEND*NSEND+DST; i++) dsts[i] += sliceSize;
+      offset += sliceSize;
     }
-    for (int i=0; i<RECV*NRECV+SRC; i++) srcs[i] += sliceSize;
-    for (int i=0; i<SEND*NSEND+DST; i++) dsts[i] += sliceSize;
-    offset += sliceSize;
   }
 
   __device__ void loadRecvConn(struct ncclConnInfo* conn, int i, T* directBuff) {