Add naive support for kernel-free and non-reduction collective operations (flagos-ai#299)

Author: MC952-arch

flagcx/core/global_comm.h

@@ -41,18 +41,18 @@ struct flagcxComm {
   flagcxInnerComm_t host_comm;
   flagcxInnerComm_t homo_comm;
   flagcxHeteroComm_t hetero_comm;
+  flagcxInnerComm_t homoInterComm;
+  flagcxInnerComm_t tunerInnerComm; // innerComm selected by tuner
   // experimental for multi-nic support
   int homoInterRootRank;
   int homoInterMyRank;
   int homoInterRanks;
   std::vector<std::vector<int>> clusterInterRankList;
-  flagcxInnerComm_t homoInterComm;
   std::vector<flagcxVendorType> clusterVendorMap;
   struct flagcxTuner *tuner;
   void *tunerContext;
   std::map<struct flagcxCommTag, flagcxInnerComm_t>
-      homoCommMap;                  // key: commTag returned by tuner
-  flagcxInnerComm_t tunerInnerComm; // innerComm selected by tuner
+      homoCommMap; // key: commTag returned by tuner
 };
 
 #endif // end include guard

flagcx/core/group.cc

@@ -196,13 +196,14 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
                   peer, comm->rank, op->args.p2pPeerSlotIdx,
                   op->args.p2pPeerOpHash);
             }
-
             // launch proxyRegister op if not yet registered
-            flagcxConnector *peerConns[] = {
-                comm->channels[op->channelId].peers[peer]->send};
-            FLAGCXCHECK(flagcxNetRegisterBuffer(
-                comm, p2p->buff, p2p->bytes, peerConns, 1, &op->args.regBufFlag,
-                &op->args.regHandle));
+            if (op->connection->transport == TRANSPORT_NET) {
+              flagcxConnector *peerConns[] = {
+                  comm->channels[op->channelId].peers[peer]->send};
+              FLAGCXCHECK(flagcxNetRegisterBuffer(
+                  comm, p2p->buff, p2p->bytes, peerConns, 1,
+                  &op->args.regBufFlag, &op->args.regHandle));
+            }
             // we don't use semaphore tracking for device func for the moment
             if (deviceAsyncLoad && deviceAsyncStore) {
               FLAGCXCHECK(deviceAdaptor->eventCreate(&op->event,
@@ -227,10 +228,10 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
             } else {
               op->args.semaphore = semaphore;
               op->event = semaphore->getEvent();
+              semaphore->counter++;
               FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
               if (launchStream == nullptr) {
                 launchStream = op->stream;
-                semaphore->counter++;
               } else {
                 FLAGCXCHECK(
                     deviceAdaptor->streamWaitEvent(launchStream, op->event));
@@ -273,13 +274,14 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
                          peer, comm->rank, op->args.p2pPeerSlotIdx,
                          op->args.p2pPeerOpHash);
             }
-
             // launch proxyRegister op if not yet registered
-            flagcxConnector *peerConns[] = {
-                comm->channels[op->channelId].peers[peer]->recv};
-            FLAGCXCHECK(flagcxNetRegisterBuffer(
-                comm, p2p->buff, p2p->bytes, peerConns, 1, &op->args.regBufFlag,
-                &op->args.regHandle));
+            if (op->connection->transport == TRANSPORT_NET) {
+              flagcxConnector *peerConns[] = {
+                  comm->channels[op->channelId].peers[peer]->recv};
+              FLAGCXCHECK(flagcxNetRegisterBuffer(
+                  comm, p2p->buff, p2p->bytes, peerConns, 1,
+                  &op->args.regBufFlag, &op->args.regHandle));
+            }
             // we don't use semaphore tracking for device func for the moment
             if (deviceAsyncLoad && deviceAsyncStore) {
               std::vector<void *> argList;
@@ -304,10 +306,10 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
             } else {
               op->args.semaphore = semaphore;
               op->event = semaphore->getEvent();
+              semaphore->counter++;
               FLAGCXCHECK(deviceAdaptor->eventRecord(op->event, op->stream));
               if (launchStream == nullptr) {
                 launchStream = op->stream;
-                semaphore->counter++;
               } else {
                 FLAGCXCHECK(
                     deviceAdaptor->streamWaitEvent(launchStream, op->event));
@@ -332,24 +334,35 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
               if (sendTasks[i]->bytes == recvTasks[j]->bytes &&
                   sendTasks[i]->dtype == recvTasks[j]->dtype) {
                 if (sendTasks[i]->buff != recvTasks[j]->buff) {
-                  flagcxEvent_t selfEvent = semaphore->getEvent();
-                  FLAGCXCHECK(deviceAdaptor->eventRecord(selfEvent,
-                                                         sendTasks[i]->stream));
+                  flagcxProxyOp *op;
+                  FLAGCXCHECK(flagcxCalloc(&op, 1));
+                  op->pattern = flagcxPatternSend;
+                  op->nbytes = sendTasks[i]->bytes;
+                  op->sendbuff = (uint8_t *)sendTasks[i]->buff;
+                  op->recvbuff = (uint8_t *)recvTasks[j]->buff;
+                  op->channelId = 0;
+                  op->root = peer;
+                  op->connection = comm->channels[op->channelId]
+                                       .peers[peer]
+                                       ->send[0]
+                                       .proxyConn.connection;
+                  op->stream = sendTasks[i]->stream;
+                  op->event = semaphore->getEvent();
+                  op->args.chunkSteps = 1; // single step
+                  op->args.semaphore = semaphore;
+                  semaphore->counter++;
+                  FLAGCXCHECK(
+                      deviceAdaptor->eventRecord(op->event, op->stream));
                   if (launchStream == nullptr) {
-                    launchStream = sendTasks[i]->stream;
+                    launchStream = op->stream;
                   } else {
                     FLAGCXCHECK(deviceAdaptor->streamWaitEvent(launchStream,
-                                                               selfEvent));
+                                                               op->event));
                   }
-
-                  FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
-                      recvTasks[j]->buff, sendTasks[i]->buff,
-                      sendTasks[i]->bytes, flagcxMemcpyDeviceToDevice,
-                      sendTasks[i]->stream, NULL));
+                  FLAGCXCHECK(flagcxProxySaveOp(comm, op));
                 }
                 free(sendTasks[i]);
                 free(recvTasks[j]);
-
                 sendTasks.erase(sendTasks.begin() + i);
                 recvTasks.erase(recvTasks.begin() + j);
                 matched = true;
@@ -377,7 +390,6 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
         }
       }
       tasks->p2pOrderSteps = newOrderSteps;
