[CRL] Refactor P2P zerocopy (flagos-ai#452)

Author: MC952-arch

flagcx/core/group.cc

@@ -283,23 +283,15 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
                          peer, comm->rank, op->args.p2pPeerSlotIdx,
                          op->args.p2pPeerOpHash);
 
-              flagcxConnector *recvConn =
-                  &comm->channels[op->channelId].peers[peer]->recv[0];
-              flagcxConnector *peerConns[] = {recvConn};
               int peerRanks[] = {peer};
               uintptr_t regOffset = 0;
               uintptr_t *peerRmtAddr = NULL;
               op->args.regBufFlag = 0;
               FLAGCXCHECK(flagcxP2pRegisterBuffer(
-                  comm, p2p->buff, p2p->bytes, peerConns, peerRanks, 1,
-                  /*isSender=*/false, &op->args.regBufFlag, &regOffset,
-                  &peerRmtAddr, op->args.p2pPeerSlotIdx));
-              if (op->args.regBufFlag) {
-                INFO(FLAGCX_REG,
-                     "flagcxGroup P2P recv reg rank %d <- %d buff %p size %zu "
-                     "offset %zu remote %p",
-                     comm->rank, peer, p2p->buff, p2p->bytes, (size_t)regOffset,
-                     peerRmtAddr ? (void *)(*peerRmtAddr) : NULL);
+                  comm, p2p->buff, p2p->bytes, peerRanks, 1,
+                  &op->args.regBufFlag, &regOffset, &peerRmtAddr));
+              if (op->args.regBufFlag && peerRmtAddr) {
+                op->args.p2pRmtAddr = (void *)peerRmtAddr;
               }
             } else if (op->connection->transport == TRANSPORT_NET) {
               op->args.chunkSize = flagcxNetChunkSize;
@@ -375,16 +367,15 @@ static flagcxResult_t groupLaunch(struct flagcxAsyncJob *job_) {
                   "peerOpHash(%ld)]",
                   peer, comm->rank, op->args.p2pPeerSlotIdx,
                   op->args.p2pPeerOpHash);
-              flagcxConnector *peerConns[] = {
-                  comm->channels[op->channelId].peers[peer]->send};
+              // Send side: register own buffer to peer's proxy for READ mode.
+              // The actual IPC address comes from SHM at proxy time.
               int peerRanks[] = {peer};
               uintptr_t regOffset = 0;
               uintptr_t *peerRmtAddr = NULL;
+              op->args.regBufFlag = 0;
               FLAGCXCHECK(flagcxP2pRegisterBuffer(
-                  comm, p2p->buff, p2p->bytes, peerConns, peerRanks, 1,
-                  /*isSender=*/true, &op->args.regBufFlag, &regOffset,
-                  &peerRmtAddr, op->args.p2pSlotIdx));
-              // peerRmtAddr is fully resolved (rmtRegAddr + peer's userOffset)
+                  comm, p2p->buff, p2p->bytes, peerRanks, 1,
+                  &op->args.regBufFlag, &regOffset, &peerRmtAddr));
               if (op->args.regBufFlag && peerRmtAddr) {
                 op->args.p2pRmtAddr = (void *)peerRmtAddr;
               }

flagcx/core/include/comm.h

