Add Pivot AllToAll algorithm for Rome model (#503)

* add a2a pivot interface

* remove debug info

* address comments

* fix bug

* remove custom script

* address comments

* fix bug

src/include/collectives.h

@@ -20,6 +20,7 @@ struct ncclDevRedOpFull {
 };
 
 #define FUNC_INDEX_P2P (ncclNumTypes+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumDevRedOps)
+#define FUNC_INDEX_ALLTOALL_PIVOT (FUNC_INDEX_P2P+1)
 #define FUNC_INDEX(func, devredop, ncclType, al, pr) ((((((func)*ncclNumDevRedOps + (devredop))*ncclNumTypes) + (ncclType))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr))
 
 #define NCCL_FUNC_NAME(func, algo, proto, devredop, type) \
@@ -93,6 +94,7 @@ DECL2(AllGather, Sum, /*undefForFloat=*/0)
 DECL(ReduceScatter)
 DECL(AllReduce)
 DECL5(SendRecv, RING, SIMPLE, Sum, int8_t)
+DECL5(AllToAllPivot, RING, SIMPLE, Sum, int8_t)
 
 extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, int8_t)(struct ncclWorkElem* args);
 extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, uint8_t)(struct ncclWorkElem* args);
@@ -126,5 +128,7 @@ extern __device__ void NCCL_ONERANK_REDUCE_NAME(PreMulSum, double)(struct ncclWo
 #define REDUCE_CHUNKSTEPS 2
 #define SENDRECV_SLICEFACTOR 1
 #define NCCL_MAX_SLICE_PER_CHUNK 2  // max value for CHUNKSTEPS/SLICESTEPS, must accord with above
+#define ALLTOALL_PIVOT_SLICESTEPS 2
+#define ALLTOALL_PIVOT_CHUNKSTEPS 4
 
 #endif

src/include/devcomm.h

@@ -26,9 +26,9 @@
 #endif
 
 
-#define NCCL_NUM_FUNCTIONS 5 // SendRecv not included for now
-typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclNumFuncs} ncclFunc_t;
-extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+1];
+#define NCCL_NUM_FUNCTIONS 5 // SendRecv and AllToAllPivot not included for now
+typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclFuncAllToAllPivot, ncclNumFuncs} ncclFunc_t;
+extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+2];
 
 #define NCCL_NUM_ALGORITHMS 3 // Tree/Ring/CollNet
 #define NCCL_ALGO_TREE 0
@@ -202,7 +202,12 @@ struct ncclWorkElem {
   union {
     struct {
       size_t count;
-      size_t lastChunkSize;
+      union {
+        size_t lastChunkSize;
+        // Pivot A2A kernel computes chunk size itself.
+        // Instead, it needs the number of bidirectional rings.
+        size_t pivotA2ANumBiRings;
+      };
       uint64_t redOpArg;
       uint16_t root;
       uint8_t bid;