Add bootstrap alltoallv (#172)

Author: liuhe-spec

flagcx/adaptor/bootstrap_adaptor.cc

@@ -151,13 +151,14 @@ flagcxResult_t bootstrapAdaptorAlltoAll(const void *sendbuff, void *recvbuff,
   return flagcxSuccess;
 }
 
-// TODO: unsupported
 flagcxResult_t
 bootstrapAdaptorAlltoAllv(const void *sendbuff, size_t *sendcounts,
                           size_t *sdispls, void *recvbuff, size_t *recvcounts,
                           size_t *rdispls, flagcxDataType_t datatype,
                           flagcxInnerComm_t comm, flagcxStream_t /*stream*/) {
-  return flagcxNotSupported;
+  FLAGCXCHECK(AlltoAllvBootstrap(comm->base, sendbuff, sendcounts, sdispls,
+                                recvbuff, recvcounts, rdispls, datatype));
+  return flagcxSuccess;
 }
 
 #define BOOTSTRAP_SEND_RECV_TAG -6767

flagcx/flagcx.cc

@@ -1363,8 +1363,53 @@ flagcxResult_t flagcxAlltoAllv(const void *sendbuff, size_t *sendcounts,
         comm->homo_comm, stream);
   } else {
     if (use_host_comm()) {
-      // TODO: to be implemented
-      return flagcxNotSupported;
+      uint64_t timers[TIMERS_COLL_COUNT] = {0};
+      timers[TIMER_COLL_TOTAL] = clockNano();
+      void *buff_in;
+      void *buff_out;
+
+      // Calculate max possible size needed for send and receive buffers
+      size_t max_send_size = 0, max_recv_size = 0 , send_size = 0 , recv_size = 0;
+      for (int i = 0; i < comm->nranks; i++) {
+        send_size = (sendcounts[i] + sdispls[i]) * getFlagcxDataTypeSize(datatype);
+        recv_size = (recvcounts[i] + rdispls[i]) * getFlagcxDataTypeSize(datatype);
+        if (send_size > max_send_size) max_send_size = send_size;
+        if (recv_size > max_recv_size) max_recv_size = recv_size;
+      }
+      timers[TIMER_COLL_ALLOC] = clockNano();
+      deviceAdaptor->deviceMalloc(&buff_in, max_send_size, flagcxMemHost, NULL);
+      deviceAdaptor->deviceMalloc(&buff_out, max_recv_size, flagcxMemHost, NULL);
+      timers[TIMER_COLL_ALLOC] = clockNano() - timers[TIMER_COLL_ALLOC];
+
+      timers[TIMER_COLL_MEM_D2H] = clockNano();
+      deviceAdaptor->deviceMemcpy(buff_in, const_cast<void *>(sendbuff), max_send_size,
+                                  flagcxMemcpyDeviceToHost, NULL, NULL);
+      timers[TIMER_COLL_MEM_D2H] = clockNano() - timers[TIMER_COLL_MEM_D2H];
+
+      timers[TIMER_COLL_COMM] = clockNano();
+      cclAdaptors[flagcxCCLAdaptorHost]->alltoAllv(
+          buff_in, sendcounts, sdispls, buff_out, recvcounts, rdispls, datatype,
+          comm->host_comm, NULL);
+      timers[TIMER_COLL_COMM] = clockNano() - timers[TIMER_COLL_COMM];
+
+      timers[TIMER_COLL_MEM_H2D] = clockNano();
+      deviceAdaptor->deviceMemcpy(recvbuff, buff_out, max_recv_size,
+                                  flagcxMemcpyHostToDevice, NULL, NULL);
+      timers[TIMER_COLL_MEM_H2D] = clockNano() - timers[TIMER_COLL_MEM_H2D];
+
+      timers[TIMER_COLL_FREE] = clockNano();
+      deviceAdaptor->deviceFree(buff_in, flagcxMemHost, NULL);
+      deviceAdaptor->deviceFree(buff_out, flagcxMemHost, NULL);
+      timers[TIMER_COLL_FREE] = clockNano() - timers[TIMER_COLL_FREE];
+
+      timers[TIMER_COLL_TOTAL] = clockNano() - timers[TIMER_COLL_TOTAL];
+      INFO(FLAGCX_COLL,
+           "Flagcx timings - %s AlltoAllv: rank %d nranks %d total %.2fms "
+           "(memory alloc %.2fms, memory free %.2fms, memory d2h %.2fms, memory h2d %.2fms, comm %.2fms)",
+           cclAdaptors[flagcxCCLAdaptorHost]->name, comm->rank, comm->nranks,
+           timers[TIMER_COLL_TOTAL] / 1e6, timers[TIMER_COLL_ALLOC] / 1e6,
+           timers[TIMER_COLL_FREE] / 1e6, timers[TIMER_COLL_MEM_D2H] / 1e6,
+           timers[TIMER_COLL_MEM_H2D] / 1e6, timers[TIMER_COLL_COMM] / 1e6);
     } else {
       // Move it into flagcxC2cPlanner workflow
       flagcxC2cPlanner planner;

flagcx/service/bootstrap.cc

@@ -997,3 +997,41 @@ flagcxResult_t bootstrapAbort(void* commState) {
   free(state);
   return flagcxSuccess;
 }
+// AlltoALlv require sendbuff and recvbuff not overlap
+flagcxResult_t AlltoAllvBootstrap(void* commState, const void* sendbuff, size_t* sendcounts, size_t* sdispls,
+    void* recvbuff, size_t* recvcounts, size_t* rdispls, flagcxDataType_t datatype){
+    struct bootstrapState* state = (struct bootstrapState*)commState;
+    int rank = state->rank;
+    int nranks = state->nranks;
+    size_t typeSize = getFlagcxDataTypeSize(datatype);
+
+    for (int i = 0; i < nranks; ++i) {
+        if (i == rank) {
+            memcpy((void *)((char*)recvbuff + rdispls[i] * typeSize),
+                   (void *)((char *)sendbuff + sdispls[i] * typeSize),
+                   sendcounts[i] * typeSize);
+        }
+        const int bootstrapTag = -9995;  // Suggest making this unique if possible
+        if (rank > i) {
+            // Send to rank i
+            FLAGCXCHECK(bootstrapSend(commState, i, bootstrapTag,
+                (void *)((char *)sendbuff + sdispls[i] * typeSize),
+                sendcounts[i] * typeSize));
+            // Recv from rank i
+            FLAGCXCHECK(bootstrapRecv(commState, i, bootstrapTag,
+                (void *)((char *)recvbuff + rdispls[i] * typeSize),
+                recvcounts[i] * typeSize));
+        } else if (rank < i) {
+            // Receive from rank i
+            FLAGCXCHECK(bootstrapRecv(commState, i, bootstrapTag,
+                (void *)((char *)recvbuff + rdispls[i] * typeSize),
+                recvcounts[i] * typeSize));
+            // Send to rank i
+            FLAGCXCHECK(bootstrapSend(commState, i, bootstrapTag,
+                (void *)((char *)sendbuff + sdispls[i] * typeSize),
+                sendcounts[i] * typeSize));
+        }
+    }
+    return flagcxSuccess;
+}
+

flagcx/service/bootstrap.h

@@ -151,6 +151,25 @@ flagcxResult_t ReduceScatterBootstrap(void* commState, const void* sendbuff, voi
  */
 flagcxResult_t AlltoAllBootstrap(void* commState, const void* sendbuff, void* recvbuff, size_t count,
                                  flagcxDataType_t datatype);
+
+/*
+ * All-to-all with variable block sizes
+ *
+ * Every rank sends j-th block of its own sendbuff to the j-th rank of the communicator.
+ * Meanwhile, every rank receives j-th block of its own recvbuff from j-th rank.
+ * 
+ * Each block can have different sizes:
+ * - sendcounts[j] specifies the number of elements to send to rank j
+ * - sdispls[j] specifies the offset in sendbuff for the j-th block
+ * - recvcounts[j] specifies the number of elements to receive from rank j
+ * - rdispls[j] specifies the offset in recvbuff for the j-th block
+ *
+ * In-place operations will happen if sendbuff == recvbuff.
+ */
+flagcxResult_t AlltoAllvBootstrap(void* commState, const void* sendbuff,  size_t* sendcounts,  size_t* sdispls,
+                                 void* recvbuff, size_t* recvcounts,  size_t* rdispls,
+                                 flagcxDataType_t datatype);
+
 #ifdef __cplusplus
 } // end extern "C"
 #endif