Adding explicit HDP flush when using RDMA via Infiniband (#95)

* Adding explicit HDP flush when using RDMA via Infiniband

src/collectives/device/all_gather.h

@@ -31,7 +31,7 @@ __device__ void ncclAllGatherKernel(struct CollectiveArgs* args) {
   WaitFlag waitDoneFromNext(ring->send.conn.head, ALLGATHER_BUFCHUNKS*ALLGATHER_SUBSTEPS);
   WaitFlag waitReadyFromPrev(ring->recv.conn.tail, ALLGATHER_SUBSTEPS);
   PostFlag postDoneToPrev(ring->recv.conn.head, ALLGATHER_SUBSTEPS, NULL, 0);
-  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, ALLGATHER_BUFCHUNKS*ALLGATHER_SUBSTEPS, ring->hdp_reg);
+  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, ALLGATHER_BUFCHUNKS*ALLGATHER_SUBSTEPS, ring->next_hdp_reg);
 
   typedef Primitives<UNROLL, ALLGATHER_SUBSTEPS, T> Prims;
 

src/collectives/device/all_reduce.h

@@ -31,7 +31,7 @@ __device__ void ncclAllReduceKernel(struct CollectiveArgs* args) {
   WaitFlag waitDoneFromNext(ring->send.conn.head, ALLREDUCE_BUFCHUNKS*ALLREDUCE_SUBSTEPS);
   WaitFlag waitReadyFromPrev(ring->recv.conn.tail, ALLREDUCE_SUBSTEPS);
   PostFlag postDoneToPrev(ring->recv.conn.head, ALLREDUCE_SUBSTEPS, NULL, 0);
-  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, ALLREDUCE_BUFCHUNKS*ALLREDUCE_SUBSTEPS, ring->hdp_reg);
+  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, ALLREDUCE_BUFCHUNKS*ALLREDUCE_SUBSTEPS, ring->next_hdp_reg);
 
   typedef Primitives<UNROLL, ALLREDUCE_SUBSTEPS, T, FUNC> Prims;
 

src/collectives/device/broadcast.h

@@ -30,7 +30,7 @@ __device__ void ncclBroadcastKernel(struct CollectiveArgs* args) {
   WaitFlag waitDoneFromNext(ring->send.conn.head, (BROADCAST_BUFCHUNKS-1)*BROADCAST_SUBSTEPS);
   WaitFlag waitReadyFromPrev(ring->recv.conn.tail, 0);
   PostFlag postDoneToPrev(ring->recv.conn.head, 0, NULL, 0);
-  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, BROADCAST_BUFCHUNKS*BROADCAST_SUBSTEPS, ring->hdp_reg);
+  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, BROADCAST_BUFCHUNKS*BROADCAST_SUBSTEPS, ring->next_hdp_reg);
 
   typedef Primitives<UNROLL, BROADCAST_SUBSTEPS, T> Prims;
 

src/collectives/device/reduce.h

@@ -27,7 +27,7 @@ __device__ void ncclReduceKernel(struct CollectiveArgs* args) {
   WaitFlag waitDoneFromNext(ring->send.conn.head, (REDUCE_BUFCHUNKS-1)*REDUCE_SUBSTEPS);
   WaitFlag waitReadyFromPrev(ring->recv.conn.tail, 0);
   PostFlag postDoneToPrev(ring->recv.conn.head, 0, NULL, 0);
-  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, REDUCE_BUFCHUNKS*REDUCE_SUBSTEPS, ring->hdp_reg);
+  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, REDUCE_BUFCHUNKS*REDUCE_SUBSTEPS, ring->next_hdp_reg);
 
   typedef Primitives<UNROLL, REDUCE_SUBSTEPS, T, FUNC> Prims;
 

src/collectives/device/reduce_scatter.h

@@ -28,7 +28,7 @@ __device__ void ncclReduceScatterKernel(struct CollectiveArgs* args) {
   WaitFlag waitDoneFromNext(ring->send.conn.head, REDUCESCATTER_BUFCHUNKS*REDUCESCATTER_SUBSTEPS);
   WaitFlag waitReadyFromPrev(ring->recv.conn.tail, REDUCESCATTER_SUBSTEPS);
   PostFlag postDoneToPrev(ring->recv.conn.head, REDUCESCATTER_SUBSTEPS, NULL, 0);
-  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, REDUCESCATTER_BUFCHUNKS*REDUCESCATTER_SUBSTEPS, ring->hdp_reg);
+  PostFlag postReadyToNext(ring->send.conn.tail, 0, ring->send.conn.fifo, REDUCESCATTER_BUFCHUNKS*REDUCESCATTER_SUBSTEPS, ring->next_hdp_reg);
 
   typedef Primitives<UNROLL, REDUCESCATTER_SUBSTEPS, T, FUNC> Prims;
 

src/include/core.h

