From 8ae1bce3bbcfaaa8d4aa743b75469d61f36c79e7 Mon Sep 17 00:00:00 2001
From: Gilbert Lee <gilbert.lee@amd.com>
Date: Thu, 10 Oct 2019 22:39:45 -0700
Subject: [PATCH] Fix for GenericOp device primitive bug

---
 src/collectives/device/primitives.h | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/collectives/device/primitives.h b/src/collectives/device/primitives.h
index 81a4d4cb7f..8c63a6a4f1 100644
--- a/src/collectives/device/primitives.h
+++ b/src/collectives/device/primitives.h
@@ -193,10 +193,11 @@ class ncclPrimitives {
       if (SEND) __threadfence_system();
       if (tid == 0) FOR_SEND(postSend);
       if (tid == 0) 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) {