-
       comm = comm->groupNext;
     } while (comm != nullptr);
   }
@@ -409,8 +421,10 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
       argsQueue.pop();
     }
   } else {
-    FLAGCXCHECK(deviceAdaptor->launchHostFunc(launchStream, cpuAsyncKernel,
-                                              (void *)semaphore.get()));
+    if (launchStream != nullptr) {
+      FLAGCXCHECK(deviceAdaptor->launchHostFunc(launchStream, cpuAsyncKernel,
+                                                (void *)semaphore.get()));
+    }
   }
   // deprecated code path for host func, since the previous
   // hang issue may be walked around by using zero copy

flagcx/core/p2p.cc

@@ -11,7 +11,7 @@ static std::map<int, std::pair<int, int>>
 void setP2pSlotInfo(int rank, int peerRank, size_t size, flagcxDataType_t dtype,
                     int isRecv, int *opHash, size_t *slotIdx) {
   // TODO: try a better hash function to reduce collisions
-  int key = rank * 1000 + int(size >> 12) + dtype * 100 + peerRank;
+  int key = rank * 1000 + int(size >> 12) + dtype * 10 + peerRank * 100;
   int opHashCounter;
   auto it = p2pOpHashMap.find(key);
   if (it != p2pOpHashMap.end()) {
@@ -231,6 +231,51 @@ flagcxResult_t flagcxP2pProxyRecv(struct flagcxP2pResources *resources,
   return flagcxSuccess;
 }
 
+flagcxResult_t flagcxP2pProxySelfCopy(struct flagcxP2pResources *resources,
+                                      void *sendData, void *recvData,
+                                      size_t size,
+                                      struct flagcxProxyArgs *args) {
+  // Make sure data is valid
+  if (!args->semaphore->pollStart())
+    return flagcxSuccess;
+
+  if (args->transmitted < args->chunkSteps) {
+    // Perform single copy step
+    if (args->copied < args->chunkSteps) {
+      FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
+          recvData, sendData, size, flagcxMemcpyDeviceToDevice,
+          resources->proxyInfo.stream, NULL));
+      FLAGCXCHECK(
+          deviceAdaptor->eventRecord(resources->proxyInfo.events[args->copied],
+                                     resources->proxyInfo.stream));
+      args->copied++;
+    }
+
+    // Check for completed copy step
+    if (args->transmitted < args->copied) {
+      flagcxResult_t res = deviceAdaptor->eventQuery(
+          resources->proxyInfo.events[args->transmitted]);
+      if (res == flagcxSuccess) {
+        args->transmitted++;
+      }
+    }
+  } else {
+    if (args->done != 1) {
+      args->semaphore->signalCounter(1);
+      // Deprecated device func handling
+      if (deviceAsyncLoad && deviceAsyncStore) {
+        if (args->deviceFuncRelaxedOrdering == 1) {
+          FLAGCXCHECK(deviceAdaptor->deviceMemcpy(
+              args->dlArgs, (void *)&args->hlArgs, sizeof(bool),
+              flagcxMemcpyHostToDevice, resources->proxyInfo.stream, NULL));
+        }
+      }
+      args->done = 1;
+    }
+  }
+  return flagcxSuccess;
+}
+
 flagcxResult_t flagcxP2pSendProxySetup(struct flagcxProxyConnection *connection,
                                        struct flagcxProxyState *proxyState,
                                        void *reqBuff, int reqSize,
@@ -249,8 +294,8 @@ flagcxResult_t flagcxP2pSendProxySetup(struct flagcxProxyConnection *connection,
 
   // Allocate shared memory and store in resources->proxyInfo
   size_t shmSize = sizeof(struct flagcxP2pShm);
-  INFO(FLAGCX_INIT,
-       "flagcxP2pSendProxySetup: Allocating shared memory size=%zu", shmSize);
+  INFO(FLAGCX_P2P, "flagcxP2pSendProxySetup: Allocating shared memory size=%zu",
+       shmSize);
   FLAGCXCHECK(flagcxShmAllocateShareableBuffer(
       shmSize, &resources->proxyInfo.desc, (void **)&resources->proxyInfo.shm,
       NULL));
@@ -264,12 +309,12 @@ flagcxResult_t flagcxP2pSendProxySetup(struct flagcxProxyConnection *connection,
     resources->proxyInfo.shm->slots[i].peerDone = 1; // 1 = slot is free
   }
 
-  INFO(FLAGCX_INIT, "flagcxP2pSendProxySetup: Copying response, shm=%p",
+  INFO(FLAGCX_P2P, "flagcxP2pSendProxySetup: Copying response, shm=%p",
        resources->proxyInfo.shm);
   memcpy(respBuff, &resources->proxyInfo, sizeof(struct flagcxP2pShmProxyInfo));
   *done = 1;
 
-  INFO(FLAGCX_INIT, "flagcxP2pSendProxySetup: Completed successfully");
+  INFO(FLAGCX_P2P, "flagcxP2pSendProxySetup: Completed successfully");
   return flagcxSuccess;
 }
 
@@ -278,7 +323,7 @@ flagcxResult_t flagcxP2pRecvProxySetup(struct flagcxProxyConnection *connection,
                                        void *reqBuff, int reqSize,
                                        void *respBuff, int respSize,
                                        int *done) {
-  INFO(FLAGCX_INIT,
+  INFO(FLAGCX_P2P,
        "flagcxP2pRecvProxySetup: reqSize=%d respSize=%d expectedReqSize=%zu "
        "expectedRespSize=%zu",
        reqSize, respSize, sizeof(struct flagcxP2pRequest),
@@ -334,7 +379,7 @@ flagcxP2pSendProxyConnect(struct flagcxProxyConnection *connection,
   }
 
   *done = 1;
-  INFO(FLAGCX_INIT, "flagcxP2pSendProxyConnect: Completed, recvFifo=%p",
+  INFO(FLAGCX_P2P, "flagcxP2pSendProxyConnect: Completed, recvFifo=%p",
        resources->proxyInfo.recvFifo);
   return flagcxSuccess;
 }
@@ -361,7 +406,7 @@ flagcxP2pRecvProxyConnect(struct flagcxProxyConnection *connection,
   }
 
   *done = 1;
-  INFO(FLAGCX_INIT, "flagcxP2pRecvProxyConnect: Completed");
+  INFO(FLAGCX_P2P, "flagcxP2pRecvProxyConnect: Completed");
   return flagcxSuccess;
 }
 

flagcx/core/p2p.h

@@ -92,6 +92,11 @@ flagcxResult_t flagcxP2pProxyRecv(struct flagcxP2pResources *resources,
                                   void *data, size_t size,
                                   struct flagcxProxyArgs *args);
 
+flagcxResult_t flagcxP2pProxySelfCopy(struct flagcxP2pResources *resources,
+                                      void *sendData, void *recvData,
+                                      size_t size,
+                                      struct flagcxProxyArgs *args);
+
 flagcxResult_t flagcxP2pSendProxySetup(struct flagcxProxyConnection *connection,
                                        struct flagcxProxyState *proxyState,
                                        void *reqBuff, int reqSize,

flagcx/core/proxy.cc

@@ -122,14 +122,18 @@ flagcxResult_t flagcxProxySaveOp(struct flagcxHeteroComm *comm,
     *justInquire = false;
   switch (op->pattern) {
     case flagcxPatternSend:
-    case flagcxPatternRecv: {
+      // Self-copy will be saved as a send operation
+      if (op->root == comm->rank)
+        op->selfCopy = 1;
+      FLAGCXCHECK(
+          SaveProxy(comm, channel, proxySend, op->root, op, 0, justInquire));
+      break;
+    case flagcxPatternRecv:
       if (op->root == comm->rank)
         return flagcxSuccess;
       FLAGCXCHECK(
-          SaveProxy(comm, channel,
-                    op->pattern == flagcxPatternSend ? proxySend : proxyRecv,
-                    op->root, op, 0, justInquire));
-    } break;
+          SaveProxy(comm, channel, proxyRecv, op->root, op, 0, justInquire));
+      break;
   }
   return flagcxSuccess;
 }
@@ -207,8 +211,13 @@ static flagcxResult_t progressOps(struct flagcxProxyState *proxyState,
             } else if (op->connection->transport == TRANSPORT_P2P) {
               struct flagcxP2pResources *resources =
                   (flagcxP2pResources *)op->connection->transportResources;
-              flagcxP2pProxySend(resources, op->recvbuff, op->nbytes,
-                                 &op->args);
+              if (op->selfCopy == 0) {
+                flagcxP2pProxySend(resources, op->recvbuff, op->nbytes,
+                                   &op->args);
+              } else {
+                flagcxP2pProxySelfCopy(resources, op->sendbuff, op->recvbuff,
+                                       op->nbytes, &op->args);
+              }
               if (deviceAsyncLoad && deviceAsyncStore) {
                 if (op->args.done == 1 && op->args.eventRecorded) {
                   if (deviceAdaptor->eventQuery(op->event) == flagcxSuccess) {
@@ -407,8 +416,10 @@ static flagcxResult_t expectedProxyResponseStore(struct flagcxProxyState *state,
         return flagcxInternalError;
       }
 
-      memcpy(elem->respBuff, respBuff, respSize);
-      free(respBuff);
+      if (respSize > 0 && respBuff != NULL) {
+        memcpy(elem->respBuff, respBuff, respSize);
+        free(respBuff);
+      }
       elem->done = true;
       elem->res = res;
       return flagcxSuccess;
@@ -731,7 +742,8 @@ static flagcxResult_t proxyProgressAsync(flagcxProxyAsyncOp **opHead,
           resources->netAdaptor->deregMr(resources->netRecvComm, handle));
     }
     done = 1;
-  } else if (op->type == flagcxProxyMsgSetup) {
+  } else if (op->type == flagcxProxyMsgSetup &&
+             op->connection->transport == TRANSPORT_P2P) {
     if (op->connection->send) {
       // P2P Send side setup
       INFO(FLAGCX_PROXY, "Calling flagcxP2pSendProxySetup");
@@ -753,8 +765,11 @@ static flagcxResult_t proxyProgressAsync(flagcxProxyAsyncOp **opHead,
          "proxyProgressAsync opId=%p op.type=%d op.reqBuff=%p op.respSize=%d "
          "done",
          op->opId, op->type, op->reqBuff, op->respSize);
-    if (op->type == flagcxProxyMsgConnect)
-      __atomic_store_n(&op->connection->state, connConnected, __ATOMIC_RELEASE);
+    if (op->connection->transport == TRANSPORT_NET) {
+      if (op->type == flagcxProxyMsgConnect)
+        __atomic_store_n(&op->connection->state, connConnected,
+                         __ATOMIC_RELEASE);
+    }
 
     /* if setup or connect is done, we should not return any error at this point
      * since flagcxSocketSend might already send the respBuff to the requester.

flagcx/core/proxy.h

@@ -96,12 +96,12 @@ struct flagcxProxyArgs {
   size_t chunkSize;
   size_t stepSize;
   void *stepBuff;
-  int waitCopy;
-  int posted;
-  int copied;
-  int postFlush;
-  int flushed;
-  int transmitted;
+  int waitCopy = 0;
+  int posted = 0;
+  int copied = 0;
+  int postFlush = 0;
+  int flushed = 0;
+  int transmitted = 0;
   int sendStepMask;
   size_t totalCopySize;
   size_t totalPostSize;
@@ -186,6 +186,7 @@ struct flagcxProxyOp {
   flagcxProxyArgs args;
   flagcxStream_t stream;
   flagcxEvent_t event; // used to record host/device func
+  int selfCopy = 0;
 };
 
 #define FLAGCX_MAX_NETDEVS 128