Merge remote-tracking branch 'nccl/master' into develop

src/transport/coll_net.cc

@@ -1,6 +1,6 @@
 /*************************************************************************
- * Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -8,11 +8,15 @@
 #include "comm.h"
 #include "coll_net.h"
 #include "graph.h"
+#include "proxy.h"
+#include "gdrwrap.h"
 
-#define COLLNET_GROUP_NSUBS 8
-#define COLLNET_MAX_GROUPS (NCCL_PROXY_MAX_SUBS/COLLNET_GROUP_NSUBS)
+int64_t ncclParamGdrCopySyncEnable();
+int64_t ncclParamGdrCopyFlushEnable();
 
 struct collNetRecvConnectInfo {
+  int rank;
+  int nranks;
   collNetHandle_t collNetHandle;
 };
 
@@ -21,132 +25,287 @@ struct collNetSendConnectInfo {
   void* reqFifo;
 };
 
+#define COLLNET_GROUP_NSUBS 8
+#define COLLNET_MAX_GROUPS (NCCL_PROXY_MAX_SUBS/COLLNET_GROUP_NSUBS)
+
+#define NCCL_NET_MAP_HOSTMEM 0
+#define NCCL_NET_MAP_DEVMEM 1
+#define NCCL_NET_MAP_SHARED_HOSTMEM 2
+#define NCCL_NET_MAP_SHARED_DEVMEM 3
+#define NCCL_NET_MAP_GDCMEM 4
+#define NCCL_NET_MAP_MEMS 5
+
+#define NCCL_NET_MAP_MASK_DEVMEM 0x40000000
+#define NCCL_NET_MAP_MASK_SHARED 0x80000000
+#define NCCL_NET_MAP_MASK_USED   0x20000000
+#define NCCL_NET_MAP_MASK_OFFSET 0x1fffffff
+
+#define NCCL_NET_MAP_OFFSET_BANK(mapStruct, offsetName) \
+  ((mapStruct)->offsets.offsetName >> 30)
+
+#define NCCL_NET_MAP_OFFSET_NULL(mapStruct, offsetName) \
+  (((mapStruct)->offsets.offsetName >> 29) == 0)
+
+#define NCCL_NET_MAP_GET_POINTER(mapStruct, cpuOrGpu, offsetName) \
+  (NCCL_NET_MAP_OFFSET_NULL(mapStruct, offsetName) ? NULL : \
+   (mapStruct)->mems[NCCL_NET_MAP_OFFSET_BANK(mapStruct, offsetName)].cpuOrGpu##Ptr + ((mapStruct)->offsets.offsetName & NCCL_NET_MAP_MASK_OFFSET))
+
+#define NCCL_NET_MAP_DEV_MEM(mapStruct, offsetName) \
+  (((mapStruct)->offsets.offsetName & NCCL_NET_MAP_MASK_DEVMEM) != 0)
+
+#define NCCL_NET_MAP_ADD_POINTER(mapStruct, shared, dev, memSize, offsetName) do { \
+    int bank = NCCL_NET_MAP_MASK_USED + (dev)*NCCL_NET_MAP_MASK_DEVMEM + (shared)*NCCL_NET_MAP_MASK_SHARED; \
+    if ((shared) == 0) { \
+      if (dev) { \
+        (mapStruct)->offsets.offsetName = bank + (mapStruct)->mems[NCCL_NET_MAP_DEVMEM].size; \
+        (mapStruct)->mems[NCCL_NET_MAP_DEVMEM].size += memSize; \
+      } else { \
+        (mapStruct)->offsets.offsetName = bank + (mapStruct)->mems[NCCL_NET_MAP_HOSTMEM].size; \
+        (mapStruct)->mems[NCCL_NET_MAP_HOSTMEM].size += memSize; \
+      } \
+    } else { \
+      (mapStruct)->offsets.offsetName = bank; \
+    } \
+} while (0);
+
+struct connectMapMem{
+  char* gpuPtr;
+  char* cpuPtr;
+  int size;
+};
+
+struct connectMap {
+  int shared;
+  // First 3 bits of offsets determine the mem bank. 001 is host mem, 011 is dev mem, 101 is shared host mem and 111 is shared dev mem.
+  struct connectMapMem mems[NCCL_NET_MAP_MEMS];
+  // Offsets. 3 MSBs indicate mem bank, 111 indicates NULL.
+  struct {
+    uint32_t sendMem;
+    uint32_t recvMem;
+    uint32_t buffs[NCCL_NUM_PROTOCOLS];
+  } offsets;
+};
+
 struct reqSlot {
   volatile void* recvBuff;
   volatile int size;
 };
 
-struct collNetSendResources {
-  struct ncclComm* comm;
+struct sendResources {
+  struct connectMap map;
   void* collNetComm;
   struct ncclSendMem* sendMem;
   struct ncclRecvMem* recvMem;
+
+  int rank;
+  int nranks;
   int netDev;
   int useGdr;
+  uint64_t* gdcSync;
+  void* gdrDesc;
   void* sendMhandles[NCCL_NUM_PROTOCOLS];
   void* recvMhandles[NCCL_NUM_PROTOCOLS];
-  struct ncclRecvMem* devRecvMem;
   uint64_t step;
-  uint64_t llLastCleaning;
   struct reqSlot (*reqFifo)[NCCL_STEPS];
   int collNetRank;
   uint32_t* curr_hdp_reg;  // Curr GPU in ring (for rdma transport use only)
 };
 
-struct collNetRecvResources {
-  struct ncclComm* comm;
+struct recvResources {
+  struct connectMap map;
   void* collNetComm;
   struct ncclSendMem* sendMem;
   struct ncclRecvMem* recvMem;
+
+  int rank;
+  int nranks;
   int netDev;
   int useGdr;
+  uint64_t* gdcSync;
+  uint64_t* gdcFlush;
+  void* gdrDesc;
   void* mhandles[NCCL_NUM_PROTOCOLS];
-  struct ncclRecvMem* devRecvMem;
   uint64_t step;
-  uint64_t llLastCleaning;
   struct reqSlot reqFifo[COLLNET_MAX_GROUPS][NCCL_STEPS];
   int collNetRank;
   uint32_t* curr_hdp_reg;  // Curr GPU in ring (for rdma transport use only)
 };
 
-struct collNetSharedResources {
-  void* collNetListenComms[MAXCHANNELS];
-  void* collNetComms[MAXCHANNELS];
-  int collNetCommRefCount[MAXCHANNELS];
-};
-
 /* Determine if we can communicate with the peer */
-ncclResult_t collNetCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTopoGraph* graph, struct ncclPeerInfo* info1, struct ncclPeerInfo* info2) {
+static ncclResult_t canConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTopoGraph* graph, struct ncclPeerInfo* info1, struct ncclPeerInfo* info2) {
   *ret = 1;
   return ncclSuccess;
 }
 
-ncclResult_t collNetSharedListen(struct ncclComm* comm, int netDev, void* collNetHandle) {
-  struct collNetSharedResources* resources = (struct collNetSharedResources*)comm->proxyState.sharedBuffs.collNetResources;
+struct setupReq {
+  int netDev;
+  int useGdr;
+};
+
+
+/* Setup send connector, and return connect information for others in the coll
+ * communicator to connect to me */
+static ncclResult_t sendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* send, int channelId, int connIndex) {
+  struct setupReq req;
+
+  int proxyRank;
+  NCCLCHECK(ncclTopoGetNetDev(comm, myInfo->rank, graph, channelId, -1, &req.netDev, &proxyRank));
+  NCCLCHECK(ncclTopoCheckGdr(comm->topo, myInfo->busId, req.netDev, 1, &req.useGdr));
+  send->conn.direct |= req.useGdr ? NCCL_DIRECT_NIC : 0;
+
+  NCCLCHECK(ncclTopoGetLocalRank(comm->topo, myInfo->rank, &send->proxyConn.localRank));
+  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_COLLNET, 1, myInfo->rank, &send->proxyConn));
+  NCCLCHECK(ncclProxyCall(&send->proxyConn, ncclProxyMsgSetup, &req, sizeof(req), NULL, 0));
+
+  INFO(NCCL_INIT|NCCL_NET,"CollNet %02d : %d [send] via COLLNET/%s/%d%s", channelId, myInfo->rank, collNetName(), req.netDev,
+      req.useGdr ? "/GDRDMA" : "");
+  return ncclSuccess;
+}
+
+static ncclResult_t recvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* recv, int channelId, int connIndex) {
+  struct setupReq req;
+
+  int proxyRank;
+  NCCLCHECK(ncclTopoGetNetDev(comm, myInfo->rank, graph, channelId, -1, &req.netDev, &proxyRank));
+  NCCLCHECK(ncclTopoCheckGdr(comm->topo, myInfo->busId, req.netDev, 0, &req.useGdr));
+  recv->conn.direct |= req.useGdr ? NCCL_DIRECT_NIC : 0;
+
+  NCCLCHECK(ncclTopoGetLocalRank(comm->topo, myInfo->rank, &recv->proxyConn.localRank));
+  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_COLLNET, 0, myInfo->rank, &recv->proxyConn));
+  struct collNetRecvConnectInfo* info = (struct collNetRecvConnectInfo*) connectInfo;
+  NCCLCHECK(ncclProxyCall(&recv->proxyConn, ncclProxyMsgSetup, &req, sizeof(req), &info->collNetHandle, sizeof(collNetHandle_t)));
+
+  INFO(NCCL_INIT|NCCL_NET,"CollNet %02d : %d [receive] via COLLNET/%s/%d%s", channelId, myInfo->rank, collNetName(), req.netDev,
+      req.useGdr ? "/GDRDMA" : "");
+  return ncclSuccess;
+}
+
+static ncclResult_t collNetDumpMap(struct connectMap* map) {
+  printf("Dump map\n");
+  struct connectMapMem *mem = map->mems+NCCL_NET_MAP_HOSTMEM;
+  printf("Mem 0: Host mem (%x B) CPU %p GPU %p\n", mem->size, mem->cpuPtr, mem->gpuPtr);
+  mem = map->mems+NCCL_NET_MAP_DEVMEM;
+  printf("Mem 1: Vid  mem CPU (%x B) %p GPU %p\n", mem->size, mem->cpuPtr, mem->gpuPtr);
+  mem = map->mems+NCCL_NET_MAP_SHARED_HOSTMEM;
+  printf("Mem 2: Shared Host mem (%x B) CPU %p GPU %p\n", mem->size, mem->cpuPtr, mem->gpuPtr);
+  mem = map->mems+NCCL_NET_MAP_SHARED_DEVMEM;
+  printf("Mem 3: Shared Vid  (%x B) mem CPU %p GPU %p\n", mem->size, mem->cpuPtr, mem->gpuPtr);
+  printf("SendMem -> Used %d Bank %d Offset %x, cpu %p gpu %p\n",
+      map->offsets.sendMem & NCCL_NET_MAP_MASK_USED ? 1 : 0,
+      NCCL_NET_MAP_OFFSET_BANK(map, sendMem), map->offsets.sendMem & NCCL_NET_MAP_MASK_OFFSET,
+      NCCL_NET_MAP_GET_POINTER(map, cpu, sendMem), NCCL_NET_MAP_GET_POINTER(map, gpu, sendMem));
+  printf("RecvMem -> Used %d Bank %d Offset %x, cpu %p gpu %p\n",
+      map->offsets.recvMem & NCCL_NET_MAP_MASK_USED ? 1 : 0,
+      NCCL_NET_MAP_OFFSET_BANK(map, recvMem), map->offsets.recvMem & NCCL_NET_MAP_MASK_OFFSET,
+      NCCL_NET_MAP_GET_POINTER(map, cpu, recvMem), NCCL_NET_MAP_GET_POINTER(map, gpu, recvMem));
+  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
+    printf("Proto %d -> Used %d Bank %d Offset %x, cpu %p, gpu %p\n", p,
+        map->offsets.buffs[p] & NCCL_NET_MAP_MASK_USED ? 1 : 0,
+        NCCL_NET_MAP_OFFSET_BANK(map, buffs[p]), map->offsets.buffs[p] & NCCL_NET_MAP_MASK_OFFSET,
+        NCCL_NET_MAP_GET_POINTER(map, cpu, buffs[p]), NCCL_NET_MAP_GET_POINTER(map, gpu, buffs[p]));
+  }
+  printf("End of dump\n");
+  return ncclSuccess;
+}
+
+struct collNetConnectArgs {
+  int rank;
+  int nranks;
+  struct ncclConnect* connectInfos;
+};
+
+static ncclResult_t sendConnect(struct ncclComm* comm, struct ncclConnect* connectInfos, int nranks, int rank, struct ncclConnector* send) {
+  // We're on the same process as the proxy. We can pass a pointer to a struct.
+  struct collNetConnectArgs args = { rank, nranks, connectInfos };
+  struct connectMap* map;
+  NCCLCHECK(ncclProxyCall(&send->proxyConn, ncclProxyMsgConnect, &args, sizeof(struct collNetConnectArgs), &map, sizeof(struct connectMap*)));
+
+  // If collnet connect failed, propagate error to fallback on regular p2p
+  if (map == NULL) return ncclSystemError;
+
+  //NCCLCHECK(collNetDumpMap(map));
+
+  struct ncclSendMem *sendMem = (struct ncclSendMem*) NCCL_NET_MAP_GET_POINTER(map, gpu, sendMem);
+  void* gdcMem = map->mems[NCCL_NET_MAP_GDCMEM].gpuPtr;
+  send->conn.head = gdcMem ? (uint64_t*)gdcMem : &sendMem->head;
+
+  struct ncclRecvMem *recvMem = (struct ncclRecvMem*) NCCL_NET_MAP_GET_POINTER(map, gpu, recvMem);
+  send->conn.tail = &recvMem->tail;
+  send->conn.sizesFifo = recvMem->sizesFifo;
+  for (int i=0; i<NCCL_STEPS; i++) send->conn.sizesFifo[i] = -1;
+  send->conn.offsFifo = recvMem->offsFifo;
+
+  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++)
+    send->conn.buffs[p] = NCCL_NET_MAP_GET_POINTER(map, gpu, buffs[p]);
+  return ncclSuccess;
+}
+
+static ncclResult_t recvConnect(struct ncclComm* comm, struct ncclConnect* connectInfos, int nranks, int rank, struct ncclConnector* recv) {
+  // We're on the same process as the proxy. We can pass a pointer to a struct.
+  struct collNetConnectArgs args = { rank, nranks, connectInfos };
+  struct connectMap* map;
+  NCCLCHECK(ncclProxyCall(&recv->proxyConn, ncclProxyMsgConnect, &args, sizeof(struct collNetConnectArgs), &map, sizeof(struct connectMap*)));
+
+  // If collnet connect failed, propagate error to fallback on regular p2p
+  if (map == NULL) return ncclSystemError;
+
+  //NCCLCHECK(collNetDumpMap(map));
+
+  struct ncclSendMem *sendMem = (struct ncclSendMem*) NCCL_NET_MAP_GET_POINTER(map, gpu, sendMem);
+  recv->conn.head = &sendMem->head;
+
+  struct ncclRecvMem *recvMem = (struct ncclRecvMem*) NCCL_NET_MAP_GET_POINTER(map, gpu, recvMem);
+  void* gdcMem = map->mems[NCCL_NET_MAP_GDCMEM].gpuPtr;
+  recv->conn.tail = gdcMem ? (uint64_t*)gdcMem : &recvMem->tail;
+  recv->conn.offsFifo = recvMem->offsFifo;
+
+  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
+    recv->conn.buffs[p] = NCCL_NET_MAP_GET_POINTER(map, gpu, buffs[p]);
+  }
+  return ncclSuccess;
+}
+
+static ncclResult_t sendFree(struct ncclConnector* send) {
+  return ncclSuccess;
+}
+
+static ncclResult_t recvFree(struct ncclConnector* recv) {
+  return ncclSuccess;
+}
+
+static ncclResult_t sendProxySetup(struct ncclProxyConnection* connection, struct ncclComm* comm, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) {
+  struct setupReq* req = (struct setupReq*)reqBuff;
+  if (reqSize != sizeof(struct setupReq)) return ncclInternalError;
+
+  struct sendResources* resources;
+  NCCLCHECK(ncclCalloc(&resources, 1));
+  connection->transportResources = resources;
+  connection->shared = 1;
+
+  resources->netDev = req->netDev;
+  resources->useGdr = req->useGdr;
+  return ncclSuccess;
+}
+
+struct sharedResources {
+  void* collNetListenComms[MAXCHANNELS];
+  void* collNetComms[MAXCHANNELS];
+  int commRefCount[NCCL_MAX_NETDEVS];
+};
+
+ncclResult_t sharedListen(struct ncclComm* comm, int netDev, void* collNetHandle) {
+  struct sharedResources* resources = (struct sharedResources*)comm->proxyState.progressState.collNet.resources;
   if (resources == NULL) {
     NCCLCHECK(ncclCalloc(&resources, 1));
-    comm->proxyState.sharedBuffs.collNetResources = resources;
+    comm->proxyState.progressState.collNet.resources = resources;
   }
   if (resources->collNetComms[netDev] == NULL)
     NCCLCHECK(collNetListen(netDev, collNetHandle, resources->collNetListenComms+netDev));
   return ncclSuccess;
 }
 
