Clique kernel support (#295) (#15)

* Adding experimental clique-based kernels (opt-in only)

Co-authored-by: Stanley Tsang <stanley.tsang@amd.com>
Co-authored-by: Gilbert Lee <gilbert.lee@amd.com>
Co-authored-by: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>

Co-authored-by: Stanley Tsang <stanley.tsang@amd.com>
Co-authored-by: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>

[ROCm/rccl commit: 3e62ceddc5]

projects/rccl/src/include/bootstrap.h

@@ -12,7 +12,7 @@
 ncclResult_t bootstrapNetInit();
 ncclResult_t bootstrapCreateRoot(ncclUniqueId* commId, bool idFromEnv);
 ncclResult_t bootstrapGetUniqueId(ncclUniqueId* out);
-ncclResult_t bootstrapInit(ncclUniqueId* id, int rank, int nranks, void** commState);
+ncclResult_t bootstrapInit(ncclUniqueId* id, int rank, int nranks, void** commState, int* rootPid); // [RCCL] Adding rootPid
 ncclResult_t bootstrapAllGather(void* commState, void* allData, int size);
 ncclResult_t bootstrapSend(void* commState, int peer, void* data, int size);
 ncclResult_t bootstrapRecv(void* commState, int peer, void* data, int size);

projects/rccl/src/include/comm.h

@@ -10,6 +10,9 @@
 
 #include "transport.h"
 #include "p2p.h"
+// [RCCL]
+#include "clique/CliqueManager.h"
+// [/RCCL]
 
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
 #else
@@ -156,8 +159,12 @@ struct ncclComm {
   int p2pSendCount;
   int p2pRecvCount;
 
-  // RCCL AllToAll/Scatter/Gather API
-  bool alltoallDisable;
+  // [RCCL]
+  bool alltoallDisable;            // RCCL AllToAll/Scatter/Gather API
+  CliqueManager* cliqueManager;    // CliqueManager handles pointer collection / distribution for clique-based kernels
+  int rootPid;                     // Process ID of root
+  // [/RCCL]
+
 };
 
 #endif

projects/rccl/src/include/devcomm.h

@@ -1,6 +1,6 @@
 /*************************************************************************
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2020 Advanced Micro Devices, Inc. All rights reserved.
+ * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -12,6 +12,9 @@
 #include "rccl_bfloat16.h"
 #include "align.h"
 #include <stdint.h>
+// [RCCL] Support for clique-based kernels
+#include "clique/CliqueCommon.h"
+// [/RCCL]
 
 // Convert volatile access to atomic
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
@@ -22,6 +25,7 @@
 #define STORE(DST, SRC) *(DST) = (SRC)
 #endif
 
+
 #define NCCL_NUM_FUNCTIONS 5 // SendRecv not included for now
 typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclFuncAllToAll, ncclFuncAllToAllv } ncclFunc_t;
 extern const char* ncclFuncStr[];
@@ -35,6 +39,7 @@ extern const char* ncclAlgoStr[NCCL_NUM_ALGORITHMS];
 #define NCCL_NUM_PROTOCOLS 3 // Simple/LL/LL128
 #define NCCL_PROTO_LL 0
 #define NCCL_PROTO_LL128 1
+#define NCCL_PROTO_CLIQUE 1  // [RCCL] Clique takes up same protocol as unused LL128
 #define NCCL_PROTO_SIMPLE 2
 extern const char* ncclProtoStr[NCCL_NUM_PROTOCOLS];
 
@@ -188,6 +193,17 @@ struct ncclWorkElem {
       uint8_t bid;
       uint8_t nChannels;
     } a2av;
+    // [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;
+      uint8_t bid;
+      uint8_t nChannels;
+    } clique;
+    // [/RCCL]
     uint64_t align[3];
   };
 };