@@ -153,11 +153,12 @@ struct ncclRing {
       int* userRanks;
       int* devUserRanks;
 
-      // Next GPU's HDP_MEM_FLUSH_ADDR: HDP Memory Coherency Flush Control. This register
+      // GPU's HDP_MEM_FLUSH_ADDR: HDP Memory Coherency Flush Control. This register
       // allows software to explicitly initiate a flush read to HDP memory. See more
       // descriptions in primitives.h.
-      uint32_t* hdp_reg;
-
+      uint32_t* next_hdp_reg;  // Next GPU in ring (for p2p transport use only)
+      uint32_t* curr_hdp_reg;  // Curr GPU in ring (for rdma transport use only)
+      
       // Operation list for aggregation
       struct ncclColl* collectives;
       struct ncclColl* devCollectives;

src/include/transport.h

@@ -110,4 +110,8 @@ inline void transportProxyIdle(int idle) {
   sched_yield();
 }
 
+// Function to get GPU's HDP_MEM_FLUSH_ADDR: HDP Memory Coherency Flush Control
+// This register allows software to explicitly initiate a flush read to HDP memory
+ncclResult_t getGpuHdpReg(int cudaDev, uint32_t** hdp);
+
 #endif

src/transport.cu

@@ -186,3 +186,56 @@ ncclResult_t transportDestroyProxy(struct ncclConnector* connector) {
   }
   return ncclSuccess;
 }
+
+
+ncclResult_t getGpuHdpReg(int cudaDev, uint32_t** hdp) {
+  auto convert_bdf = [](const char *busId) {
+    char bdf[9];
+    strncpy(bdf, busId, 4);
+    strncpy(bdf+4, busId+5, 2);
+    strncpy(bdf+6, busId+8, 2);
+    bdf[8] = '\0';
+    uint16_t id = (uint16_t)strtol(bdf, NULL, 16);
+    return id;
+  };
+
+  union find_agent_args {
+    hsa_agent_t agent;
+    uint16_t id;
+  } args;
+
+  const auto& find_agent = [](hsa_agent_t agent, void* arg) {
+    uint16_t id = ((union find_agent_args *)arg)->id;
+    hsa_device_type_t type;
+    hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, (void*)&type);
+    if(type == HSA_DEVICE_TYPE_GPU) {
+      uint16_t bdf_id = 1;
+      hsa_agent_get_info(agent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_BDFID, &bdf_id);
+      if(bdf_id == id) {
+        ((union find_agent_args *)arg)->agent=agent;
+        return HSA_STATUS_INFO_BREAK;
+      }
+    }
+    return HSA_STATUS_SUCCESS;
+  };
+
+#define PCI_BUS_ID_BUFFER_SIZE 16
+  char busId[PCI_BUS_ID_BUFFER_SIZE];
+  *hdp = NULL;
+  CUDACHECK(hipDeviceGetPCIBusId(busId, PCI_BUS_ID_BUFFER_SIZE, cudaDev));
+  args.id = convert_bdf(busId);
+  hsa_status_t err = hsa_iterate_agents(find_agent, (void*)&args);
+  if (err != HSA_STATUS_INFO_BREAK) {
+    WARN("failed to get locate HSA agent for GPU %d", cudaDev);
+    return ncclSystemError;
+  }
+  hsa_amd_hdp_flush_t hdpinfo;
+  err = hsa_agent_get_info(args.agent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_HDP_FLUSH, &hdpinfo);
+  if ((err != HSA_STATUS_SUCCESS) && (err != HSA_STATUS_INFO_BREAK)) {
+    WARN("failed to get HSA_AMD_AGENT_INFO_HDP_FLUSH for GPU %d", cudaDev);
+    return ncclSystemError;
+  }
+  *hdp = hdpinfo.HDP_MEM_FLUSH_CNTL;
+  return ncclSuccess;
+#undef PCI_BUS_ID_BUFFER_SIZE
+}

src/transport/net.cu

