Add Alltoallv RCCL kernel implementation (#269)

* Add alltoallv API and implementation

* Extend Rome P2P channel limit to multinode and alltoall kernels

* topo_expl: fix compilation and sync up with main

* gtest: use RCCL alltoallv API

* Code review changes

src/include/collectives.h

@@ -8,7 +8,7 @@
 #ifndef NCCL_COLLECTIVES_H_
 #define NCCL_COLLECTIVES_H_
 
-#define FUNC_INDEX_P2P (3+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps)
+#define FUNC_INDEX_P2P (4+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps)
 #define FUNC_INDEX(coll, redop, dtype, al, pr) ((coll >= NCCL_NUM_FUNCTIONS) \
   ? (coll-NCCL_NUM_FUNCTIONS+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps) \
   : ((((((coll)*ncclNumOps + (redop))*ncclNumTypes) + (dtype))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr)))
@@ -61,6 +61,7 @@
   DECL_COLL5(ncclGather, copy, i8) \
   DECL_COLL5(ncclScatter, copy, i8) \
   DECL_COLL5(ncclAllToAll, copy, i8) \
+  DECL_COLL5(ncclAllToAllv, copy, i8) \
   DECL_COLL5(ncclSendRecv, copy, i8) \
 
 DECL_ALL_COLLS
@@ -89,5 +90,7 @@ DECL_ALL_COLLS
 #define SCATTER_CHUNKSTEPS 4
 #define ALLTOALL_SLICESTEPS 4
 #define ALLTOALL_CHUNKSTEPS 4
+#define ALLTOALLV_SLICESTEPS 4
+#define ALLTOALLV_CHUNKSTEPS 4
 
 #endif

src/include/devcomm.h

@@ -23,8 +23,8 @@
 #endif
 
 #define NCCL_NUM_FUNCTIONS 5 // SendRecv not included for now
-typedef enum { ncclCollBroadcast, ncclCollReduce, ncclCollAllGather, ncclCollReduceScatter, ncclCollAllReduce, ncclCollGather, ncclCollScatter, ncclCollAllToAll, ncclCollSendRecv} ncclFunc_t;
-extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+3];
+typedef enum { ncclCollBroadcast, ncclCollReduce, ncclCollAllGather, ncclCollReduceScatter, ncclCollAllReduce, ncclCollGather, ncclCollScatter, ncclCollAllToAll, ncclCollAllToAllv, ncclCollSendRecv} ncclFunc_t;
+extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+4];
 
 #define NCCL_NUM_ALGORITHMS 3 // Tree/Ring/CollNet
 #define NCCL_ALGO_TREE 0
@@ -183,6 +183,13 @@ struct CollectiveArgs {
       size_t sendCount;
       size_t recvCount;
     } p2p;
+    struct {
+      uint16_t nThreads;
+      uint8_t bid;
+      uint8_t nChannels;
+      size_t count;
+      size_t* extra;
+    } a2av;
   };
 };
 struct ncclColl {
@@ -215,6 +222,7 @@ struct ncclChannel {
 
       // Operation list for aggregation
       struct ncclColl* collectives;
+      size_t* collectivesExtra;
       int collStart;
       int collCount;
       int collFifoHead; // Only used by GPU

src/include/info.h

@@ -40,6 +40,11 @@ struct ncclInfo {
   // Algorithm details
   int chunkSteps;
   int sliceSteps;
+  // For alltoallv
+  const size_t *sendcounts;
+  const size_t *sdispls;
+  const size_t *recvcounts;
+  const size_t *rdispls;
   // Computed later
   int algorithm;
   int protocol;