Merge remote-tracking branch 'nccl/master' into develop

[ROCm/rccl commit: 3a919c1f49]

projects/rccl/src/include/devcomm.h

@@ -96,8 +96,11 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
 #define NCCL_LL128_SHMEM_ELEMS_PER_THREAD 2
 #define NCCL_LL128_SHMEM_SIZE (NCCL_LL128_SHMEM_ELEMS_PER_THREAD*NCCL_LL128_MAX_NTHREADS)
 
-#define NCCL_DIRECT_GPU 0x01
-#define NCCL_DIRECT_NIC 0x10
+#define NCCL_DIRECT_WRITE 0x01
+#define NCCL_DIRECT_READ  0x02
+#define NCCL_DIRECT_NIC   0x04
+#define NCCL_IPC_WRITE    0x08
+#define NCCL_IPC_READ     0x10
 
 struct ncclConnInfo {
   // Regular comm mechanism
@@ -108,6 +111,7 @@ struct ncclConnInfo {
   int direct;         // Direct communication
   int shared;         // Buffers are shared
   void **ptrExchange; // Pointer exchange for direct communication
+  uint64_t* redOpArgExchange; // PreOp scaler exchange for direct pull case
 
   int *sizesFifo;     // Sizes fifo from GPU to proxy
   void* *ptrsFifo;      // Buffer fifo from proxy to GPU
@@ -175,7 +179,7 @@ struct ncclPeer {
 struct ncclDevComm;
 
 #pragma pack(push)  /* push current alignment to stack */
-#pragma pack(8)     /* set alignment to 4 bytes boundary */
+#pragma pack(8)     /* set alignment to 8 bytes boundary */
 #define NCCL_MAX_WORK_ELEMENTS 1
 #define NCCL_MAX_GROUPS (NCCL_MAX_NTHREADS/WARP_SIZE)
 
@@ -187,8 +191,9 @@ struct ncclWorkElem {
   struct ncclDevComm* comm;
   uint16_t nThreads;
   uint16_t funcIndex;
-  uint16_t index;
-  uint16_t active;
+  uint8_t regUsed;
+  uint8_t direct;
+  uint8_t active, redOpArgIsPtr;
 
   const void * sendbuff;
   void * recvbuff;
@@ -198,10 +203,12 @@ struct ncclWorkElem {
     struct {
       size_t count;
       size_t lastChunkSize;
-      uint32_t root;
+      uint64_t redOpArg;
+      uint16_t root;
       uint8_t bid;
       uint8_t nChannels;
-      uint8_t connIndex;
+      uint16_t connIndex;
+      uint16_t opCount;
     } coll;
     struct {
       size_t sendCount;
@@ -217,18 +224,16 @@ struct ncclWorkElem {
         };
         uint16_t padding;
       };
-    } p2p;
-    struct {
-      uint16_t padding[15];
       uint16_t opCount;
-    } op;
+    } p2p;
     // [RCCL] Clique-based arguments
     //        NOTE: Follows same field structure as coll
     //              because nChannels is accessed from "coll" struct.
     struct {
       size_t count;
       cliqueDevicePtrs_t* ptrs;
-      uint32_t unused;
+      uint64_t unused_1;
+      uint16_t unused_2;
       uint8_t bid;
       uint8_t nChannels;
     } clique;
@@ -236,11 +241,24 @@ struct ncclWorkElem {
     uint64_t align[4];
   };
 };
-struct ncclWork {
-  struct ncclWorkElem elems[NCCL_MAX_WORK_ELEMENTS];
-};
 static_assert(sizeof(struct ncclWorkElem) == (0x10*sizeof(int)), "ncclWorkElem must have a pow2 size");
 
+struct ncclWorkRegElem {
+  struct ncclWorkElem elem;
+  void* dnInputs[NCCL_MAX_DIRECT_ARITY+1];
+  void* dnOutputs[NCCL_MAX_DIRECT_ARITY+1];
+  void* upOutputs[NCCL_MAX_DIRECT_ARITY+1];
+};
+#define NCCL_REG_ELEM_FACTOR 4
+static_assert(sizeof(struct ncclWorkRegElem) == (NCCL_REG_ELEM_FACTOR*sizeof(struct ncclWorkElem)), "ncclWorkRegElem size must be pow2 times ncclWorkElem size");
+
+struct ncclWork {
+  union {
+    struct ncclWorkElem elems[NCCL_MAX_WORK_ELEMENTS];
+    struct ncclWorkRegElem regElems[NCCL_MAX_WORK_ELEMENTS/NCCL_REG_ELEM_FACTOR];
+  };
+};
+
 struct ncclChannel {
   union {
     struct {
@@ -330,6 +348,7 @@ struct ncclProf {
 typedef enum {
   ncclCollTraceNotReady,
   ncclCollTraceKernelLaunchType,
+  ncclCollTraceKernelEndType,
   ncclCollTraceCollEndType,
   ncclCollTraceAbortType,
   ncclCollTraceDataType