@@ -293,6 +293,8 @@ struct flagcxHeteroComm {
   uint64_t intraBarrierGate; // only used if this is intraComm0
 
   struct flagcxProxyState *proxyState;
+  struct flagcxProxyConnector
+      *gproxyConn;      // Array[nRanks], per-peer proxy connector for IPC reg
   int proxyRefCountOld; /* store proxy post-atomic-sub refcount */
   // Whether this communicator uses collNet
   int collNetSupport;

flagcx/core/include/device.h

@@ -120,6 +120,7 @@ struct flagcxProxyConnector {
   int tpRank;
   int tpLocalRank;
   int sameProcess;
+  bool initialized;
   struct flagcxProxyConnection *connection;
   flagcxResult_t (*proxyProgress)(
       struct flagcxProxyState *proxyState,

flagcx/core/include/launch_kernel.h

@@ -46,9 +46,11 @@ struct flagcxHostSemaphore : public flagcxSemaphore {
   std::unordered_map<int, int> stepInfo;    // opId -> sigalId
   std::vector<std::pair<int, int>> signals; // [curStep, nSteps]
   std::vector<flagcxEvent_t> events;
+  bool frozen; // true during execution phase
 
   flagcxHostSemaphore() {
     counter = 0;
+    frozen = false;
     stepInfo.reserve(FLAGCX_OPS_PER_SEMAPHORE);
     signals.reserve(FLAGCX_SIGNALS_PER_SEMAPHORE);
     events.reserve(FLAGCX_SIGNALS_PER_SEMAPHORE);
@@ -65,22 +67,27 @@ struct flagcxHostSemaphore : public flagcxSemaphore {
     return event;
   }
   void signalStart() override {
+    frozen =
+        true; // freeze: no more structural mutations until wait() completes
     for (auto it = stepInfo.begin(); it != stepInfo.end(); ++it) {
       __atomic_store_n(&signals[it->second].first, 0, __ATOMIC_RELEASE);
     }
   }
   void *getSignals() override { return nullptr; }
   void subCounter(int opId = 0) override {
-    assert(stepInfo.find(opId) != stepInfo.end());
-    __atomic_fetch_add(&signals[stepInfo[opId]].first, 1, __ATOMIC_RELEASE);
+    auto it = stepInfo.find(opId);
+    assert(it != stepInfo.end());
+    int idx = it->second;
+    __atomic_fetch_add(&signals[idx].first, 1, __ATOMIC_RELEASE);
     INFO(FLAGCX_PROXY,
          "SubCounter curStep[%d] = %d, nSteps[%d] = %d, counter %d", opId,
-         signals[stepInfo[opId]].first, opId, signals[stepInfo[opId]].second,
-         counter);
+         signals[idx].first, opId, signals[idx].second, counter);
   }
   void addCounter(int opId = 0) override {
-    if (stepInfo.find(opId) != stepInfo.end()) {
-      __atomic_fetch_add(&signals[stepInfo[opId]].second, 1, __ATOMIC_RELEASE);
+    assert(!frozen); // must not mutate during execution phase
+    auto it = stepInfo.find(opId);
+    if (it != stepInfo.end()) {
+      __atomic_fetch_add(&signals[it->second].second, 1, __ATOMIC_RELEASE);
     } else {
       signals.emplace_back(-1, 1);
       stepInfo[opId] = (int)signals.size() - 1;
@@ -89,8 +96,9 @@ struct flagcxHostSemaphore : public flagcxSemaphore {
   }
   int getCounter() override { return counter; }
   int pollStart(int opId = 0, int step = 0) override {
-    assert(stepInfo.find(opId) != stepInfo.end());
-    return (signals[stepInfo[opId]].first >= step);
+    auto it = stepInfo.find(opId);
+    assert(it != stepInfo.end());
+    return (signals[it->second].first >= step);
   }
   int pollEnd() override {
     return (__atomic_load_n(&counter, __ATOMIC_ACQUIRE) == 0);
@@ -109,6 +117,7 @@ struct flagcxHostSemaphore : public flagcxSemaphore {
       sched_yield();
     }
     __atomic_store_n(&counter, 0, __ATOMIC_RELEASE);
+    frozen = false; // unfreeze: allow addCounter for next round
   }
 };
 

flagcx/core/include/p2p.h

@@ -58,9 +58,13 @@ struct flagcxP2pSyncSlot {
 struct p2pRegInfo {
   int copyDone;    // Indicates if the copy operation is complete
   int copyStarted; // Indicates if the copy operation has started
+  // WRITE mode: recv publishes into sender's slot
+  int ipcRecvRegReady;      // 1 = ipcRecvRmtAddr valid; recv sets
+  uintptr_t ipcRecvRmtAddr; // Recv's buffer mapped in sender's address space
+  // READ mode: sender publishes into receiver's slot
+  int ipcSendRegReady; // 1 = ipcSendRmtAddr valid; sender sets
   uintptr_t
-      ipcUserOffset; // Per-slot IPC offset (recv-side writes, send-side reads)
-  int ipcRegReady;   // 1 = ipcUserOffset is valid for current op
+      ipcSendRmtAddr; // Sender's buffer mapped in receiver's address space
 };
 
 struct flagcxP2pShm {
@@ -149,12 +153,11 @@ flagcxResult_t flagcxP2pImportShareableBuffer(struct flagcxHeteroComm *comm,
                                               struct flagcxP2pIpcDesc *ipcDesc,
                                               void **devMemPtr);
 
-flagcxResult_t
-flagcxP2pRegisterBuffer(struct flagcxHeteroComm *comm, const void *userbuff,
-                        size_t buffSize, struct flagcxConnector **peerConns,
-                        int *peerRanks, int nPeers, bool isSender,
-                        int *regBufFlag, uintptr_t *offsetOut,
-                        uintptr_t **peerRmtAddrsOut, size_t shmRegSlotIdx);
+flagcxResult_t flagcxP2pRegisterBuffer(struct flagcxHeteroComm *comm,
+                                       const void *userbuff, size_t buffSize,
+                                       int *peerRanks, int nPeers,
+                                       int *regBufFlag, uintptr_t *offsetOut,
+                                       uintptr_t **peerRmtAddrsOut);
 
 flagcxResult_t flagcxP2pDeregisterBuffer(struct flagcxHeteroComm *comm,
                                          struct flagcxIpcRegInfo *info);

flagcx/core/include/proxy.h

@@ -144,7 +144,8 @@ struct flagcxProxyArgs {
   uint64_t p2pPeerOpHash = -1;
   size_t p2pSlotIdx = 0;
   size_t p2pPeerSlotIdx = 0;
-  void *p2pRmtAddr = nullptr; // Remote address for P2P zero-copy
+  void *p2pRmtAddr = nullptr; // remote addr for zero-copy P2P (send side reads,
+                              // recv side writes)
 
   union flagcxProxyOpSpecifics specifics;
 };
@@ -336,7 +337,8 @@ struct flagcxProxyState {
   pthread_mutex_t mutex;
   pthread_cond_t cond;
   union flagcxSocketAddress *peerAddresses;
-  struct flagcxSocket peerSock;
+  struct flagcxSocket *peerSocks; // Array[nRanks], indexed by rank
+  int nPeerSocks;                 // Number of allocated peerSocks entries
   struct flagcxProxyOps proxyOps[MAXCHANNELS];
 
   struct flagcxProxyOps *prodProgChannelHead; /*producer*/
@@ -374,6 +376,7 @@ enum proxyConnectState {
 struct flagcxProxyConnection {
   int send, transport, shared;
   int tpLocalRank, sameProcess;
+  int cudaDev;
   struct flagcxSocket *sock;
   struct flagcxTransportComm *tcomm;
   struct flagcxProxyArgs *proxyAppend;

flagcx/core/include/reg_pool.h

@@ -22,6 +22,7 @@ class flagcxRegPool {
   flagcxResult_t addP2pHandle(void *comm, flagcxRegItem *reg, void *handle,
                               struct flagcxProxyConnector *proxyConn);
   flagcxResult_t removeRegItemP2pHandles(void *comm, flagcxRegItem *reg);
+  flagcxResult_t removeAllP2pHandles(void *comm);
   flagcxResult_t registerBuffer(void *comm, void *data, size_t length);
   flagcxResult_t deregisterBuffer(void *comm, void *handle);
   std::map<uintptr_t, std::map<uintptr_t, flagcxRegItem *>> &getGlobalMap();

flagcx/core/include/register.h

@@ -53,6 +53,7 @@ struct flagcxIpcRegInfo {
   struct flagcxProxyConnector *ipcProxyconn;
   struct flagcxIpcImpInfo impInfo;
   bool handleReady;
+  bool sameProcess; // cached at registration time for safe deregister
 };
 
 struct flagcxRegItem {

flagcx/core/init.cc

@@ -12,6 +12,7 @@
 #include "group.h"
 #include "net.h"
 #include "p2p.h"
+#include "reg_pool.h"
 #include "topo.h"
 #include "transport.h"
 #include "type.h"
@@ -294,8 +295,19 @@ static flagcxResult_t flagcxCommInitRankFunc(struct flagcxAsyncJob *job_) {
     int nranks = comm->nRanks;
     for (int i = 0; i < MAXCHANNELS; i++) {
       FLAGCXCHECK(flagcxCalloc(&comm->channels[i].peers, nranks));
-      for (int r = 0; r < nranks; r++)
+      for (int r = 0; r < nranks; r++) {
         FLAGCXCHECK(flagcxCalloc(&comm->channels[i].peers[r], nranks));
+      }
+    }
+    // Set tpRank = comm->rank for all channel connectors so local RPCs
+    // route through peerSocks[myRank] to the local service thread
+    for (int i = 0; i < MAXCHANNELS; i++) {
+      for (int r = 0; r < nranks; r++) {
+        for (int c = 0; c < FLAGCX_MAX_CONNS; c++) {
+          comm->channels[i].peers[r]->send[c].proxyConn.tpRank = comm->rank;
+          comm->channels[i].peers[r]->recv[c].proxyConn.tpRank = comm->rank;
+        }
+      }
     }
     FLAGCXCHECK(flagcxCalloc(&comm->connectSend, nranks));
     FLAGCXCHECK(flagcxCalloc(&comm->connectRecv, nranks));
@@ -339,6 +351,33 @@ static flagcxResult_t flagcxCommInitRankFunc(struct flagcxAsyncJob *job_) {
     INFO(FLAGCX_INIT, "Flagcx RuntimeProxy flag set to %d", runtimeProxy);
     if (!runtimeProxy) {
       FLAGCXCHECK(flagcxProxyInit(comm));
+
+      // Allocate gproxyConn array and populate peerAddresses for peer proxy
+      // connections
+      FLAGCXCHECK(flagcxCalloc(&comm->gproxyConn, comm->nRanks));
+      FLAGCXCHECK(flagcxCalloc(&comm->proxyState->peerAddresses, comm->nRanks));
+      comm->proxyState->peerAddresses[comm->rank] =
+          comm->proxyState->listenSock.addr;
+      FLAGCXCHECK(bootstrapAllGather(comm->bootstrap,
+                                     comm->proxyState->peerAddresses,
+                                     sizeof(union flagcxSocketAddress)));
+
+      // Pre-connect all peer sockets (including self)
+      FLAGCXCHECK(flagcxCalloc(&comm->proxyState->peerSocks, comm->nRanks));
+      comm->proxyState->nPeerSocks = comm->nRanks;
+      for (int i = 0; i < comm->nRanks; i++) {
+        FLAGCXCHECK(flagcxSocketSetFd(-1, &comm->proxyState->peerSocks[i]));
+      }
+      for (int i = 0; i < comm->nRanks; i++) {
+        struct flagcxSocket *sock = &comm->proxyState->peerSocks[i];
+        FLAGCXCHECK(flagcxSocketInit(sock, comm->proxyState->peerAddresses + i,
+                                     comm->magic, flagcxSocketTypeProxy));
+        FLAGCXCHECK(flagcxSocketConnect(sock));
+        int ready = 0;
+        while (!ready) {
+          FLAGCXCHECK(flagcxSocketReady(sock, &ready));
+        }
+      }
     }
   }
 
@@ -461,6 +500,8 @@ flagcxResult_t flagcxHeteroCommUserRank(const flagcxHeteroComm_t comm,
 
 flagcxResult_t flagcxHeteroCommDestroy(flagcxHeteroComm_t comm) {
   FLAGCXCHECK(flagcxHeteroRmaProxyStop(comm));
+  // Clean up P2P IPC handles while proxy is still alive and peerSocks valid
+  FLAGCXCHECK(globalRegPool.removeAllP2pHandles(comm));
   flagcxProxyDestroy(comm);
   for (int i = 0; i < MAXCHANNELS; i++) {
     for (int r = 0; r < comm->nRanks; r++) {
@@ -477,6 +518,19 @@ flagcxResult_t flagcxHeteroCommDestroy(flagcxHeteroComm_t comm) {
 
   free(comm->connectSend);
   free(comm->connectRecv);
+  if (comm->gproxyConn) {
+    // gproxyConn[i].connection is an opaque handle pointing to a
+    // flagcxProxyConnection allocated and owned by the peer's service thread.
+    // Do NOT free it here — the peer frees it when its service thread exits.
+    free(comm->gproxyConn);
+  }
+  free(comm->proxyState->peerAddresses);
+  if (comm->proxyState->peerSocks != NULL) {
+    for (int i = 0; i < comm->proxyState->nPeerSocks; i++) {
+      flagcxSocketClose(&comm->proxyState->peerSocks[i]);
+    }
+    free(comm->proxyState->peerSocks);
+  }
   free(comm->proxyState);
   free(comm->tasks.peers);
   free(comm->tasks.p2pOrder);