[CRL] Shard P2P QPs across per-worker CQs and pin worker threads to NIC-local CPUs (#504)

Author: leoda1

flagcx/adaptor/device/ducuda_adaptor.cc

@@ -301,14 +301,12 @@ flagcxResult_t ducudaAdaptorDmaSupport(bool *dmaBufferSupport) {
   return flagcxSuccess;
 }
 
-flagcxResult_t
-ducudaAdaptorMemGetHandleForAddressRange(void *handleOut, void *buffer,
-                                       size_t size, unsigned long long flags) {
-  //unsupportted on dcu
+flagcxResult_t ducudaAdaptorMemGetHandleForAddressRange(
+    void *handleOut, void *buffer, size_t size, unsigned long long flags) {
+  // unsupportted on dcu
   return flagcxNotSupported;
 }
 
-
 flagcxResult_t ducudaAdaptorGetDeviceProperties(struct flagcxDevProps *props,
                                                 int dev) {
   if (props == NULL) {
@@ -348,7 +346,7 @@ flagcxResult_t ducudaAdaptorGetDeviceByPciBusId(int *dev,
 }
 
 flagcxResult_t ducudaAdaptorStreamWaitValue64(flagcxStream_t stream, void *addr,
-                                            uint64_t value, int flags) {
+                                              uint64_t value, int flags) {
   (void)flags;
   if (stream == NULL || addr == NULL)
     return flagcxInvalidArgument;
@@ -357,8 +355,9 @@ flagcxResult_t ducudaAdaptorStreamWaitValue64(flagcxStream_t stream, void *addr,
                                      CU_STREAM_WAIT_VALUE_GEQ);
   return (err == CUDA_SUCCESS) ? flagcxSuccess : flagcxUnhandledDeviceError;
 }
-flagcxResult_t ducudaAdaptorStreamWriteValue64(flagcxStream_t stream, void *addr,
-                                            uint64_t value, int flags) {
+flagcxResult_t ducudaAdaptorStreamWriteValue64(flagcxStream_t stream,
+                                               void *addr, uint64_t value,
+                                               int flags) {
   (void)flags;
   if (stream == NULL || addr == NULL)
     return flagcxInvalidArgument;
@@ -368,7 +367,7 @@ flagcxResult_t ducudaAdaptorStreamWriteValue64(flagcxStream_t stream, void *addr
   return (err == CUDA_SUCCESS) ? flagcxSuccess : flagcxUnhandledDeviceError;
 }
 flagcxResult_t ducudaAdaptorEventElapsedTime(float *ms, flagcxEvent_t start,
-                                           flagcxEvent_t end) {
+                                             flagcxEvent_t end) {
   if (ms == NULL || start == NULL || end == NULL) {
     return flagcxInvalidArgument;
   }
@@ -394,8 +393,10 @@ flagcxResult_t ducudaAdaptorHostUnregister(void *ptr) {
 
 // Symmetric memory VMM stubs (not supported)
 flagcxResult_t ducudaAdaptorSymPhysAlloc(void *ptr, size_t size,
-                                       void **physHandle, void *shareableHandle,
-                                       size_t *handleSize, size_t *allocSize) {
+                                         void **physHandle,
+                                         void *shareableHandle,
+                                         size_t *handleSize,
+                                         size_t *allocSize) {
   if (ptr == NULL || physHandle == NULL || shareableHandle == NULL ||
       handleSize == NULL || allocSize == NULL)
     return flagcxInvalidArgument;
@@ -425,7 +426,7 @@ flagcxResult_t ducudaAdaptorSymPhysAlloc(void *ptr, size_t size,
   *physHandle = cuHandle;
   return flagcxSuccess;
 }
-flagcxResult_t ducudaAdaptorSymPhysFree(void *physHandle) { 
+flagcxResult_t ducudaAdaptorSymPhysFree(void *physHandle) {
   if (physHandle == NULL)
     return flagcxSuccess;
   CUmemGenericAllocationHandle *cuHandle =
@@ -435,8 +436,8 @@ flagcxResult_t ducudaAdaptorSymPhysFree(void *physHandle) {
   return flagcxSuccess;
 }
 flagcxResult_t ducudaAdaptorSymFlatMap(void *peerHandles[], int nPeers,
-                                     int selfIndex, void *selfPhysHandle,
-                                     size_t allocSize, void **flatBase) {
+                                       int selfIndex, void *selfPhysHandle,
+                                       size_t allocSize, void **flatBase) {
   if (peerHandles == NULL || selfPhysHandle == NULL || flatBase == NULL ||
       nPeers <= 0 || allocSize == 0)
     return flagcxInvalidArgument;
@@ -481,7 +482,7 @@ flagcxResult_t ducudaAdaptorSymFlatMap(void *peerHandles[], int nPeers,
   return flagcxSuccess;
 }
 flagcxResult_t ducudaAdaptorSymFlatUnmap(void *flatBase, size_t allocSize,
-                                       int nPeers) {
+                                         int nPeers) {
   if (flatBase == NULL)
     return flagcxSuccess;
   CUdeviceptr base = (CUdeviceptr)flatBase;
@@ -494,16 +495,16 @@ flagcxResult_t ducudaAdaptorSymMulticastSupported(int *supported) {
   // not supported on dcu
   if (supported == NULL)
     return flagcxInvalidArgument;
-  
+
   if (supported)
     *supported = 0;
   return flagcxSuccess;
 }
 flagcxResult_t ducudaAdaptorSymMulticastCreate(size_t allocSize,
-                                             int nLocalDevices,
-                                             const int *localDeviceOrdinals,
-                                             void **mcHandle,
-                                             int *shareableFd) {
+                                               int nLocalDevices,
+                                               const int *localDeviceOrdinals,
+                                               void **mcHandle,
+                                               int *shareableFd) {
   // not supported on dcu
   if (mcHandle)
     *mcHandle = NULL;
@@ -514,9 +515,9 @@ flagcxResult_t ducudaAdaptorSymMulticastCreate(size_t allocSize,
   return flagcxNotSupported;
 }
 flagcxResult_t ducudaAdaptorSymMulticastBind(void *mcHandle, int importFd,
-                                           void *physHandle, size_t allocSize,
-                                           int localRank, int nLocalDevices,
-                                           void **mcBase, size_t *mcMapSize) {
+                                             void *physHandle, size_t allocSize,
+                                             int localRank, int nLocalDevices,
+                                             void **mcBase, size_t *mcMapSize) {
   // not supported on dcu
   if (mcBase)
     *mcBase = NULL;
@@ -526,7 +527,8 @@ flagcxResult_t ducudaAdaptorSymMulticastBind(void *mcHandle, int importFd,
 
   return flagcxNotSupported;
 }
-flagcxResult_t ducudaAdaptorSymMulticastTeardown(void *mcBase, size_t mcMapSize) {
+flagcxResult_t ducudaAdaptorSymMulticastTeardown(void *mcBase,
+                                                 size_t mcMapSize) {
   // not supported on dcu
   return flagcxSuccess;
 }
@@ -590,9 +592,13 @@ struct flagcxDeviceAdaptor ducudaAdaptor {
                                         // *dev, const char *pciBusId);
       ducudaAdaptorLaunchHostFunc,
       // DMA buffer
-      ducudaAdaptorDmaSupport, // flagcxResult_t (*dmaSupport)(bool *dmaBufferSupport);
-      ducudaAdaptorMemGetHandleForAddressRange, // flagcxResult_t (*memGetHandleForAddressRange)(void *handleOut,
-            // void *buffer, size_t size, unsigned long long flags);
+      ducudaAdaptorDmaSupport, // flagcxResult_t (*dmaSupport)(bool
+                               // *dmaBufferSupport);
+      ducudaAdaptorMemGetHandleForAddressRange, // flagcxResult_t
+                                                // (*memGetHandleForAddressRange)(void
+                                                // *handleOut, void *buffer,
+                                                // size_t size, unsigned long
+                                                // long flags);
       ducudaAdaptorHostRegister,   // flagcxResult_t (*hostRegister)(void *,
                                    // size_t);
       ducudaAdaptorHostUnregister, // flagcxResult_t (*hostUnregister)(void *);

flagcx/adaptor/net/ibrc_p2p_adaptor.cc

@@ -22,7 +22,6 @@
 #include <pthread.h>
 #include <stdint.h>
 #include <string.h>
-#include <string>
 #include <thread>
 #include <unistd.h>
 #include <vector>
@@ -629,7 +628,6 @@ flagcxP2pBuildSingleSliceReq(struct flagcxP2pSendComm *comm, uint64_t localVa,
   req->slice.lkey = lkey;
   req->slice.rkey = rkey;
   req->slice.opcode = opcode;
-  req->slice.peerNicPath = std::string();
   req->slice.task = &req->task;
   req->slice.qpDepth = NULL;
   req->task.sliceList.push_back(&req->slice);
@@ -710,7 +708,6 @@ flagcxP2pIgetBatch(void *sendComm, int count, const uint64_t *srcOffs,
                               dst->lkey,
                               src->rkey,
                               FLAGCX_SLICE_OP_READ,
-                              std::string(),
                               &req->task,
                               NULL};
     req->task.sliceList.push_back(s);
@@ -745,8 +742,11 @@ static inline enum ibv_wr_opcode flagcxSliceOpcodeToVerbs(uint8_t op) {
 }
 
 extern "C" flagcxResult_t flagcxP2pSliceBatch(void *sendComm, struct ibv_qp *qp,
-                                              int count, FlagcxSlice **slices) {
+                                              int count, FlagcxSlice **slices,
+                                              int *failedCount) {
   struct flagcxP2pSendComm *comm = (struct flagcxP2pSendComm *)sendComm;
+  if (failedCount != NULL)
+    *failedCount = 0;
   const char *opLabel = (slices != NULL && count > 0 && slices[0] != NULL &&
                          slices[0]->opcode == FLAGCX_SLICE_OP_READ)
                             ? "READ"
@@ -757,23 +757,50 @@ extern "C" flagcxResult_t flagcxP2pSliceBatch(void *sendComm, struct ibv_qp *qp,
     WARN("NET/IB_P2P : invalid sliceBatch arguments (op=%s, count=%d, qp=%p, "
          "max=%d)",
          opLabel, count, (void *)qp, maxWrPerPost);
+    int failed = 0;
+    if (slices != NULL && count > 0) {
+      for (int i = 0; i < count; i++) {
+        if (slices[i] != NULL) {
+          if (slices[i]->qpDepth != NULL)
+            __sync_fetch_and_sub(slices[i]->qpDepth, 1);
+          slices[i]->markFailed();
+          failed++;
+        }
+      }
+    }
+    if (failedCount != NULL)
+      *failedCount = failed;
     return flagcxInternalError;
   }
 
-  // count can be up to flagcxP2pGlobalConfig().maxWrPerPost (default 256,
-  // bounded at 1024). Heap-allocate to keep the stack small.
-  std::vector<struct ibv_send_wr> wrs(count);
-  std::vector<struct ibv_sge> sges(count);
+  static thread_local std::vector<struct ibv_send_wr> wrScratch;
+  static thread_local std::vector<struct ibv_sge> sgeScratch;
+  if ((int)wrScratch.size() < maxWrPerPost) {
+    wrScratch.resize(maxWrPerPost);
+    sgeScratch.resize(maxWrPerPost);
+  }
+  struct ibv_send_wr *wrs = wrScratch.data();
+  struct ibv_sge *sges = sgeScratch.data();
+  memset(wrs, 0, sizeof(*wrs) * count);
 
   for (int i = 0; i < count; i++) {
     FlagcxSlice *s = slices[i];
     if (s == NULL) {
       WARN("NET/IB_P2P : sliceBatch slice[%d] is NULL", i);
-      for (int k = 0; k < i; k++)
+      for (int k = 0; k < i; k++) {
+        if (slices[k]->qpDepth != NULL)
+          __sync_fetch_and_sub(slices[k]->qpDepth, 1);
         slices[k]->markFailed();
-      for (int k = i; k < count; k++)
-        if (slices[k])
+      }
+      for (int k = i; k < count; k++) {
+        if (slices[k]) {
+          if (slices[k]->qpDepth != NULL)
+            __sync_fetch_and_sub(slices[k]->qpDepth, 1);
           slices[k]->markFailed();
+        }
+      }
+      if (failedCount != NULL)
+        *failedCount = count;
       return flagcxInternalError;
     }
 
@@ -792,11 +819,11 @@ extern "C" flagcxResult_t flagcxP2pSliceBatch(void *sendComm, struct ibv_qp *qp,
   }
 
   struct ibv_send_wr *bad_wr = NULL;
-  flagcxResult_t res = flagcxWrapIbvPostSend(qp, &wrs[0], &bad_wr);
+  flagcxResult_t res = flagcxWrapIbvPostSend(qp, wrs, &bad_wr);
   if (res != flagcxSuccess) {
     int failedFrom = 0;
     if (bad_wr != NULL) {
-      ptrdiff_t off = bad_wr - &wrs[0];
+      ptrdiff_t off = bad_wr - wrs;
       if (off >= 0 && off < count)
         failedFrom = (int)off;
     }
@@ -807,6 +834,8 @@ extern "C" flagcxResult_t flagcxP2pSliceBatch(void *sendComm, struct ibv_qp *qp,
         __sync_fetch_and_sub(slices[k]->qpDepth, 1);
       slices[k]->markFailed();
     }
+    if (failedCount != NULL)
+      *failedCount = count - failedFrom;
     WARN("NET/IB_P2P : sliceBatch ibv_post_send failed (op=%s, count=%d, "
          "failedFrom=%d)",
          opLabel, count, failedFrom);