@@ -290,6 +290,13 @@ ncclResult_t netRecvSetup(ncclTinfo_t* myOpaqueInfo, ncclTinfo_t* peerOpaqueInfo
   resources->netDev = getDev(ring->id, myInfo->ndev, myInfo->distances);
   NCCLCHECK(netGetGdrSupport(resources->netDev, myInfo->distances[resources->netDev], 0, &resources->useGdr));
 
+  if (resources->useGdr) {
+    // Collect HDR register for local GPU to initiate flush after receive
+    int cudaDev;
+    hipGetDevice(&cudaDev);
+    NCCLCHECK(getGpuHdpReg(cudaDev, &ring->curr_hdp_reg));
+  }
+
   int sendSize = sizeof(struct ncclSendMem);
   NCCLCHECK(ncclCudaHostAlloc((void**)&resources->hostSendMem, (void**)&resources->devHostSendMem, sendSize));
 
@@ -533,7 +540,13 @@ ncclResult_t netRecvProxy(struct ncclProxyArgs* args) {
         if (nextBuff) memcpy(nextBuff+slot*sliceSize, localBuff+slot*sliceSize, size);
         head++;
         if (llMode == 0) {
-          if (ptrType == NCCL_PTR_CUDA) ncclNetFlush(resources->netRecvComm, localBuff+slot*sliceSize, size);
+          if (ptrType == NCCL_PTR_CUDA) {
+              ncclNetFlush(resources->netRecvComm, localBuff+slot*sliceSize, size);
+
+              // Flush local HDP register after local read-back finishes
+              STORE(ring->curr_hdp_reg, 0x1);
+              TRACE(NCCL_NET, "Flushing GPU memory via HDP %p", ring->curr_hdp_reg);
+          }
           //TRACE(NCCL_NET,"head %d tail %d slot %d size %d ptrType %d", head, tail, slot, size, ptrType);
           STORE(nextTail, head);
         }

src/transport/p2p.cu

@@ -482,56 +482,6 @@ end:
   } while (0)
 
 /* Send: Create and return connect structures for this peer to connect to me */
-static ncclResult_t getGpuHdpReg(int cudaDev, uint32_t** hdp) {
-  auto convert_bdf = [](const char *busId) {
-    char bdf[9];
-    strncpy(bdf, busId, 4);
-    strncpy(bdf+4, busId+5, 2);
-    strncpy(bdf+6, busId+8, 2);
-    bdf[8] = '\0';
-    uint16_t id = (uint16_t)strtol(bdf, NULL, 16);
-    return id;
-  };
-
-  union find_agent_args {
-    hsa_agent_t agent;
-    uint16_t id;
-  } args;
-
-  const auto& find_agent = [](hsa_agent_t agent, void* arg) {
-    uint16_t id = ((union find_agent_args *)arg)->id;
-    hsa_device_type_t type;
-    hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, (void*)&type);
-    if(type == HSA_DEVICE_TYPE_GPU) {
-      uint16_t bdf_id = 1;
-      hsa_agent_get_info(agent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_BDFID, &bdf_id);
-      if(bdf_id == id) {
-        ((union find_agent_args *)arg)->agent=agent;
-        return HSA_STATUS_INFO_BREAK;
-      }
-    }
-    return HSA_STATUS_SUCCESS;
-  };
-
-  char busId[NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
-  *hdp = NULL;
-  CUDACHECK(hipDeviceGetPCIBusId(busId, NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE, cudaDev));
-  args.id = convert_bdf(busId);
-  hsa_status_t err = hsa_iterate_agents(find_agent, (void*)&args);
-  if (err != HSA_STATUS_INFO_BREAK) {
-    WARN("failed to get locate HSA agent for GPU %d", cudaDev);
-    return ncclSystemError;
-  }
-  hsa_amd_hdp_flush_t hdpinfo;
-  err = hsa_agent_get_info(args.agent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_HDP_FLUSH, &hdpinfo);
-  if ((err != HSA_STATUS_SUCCESS) && (err != HSA_STATUS_INFO_BREAK)) {
-    WARN("failed to get HSA_AMD_AGENT_INFO_HDP_FLUSH for GPU %d", cudaDev);
-    return ncclSystemError;
-  }
-  *hdp = hdpinfo.HDP_MEM_FLUSH_CNTL;
-  return ncclSuccess;
-}
-
 ncclResult_t p2pSendSetup(ncclTinfo_t* myOpaqueInfo, ncclTinfo_t* peerOpaqueInfo, struct ncclConnect* connectInfo, struct ncclRing* ring) {
   struct p2pInfo* myInfo = (struct p2pInfo*)myOpaqueInfo;
   struct p2pInfo* peerInfo = (struct p2pInfo*)peerOpaqueInfo;
@@ -542,8 +492,8 @@ ncclResult_t p2pSendSetup(ncclTinfo_t* myOpaqueInfo, ncclTinfo_t* peerOpaqueInfo
     return ncclInternalError;
   }
   if (linktype != HSA_AMD_LINK_INFO_TYPE_XGMI) {
-    NCCLCHECK(getGpuHdpReg(peerInfo->cudaDev, &ring->hdp_reg));
-    TRACE(NCCL_INIT|NCCL_P2P,"Ring %02d : %d -> %d HDP %p", ring->id, myInfo->rank, peerInfo->rank, ring->hdp_reg);
+    NCCLCHECK(getGpuHdpReg(peerInfo->cudaDev, &ring->next_hdp_reg));
+    TRACE(NCCL_INIT|NCCL_P2P,"Ring %02d : %d -> %d HDP %p", ring->id, myInfo->rank, peerInfo->rank, ring->next_hdp_reg);
   }
   if (myInfo->pidHash == peerInfo->pidHash) {
     info.direct = 1;