-/* Setup send connector, and return connect information for others in the coll communicator to connect to me */
-ncclResult_t collNetSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* send, int channelId, int connIndex) {
-  struct collNetSendResources* resources;
-  NCCLCHECK(ncclCalloc(&resources, 1));
-  send->transportResources = resources;
-  send->conn.shared = 1;
-  resources->comm = comm;
-
-  NCCLCHECK(ncclTopoGetNetDev(comm->topo, myInfo->rank, graph, channelId, 0, &resources->netDev));
-  NCCLCHECK(ncclTopoCheckGdr(comm->topo, myInfo->busId, resources->netDev, 1, &resources->useGdr));
-
-  send->proxyAppendPtr = comm->proxyState.sharedBuffs.proxyAppendCollNet+2*resources->netDev+1;
-
-  NCCLCHECK(ncclCudaHostCalloc(&resources->sendMem, 1));
-
-  int recvSize = offsetof(struct ncclRecvMem, buff);
-  // Simple uses shared buffers and we don't support LL128
-  recvSize += send->comm->buffSizes[NCCL_PROTO_LL];
-
-  if (resources->useGdr) {
-    NCCLCHECK(ncclCudaCalloc((char**)(&resources->devRecvMem), recvSize, resources->useGdr));
-    CUDACHECK(hipDeviceGetAttribute((int*)&resources->curr_hdp_reg, hipDeviceAttributeHdpMemFlushCntl, myInfo->cudaDev));
-    send->conn.curr_hdp_reg = resources->curr_hdp_reg;
-  }
-  NCCLCHECK(ncclCudaHostCalloc((char**)&resources->recvMem, recvSize));
-
-  INFO(NCCL_INIT|NCCL_NET,"CollNet %02d : %d [send] via COLLNET/%s/%d%s comm %p nRanks %02d", channelId, myInfo->rank, collNetName(), resources->netDev,
-      resources->useGdr ? "/GDRDMA" : "", comm, comm->nRanks);
-  return ncclSuccess;
-}
-
-/* Setup recv connector */
-ncclResult_t collNetRecvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclConnect* connectInfo, struct ncclConnector* recv, int channelId, int connIndex) {
-  struct collNetRecvResources* resources;
-  NCCLCHECK(ncclCalloc(&resources, 1));
-  recv->transportResources = resources;
-  recv->conn.shared = 1;
-  resources->comm = comm;
-
-  NCCLCHECK(ncclTopoGetNetDev(comm->topo, myInfo->rank, graph, channelId, 0, &resources->netDev));
-  NCCLCHECK(ncclTopoCheckGdr(comm->topo, myInfo->busId, resources->netDev, 0, &resources->useGdr));
-
-  recv->proxyAppendPtr = comm->proxyState.sharedBuffs.proxyAppendCollNet+2*resources->netDev;
-
-  NCCLCHECK(ncclCudaHostCalloc(&resources->sendMem, 1));
-
-  int recvSize = offsetof(struct ncclRecvMem, buff);
-  // Simple uses shared buffers and we don't support LL128
-  recvSize += recv->comm->buffSizes[NCCL_PROTO_LL];
-
-  if (resources->useGdr) {
-    NCCLCHECK(ncclCudaCalloc((char**)(&resources->devRecvMem), recvSize, resources->useGdr));
-  }
-  NCCLCHECK(ncclCudaHostCalloc((char**)&resources->recvMem, recvSize));
-
-  INFO(NCCL_INIT|NCCL_NET,"CollNet %02d : %d [receive] via COLLNET/%s/%d%s comm %p nRanks %02d", channelId, myInfo->rank, collNetName(), resources->netDev,
-      resources->useGdr ? "/GDRDMA" : "", comm, comm->nRanks);
-  struct collNetRecvConnectInfo* info = (struct collNetRecvConnectInfo*) connectInfo;
-
-  NCCLCHECK(collNetSharedListen(comm, resources->netDev, &info->collNetHandle));
-  return ncclSuccess;
-}
-
-ncclResult_t collNetSharedConnect(struct ncclComm* comm, int netDev, struct ncclConnect* connectInfos, int nranks, int rank, void** collNetComm) {
-  struct collNetSharedResources* resources = (struct collNetSharedResources*)comm->proxyState.sharedBuffs.collNetResources;
+static ncclResult_t sharedConnect(struct ncclComm* comm, int netDev, struct ncclConnect* connectInfos, int nranks, int rank, void** collNetComm) {
+  struct sharedResources* resources = (struct sharedResources*)comm->proxyState.progressState.collNet.resources;
   if (resources->collNetComms[netDev] == NULL) {
     // Connect to coll comm
     collNetHandle_t** handlePtrs = NULL;
@@ -159,157 +318,255 @@ ncclResult_t collNetSharedConnect(struct ncclComm* comm, int netDev, struct nccl
           resources->collNetListenComms[netDev],
           resources->collNetComms+netDev);
     free(handlePtrs);
-    NCCLCHECK(ret);
-    // Close listen comm
-    NCCLCHECK(collNetCloseListen(resources->collNetListenComms[netDev]));
+    if (ret == ncclSuccess) {
+      // Close listen comm
+      NCCLCHECK(collNetCloseListen(resources->collNetListenComms[netDev]));
+    } else {
+      resources->collNetListenComms[netDev] = NULL;
+    }
   }
   *collNetComm = resources->collNetComms[netDev];
-  resources->collNetCommRefCount[netDev]++;
+  if (*collNetComm) resources->commRefCount[netDev]++;
   return ncclSuccess;
 }
 
-ncclResult_t collNetSendConnect(struct ncclComm* comm, struct ncclConnect* connectInfos, int nranks, int rank, struct ncclConnector* send) {
-  // Setup device pointers
-  struct collNetSendResources* resources = (struct collNetSendResources*)send->transportResources;
-  struct collNetSendConnectInfo* info = (struct collNetSendConnectInfo*)(connectInfos+rank);
+static ncclResult_t sharedFree(struct ncclComm* comm, int netDev) {
+  struct sharedResources* resources = (struct sharedResources*)comm->proxyState.progressState.collNet.resources;
+  resources->commRefCount[netDev]--;
+  if (resources->commRefCount[netDev] == 0) {
+    NCCLCHECK(collNetCloseColl(resources->collNetComms[netDev]));
+  }
+  for (int n=0; n<NCCL_MAX_NETDEVS; n++) if (resources->commRefCount[n]) return ncclSuccess;
+  comm->proxyState.progressState.collNet.resources = NULL;
+  free(resources);
+  return ncclSuccess;
+}
 
-  // Intermediate buffering on GPU for GPU Direct RDMA, but LL buffer is always on host
-  send->conn.buffs[NCCL_PROTO_LL] = resources->recvMem->buff;
-  send->conn.buffs[NCCL_PROTO_LL128] = send->conn.buffs[NCCL_PROTO_SIMPLE] = NULL;
-  send->conn.direct |= resources->useGdr ? NCCL_DIRECT_NIC : 0;
+static ncclResult_t sharedBuffersInit(struct ncclComm* comm, int cuda, char** gpuPtr, char** cpuPtr, int* size) {
+  struct ncclProxySharedCollNet* state = &comm->proxyState.progressState.collNet;
+  if (state->size == 0) {
+    state->size = 2*comm->nChannels*comm->buffSizes[NCCL_PROTO_SIMPLE];
+  }
 
-  // Head/Tail/Opcount/Fifos are always on host
-  send->conn.tail = &resources->recvMem->tail;
-  send->conn.sizesFifo = resources->recvMem->sizesFifo;
-  send->conn.ptrsFifo = resources->recvMem->ptrsFifo;
-  send->conn.head = &resources->sendMem->head;
-  resources->sendMem->head = -NCCL_STEPS; // Don't give any credit yet when sharing buffers
-  for (int i=0; i<NCCL_STEPS; i++) send->conn.sizesFifo[i] = -1;
+  *size = state->size;
+
+  if (cuda && state->cudaBuff == NULL) {
+    NCCLCHECK(ncclCudaCalloc(&state->cudaBuff, *size, cuda));
+  }
+  if (!cuda && state->hostBuff == NULL) {
+    NCCLCHECK(ncclCudaHostCalloc(&state->hostBuff, *size));
+  }
+  *gpuPtr = *cpuPtr = cuda ? state->cudaBuff : state->hostBuff;
+  return ncclSuccess;
+}
+
+static ncclResult_t sharedBuffersGet(struct ncclComm* comm, int type, int slot, int channel, int* offset) {
+  // Use different pools for different channels and also separate send/recv.
+  int slotSize = comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS;
+  int globalSlot = (type*NCCL_STEPS+slot)*comm->nChannels+channel;
+  *offset = slotSize * globalSlot;
+  return ncclSuccess;
+}
+
+static ncclResult_t sharedBuffersDestroy(struct ncclComm* comm) {
+  struct ncclProxySharedCollNet* state = &comm->proxyState.progressState.collNet;
+  if (state->size == 0) return ncclSuccess;
+  CUDACHECK(hipFree(state->cudaBuff));
+  NCCLCHECK(ncclCudaHostFree(state->hostBuff));
+  // This will be called multiple times, with multiple channels and send/recv. Make sure we only do it once.
+  state->size = 0;
+  return ncclSuccess;
+}
+
+static ncclResult_t recvProxySetup(struct ncclProxyConnection* connection, struct ncclComm* comm, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) {
+  struct setupReq* req = (struct setupReq*)reqBuff;
+  if (reqSize != sizeof (struct setupReq)) return ncclInternalError;
+
+  struct recvResources* resources;
+  NCCLCHECK(ncclCalloc(&resources, 1));
+  connection->transportResources = resources;
+  connection->shared = 1;
+
+  resources->netDev = req->netDev;
+  resources->useGdr = req->useGdr;
+
+  collNetHandle_t* netHandle = (collNetHandle_t*) respBuff;
+  if (respSize != sizeof(collNetHandle_t)) return ncclInternalError;
+
+  NCCLCHECK(sharedListen(comm, req->netDev, netHandle));
+  return ncclSuccess;
+}
+
+static ncclResult_t sendProxyConnect(struct ncclProxyConnection* connection, struct ncclComm* comm, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) {
+  if (reqSize != sizeof(struct collNetConnectArgs)) { WARN("sendProxyConnect: reqSize is %d != %ld\n", reqSize, sizeof(struct collNetConnectArgs)); return ncclInternalError; }
+  struct collNetConnectArgs* args = (struct collNetConnectArgs*)reqBuff;
+  struct collNetSendConnectInfo* info = (struct collNetSendConnectInfo*)(args->connectInfos+args->rank);
+
+  struct sendResources* resources = (struct sendResources*)(connection->transportResources);
 
   // Get info from recv side
-  resources->collNetRank = rank;
+  resources->collNetRank = args->rank;
   resources->reqFifo = (struct reqSlot (*)[NCCL_STEPS])(info->reqFifo);
 
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++)
     resources->recvMhandles[p] = info->mhandles[p];
 
-  NCCLCHECK(collNetSharedConnect(comm, resources->netDev, connectInfos, nranks, rank, &resources->collNetComm));
+  NCCLCHECK(sharedConnect(comm, resources->netDev, args->connectInfos, args->nranks, args->rank, &resources->collNetComm));
+
+  // Collnet connect is allowed to fail. Gracefully handle that case by returning NULL to the caller.
+  if (respSize != sizeof(struct connectMap*)) { WARN("sendProxyConnect: respSize is %d != %ld\n", respSize, sizeof(void*)); return ncclInternalError; }
+  if (resources->collNetComm == NULL) {
+    *((struct connectMap**)respBuff) = NULL;
+    return ncclSuccess;
+  }
+  connection->proxyAppendPtr = comm->proxyState.progressState.collNet.proxyAppend+2*resources->netDev;
+
+  struct connectMap* map = &resources->map;
+
+  NCCL_NET_MAP_ADD_POINTER(map, 0, 0, sizeof(struct ncclSendMem), sendMem);
+  NCCL_NET_MAP_ADD_POINTER(map, 0, 0, sizeof(struct ncclRecvMem), recvMem);
+
+  NCCLCHECK(ncclCudaHostCalloc(&map->mems[NCCL_NET_MAP_HOSTMEM].cpuPtr, map->mems[NCCL_NET_MAP_HOSTMEM].size));
+  map->mems[NCCL_NET_MAP_HOSTMEM].gpuPtr = map->mems[NCCL_NET_MAP_HOSTMEM].cpuPtr;
+  if (ncclGdrCopy && ncclParamGdrCopySyncEnable()) {
+    uint64_t *cpuPtr, *gpuPtr;
+    NCCLCHECK(ncclGdrCudaCalloc(&cpuPtr, &gpuPtr, 1, &resources->gdrDesc));
+
+    resources->gdcSync = cpuPtr;
+    struct connectMapMem* gdcMem = map->mems+NCCL_NET_MAP_GDCMEM;
+    gdcMem->cpuPtr = (char*)cpuPtr;
+    gdcMem->gpuPtr = (char*)gpuPtr;
+    gdcMem->size = sizeof(uint64_t); // sendMem->head
+  }
+
+  resources->sendMem = (struct ncclSendMem*) NCCL_NET_MAP_GET_POINTER(map, cpu, sendMem);
+  resources->recvMem = (struct ncclRecvMem*) NCCL_NET_MAP_GET_POINTER(map, cpu, recvMem);
+  // Don't give credits yet in shared mode.
+  resources->sendMem->head = -NCCL_STEPS;
 
-  int size;
-  char* ptr;
   // Allocate & Register shared buffers for the Simple protocol
-  NCCLCHECK(ncclProxySharedBuffersInit(send->comm, resources->useGdr, &size, &ptr));
-  NCCLCHECK(collNetRegMr(resources->collNetComm, ptr, size,
+  int bank = resources->useGdr ? NCCL_NET_MAP_SHARED_DEVMEM : NCCL_NET_MAP_SHARED_HOSTMEM;
+  struct connectMapMem* mapMem = map->mems+bank;
+  NCCLCHECK(sharedBuffersInit(comm, resources->useGdr, &mapMem->gpuPtr, &mapMem->cpuPtr, &mapMem->size));
+  NCCL_NET_MAP_ADD_POINTER(map, 1, resources->useGdr, mapMem->size, buffs[NCCL_PROTO_SIMPLE]);
+
+  NCCLCHECK(collNetRegMr(resources->collNetComm, mapMem->cpuPtr, mapMem->size,
         resources->useGdr ? NCCL_PTR_CUDA : NCCL_PTR_HOST,
         &resources->sendMhandles[NCCL_PROTO_SIMPLE]));
 
-  // Allocate & Register shared buffers for the LL protocol
-  NCCLCHECK(ncclProxySharedBuffersInit(send->comm, 0, &size, &ptr));
-  NCCLCHECK(collNetRegMr(resources->collNetComm, ptr, size,
-        NCCL_PTR_HOST,
-        &resources->sendMhandles[NCCL_PROTO_LL]));
+  *((struct connectMap**)respBuff) = &resources->map;
   return ncclSuccess;
 }
 
-ncclResult_t collNetRecvConnect(struct ncclComm* comm, struct ncclConnect* connectInfos, int nranks, int rank, struct ncclConnector* recv) {
-  // Setup device pointers
-  struct collNetRecvResources* resources = (struct collNetRecvResources*)recv->transportResources;
-  struct collNetSendConnectInfo* info = (struct collNetSendConnectInfo*)(connectInfos+rank);
-  resources->collNetRank = rank;
+static ncclResult_t recvProxyConnect(struct ncclProxyConnection* connection, struct ncclComm* comm, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) {
+  if (reqSize != sizeof(struct collNetConnectArgs)) { WARN("recvProxyConnect: reqSize is %d != %ld\n", reqSize, sizeof(struct collNetConnectArgs)); return ncclInternalError; }
+  struct collNetConnectArgs* args = (struct collNetConnectArgs*)reqBuff;
 
-  // Intermediate buffering on GPU for GPU Direct RDMA
-  struct ncclRecvMem* recvMem = resources->useGdr ? resources->devRecvMem : resources->recvMem;
-  int offset = 0;
-  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
-    recv->conn.buffs[p] = (p == NCCL_PROTO_LL ? resources->recvMem->buff : recvMem->buff) + offset;
-    offset += recv->comm->buffSizes[p];
+  struct recvResources* resources = (struct recvResources*)(connection->transportResources);
+  struct collNetSendConnectInfo* info = (struct collNetSendConnectInfo*)(args->connectInfos+args->rank);
+  resources->collNetRank = args->rank;
+
+  NCCLCHECK(sharedConnect(comm, resources->netDev, args->connectInfos, args->nranks, args->rank, &resources->collNetComm));
+
+  // Collnet connect is allowed to fail. Gracefully handle that case by returning NULL to the caller.
+  if (respSize != sizeof(struct connectMap*)) { WARN("sendProxyConnect: respSize is %d != %ld\n", respSize, sizeof(void*)); return ncclInternalError; }
+  if (resources->collNetComm == NULL) {
+    *((struct connectMap**)respBuff) = NULL;
+    return ncclSuccess;
   }
-  recv->conn.direct |= resources->useGdr ? NCCL_DIRECT_NIC : 0;
+  connection->proxyAppendPtr = comm->proxyState.progressState.collNet.proxyAppend+2*resources->netDev+1;
 
-  // Head/Tail/Opcount are always on host
-  recv->conn.tail = &resources->recvMem->tail;
-  recv->conn.ptrsFifo = resources->recvMem->ptrsFifo;
-  recv->conn.head = &resources->sendMem->head;
+  struct connectMap* map = &resources->map;
 
-  NCCLCHECK(collNetSharedConnect(comm, resources->netDev, connectInfos, nranks, rank, &resources->collNetComm));
+  NCCL_NET_MAP_ADD_POINTER(map, 0, 0, sizeof(struct ncclSendMem), sendMem);
+  NCCL_NET_MAP_ADD_POINTER(map, 0, 0, sizeof(struct ncclRecvMem), recvMem);
 
-  int size;
-  char* ptr;
+  NCCLCHECK(ncclCudaHostCalloc(&map->mems[NCCL_NET_MAP_HOSTMEM].cpuPtr, map->mems[NCCL_NET_MAP_HOSTMEM].size));
+  map->mems[NCCL_NET_MAP_HOSTMEM].gpuPtr = map->mems[NCCL_NET_MAP_HOSTMEM].cpuPtr;
+  if (ncclGdrCopy) {
+    uint64_t *cpuPtr, *gpuPtr;
+    NCCLCHECK(ncclGdrCudaCalloc(&cpuPtr, &gpuPtr, 2, &resources->gdrDesc));
+
+    if (ncclParamGdrCopySyncEnable()) {
+      resources->gdcSync = cpuPtr;
+      struct connectMapMem* gdcMem = map->mems+NCCL_NET_MAP_GDCMEM;
+      gdcMem->cpuPtr = (char*)cpuPtr;
+      gdcMem->gpuPtr = (char*)gpuPtr;
+      gdcMem->size = sizeof(uint64_t);
+    }
+    if (ncclParamGdrCopyFlushEnable()) resources->gdcFlush = cpuPtr + 1;
+  }
+
+  resources->sendMem = (struct ncclSendMem*) NCCL_NET_MAP_GET_POINTER(map, cpu, sendMem);
+  resources->recvMem = (struct ncclRecvMem*) NCCL_NET_MAP_GET_POINTER(map, cpu, recvMem);
 
   // Allocate & Register shared buffers for the Simple protocol
-  NCCLCHECK(ncclProxySharedBuffersInit(recv->comm, resources->useGdr, &size, &ptr));
-  NCCLCHECK(collNetRegMr(resources->collNetComm, ptr, size,
+  int bank = resources->useGdr ? NCCL_NET_MAP_SHARED_DEVMEM : NCCL_NET_MAP_SHARED_HOSTMEM;
+  struct connectMapMem* mapMem = map->mems+bank;
+  NCCLCHECK(sharedBuffersInit(comm, resources->useGdr, &mapMem->gpuPtr, &mapMem->cpuPtr, &mapMem->size));
+  NCCL_NET_MAP_ADD_POINTER(map, 1, resources->useGdr, mapMem->size, buffs[NCCL_PROTO_SIMPLE]);
+
+  NCCLCHECK(collNetRegMr(resources->collNetComm, mapMem->cpuPtr, mapMem->size,
         resources->useGdr ? NCCL_PTR_CUDA : NCCL_PTR_HOST,
         &resources->mhandles[NCCL_PROTO_SIMPLE]));
 
-  // Allocate & Register shared buffers for the LL protocol
-  NCCLCHECK(ncclProxySharedBuffersInit(recv->comm, 0, &size, &ptr));
-  NCCLCHECK(collNetRegMr(resources->collNetComm, ptr, size,
-        NCCL_PTR_HOST,
-        &resources->mhandles[NCCL_PROTO_LL]));
-
   // Pass info to send side
   info->reqFifo = resources->reqFifo;
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++)
     info->mhandles[p] = resources->mhandles[p];
 
+  if (respSize != sizeof(struct connectMap*)) { WARN("recvProxyConnect: respSize is %d != %ld\n", respSize, sizeof(void*)); return ncclInternalError; }
+  *((struct connectMap**)respBuff) = &resources->map;
   return ncclSuccess;
 }
 
-ncclResult_t collNetSharedFree(struct ncclComm* comm, int netDev) {
-  struct collNetSharedResources* resources = (struct collNetSharedResources*)comm->proxyState.sharedBuffs.collNetResources;
-  resources->collNetCommRefCount[netDev]--;
-  if (resources->collNetCommRefCount[netDev] == 0) {
-    NCCLCHECK(collNetCloseColl(resources->collNetComms[netDev]));
+static ncclResult_t sendProxyFree(struct ncclProxyConnection* connection, struct ncclComm* comm) {
+  struct sendResources* resources = (struct sendResources*)(connection->transportResources);
+  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
+    if (resources->sendMhandles[p]) {
+      NCCLCHECK(collNetDeregMr(resources->collNetComm, resources->sendMhandles[p]));
+    }
   }
-  for (int c=0; c<MAXCHANNELS; c++) if (resources->collNetCommRefCount[c]) return ncclSuccess;
-  comm->proxyState.sharedBuffs.collNetResources = NULL;
-  free(resources);
+  struct connectMapMem* mems = resources->map.mems;
+  NCCLCHECK(ncclCudaHostFree(mems[NCCL_NET_MAP_HOSTMEM].cpuPtr));
+  CUDACHECK(hipFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
+  if (mems[NCCL_NET_MAP_GDCMEM].cpuPtr) NCCLCHECK(ncclGdrCudaFree(resources->gdrDesc));
+  NCCLCHECK(sharedBuffersDestroy(comm));
+  NCCLCHECK(sharedFree(comm, resources->netDev));
+  free(connection->transportResources);
   return ncclSuccess;
 }
 
-ncclResult_t collNetSendFree(void* sendTransportResources) {
-  struct collNetSendResources* resources = (struct collNetSendResources*)sendTransportResources;
-  NCCLCHECK(ncclCudaHostFree(resources->sendMem));
-  NCCLCHECK(ncclCudaHostFree(resources->recvMem));
-  if (resources->collNetComm) {
-    NCCLCHECK(collNetDeregMr(resources->collNetComm, resources->sendMhandles[NCCL_PROTO_LL]));
-    NCCLCHECK(collNetDeregMr(resources->collNetComm, resources->sendMhandles[NCCL_PROTO_SIMPLE]));
+static ncclResult_t recvProxyFree(struct ncclProxyConnection* connection, struct ncclComm* comm) {
+  struct recvResources* resources = (struct recvResources*)(connection->transportResources);
+  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
+    if (resources->mhandles[p]) {
+      NCCLCHECK(collNetDeregMr(resources->collNetComm, resources->mhandles[p]));
+    }
   }
-  if (resources->useGdr) CUDACHECK(hipFree(resources->devRecvMem));
-
-  NCCLCHECK(collNetSharedFree(resources->comm, resources->netDev));
-  free(resources);
+  struct connectMapMem* mems = resources->map.mems;
+  NCCLCHECK(ncclCudaHostFree(mems[NCCL_NET_MAP_HOSTMEM].cpuPtr));
+  CUDACHECK(hipFree(mems[NCCL_NET_MAP_DEVMEM].cpuPtr));
+  if (mems[NCCL_NET_MAP_GDCMEM].cpuPtr) NCCLCHECK(ncclGdrCudaFree(resources->gdrDesc));
+  NCCLCHECK(sharedBuffersDestroy(comm));
+  NCCLCHECK(sharedFree(comm, resources->netDev));
+  free(connection->transportResources);
   return ncclSuccess;
 }
 
-ncclResult_t collNetRecvFree(void* recvTransportResources) {
-  struct collNetRecvResources* resources = (struct collNetRecvResources*)recvTransportResources;
-  NCCLCHECK(ncclCudaHostFree(resources->sendMem));
-  NCCLCHECK(ncclCudaHostFree(resources->recvMem));
-  if (resources->collNetComm) {
-    NCCLCHECK(collNetDeregMr(resources->collNetComm, resources->mhandles[NCCL_PROTO_LL]));
-    NCCLCHECK(collNetDeregMr(resources->collNetComm, resources->mhandles[NCCL_PROTO_SIMPLE]));
-  }
-  if (resources->useGdr) CUDACHECK(hipFree(resources->devRecvMem));
-
-  NCCLCHECK(collNetSharedFree(resources->comm, resources->netDev));
-  free(resources);
-  return ncclSuccess;
-}
 
 #define LAST_OF_GROUP(s) \
   (s % COLLNET_GROUP_NSUBS == COLLNET_GROUP_NSUBS-1 || s == args->nsubs-1)
 
-ncclResult_t collNetSendProxy(struct ncclProxyArgs* args) {
-  if (args->protocol == NCCL_PROTO_LL128) {
-    WARN("CollNet does not support LL128");
+static ncclResult_t sendProxyProgress(struct ncclComm* comm, struct ncclProxyArgs* args) {
+  if (args->protocol != NCCL_PROTO_SIMPLE) {
+    WARN("CollNet does not support LL/LL128");
     return ncclInternalError;
   }
   if (args->state == ncclProxyOpReady) {
     for (int s=0; s<args->nsubs; s++) {
       struct ncclProxySubArgs* sub = args->subs+s;
-      struct collNetSendResources* resources = (struct collNetSendResources*) (sub->connector->transportResources);
+      struct sendResources* resources = (struct sendResources*) (sub->connection->transportResources);
       // Round to next multiple of sliceSteps
       sub->base = ROUNDUP(resources->step, args->chunkSteps);
       sub->posted = sub->received = sub->transmitted = sub->done = 0;
@@ -325,23 +582,21 @@ ncclResult_t collNetSendProxy(struct ncclProxyArgs* args) {
     int perGroupSteps = NCCL_STEPS / nGroups;
     for (int s=0; s<args->nsubs; s++) {
       struct ncclProxySubArgs* sub = args->subs+s;
-      struct collNetSendResources* resources = (struct collNetSendResources*) (sub->connector->transportResources);
+      struct sendResources* resources = (struct sendResources*) (sub->connection->transportResources);
       void* sendMhandle = resources->sendMhandles[p];
       void* recvMhandle = resources->recvMhandles[p];
-      int stepSize = sub->connector->comm->buffSizes[p] / NCCL_STEPS;
       auto reqFifo = resources->reqFifo;
       if (sub->posted < sub->nsteps && sub->posted < sub->done + NCCL_STEPS) {
         int buffSlot = (sub->base+sub->posted)%NCCL_STEPS;
-        if (p == NCCL_PROTO_SIMPLE) {
-          char* ptr;
-          int sharedBuffSlot = sub->posted%NCCL_STEPS;
-          NCCLCHECK(ncclProxySharedBuffersGetCollNet(sub->connector->comm, resources->useGdr, 0, sharedBuffSlot, 0, &ptr));
-          resources->recvMem->ptrsFifo[buffSlot] = ptr + s*args->chunkSize;
-          __sync_synchronize();
-        }
-        volatile uint64_t* sendHead = &resources->sendMem->head;
+        int sharedBuffSlot = sub->posted%NCCL_STEPS;
+        int offset;
+        NCCLCHECK(sharedBuffersGet(comm, 0, sharedBuffSlot, 0, &offset));
+        resources->recvMem->offsFifo[buffSlot] = offset + s*args->chunkSize;
+        __sync_synchronize();
+        volatile uint64_t* sendHead = resources->gdcSync ? resources->gdcSync : &resources->sendMem->head;
         sub->posted += args->sliceSteps;
         *sendHead = sub->base + sub->posted - NCCL_STEPS;
+        if (resources->gdcSync) wc_store_fence(); // Flush out WC write
       }
       // Enforce sync between operations of the same group.
       bool groupSync = (((s == 0) && ((sub+args->nsubs-1)->received == sub->received)) || (s && (sub-1)->received > sub->received));
@@ -350,30 +605,15 @@ ncclResult_t collNetSendProxy(struct ncclProxyArgs* args) {
         int sharedBuffSlot = sub->received%NCCL_STEPS;
         volatile int* sizesFifo = resources->recvMem->sizesFifo;
         volatile uint64_t* recvTail = &resources->recvMem->tail;
-        if (sizesFifo[buffSlot] != -1 && ((*recvTail > (sub->base+sub->received)) || p == NCCL_PROTO_LL)) {
+        char* localBuff = NCCL_NET_MAP_GET_POINTER(&resources->map, gpu, buffs[p]);
+        if (sizesFifo[buffSlot] != -1 && ((*recvTail > (sub->base+sub->received)))) {
           // We have something to receive, let's check whether data is ready.
-          int size = sizesFifo[buffSlot];
           int ready = 1;
           if (s == 0) {
-            NCCLCHECK(ncclProxySharedBuffersGetCollNet(sub->connector->comm, p == NCCL_PROTO_SIMPLE ? resources->useGdr : 0, 0, sharedBuffSlot, 0, &args->sharedBuff[sharedBuffSlot]));
-            args->sharedSize[sharedBuffSlot] = p == NCCL_PROTO_SIMPLE ? args->chunkSize : size/2;
-          }
-          if (p == NCCL_PROTO_LL) {
-            char* localBuff = sub->connector->conn.buffs[p];
-            uint32_t flag = NCCL_LL_FLAG(sub->base + sub->received + 1);
-            int nFifoLines = size / sizeof(union ncclLLFifoLine);
-            union ncclLLFifoLine* lines = (union ncclLLFifoLine*)(localBuff+buffSlot*stepSize);
-            // Pack data into the shared buffer
-            uint32_t* sendBuff = (uint32_t*)(args->sharedBuff[sharedBuffSlot]+args->sharedSize[sharedBuffSlot]*s);
-            for (int i=0; i<nFifoLines; i++) {
-              volatile uint32_t *f1 = &lines[i].flag1;
-              volatile uint32_t *d1 = &lines[i].data1;
-              volatile uint32_t *f2 = &lines[i].flag2;
-              volatile uint32_t *d2 = &lines[i].data2;
-              if (f1[0] != flag || f2[0] != flag) { ready = 0; break; }
-              sendBuff[2*i] = d1[0];
-              sendBuff[2*i+1] = d2[0];
-            }
+            int offset;
+            NCCLCHECK(sharedBuffersGet(comm, 0, sharedBuffSlot, 0, &offset));
+            args->sharedBuff[sharedBuffSlot] = localBuff + offset;
+            args->sharedSize[sharedBuffSlot] = args->chunkSize;
           }
           if (ready) {
             sizesFifo[buffSlot] = -1;
@@ -437,15 +677,15 @@ ncclResult_t collNetSendProxy(struct ncclProxyArgs* args) {
   return ncclSuccess;
 }
 
-ncclResult_t collNetRecvProxy(struct ncclProxyArgs* args) {
-  if (args->protocol == NCCL_PROTO_LL128) {
-    WARN("CollNet does not support LL128");
+static ncclResult_t recvProxyProgress(struct ncclComm* comm, struct ncclProxyArgs* args) {
+  if (args->protocol != NCCL_PROTO_SIMPLE) {
+    WARN("CollNet does not support LL/LL128");
     return ncclInternalError;
   }
   if (args->state == ncclProxyOpReady) {
     for (int s=0; s<args->nsubs; s++) {
       struct ncclProxySubArgs* sub = args->subs+s;
-      struct collNetRecvResources* resources = (struct collNetRecvResources*) (sub->connector->transportResources);
+      struct recvResources* resources = (struct recvResources*) (sub->connection->transportResources);
       // Round to next multiple of sliceSteps
       sub->base = ROUNDUP(resources->step, args->chunkSteps);
       sub->posted = sub->received = sub->flushed = sub->transmitted = sub->done = 0;
@@ -460,19 +700,20 @@ ncclResult_t collNetRecvProxy(struct ncclProxyArgs* args) {
     int perGroupSteps = NCCL_STEPS / nGroups;
     for (int s=0; s<args->nsubs; s++) {
       struct ncclProxySubArgs* sub = args->subs+s;
-      struct collNetRecvResources* resources = (struct collNetRecvResources*) (sub->connector->transportResources);
+      struct recvResources* resources = (struct recvResources*) (sub->connection->transportResources);
       void* mhandle = resources->mhandles[p];
-      int stepSize = sub->connector->comm->buffSizes[p] / NCCL_STEPS;
       auto reqFifo = resources->reqFifo;
+      char* localBuff = NCCL_NET_MAP_GET_POINTER(&resources->map, cpu, buffs[p]);
+
       // Enforce sync between operations of the same group.
       if (LAST_OF_GROUP(s) && (sub->posted < sub->done + perGroupSteps) && (sub->posted < sub->nsteps)) {
         int group = s / COLLNET_GROUP_NSUBS;
         int buffSlot = (sub->base+sub->posted)%NCCL_STEPS;
-        char* ptr;
         int sharedBuffSlot = sub->posted%NCCL_STEPS;
         int startChannel = group*COLLNET_GROUP_NSUBS;
-        NCCLCHECK(ncclProxySharedBuffersGetCollNet(sub->connector->comm, p == NCCL_PROTO_SIMPLE ? resources->useGdr : 0, 1, sharedBuffSlot, startChannel, &ptr));
-        reqFifo[group][buffSlot].recvBuff = ptr;
+        int offset;
+        NCCLCHECK(sharedBuffersGet(comm, 1, sharedBuffSlot, startChannel, &offset));
+        reqFifo[group][buffSlot].recvBuff = localBuff + offset;
         TRACE(NCCL_NET, "recvProxy [%lu/%d/%d] posted buffer %p", sub->posted, group, buffSlot, reqFifo[group][buffSlot].recvBuff);
         sub->posted += args->sliceSteps;
         args->idle = 0;
@@ -487,11 +728,24 @@ ncclResult_t collNetRecvProxy(struct ncclProxyArgs* args) {
           int totalSize = args->sharedSize[sharedBuffSlot]*(s-group*COLLNET_GROUP_NSUBS+1);
           TRACE(NCCL_NET, "recvProxy [%lu/%d/%d] received, size %d", sub->received, group, buffSlot, totalSize);
           sub->received += args->sliceSteps;
-          if (reqFifo[group][buffSlot].size > 0 && p == NCCL_PROTO_SIMPLE && resources->useGdr) {
-            int startChannel = group*COLLNET_GROUP_NSUBS;
-            char* groupRecvAddress;
-            NCCLCHECK(ncclProxySharedBuffersGetCollNet(sub->connector->comm, 1, 1, sharedBuffSlot, startChannel, &groupRecvAddress));
-            NCCLCHECK(collNetIflush(resources->collNetComm, groupRecvAddress, totalSize, mhandle, sub->requests+buffSlot));
+          sub->requests[buffSlot] = NULL;
+          if (reqFifo[group][buffSlot].size > 0 && resources->useGdr) {
+            // GDRCOPY support
+            if (resources->gdcFlush) {
+#if defined (__x86_64__)
+              // Force a PCI-E read from GPU memory
+              asm volatile ("mov (%0), %%eax" :: "l"(resources->gdcFlush) : "%eax");
+#else
+              WARN("NET: GDR Flush only supported on x86_64");
+              return ncclInternalError;
+#endif
+              sub->requests[buffSlot] = NULL;
+            } else {
+              int startChannel = group*COLLNET_GROUP_NSUBS;
+              int offset;
+              NCCLCHECK(sharedBuffersGet(comm, 1, sharedBuffSlot, startChannel, &offset));
+              NCCLCHECK(collNetIflush(resources->collNetComm, localBuff + offset, totalSize, mhandle, sub->requests+buffSlot));
+            }
           } else {
             for (int i=group*COLLNET_GROUP_NSUBS; i<=s; i++) args->subs[i].flushed += args->sliceSteps;
           }
@@ -517,27 +771,14 @@ ncclResult_t collNetRecvProxy(struct ncclProxyArgs* args) {
         int buffSlot = (sub->base + sub->transmitted)%NCCL_STEPS;
         int sharedBuffSlot = sub->transmitted%NCCL_STEPS;
         int startChannel = group*COLLNET_GROUP_NSUBS;
-        char* groupRecvAddress;
-        NCCLCHECK(ncclProxySharedBuffersGetCollNet(sub->connector->comm, p == NCCL_PROTO_SIMPLE ? resources->useGdr : 0, 1, sharedBuffSlot, startChannel, &groupRecvAddress));
-        char* ptr = groupRecvAddress + (s%COLLNET_GROUP_NSUBS)*args->sharedSize[sharedBuffSlot];
-        if (p == NCCL_PROTO_SIMPLE) {
-          volatile void** ptrsFifo = (volatile void**)resources->recvMem->ptrsFifo;
-          ptrsFifo[buffSlot] = ptr;
-          __sync_synchronize();
-          resources->recvMem->tail = sub->base + sub->flushed;
-        }
-        if (p == NCCL_PROTO_LL) { // ll
-          // re-attach flag
-          char* localBuff = sub->connector->conn.buffs[p];
-          uint32_t flag = NCCL_LL_FLAG(sub->base + sub->transmitted + 1);
-          union ncclLLFifoLine* lines = (union ncclLLFifoLine*)(localBuff+buffSlot*stepSize);
-          uint32_t* recvData = (uint32_t*)ptr;
-          int nFifoLines = DIVUP(args->sharedSize[sharedBuffSlot], 2*sizeof(uint32_t));
-          for (int i=0; i<nFifoLines; i++) {
-            lines[i].v[0] = ((uint64_t)flag << 32) + recvData[2*i];
-            lines[i].v[1] = ((uint64_t)flag << 32) + recvData[2*i+1];
-          }
-        }
+        int offset;
+        NCCLCHECK(sharedBuffersGet(comm, 1, sharedBuffSlot, startChannel, &offset));
+        volatile int* offsFifo = (volatile int*)resources->recvMem->offsFifo;
+        offsFifo[buffSlot] = offset + (s%COLLNET_GROUP_NSUBS)*args->chunkSize;
+        __sync_synchronize();
+        volatile uint64_t* recvTail = resources->gdcSync ? resources->gdcSync : &resources->recvMem->tail;
+        *recvTail = sub->base + sub->flushed;
+        if (resources->gdcSync) wc_store_fence(); // Flush out WC write
         sub->transmitted += args->sliceSteps;
         args->idle = 0;
         continue;
@@ -562,7 +803,7 @@ ncclResult_t collNetRecvProxy(struct ncclProxyArgs* args) {
 
 struct ncclTransport collNetTransport = {
   "COL",
-  collNetCanConnect,
-  { collNetSendSetup, collNetSendConnect, collNetSendFree, collNetSendProxy },
-  { collNetRecvSetup, collNetRecvConnect, collNetRecvFree, collNetRecvProxy }
+  canConnect,
+  { sendSetup, sendConnect, sendFree, NULL, sendProxySetup, sendProxyConnect, sendProxyFree, sendProxyProgress },
+  { recvSetup, recvConnect, recvFree, NULL, recvProxySetup, recvProxyConnect, recvProxyFree, recvProxyProgress }
 };

src/transport/net_socket.cc

@@ -1,5 +1,5 @@
 /*************************************************************************
- * Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -19,7 +19,7 @@
 /* Init functions */
 static int ncclNetIfs = -1;
 struct ncclSocketDev {
-  union socketAddress addr;
+  union ncclSocketAddress addr;
   char devName[MAX_IF_NAME_SIZE];
   char* pciPath;
 };
@@ -40,8 +40,8 @@ ncclResult_t ncclSocketInit(ncclDebugLogger_t logFunction) {
     pthread_mutex_lock(&ncclSocketLock);
     if (ncclNetIfs == -1) {
       char names[MAX_IF_NAME_SIZE*MAX_IFS];
-      union socketAddress addrs[MAX_IFS];
-      ncclNetIfs = findInterfaces(names, addrs, MAX_IF_NAME_SIZE, MAX_IFS);
+      union ncclSocketAddress addrs[MAX_IFS];
+      ncclNetIfs = ncclFindInterfaces(names, addrs, MAX_IF_NAME_SIZE, MAX_IFS);
       if (ncclNetIfs <= 0) {
         WARN("NET/Socket : no interface found");
         return ncclInternalError;
@@ -53,10 +53,10 @@ ncclResult_t ncclSocketInit(ncclDebugLogger_t logFunction) {
         addrline[SOCKET_NAME_MAXLEN] = '\0';
         for (int i=0; i<ncclNetIfs; i++) {
           strcpy(ncclSocketDevs[i].devName, names+i*MAX_IF_NAME_SIZE);
-          memcpy(&ncclSocketDevs[i].addr, addrs+i, sizeof(union socketAddress));
+          memcpy(&ncclSocketDevs[i].addr, addrs+i, sizeof(union ncclSocketAddress));
           NCCLCHECK(ncclSocketGetPciPath(ncclSocketDevs[i].devName, &ncclSocketDevs[i].pciPath));
           snprintf(line+strlen(line), MAX_LINE_LEN-strlen(line), " [%d]%s:%s", i, names+i*MAX_IF_NAME_SIZE,
-              socketToString(&addrs[i], addrline));
+              ncclSocketToString(&addrs[i], addrline));
         }
         line[MAX_LINE_LEN] = '\0';
         INFO(NCCL_INIT|NCCL_NET,"NET/Socket : Using%s", line);
@@ -97,12 +97,14 @@ ncclResult_t ncclSocketGetProperties(int dev, ncclNetProperties_t* props) {
   props->guid = dev;
   props->ptrSupport = NCCL_PTR_HOST;
   NCCLCHECK(ncclSocketGetSpeed(props->name, &props->speed));
+  props->latency = 0; // Not set
   props->port = 0;
   props->maxComms = 65536;
+  props->maxRecvs = 1;
   return ncclSuccess;
 }
 
-ncclResult_t GetSocketAddr(int dev, union socketAddress* addr) {
+ncclResult_t GetSocketAddr(int dev, union ncclSocketAddress* addr) {
   if (dev >= ncclNetIfs) return ncclInternalError;
   memcpy(addr, &ncclSocketDevs[dev].addr, sizeof(*addr));
   return ncclSuccess;
@@ -118,18 +120,33 @@ ncclResult_t GetSocketAddr(int dev, union socketAddress* addr) {
 NCCL_PARAM(SocketNsocksPerThread, "NSOCKS_PERTHREAD", -2);
 NCCL_PARAM(SocketNthreads, "SOCKET_NTHREADS", -2);
 
+enum ncclSocketCommState {
+  ncclSocketCommStateStart = 0,
+  ncclSocketCommStateConnect = 1,
+  ncclSocketCommStateAccept = 3,
+  ncclSocketCommStateSend = 4,
+  ncclSocketCommStateRecv = 5,
+};
+
+struct ncclSocketCommStage {
+  enum ncclSocketCommState state;
+  uint8_t iteration;
+  struct ncclSocket* sock;
+  struct ncclSocketComm* comm;
+};
+
 struct ncclSocketHandle {
-  union socketAddress connectAddr;
+  union ncclSocketAddress connectAddr;
   int nSocks;
   int nThreads;
+  struct ncclSocketCommStage stage;
 };
 
 struct ncclSocketTask {
   int op;
   void* data;
   int size;
-  int fd;
-  union socketAddress *addr;
+  struct ncclSocket* sock;
   int offset;
   int used;
   ncclResult_t result;
@@ -139,8 +156,7 @@ struct ncclSocketRequest {
   int op;
   void* data;
   int size;
-  int ctrlFd;
-  union socketAddress *addr;
+  struct ncclSocket* ctrlSock;
   int offset;
   int used;
   struct ncclSocketComm* comm;
@@ -154,29 +170,30 @@ struct ncclSocketTaskQueue {
   struct ncclSocketTask* tasks;
 };
 
-enum threadState {start, stop};
-
 struct ncclSocketThreadResources {
   struct ncclSocketTaskQueue threadTaskQueue;
-  enum threadState state;
+  int stop;
   struct ncclSocketComm* comm;
   pthread_mutex_t threadLock;
   pthread_cond_t  threadCond;
 };
 
 struct ncclSocketListenComm {
-  int fd;
+  struct ncclSocket sock;
+  struct ncclSocketCommStage stage;
   int nSocks;
   int nThreads;
+  int dev;
 };
 
 struct ncclSocketComm {
-  int ctrlFd;
-  union socketAddress addr;
-  int fds[MAX_SOCKETS];
+  struct ncclSocket ctrlSock;
+  struct ncclSocket socks[MAX_SOCKETS];
+  int dev;
+  int cudaDev;
   int nSocks;
   int nThreads;
-  int nextFd;
+  int nextSock;
   struct ncclSocketRequest requests[MAX_REQUESTS];
   pthread_t helperThread[MAX_THREADS];
   struct ncclSocketThreadResources threadResources[MAX_THREADS];
@@ -185,7 +202,6 @@ struct ncclSocketComm {
 void* persistentSocketThread(void *args_) {
   struct ncclSocketThreadResources* resource = (struct ncclSocketThreadResources*)args_;
   struct ncclSocketComm* comm = resource->comm;
-  volatile enum threadState* state = &resource->state;
   struct ncclSocketTaskQueue* myQueue = &resource->threadTaskQueue;
   int nSocksPerThread = comm->nSocks / comm->nThreads;
   while (1) {
@@ -198,7 +214,7 @@ void* persistentSocketThread(void *args_) {
         for (int j=0; j<nSocksPerThread; j++) {
           struct ncclSocketTask* r = myQueue->tasks+i+j;
           if (r != NULL && r->used == 1 && r->offset < r->size) {
-            r->result = socketProgress(r->op, r->fd, r->addr, r->data, r->size, &r->offset);
+            r->result = ncclSocketProgress(r->op, r->sock, r->data, r->size, &r->offset);
             if (r->result != ncclSuccess) {
               WARN("NET/Socket : socket progress error");
               return NULL;
@@ -211,12 +227,12 @@ void* persistentSocketThread(void *args_) {
     }
     if (idle) {
       pthread_mutex_lock(&resource->threadLock);
-      while (mark == myQueue->next && *state != stop) { // no new tasks, wait
+      while (mark == myQueue->next && resource->stop == 0) { // no new tasks, wait
         pthread_cond_wait(&resource->threadCond, &resource->threadLock);
       }
       pthread_mutex_unlock(&resource->threadLock);
     }
-    if (*state == stop) return NULL;
+    if (resource->stop) return NULL;
   }
 }
 
@@ -271,17 +287,17 @@ end:
 
 ncclResult_t ncclSocketNewListenComm(struct ncclSocketListenComm** comm) {
   NCCLCHECK(ncclCalloc(comm, 1));
-  (*comm)->fd = -1;
+  (*comm)->sock.fd = -1;
   return ncclSuccess;
 }
 
 ncclResult_t ncclSocketNewComm(struct ncclSocketComm** comm) {
   NCCLCHECK(ncclCalloc(comm, 1));
-  (*comm)->ctrlFd = -1;
+  (*comm)->ctrlSock.fd = -1;
   for (int i=0; i < MAX_SOCKETS; i++) {
-    (*comm)->fds[i] = -1;
+    (*comm)->socks[i].fd = -1;
   }
-  (*comm)->nextFd = 0;
+  (*comm)->nextSock = 0;
   return ncclSuccess;
 }
 
@@ -290,14 +306,18 @@ ncclResult_t ncclSocketListen(int dev, void* opaqueHandle, void** listenComm) {
     return ncclInternalError;
   }
   struct ncclSocketHandle* handle = (struct ncclSocketHandle*) opaqueHandle;
-  static_assert(sizeof(struct ncclSocketHandle) < NCCL_NET_HANDLE_MAXSIZE, "ncclSocketHandle size too large");
+  memset(handle, 0, sizeof(struct ncclSocketHandle));
+  static_assert(sizeof(struct ncclSocketHandle) <= NCCL_NET_HANDLE_MAXSIZE, "ncclSocketHandle size too large");
   struct ncclSocketListenComm* comm;
   NCCLCHECK(ncclSocketNewListenComm(&comm));
-  NCCLCHECK(GetSocketAddr(dev, &handle->connectAddr));
-  NCCLCHECK(createListenSocket(&comm->fd, &handle->connectAddr));
+  NCCLCHECK(GetSocketAddr(dev, &comm->sock.addr));
+  NCCLCHECK(ncclSocketListen(&comm->sock));
+  memcpy(&handle->connectAddr, &comm->sock.addr, sizeof(union ncclSocketAddress));
   NCCLCHECK(ncclSocketGetNsockNthread(dev, &comm->nSocks, &comm->nThreads));
   handle->nSocks = comm->nSocks;
   handle->nThreads = comm->nThreads;
+  comm->sock.asyncFlag = 1;
+  comm->dev = dev;
   *listenComm = comm;
   return ncclSuccess;
 }
@@ -306,38 +326,99 @@ ncclResult_t ncclSocketConnect(int dev, void* opaqueHandle, void** sendComm) {
   if (dev < 0) { // data transfer socket is based on specified dev
     return ncclInternalError;
   }
-  struct ncclSocketComm* comm;
-  NCCLCHECK(ncclSocketNewComm(&comm));
+
+  enum ncclSocketState conState;
   struct ncclSocketHandle* handle = (struct ncclSocketHandle*) opaqueHandle;
+  struct ncclSocketCommStage* stage = &handle->stage;
+  struct ncclSocketComm* comm = stage->comm;
+  uint8_t i = stage->iteration;
+  struct ncclSocket* sock = stage->sock;
+  *sendComm = NULL;
+
+  if (stage->state == ncclSocketCommStateConnect) goto socket_connect_check;
+  if (stage->state == ncclSocketCommStateSend) goto socket_send;
+
+  NCCLCHECK(ncclSocketNewComm(&comm));
+  stage->comm = comm;
   comm->nSocks = handle->nSocks;
   comm->nThreads = handle->nThreads;
-  for (int i=0; i<comm->nSocks+1; i++) {
-    int tmpFd, offset=0;
-    NCCLCHECK(connectAddress(&tmpFd, &handle->connectAddr));
-    NCCLCHECK(socketWait(NCCL_SOCKET_SEND, tmpFd, &handle->connectAddr, &i, sizeof(int), &offset));
-    if (i == comm->nSocks) comm->ctrlFd = tmpFd;
-    else comm->fds[i] = tmpFd;
+  comm->dev = dev;
+  CUDACHECK(hipGetDevice(&comm->cudaDev));
+  for (; i<comm->nSocks+1; i++) {
+    sock = i == comm->nSocks ? &comm->ctrlSock : comm->socks+i;
+    NCCLCHECK(ncclSocketInit(sock, &handle->connectAddr, NULL, 1));
+
+    stage->sock = sock;
+    stage->state = ncclSocketCommStateConnect;
+    stage->iteration = i;
+    NCCLCHECK(ncclSocketConnect(sock));
+
+socket_connect_check:
+    NCCLCHECK(ncclGetSocketState(sock, &conState));
+    if (conState == ncclSocketConnecting) {
+      /* expect user to call again */
+      return ncclSuccess;
+    } else if (conState == ncclSocketError) {
+      return ncclSystemError;
+    }
+    stage->state = ncclSocketCommStateSend;
+
+socket_send:
+    int done = 0;
+    NCCLCHECK(ncclSocketProgress(NCCL_SOCKET_SEND, sock, &i, sizeof(uint8_t), &done));
+    if (done == 0) return ncclSuccess;
   }
   *sendComm = comm;
-  comm->addr = handle->connectAddr;
   return ncclSuccess;
 }
 
 ncclResult_t ncclSocketAccept(void* listenComm, void** recvComm) {
   struct ncclSocketListenComm* lComm = (struct ncclSocketListenComm*)listenComm;
-  struct ncclSocketComm* rComm;
+  struct ncclSocketCommStage* stage = &lComm->stage;
+  struct ncclSocketComm* rComm = stage->comm;
+  uint8_t i = stage->iteration;
+  struct ncclSocket* sock = stage->sock;
+
+  *recvComm = NULL;
+  if (stage->state == ncclSocketCommStateAccept) goto socket_accept;
+  if (stage->state == ncclSocketCommStateRecv) goto socket_recv;
+
   NCCLCHECK(ncclSocketNewComm(&rComm));
+  stage->comm = rComm;
   rComm->nSocks = lComm->nSocks;
   rComm->nThreads = lComm->nThreads;
-  for (int i=0; i<rComm->nSocks+1; i++) {
-    int tmpFd, sendSockIdx, offset=0;
-    socklen_t socklen = sizeof(union socketAddress);
-    SYSCHECKVAL(accept(lComm->fd, &rComm->addr.sa, &socklen), "accept", tmpFd);
-    NCCLCHECK(socketWait(NCCL_SOCKET_RECV, tmpFd, &rComm->addr, &sendSockIdx, sizeof(int), &offset));
-    if (sendSockIdx == rComm->nSocks) rComm->ctrlFd = tmpFd;
-    else rComm->fds[sendSockIdx] = tmpFd;
+  rComm->dev = lComm->dev;
+  CUDACHECK(hipGetDevice(&rComm->cudaDev));
+  lComm->sock.asyncFlag = 1;
+  for (; i<rComm->nSocks+1; i++) {
+    uint8_t sendSockIdx;
+    ncclCalloc(&sock, 1);
+    NCCLCHECK(ncclSocketInit(sock, NULL, NULL, 1));
+    stage->sock = sock;
+    stage->state = ncclSocketCommStateAccept;
+    stage->iteration = i;
+socket_accept:
+    NCCLCHECK(ncclSocketAccept(sock, &lComm->sock));
+    if (sock->fd == -1) return ncclSuccess;
+
+    stage->state = ncclSocketCommStateRecv;
+socket_recv:
+    int done = 0;
+    NCCLCHECK(ncclSocketProgress(NCCL_SOCKET_RECV, sock, &sendSockIdx, sizeof(uint8_t), &done));
+    if (done == 0) return ncclSuccess;
+
+    if (sendSockIdx == rComm->nSocks) memcpy(&rComm->ctrlSock, sock, sizeof(struct ncclSocket));
+    else memcpy(rComm->socks+sendSockIdx, sock, sizeof(struct ncclSocket));
+
+    free(sock);
   }
   *recvComm = rComm;
+
+  /* reset lComm state */
+  stage->state = ncclSocketCommStateStart;
+  stage->iteration = 0;
+  stage->sock = NULL;
+  stage->comm = NULL;
   return ncclSuccess;
 }
 
@@ -348,8 +429,7 @@ ncclResult_t ncclSocketGetRequest(struct ncclSocketComm* comm, int op, void* dat
       r->op = op;
       r->data = data;
       r->size = size;
-      r->ctrlFd = comm->ctrlFd;
-      r->addr = &comm->addr;
+      r->ctrlSock = &comm->ctrlSock;
       r->used = 1;
       r->comm = comm;
       r->nSubs = 0;
@@ -362,7 +442,7 @@ ncclResult_t ncclSocketGetRequest(struct ncclSocketComm* comm, int op, void* dat
 }
 
 ncclResult_t ncclSocketGetTask(struct ncclSocketComm* comm, int op, void* data, int size, struct ncclSocketTask** req) {
-  int tid = comm->nextFd % comm->nThreads;
+  int tid = comm->nextSock % comm->nThreads;
   struct ncclSocketThreadResources* res = comm->threadResources+tid;
   struct ncclSocketTaskQueue* queue = &res->threadTaskQueue;
   // create helper threads and prepare per-thread task queue
@@ -377,22 +457,21 @@ ncclResult_t ncclSocketGetTask(struct ncclSocketComm* comm, int op, void* data,
     pthread_mutex_init(&res->threadLock, NULL);
     pthread_cond_init(&res->threadCond, NULL);
     pthread_create(comm->helperThread+tid, NULL, persistentSocketThread, res);
+    ncclSetThreadName(comm->helperThread[tid], "NCCL Sock%c%1u%2u%2u", op == NCCL_SOCKET_SEND ? 'S' : 'R', comm->dev, tid, comm->cudaDev);
   }
   struct ncclSocketTask* r = queue->tasks+queue->next;
   if (r->used == 0) {
     r->op = op;
     r->data = data;
     r->size = size;
-    r->fd = comm->fds[comm->nextFd];
-    r->addr = &comm->addr;
+    r->sock = comm->socks+comm->nextSock;
     r->offset = 0;
     r->result = ncclSuccess;
-    comm->nextFd = (comm->nextFd + 1) % comm->nSocks;
+    comm->nextSock = (comm->nextSock + 1) % comm->nSocks;
     r->used = 1;
     *req = r;
     pthread_mutex_lock(&res->threadLock);
     queue->next = (queue->next+1)%queue->len;
-    res->state = start;
     pthread_cond_signal(&res->threadCond);
     pthread_mutex_unlock(&res->threadLock);
     return ncclSuccess;
@@ -411,18 +490,20 @@ ncclResult_t ncclSocketTest(void* request, int* done, int* size) {
   if (r->used == 1) { /* try to send/recv size */
     int data = r->size;
     int offset = 0;
-    NCCLCHECK(socketProgress(r->op, r->ctrlFd, r->addr, &data, sizeof(int), &offset));
+    NCCLCHECK(ncclSocketProgress(r->op, r->ctrlSock, &data, sizeof(int), &offset));
 
     if (offset == 0) return ncclSuccess; /* Not ready -- retry later */
 
     // Not sure we could ever receive less than 4 bytes, but just in case ...
-    if (offset < sizeof(int)) NCCLCHECK(socketWait(r->op, r->ctrlFd, r->addr, &data, sizeof(int), &offset));
+    if (offset < sizeof(int)) NCCLCHECK(ncclSocketWait(r->op, r->ctrlSock, &data, sizeof(int), &offset));
 
     // Check size is less or equal to the size provided by the user
     if (r->op == NCCL_SOCKET_RECV && data > r->size) {
       char line[SOCKET_NAME_MAXLEN+1];
-      WARN("NET/Socket : peer %s message truncated : receiving %d bytes instead of %d", socketToString(r->addr, line), data, r->size);
-      return ncclInternalError;
+      WARN("NET/Socket : peer %s message truncated : receiving %d bytes instead of %d. If you believe your socket network is in healthy state, \
+          there may be a mismatch in collective sizes or environment settings (e.g. NCCL_PROTO, NCCL_ALGO) between ranks",
+          ncclSocketToString(&r->ctrlSock->addr, line), data, r->size);
+      return ncclInvalidUsage;
     }
     r->size = data;
     r->offset = 0;
@@ -459,7 +540,7 @@ ncclResult_t ncclSocketTest(void* request, int* done, int* size) {
       }
     } else { // progress request using main thread
       if (r->offset < r->size) {
-        NCCLCHECK(socketProgress(r->op, r->ctrlFd, r->addr, r->data, r->size, &r->offset));
+        NCCLCHECK(ncclSocketProgress(r->op, r->ctrlSock, r->data, r->size, &r->offset));
       }
       if (r->offset == r->size) {
         if (size) *size = r->size;
@@ -476,19 +557,20 @@ ncclResult_t ncclSocketRegMr(void* comm, void* data, int size, int type, void**
 }
 ncclResult_t ncclSocketDeregMr(void* comm, void* mhandle) { return ncclSuccess; }
 
-ncclResult_t ncclSocketIsend(void* sendComm, void* data, int size, void* mhandle, void** request) {
+ncclResult_t ncclSocketIsend(void* sendComm, void* data, int size, int tag, void* mhandle, void** request) {
   struct ncclSocketComm* comm = (struct ncclSocketComm*)sendComm;
   NCCLCHECK(ncclSocketGetRequest(comm, NCCL_SOCKET_SEND, data, size, (struct ncclSocketRequest**)request));
   return ncclSuccess;
 }
 
-ncclResult_t ncclSocketIrecv(void* recvComm, void* data, int size, void* mhandle, void** request) {
+ncclResult_t ncclSocketIrecv(void* recvComm, int n, void** data, int* sizes, int* tags, void** mhandles, void** request) {
   struct ncclSocketComm* comm = (struct ncclSocketComm*)recvComm;
-  NCCLCHECK(ncclSocketGetRequest(comm, NCCL_SOCKET_RECV, data, size, (struct ncclSocketRequest**)request));
+  if (n != 1) return ncclInternalError;
+  NCCLCHECK(ncclSocketGetRequest(comm, NCCL_SOCKET_RECV, data[0], sizes[0], (struct ncclSocketRequest**)request));
   return ncclSuccess;
 }
 
-ncclResult_t ncclSocketIflush(void* recvComm, void* data, int size, void* mhandle, void** request) {
+ncclResult_t ncclSocketIflush(void* recvComm, int n, void** data, int* sizes, void** mhandles, void** request) {
   // We don't support CUDA pointers, so we don't need a flush operation
   return ncclInternalError;
 }
@@ -496,7 +578,7 @@ ncclResult_t ncclSocketIflush(void* recvComm, void* data, int size, void* mhandl
 ncclResult_t ncclSocketCloseListen(void* opaqueComm) {
   struct ncclSocketListenComm* comm = (struct ncclSocketListenComm*)opaqueComm;
   if (comm) {
-    if (comm->fd != -1) close(comm->fd);
+    if (comm->sock.fd != -1) close(comm->sock.fd);
     free(comm);
   }
   return ncclSuccess;
@@ -509,16 +591,16 @@ ncclResult_t ncclSocketClose(void* opaqueComm) {
       struct ncclSocketThreadResources* res = comm->threadResources+i;
       if (comm->helperThread[i]) {
         pthread_mutex_lock(&res->threadLock);
-        res->state = stop;
+        res->stop = 1;
         pthread_cond_signal(&res->threadCond);
         pthread_mutex_unlock(&res->threadLock);
         pthread_join(comm->helperThread[i], NULL);
       }
       free(res->threadTaskQueue.tasks);
     }
-    if (comm->ctrlFd != -1) close(comm->ctrlFd);
+    if (comm->ctrlSock.fd != -1) close(comm->ctrlSock.fd);
     for (int i=0; i<comm->nSocks; i++) {
-      if (comm->fds[i] != -1) close(comm->fds[i]);
+      if (comm->socks[i].fd != -1) close(comm->socks[i].fd);
     }
     free(comm);
   }

src/transport/p2p.cc

@@ -1,6 +1,6 @@
 /*************************************************************************
- * Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -8,32 +8,30 @@
 #include "comm.h"
 #include "graph.h"
 #include "utils.h"
-#include "bootstrap.h"
 
-struct p2pConnectInfo {
-  int rank;
-  int read;
+struct ncclP2pBuff {
   void* directPtr;
   hipIpcMemHandle_t devIpc;
 };
 
+struct p2pConnectInfo {
+  int rank;
+  int read;
+  struct ncclP2pBuff p2pBuff;
+};
+static_assert(sizeof(p2pConnectInfo) <= CONNECT_SIZE, "P2P Connect info is too large");
+
 struct p2pSendResources {
   struct ncclSendMem* devMem;
-  void* ipcPtr;
   uint32_t* next_hdp_reg;  // Next GPU in ring (for p2p transport use only)
-  int remoteId;
-  int memRank;
-  void* remIpcPtr;
-  void* bootstrap;
+  void* sendMemIpc;
+  void* recvMemIpc;
 };
 
 struct p2pRecvResources {
   struct ncclRecvMem* devMem;
-  void* ipcPtr;
-  int remoteId;
-  int memRank;
-  void* remIpcPtr;
-  void* bootstrap;
+  void* sendMemIpc;
+  void* recvMemIpc;
 };
 
 #include <sys/types.h>
@@ -103,15 +101,22 @@ ncclResult_t p2pCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTop
 #else
   // Check that legacy IPC support is available
   if (p2p != 0) {
+    // Cached result of the legacyIPC detection
+    static int legacyIPC = -1;
+    if (legacyIPC >= 0) {
+      *ret = legacyIPC;
+      return ncclSuccess;
+    }
+    // Check that legacy IPC support is available (WSL WAR)
     char *dummy;
-    cudaIpcMemHandle_t ipc;
+    hipIpcMemHandle_t ipc;
     NCCLCHECK(ncclCudaCalloc(&dummy, CUDA_IPC_MIN));
-    if (cudaIpcGetMemHandle(&ipc, dummy) != cudaSuccess) {
-      INFO(NCCL_INIT|NCCL_P2P,"Legacy IPC not supported on dev %d(=%lx)",
-           cudaDev1, info1->busId);
+    if (hipIpcGetMemHandle(&ipc, dummy) != hipSuccess) {
+      INFO(NCCL_INIT|NCCL_P2P,"Legacy IPC not supported");
       *ret = 0;
     }
-    CUDACHECK(cudaFree(dummy));
+    CUDACHECK(hipFree(dummy));
+    legacyIPC = *ret;
     return ncclSuccess;
   }
 #endif
@@ -132,6 +137,7 @@ ncclResult_t p2pCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTop
     TRACE(P2P,"IPC: %016lx %016lx %016lx %016lx", devIpc[4], devIpc[5], devIpc[6], devIpc[7]); \
   } while (0)
 
+
 // Setting this to non zero causes P2P to use Reads rather than Writes
 NCCL_PARAM(P2pReadEnable, "P2P_READ_ENABLE", -2);
 
@@ -146,7 +152,7 @@ static ncclResult_t p2pGetInfo(struct ncclTopoSystem* topo, struct ncclPeerInfo*
   return ncclSuccess;
 }
 
-static ncclResult_t p2pMap(struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct p2pConnectInfo* p2pInfo, void** devMem, void** ipcPtr) {
+static ncclResult_t p2pMap(struct ncclPeerInfo* myInfo, struct ncclPeerInfo* peerInfo, struct ncclP2pBuff* p2pBuff, void** devMem, void** ipcPtr) {
   if (myInfo->pidHash == peerInfo->pidHash) {
     if (peerInfo->cudaDev != myInfo->cudaDev) {
       // Enable P2P access
@@ -159,10 +165,10 @@ static ncclResult_t p2pMap(struct ncclPeerInfo* myInfo, struct ncclPeerInfo* pee
         return ncclInternalError;
       }
     }
-    *devMem = p2pInfo->directPtr;
+    *devMem = p2pBuff->directPtr;
     *ipcPtr = NULL;
   } else {
-    CUDACHECK(hipIpcOpenMemHandle(devMem, p2pInfo->devIpc, hipIpcMemLazyEnablePeerAccess));
+    CUDACHECK(hipIpcOpenMemHandle(devMem, p2pBuff->devIpc, hipIpcMemLazyEnablePeerAccess));
     *ipcPtr = *devMem;
   }
   return ncclSuccess;
@@ -188,44 +194,40 @@ ncclResult_t p2pSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, st
     TRACE(NCCL_INIT|NCCL_P2P,"Ring %02d : %d -> %d HDP %p", channelId, myInfo->rank, peerInfo->rank, resources->next_hdp_reg);
   }
 
-  struct p2pConnectInfo info;
-  // For CollNet, we use write for scatter-reduce (conn 1), read for broadcast-gather (conn 0)
-  info.read = (connIndex == 0) ? useRead : 0;
-  const char* useReadStr = info.read ? "/read" : "";
+  static_assert(sizeof(struct p2pConnectInfo) <= sizeof(struct ncclConnect), "p2p Connect Info is too big");
+  struct p2pConnectInfo* info = (struct p2pConnectInfo*)connectInfo;
+  info->read = useRead;
+  // For CollNet, use write for scatter-reduce (conn 1), read for broadcast-gather (conn 0)
+  if (graph && connIndex == 1) info->read = 0;
+  const char* useReadStr = info->read ? "/read" : "";
 
   int sendSize = sizeof(struct ncclSendMem);
   // For P2P Read the SIMPLE buffer is tagged on the end of the ncclSendMem structure
-  if (info.read) sendSize += send->comm->buffSizes[NCCL_PROTO_SIMPLE];
+  if (info->read) sendSize += send->comm->buffSizes[NCCL_PROTO_SIMPLE];
   ALIGN_SIZE(sendSize, CUDA_IPC_MIN);
 
-  resources->remoteId = -1;
-  resources->bootstrap = comm->bootstrap;
   if (intermediateRank == -1) {
-    NCCLCHECK(ncclCudaCalloc((char**)&info.directPtr, sendSize, true));
-    info.rank = myInfo->rank;
+    info->rank = myInfo->rank;
     if (myInfo->pidHash == peerInfo->pidHash) {
-      send->conn.direct |= info.read ? NCCL_DIRECT_READ : NCCL_DIRECT_WRITE;
+      send->conn.direct |= info->read ? NCCL_DIRECT_READ : NCCL_DIRECT_WRITE;
       INFO(NCCL_INIT|NCCL_P2P, "Channel %02d : %d[%lx] -> %d[%lx] via P2P/direct pointer%s",
           channelId, myInfo->rank, myInfo->busId, peerInfo->rank, peerInfo->busId, useReadStr);
     } else {
-      send->conn.direct |= info.read ? NCCL_IPC_READ : NCCL_IPC_WRITE;
-      CUDACHECK(hipIpcGetMemHandle(&info.devIpc, info.directPtr));
+      send->conn.direct |= info->read ? NCCL_IPC_READ : NCCL_IPC_WRITE;
       INFO(NCCL_INIT|NCCL_P2P,"Channel %02d : %d[%lx] -> %d[%lx] via P2P/IPC%s",
           channelId, myInfo->rank, myInfo->busId, peerInfo->rank, peerInfo->busId, useReadStr);
     }
   } else {
-    NCCLCHECK(bootstrapRemAlloc(sendSize, intermediateRank, resources->bootstrap, &resources->remoteId, &info.devIpc, &info.directPtr));
-    info.rank = intermediateRank;
-    INFO(NCCL_INIT|NCCL_P2P, "Channel %02d : %d[%lx] -> %d[%lx] via P2P/indirect/%d[%lx]%s comm %p nRanks %02d",
+    info->rank = intermediateRank;
+    INFO(NCCL_INIT|NCCL_P2P, "Channel %02d : %d[%lx] -> %d[%lx] via P2P/indirect/%d[%lx]%s",
         channelId, myInfo->rank, myInfo->busId, peerInfo->rank, peerInfo->busId, intermediateRank,
-	comm->peerInfo[intermediateRank].busId, useReadStr, comm, comm->nRanks);
+	comm->peerInfo[intermediateRank].busId, useReadStr);
   }
-  resources->memRank = info.rank;
 
-  NCCLCHECK(p2pMap(myInfo, comm->peerInfo+info.rank, &info, (void**)&resources->devMem, &resources->ipcPtr));
+  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_P2P, 1, info->rank, &send->proxyConn));
+  NCCLCHECK(ncclProxyCall(&send->proxyConn, ncclProxyMsgSetup, &sendSize, sizeof(int), &info->p2pBuff, sizeof(struct ncclP2pBuff)));
 
-  static_assert(sizeof(struct p2pConnectInfo) <= sizeof(struct ncclConnect), "p2p Connect Info is too big");
-  memcpy(connectInfo, &info, sizeof(struct p2pConnectInfo));
+  NCCLCHECK(p2pMap(myInfo, comm->peerInfo+info->rank, &info->p2pBuff, (void**)&resources->devMem, &resources->sendMemIpc));
   return ncclSuccess;
 }
 
@@ -238,36 +240,32 @@ ncclResult_t p2pRecvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, st
   int useRead, intermediateRank;
   NCCLCHECK(p2pGetInfo(comm->topo, myInfo, peerInfo, &useRead, &intermediateRank));
 
-  struct p2pConnectInfo info;
-  // For CollNet, we use write for scatter-reduce (conn 1), read for broadcast-gather (conn 0)
-  info.read = (connIndex == 0) ? useRead : 0;
+  static_assert(sizeof(struct p2pConnectInfo) <= sizeof(struct ncclConnect), "p2p Connect Info is too big");
+  struct p2pConnectInfo* info = (struct p2pConnectInfo*)connectInfo;
+  info->read = useRead;
+  // For CollNet, use write for scatter-reduce (conn 1), read for broadcast-gather (conn 0)
+  if (graph && connIndex == 1) info->read = 0;
 
-  int recvSize = offsetof(struct ncclRecvMem, buff);
+  int recvSize = sizeof(struct ncclRecvMem);
   // For P2P Read the SIMPLE buffer is tagged on the end of the ncclSendMem structure
-  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) if (!(info.read && p == NCCL_PROTO_SIMPLE)) recvSize += recv->comm->buffSizes[p];
+  for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) if (!(info->read && p == NCCL_PROTO_SIMPLE)) recvSize += recv->comm->buffSizes[p];
   ALIGN_SIZE(recvSize, CUDA_IPC_MIN);
 
-  resources->remoteId = -1;
-  resources->bootstrap = comm->bootstrap;
   if (intermediateRank == -1) {
-    NCCLCHECK(ncclCudaCalloc((char**)&info.directPtr, recvSize, true));
-    info.rank = myInfo->rank;
+    info->rank = myInfo->rank;
     if (myInfo->pidHash == peerInfo->pidHash) {
-      recv->conn.direct |= info.read ? NCCL_DIRECT_READ : NCCL_DIRECT_WRITE;
+      recv->conn.direct |= info->read ? NCCL_DIRECT_READ : NCCL_DIRECT_WRITE;
     } else {
-      recv->conn.direct |= info.read ? NCCL_IPC_READ : NCCL_IPC_WRITE;
-      CUDACHECK(hipIpcGetMemHandle(&info.devIpc, info.directPtr));
+      recv->conn.direct |= info->read ? NCCL_IPC_READ : NCCL_IPC_WRITE;
     }
   } else {
-    NCCLCHECK(bootstrapRemAlloc(recvSize, intermediateRank, resources->bootstrap, &resources->remoteId, &info.devIpc, &info.directPtr));
-    info.rank = intermediateRank;
+    info->rank = intermediateRank;
   }
-  resources->memRank = info.rank;
 
-  NCCLCHECK(p2pMap(myInfo, comm->peerInfo+info.rank, &info, (void**)&resources->devMem, &resources->ipcPtr));
+  NCCLCHECK(ncclProxyConnect(comm, TRANSPORT_P2P, 0, info->rank, &recv->proxyConn));
+  NCCLCHECK(ncclProxyCall(&recv->proxyConn, ncclProxyMsgSetup, &recvSize, sizeof(int), &info->p2pBuff, sizeof(struct ncclP2pBuff)));
 
-  static_assert(sizeof(struct p2pConnectInfo) <= sizeof(struct ncclConnect), "p2p Connect Info is too big");
-  memcpy(connectInfo, &info, sizeof(struct p2pConnectInfo));
+  NCCLCHECK(p2pMap(myInfo, comm->peerInfo+info->rank, &info->p2pBuff, (void**)&resources->devMem, &resources->recvMemIpc));
   return ncclSuccess;
 }
 
@@ -277,16 +275,16 @@ static ncclResult_t p2pSendConnect(struct ncclComm* comm, struct ncclConnect* co
   struct ncclRecvMem* remDevMem;
   struct p2pConnectInfo* info = (struct p2pConnectInfo*)connectInfo;
 
-  NCCLCHECK(p2pMap(comm->peerInfo+rank, comm->peerInfo+info->rank, info, (void**)&remDevMem, &resources->remIpcPtr));
+  NCCLCHECK(p2pMap(comm->peerInfo+rank, comm->peerInfo+info->rank, &info->p2pBuff, (void**)&remDevMem, &resources->recvMemIpc));
 
-  int offset = 0;
+  char* buff = (char*)(remDevMem+1);
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
     if (info->read && p == NCCL_PROTO_SIMPLE) {
       /* For P2P Read the SIMPLE buffer is local (ncclSendMem) */
-      send->conn.buffs[p] = resources->devMem->buff;
+      send->conn.buffs[p] = (char*)(resources->devMem+1);
     } else {
-      send->conn.buffs[p] = remDevMem->buff + offset;
-      offset += send->comm->buffSizes[p];
+      send->conn.buffs[p] = buff;
+      buff += send->comm->buffSizes[p];
     }
   }
   send->conn.tail = &remDevMem->tail;
@@ -303,16 +301,16 @@ ncclResult_t p2pRecvConnect(struct ncclComm* comm, struct ncclConnect* connectIn
   struct ncclSendMem* remDevMem;
   struct p2pConnectInfo* info = (struct p2pConnectInfo*)connectInfo;
 
-  NCCLCHECK(p2pMap(comm->peerInfo+rank, comm->peerInfo+info->rank, info, (void**)&remDevMem, &resources->remIpcPtr));
+  NCCLCHECK(p2pMap(comm->peerInfo+rank, comm->peerInfo+info->rank, &info->p2pBuff, (void**)&remDevMem, &resources->sendMemIpc));
 
-  int offset = 0;
+  char* buff = (char*)(resources->devMem+1);
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
     if (info->read && p == NCCL_PROTO_SIMPLE) {
       /* For P2P Read the SIMPLE buffer is remote (ncclSendMem) */
-      recv->conn.buffs[p] = remDevMem->buff;
+      recv->conn.buffs[p] = (char*)(remDevMem+1);
     } else {
-      recv->conn.buffs[p] = resources->devMem->buff + offset;
-      offset += recv->comm->buffSizes[p];
+      recv->conn.buffs[p] = buff;
+      buff += recv->comm->buffSizes[p];
     }
   }
   recv->conn.tail = &resources->devMem->tail;
@@ -322,39 +320,49 @@ ncclResult_t p2pRecvConnect(struct ncclComm* comm, struct ncclConnect* connectIn
   return ncclSuccess;
 }
 
-ncclResult_t p2pSendFree(void* resources) {
-  struct p2pSendResources* sendRes = (struct p2pSendResources*)resources;
-  if (sendRes->ipcPtr)
-    CUDACHECK(hipIpcCloseMemHandle(sendRes->ipcPtr));
-  if (sendRes->remIpcPtr)
-    CUDACHECK(hipIpcCloseMemHandle(sendRes->remIpcPtr));
-  if (sendRes->remoteId != -1) {
-    NCCLCHECK(bootstrapRemFree(sendRes->remoteId, sendRes->memRank, sendRes->bootstrap));
-    sendRes->devMem = NULL;
-  }
-  CUDACHECK(hipFree(sendRes->devMem));
-  free(sendRes);
+ncclResult_t p2pSendFree(struct ncclConnector* send) {
+  struct p2pSendResources* resources = (struct p2pSendResources*)send->transportResources;
+  if (resources->sendMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->sendMemIpc));
+  if (resources->recvMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->recvMemIpc));
+  free(resources);
   return ncclSuccess;
 }
 
-ncclResult_t p2pRecvFree(void* resources) {
-  struct p2pRecvResources* recvRes = (struct p2pRecvResources*)resources;
-  if (recvRes->ipcPtr)
-    CUDACHECK(hipIpcCloseMemHandle(recvRes->ipcPtr));
-  if (recvRes->remIpcPtr)
-    CUDACHECK(hipIpcCloseMemHandle(recvRes->remIpcPtr));
-  if (recvRes->remoteId != -1) {
-    NCCLCHECK(bootstrapRemFree(recvRes->remoteId, recvRes->memRank, recvRes->bootstrap));
-    recvRes->devMem = NULL;
+ncclResult_t p2pRecvFree(struct ncclConnector* recv) {
+  struct p2pRecvResources* resources = (struct p2pRecvResources*)recv->transportResources;
+  if (resources->sendMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->sendMemIpc));
+  if (resources->recvMemIpc) CUDACHECK(hipIpcCloseMemHandle(resources->recvMemIpc));
+  free(resources);
+  return ncclSuccess;
+}
+
+static ncclResult_t p2pProxySetup(struct ncclProxyConnection* connection, struct ncclComm* comm, void* reqBuff, int reqSize, void* respBuff, int respSize, int* done) {
+  if (reqSize != sizeof(int)) return ncclInternalError;
+  int size = *((int*)reqBuff);
+  if (respSize != sizeof(struct ncclP2pBuff)) return ncclInternalError;
+  struct ncclP2pBuff* p2pBuff = (struct ncclP2pBuff*)respBuff;
+  NCCLCHECK(ncclCudaCalloc((char**)&p2pBuff->directPtr, size, true));
+  connection->transportResources = p2pBuff->directPtr;
+  hipError_t res = hipIpcGetMemHandle(&p2pBuff->devIpc, p2pBuff->directPtr);
+  if (res != hipSuccess) {
+    WARN("hipIpcGetMemHandle failed : %s", hipGetErrorString(res));
+    hipFree(p2pBuff->directPtr);
+    free(p2pBuff);
+    CUDACHECK(res);
   }
-  CUDACHECK(hipFree(recvRes->devMem));
-  free(recvRes);
+  *done = 1;
+  return ncclSuccess;
+}
+
+static ncclResult_t p2pProxyFree(struct ncclProxyConnection* connection, struct ncclComm* comm) {
+  // Do not check return code as CUDA may have already shut down
+  hipFree(connection->transportResources);
   return ncclSuccess;
 }
 
 struct ncclTransport p2pTransport = {
   "P2P",
   p2pCanConnect,
-  { p2pSendSetup, p2pSendConnect, p2pSendFree, NULL },
-  { p2pRecvSetup, p2pRecvConnect, p2pRecvFree, NULL }
+  { p2pSendSetup, p2pSendConnect, p2pSendFree, NULL, p2pProxySetup, NULL, p2pProxyFree, NULL },
+  { p2pRecvSetup, p2pRecvConnect, p2pRecvFree, NULL, p2pProxySetup, NULL, p2pProxyFree, NULL }
 };

src/transport/shm.cc

@@ -1,6 +1,5 @@
 /*************************************************************************
- * Copyright (c) 2016-2021, NVIDIA CORPORATION. All rights reserved.
- * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2016-2022, NVIDIA CORPORATION. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -9,12 +8,10 @@
 #include "shm.h"
 
 struct shmConnectInfo {
-  uint64_t pidHash;
-  int id;
-  int sendRank;
-  int recvRank;
+  char shmName[7];
   int shmSize;
 };
+static_assert(sizeof(shmConnectInfo) <= CONNECT_SIZE, "SHM Connect info is too large");
 
 struct shmSendResources {
   int remShmSize;
@@ -63,22 +60,17 @@ ncclResult_t shmSendSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, st
   NCCLCHECK(ncclCalloc(&resources, 1));
   send->transportResources = resources;
 
-  struct shmConnectInfo info;
-  info.id = channelId;
-  info.pidHash = myInfo->pidHash;
-  info.sendRank = myInfo->rank;
-  info.recvRank = peerInfo->rank;
+  static_assert(sizeof(struct shmConnectInfo) <= sizeof(struct ncclConnect), "shm Connect Info is too big");
+  struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;
 
-  char shmName[MAX_SHM_NAME_LEN];
-  sprintf(shmName, "nccl-shm-send-%lx-%d-%d-%d", info.pidHash, info.id, info.sendRank, info.recvRank);
-  info.shmSize = resources->shmSize = sizeof(struct ncclSendMem);
-  TRACE(NCCL_SHM,"Open shmName %s shmSize %d", shmName, info.shmSize);
-  NCCLCHECK(shmOpen(shmName, resources->shmSize, (void**)&resources->hostMem, (void**)&resources->devHostMem, 1));
+  char shmPath[PATH_MAX];
+  shmPath[0] = '\0';
+  info->shmSize = resources->shmSize = sizeof(struct ncclSendMem);
+  NCCLCHECK(ncclShmOpen(shmPath, resources->shmSize, (void**)&resources->hostMem, (void**)&resources->devHostMem, 1));
+  TRACE(NCCL_SHM,"Opened shmName %s shmSize %d", shmPath, info->shmSize);
+  memcpy(info->shmName, shmPath+sizeof("/dev/shm/nccl-")-1, sizeof(info->shmName));
 
-  INFO(NCCL_INIT|NCCL_SHM,"Channel %02d : %d[%lx] -> %d[%lx] via direct shared memory comm %p nRanks %02d", channelId, myInfo->rank, 
-		  myInfo->busId, peerInfo->rank, peerInfo->busId, comm, comm->nRanks);
-  static_assert(sizeof(struct shmConnectInfo) <= sizeof(struct ncclConnect), "shm Connect Recv Info is too big");
-  memcpy(connectInfo, &info, sizeof(struct shmConnectInfo));
+  INFO(NCCL_INIT|NCCL_SHM,"Channel %02d : %d[%lx] -> %d[%lx] via direct shared memory", channelId, myInfo->rank, myInfo->busId, peerInfo->rank, peerInfo->busId);
   return ncclSuccess;
 }
 
@@ -87,22 +79,18 @@ ncclResult_t shmRecvSetup(struct ncclComm* comm, struct ncclTopoGraph* graph, st
   NCCLCHECK(ncclCalloc(&resources, 1));
   recv->transportResources = resources;
 
-  struct shmConnectInfo info;
-  info.id = channelId;
-  info.pidHash = myInfo->pidHash;
-  info.sendRank = peerInfo->rank;
-  info.recvRank = myInfo->rank;
+  static_assert(sizeof(struct shmConnectInfo) <= sizeof(struct ncclConnect), "shm Connect Info is too big");
+  struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;
 
-  char shmName[MAX_SHM_NAME_LEN];
-  sprintf(shmName, "nccl-shm-recv-%lx-%d-%d-%d", info.pidHash, info.id, info.sendRank, info.recvRank);
-  int shmSize = offsetof(struct ncclRecvMem, buff);
+  char shmPath[PATH_MAX];
+  shmPath[0] = '\0';
+  int shmSize = sizeof(struct ncclRecvMem);
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) shmSize += recv->comm->buffSizes[p];
-  info.shmSize = resources->shmSize = shmSize;
-  TRACE(NCCL_SHM,"Open shmName %s shmSize %d", shmName, info.shmSize);
-  NCCLCHECK(shmOpen(shmName, resources->shmSize, (void**)&resources->hostMem, (void**)&resources->devHostMem, 1));
+  info->shmSize = resources->shmSize = shmSize;
+  NCCLCHECK(ncclShmOpen(shmPath, resources->shmSize, (void**)&resources->hostMem, (void**)&resources->devHostMem, 1));
+  TRACE(NCCL_SHM,"Opened shmName %s shmSize %d", shmPath, info->shmSize);
+  memcpy(info->shmName, shmPath+sizeof("/dev/shm/nccl-")-1, sizeof(info->shmName));
 
-  static_assert(sizeof(struct shmConnectInfo) <= sizeof(struct ncclConnect), "shm Connect Send Info is too big");
-  memcpy(connectInfo, &info, sizeof(struct shmConnectInfo));
   return ncclSuccess;
 }
 
@@ -112,18 +100,18 @@ ncclResult_t shmSendConnect(struct ncclComm* comm, struct ncclConnect* connectIn
   struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;
   struct shmSendResources* resources = (struct shmSendResources*)send->transportResources;
 
-  char shmName[MAX_SHM_NAME_LEN];
-  sprintf(shmName, "nccl-shm-recv-%lx-%d-%d-%d", info->pidHash, info->id, info->sendRank, info->recvRank);
+  char shmPath[PATH_MAX];
+  sprintf(shmPath, "/dev/shm/nccl-%s", info->shmName);
   resources->remShmSize = info->shmSize;
-  TRACE(NCCL_SHM,"Open shmName %s shmSize %d", shmName, info->shmSize);
-  NCCLCHECK(shmOpen(shmName, resources->remShmSize, (void**)&resources->remHostMem, (void**)&resources->devRemHostMem, 0));
+  TRACE(NCCL_SHM,"Open shmName %s shmSize %d", shmPath, info->shmSize);
+  NCCLCHECK(ncclShmOpen(shmPath, resources->remShmSize, (void**)&resources->remHostMem, (void**)&resources->devRemHostMem, 0));
   // Remove the file to ensure proper clean-up
-  NCCLCHECK(shmUnlink(shmName));
+  NCCLCHECK(ncclShmUnlink(shmPath));
 
   send->transportResources = resources;
   int offset = 0;
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
-    send->conn.buffs[p] = resources->devRemHostMem->buff + offset;
+    send->conn.buffs[p] = (char*)(resources->devRemHostMem+1) + offset;
     offset += send->comm->buffSizes[p];
   }
   send->conn.tail = &resources->devRemHostMem->tail;
@@ -137,35 +125,35 @@ ncclResult_t shmRecvConnect(struct ncclComm* comm, struct ncclConnect* connectIn
   struct shmRecvResources* resources = (struct shmRecvResources*)recv->transportResources;
   struct shmConnectInfo* info = (struct shmConnectInfo*)connectInfo;
 
-  char shmName[MAX_SHM_NAME_LEN];
-  sprintf(shmName, "nccl-shm-send-%lx-%d-%d-%d", info->pidHash, info->id, info->sendRank, info->recvRank);
+  char shmPath[PATH_MAX];
+  sprintf(shmPath, "/dev/shm/nccl-%s", info->shmName);
   resources->remShmSize = info->shmSize;
-  TRACE(NCCL_SHM,"Open shmName %s shmSize %d", shmName, info->shmSize);
-  NCCLCHECK(shmOpen(shmName, resources->remShmSize, (void**)&resources->remHostMem, (void**)&resources->devRemHostMem, 0));
-  NCCLCHECK(shmUnlink(shmName));
+  TRACE(NCCL_SHM,"Open shmName %s shmSize %d", shmPath, info->shmSize);
+  NCCLCHECK(ncclShmOpen(shmPath, resources->remShmSize, (void**)&resources->remHostMem, (void**)&resources->devRemHostMem, 0));
+  NCCLCHECK(ncclShmUnlink(shmPath));
   recv->conn.head = &resources->devRemHostMem->head;
 
   int offset = 0;
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
-    recv->conn.buffs[p] = resources->devHostMem->buff + offset;
+    recv->conn.buffs[p] = (char*)(resources->devHostMem+1) + offset;
     offset += recv->comm->buffSizes[p];
   }
   recv->conn.tail = &resources->devHostMem->tail;
   return ncclSuccess;
 }
 
-ncclResult_t shmSendFree(void* transportResources) {
-  struct shmSendResources* resources = (struct shmSendResources*)transportResources;
-  NCCLCHECK(shmClose(resources->hostMem, resources->devHostMem, resources->shmSize));
-  NCCLCHECK(shmClose(resources->remHostMem, resources->devRemHostMem, resources->remShmSize));
+ncclResult_t shmSendFree(struct ncclConnector* send) {
+  struct shmRecvResources* resources = (struct shmRecvResources*)send->transportResources;
+  NCCLCHECK(ncclShmClose(resources->hostMem, resources->devHostMem, resources->shmSize));
+  NCCLCHECK(ncclShmClose(resources->remHostMem, resources->devRemHostMem, resources->remShmSize));
   free(resources);
   return ncclSuccess;
 }
 
-ncclResult_t shmRecvFree(void* transportResources) {
-  struct shmRecvResources* resources = (struct shmRecvResources*)transportResources;
-  NCCLCHECK(shmClose(resources->hostMem, resources->devHostMem, resources->shmSize));
-  NCCLCHECK(shmClose(resources->remHostMem, resources->devRemHostMem, resources->remShmSize));
+ncclResult_t shmRecvFree(struct ncclConnector* recv) {
+  struct shmRecvResources* resources = (struct shmRecvResources*)recv->transportResources;
+  NCCLCHECK(ncclShmClose(resources->hostMem, resources->devHostMem, resources->shmSize));
+  NCCLCHECK(ncclShmClose(resources->remHostMem, resources->devRemHostMem, resources->remShmSize));
   free(resources);
   return ncclSuccess;
 }
@@ -173,6 +161,6 @@ ncclResult_t shmRecvFree(void* transportResources) {
 struct ncclTransport shmTransport = {
   "SHM",
   shmCanConnect,
-  { shmSendSetup, shmSendConnect, shmSendFree, NULL },
-  { shmRecvSetup, shmRecvConnect, shmRecvFree, NULL }
+  { shmSendSetup, shmSendConnect, shmSendFree, NULL, NULL, NULL, NULL, NULL },
+  { shmRecvSetup, shmRecvConnect, shmRecvFree, NULL, NULL, NULL, NULL, NULL }
 };