Port alltoall[v]

src/include/collectives.h

@@ -8,8 +8,10 @@
 #ifndef NCCL_COLLECTIVES_H_
 #define NCCL_COLLECTIVES_H_
 
-#define FUNC_INDEX_P2P 1800
-#define FUNC_INDEX(func, redop, ncclType, al, pr) ((((((func)*ncclNumOps + (redop))*ncclNumTypes) + (ncclType))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr))
+#define FUNC_INDEX_P2P (NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps)
+#define FUNC_INDEX(func, redop, ncclType, al, pr) ((func >= NCCL_NUM_FUNCTIONS) \
+  ? (func-NCCL_NUM_FUNCTIONS+NCCL_NUM_FUNCTIONS*NCCL_NUM_ALGORITHMS*NCCL_NUM_PROTOCOLS*ncclNumTypes*ncclNumOps) \
+  : ((((((func)*ncclNumOps + (redop))*ncclNumTypes) + (ncclType))*NCCL_NUM_ALGORITHMS+(al))*NCCL_NUM_PROTOCOLS+(pr)))
 
 #define NCCL_FUNC_NAME(func, algo, proto, redop, type) \
   ncclFunction_##func##_##algo##_##proto##_##redop##_##type
@@ -60,6 +62,8 @@
   DECL(ReduceScatter) \
   DECL(AllReduce) \
   DECL5(SendRecv, RING, SIMPLE, Sum, int8_t) \
+  DECL5(AllToAll, RING, SIMPLE, Sum, int8_t) \
+  DECL5(AllToAllv, RING, SIMPLE, Sum, int8_t) \
 
 DECL_ALL
 
@@ -81,4 +85,12 @@ DECL_ALL
 #define REDUCE_SLICESTEPS 1
 #define REDUCE_CHUNKSTEPS 1
 #define SENDRECV_SLICEFACTOR 1
+#define GATHER_SLICESTEPS 4
+#define GATHER_CHUNKSTEPS 4
+#define SCATTER_SLICESTEPS 4
+#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 { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv} ncclFunc_t;
-extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS];
+typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclFuncAllToAll, ncclFuncAllToAllv } ncclFunc_t;
+extern const char* ncclFuncStr[];
 
 #define NCCL_NUM_ALGORITHMS 3 // Tree/Ring/CollNet
 #define NCCL_ALGO_TREE 0
@@ -183,6 +183,11 @@ struct ncclWorkElem {
       int32_t delta;
       uint16_t nThreads;
     } p2p;
+    struct {
+      size_t count;
+      uint8_t bid;
+      uint8_t nChannels;
+    } a2av;
     uint64_t align[3];
   };
 };
@@ -208,6 +213,7 @@ struct ncclChannel {
       struct ncclWork* workFifo;
       int workCount;
       uint64_t workFifoTail; // Only used by CPU
+      size_t* a2avParams;
 
 #ifdef ENABLE_PROFILING
       struct timeval tvs;

src/include/info.h

@@ -20,7 +20,8 @@ typedef enum {
   ncclPatternTreeDown,
   ncclPatternTreeUpDown,
   ncclPatternCollTreeUp,
-  ncclPatternCollTreeDown
+  ncclPatternCollTreeDown,
+  ncclPatternAll
 } ncclPattern_t;
 
 // Used to pass NCCL call information between functions
@@ -39,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;