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

src/include/align.h

@@ -16,4 +16,29 @@
 #define ALIGN_SIZE(size, align) \
   size = ((size + (align) - 1) / (align)) * (align);
 
+#if !__CUDA_ARCH__
+  #ifndef __host__
+    #define __host__
+  #endif
+  #ifndef __device__
+    #define __device__
+  #endif
+#endif
+
+template<typename X, typename Y, typename Z = decltype(X()+Y())>
+__host__ __device__ constexpr Z divUp(X x, Y y) {
+  return (x+y-1)/y;
+}
+
+template<typename X, typename Y, typename Z = decltype(X()+Y())>
+__host__ __device__ constexpr Z roundUp(X x, Y y) {
+  return (x+y-1) - (x+y-1)%y;
+}
+
+// assumes second argument is a power of 2
+template<typename X, typename Z = decltype(X()+int())>
+__host__ __device__ constexpr Z alignUp(X x, int a) {
+  return (x+a-1) & Z(-a);
+}
+
 #endif

src/include/alloc.h

@@ -14,12 +14,13 @@
 #include <sys/mman.h>
 
 template <typename T>
-static ncclResult_t ncclCudaHostCalloc(T** ptr, size_t nelem) {
+static ncclResult_t ncclCudaHostCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
   CUDACHECK(hipHostMalloc(ptr, nelem*sizeof(T), hipHostMallocMapped));
   memset(*ptr, 0, nelem*sizeof(T));
-  INFO(NCCL_ALLOC, "Cuda Host Alloc Size %ld pointer %p", nelem*sizeof(T), *ptr);
+  INFO(NCCL_ALLOC, "%s:%d Cuda Host Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return ncclSuccess;
 }
+#define ncclCudaHostCalloc(...) ncclCudaHostCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
 static inline ncclResult_t ncclCudaHostFree(void* ptr) {
   CUDACHECK(hipHostFree(ptr));
@@ -27,7 +28,7 @@ static inline ncclResult_t ncclCudaHostFree(void* ptr) {
 }
 
 template <typename T>
-static ncclResult_t ncclCalloc(T** ptr, size_t nelem) {
+static ncclResult_t ncclCallocDebug(T** ptr, size_t nelem, const char *filefunc, int line) {
   void* p = malloc(nelem*sizeof(T));
   if (p == NULL) {
     WARN("Failed to malloc %ld bytes", nelem*sizeof(T));
@@ -35,9 +36,10 @@ static ncclResult_t ncclCalloc(T** ptr, size_t nelem) {
   }
   memset(p, 0, nelem*sizeof(T));
   *ptr = (T*)p;
-  INFO(NCCL_ALLOC, "Mem Alloc Size %ld pointer %p", nelem*sizeof(T), *ptr);
+  INFO(NCCL_ALLOC, "%s:%d Mem Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   return ncclSuccess;
 }
+#define ncclCalloc(...) ncclCallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
 struct __attribute__ ((aligned(64))) allocationTracker {
   union {
@@ -53,7 +55,7 @@ static_assert(sizeof(struct allocationTracker) == 64, "allocationTracker must be
 extern struct allocationTracker allocTracker[];
 
 template <typename T>
-static ncclResult_t ncclCudaCalloc(T** ptr, size_t nelem, bool isFineGrain = false) {
+static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr, size_t nelem, bool isFineGrain = false) {
   // Need async stream for P2P pre-connect + CUDA Graph
   hipStream_t stream;
   CUDACHECK(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
@@ -64,7 +66,7 @@ static ncclResult_t ncclCudaCalloc(T** ptr, size_t nelem, bool isFineGrain = fal
   CUDACHECK(hipMemsetAsync(*ptr, 0, nelem*sizeof(T), stream));
   CUDACHECK(hipStreamSynchronize(stream));
   CUDACHECK(hipStreamDestroy(stream));
-  INFO(NCCL_ALLOC, "Cuda Alloc Size %ld pointer %p", nelem*sizeof(T), *ptr);
+  INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   int dev;
   CUDACHECK(hipGetDevice(&dev));
   if (dev < MAX_ALLOC_TRACK_NGPU) {
@@ -73,6 +75,7 @@ static ncclResult_t ncclCudaCalloc(T** ptr, size_t nelem, bool isFineGrain = fal
   }
   return ncclSuccess;
 }
+#define ncclCudaCalloc(...) ncclCudaCallocDebug(__FILE__, __LINE__, __VA_ARGS__)
 
 template <typename T>
 static ncclResult_t ncclCudaMemcpy(T* dst, T* src, size_t nelem) {
@@ -93,7 +96,7 @@ static bool hasFineGrainVramPcie() {
 // Allocate memory to be potentially ibv_reg_mr'd. This needs to be
 // allocated on separate pages as those pages will be marked DONTFORK
 // and if they are shared, that could cause a crash in a child process
-static ncclResult_t ncclIbMalloc(void** ptr, size_t size) {
+static ncclResult_t ncclIbMallocDebug(void** ptr, size_t size, const char *filefunc, int line) {
   size_t page_size = sysconf(_SC_PAGESIZE);
   void* p;
   int size_aligned = ROUNDUP(size, page_size);
@@ -101,8 +104,9 @@ static ncclResult_t ncclIbMalloc(void** ptr, size_t size) {
   if (ret != 0) return ncclSystemError;
   memset(p, 0, size);
   *ptr = p;
-  INFO(NCCL_ALLOC, "Ib Alloc Size %ld pointer %p", size, *ptr);
+  INFO(NCCL_ALLOC, "%s:%d Ib Alloc Size %ld pointer %p", filefunc, line, size, *ptr);
   return ncclSuccess;
 }
+#define ncclIbMalloc(...) ncclIbMallocDebug(__VA_ARGS__, __FILE__, __LINE__)
 
 #endif

src/include/bootstrap.h

@@ -1,5 +1,6 @@
 /*************************************************************************
  * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
+ * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
  ************************************************************************/
@@ -16,6 +17,7 @@ ncclResult_t bootstrapInit(ncclUniqueId* id, int rank, int nranks, void** commSt
 ncclResult_t bootstrapAllGather(void* commState, void* allData, int size);
 ncclResult_t bootstrapSend(void* commState, int peer, int tag, void* data, int size);
 ncclResult_t bootstrapRecv(void* commState, int peer, int tag, void* data, int size);
+ncclResult_t bootstrapBarrier(void* commState, int *ranks, int tag, int rank, int nranks);
 ncclResult_t bootstrapRemAlloc(size_t size, int rank, void* commState, int* id, hipIpcMemHandle_t* ipc, void** ptr);
 ncclResult_t bootstrapRemFree(int id, int rank, void* commState);
 ncclResult_t bootstrapClose(void* commState);

src/include/collectives.h

@@ -1,5 +1,5 @@
 /*************************************************************************
- * Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
  * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
  *
  * See LICENSE.txt for license information
@@ -25,6 +25,9 @@
   extern __device__ __attribute__((noinline)) void NCCL_FUNC_NAME(func, algo, proto, redop, type)(struct ncclWorkElem* args); \
   extern __global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(struct ncclWorkElem first); \
 
+  //extern __device__ void NCCL_FUNC_NAME(func, algo, proto, redop, type)(); \
+  //extern __global__ void NCCL_KERN_NAME(func, algo, proto, redop, type)(ncclWorkElem c); \
+
 #define DECL4(func, algo, redop, type) \
   DECL5(func, algo, SIMPLE, redop, type) \
   DECL5(func, algo, LL,     redop, type) \
@@ -35,6 +38,19 @@
   DECL4(func, TREE,    redop, type) \
   DECL4(func, COLLNET, redop, type)
 
+#if defined(__CUDA_BF16_TYPES_EXIST__)
+#define DECL2(func, redop) \
+  DECL3(func, redop, int8_t) \
+  DECL3(func, redop, uint8_t) \
+  DECL3(func, redop, int32_t) \
+  DECL3(func, redop, uint32_t) \
+  DECL3(func, redop, int64_t) \
+  DECL3(func, redop, uint64_t) \
+  DECL3(func, redop, half) \
+  DECL3(func, redop, float) \
+  DECL3(func, redop, double) \
+  DECL3(func, redop, __nv_bfloat16)
+#else
 #define DECL2(func, redop) \
   DECL3(func, redop, int8_t) \
   DECL3(func, redop, uint8_t) \
@@ -46,12 +62,14 @@
   DECL3(func, redop, float) \
   DECL3(func, redop, double) \
   DECL3(func, redop, rccl_bfloat16)
+#endif
 
 #define DECL(func) \
   DECL2(func, Sum) \
   DECL2(func, Prod) \
   DECL2(func, Min) \
-  DECL2(func, Max)
+  DECL2(func, Max) \
+  DECL2(func, Avg) \
 
 #define DECL_ALL \
   DECL2(Broadcast, Sum) \

src/include/comm.h

@@ -82,6 +82,7 @@ struct ncclComm {
   int nRanks;  // number of GPUs in communicator
   int cudaDev; // my cuda device index
   int64_t busId;   // my PCI bus ID in int format
+  cpu_set_t cpuAffinity; // CPU affinity of the GPU
 
   int node;
   int nNodes;
@@ -162,11 +163,13 @@ struct ncclComm {
   struct ncclInfo* asyncOps;
   int asyncOpCount;
   size_t asyncTotalSize;
+  ssize_t channelSize;
   int lastChannel;
+  enum { ROUND_ROBIN, SHORTEST_QUEUE } asyncAllocMode;
 
   //list of async p2p operation queued in a group semantics
-  struct ncclP2Plist* p2pSends;
-  struct ncclP2Plist* p2pRecvs;
+  ncclP2Plist** p2pSends;
+  ncclP2Plist** p2pRecvs;
   int p2pSendCount;
   int p2pRecvCount;
 

src/include/core.h

@@ -1,6 +1,10 @@
 /*************************************************************************
+<<<<<<< HEAD
  * Copyright (c) 2015-2020, NVIDIA CORPORATION. All rights reserved.
  * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved.
+=======
+ * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved.
+>>>>>>> nccl/master
  *
  * See LICENSE.txt for license information
  ************************************************************************/

src/include/devcomm.h

@@ -27,7 +27,7 @@
 
 
 #define NCCL_NUM_FUNCTIONS 5 // SendRecv not included for now
-typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv} ncclFunc_t;
+typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclNumFuncs} ncclFunc_t;
 extern const char* ncclFuncStr[NCCL_NUM_FUNCTIONS+1];
 
 #define NCCL_NUM_ALGORITHMS 3 // Tree/Ring/CollNet
@@ -142,6 +142,8 @@ struct ncclRing {
   // devices. Ordered from current device.
   int* userRanks;
   int* devUserRanks;
+
+  int index; // This rank's index in the ring
 };
 
 
@@ -173,7 +175,7 @@ struct ncclPeer {
 struct ncclDevComm;
 
 #pragma pack(push)  /* push current alignment to stack */
-#pragma pack(4)     /* set alignment to 4 bytes boundary */
+#pragma pack(8)     /* set alignment to 4 bytes boundary */
 #define NCCL_MAX_WORK_ELEMENTS 1
 #define NCCL_MAX_GROUPS (NCCL_MAX_NTHREADS/WARP_SIZE)
 
@@ -255,6 +257,7 @@ struct ncclChannel {
       // Operation list for aggregation
       struct ncclWork* workFifo;
       int workCount;
+      size_t totalSize;
       uint64_t workFifoTail; // Only used by CPU
 
 #ifdef ENABLE_PROFILING
@@ -285,7 +288,6 @@ struct ncclProf {
     struct {
       uint64_t total_cycle;
       uint64_t wait_cycle[MAXCHANNELS];      // total wait cycle
-      uint64_t wait_recv_cycle[MAXCHANNELS]; // recv wait cycle
       // primtive cycles
       uint64_t send_cycle;
       uint64_t directSend_cycle;
@@ -376,4 +378,9 @@ struct ncclDevComm {
 #endif
 };
 
+struct ncclDevCommAndChannels {
+  ncclDevComm comm;
+  ncclChannel channels[MAXCHANNELS];
+};
+
 #endif

src/include/enqueue.h

@@ -11,6 +11,9 @@
 #include "group.h"
 #include "collectives.h"
 
+#define NCCL_MIN_CHANNEL_SIZE (NCCL_LL_THREAD_THRESHOLD*64)
+#define NCCL_AGG_CHANNEL_SIZE (1LL << 21) /* 2 MiB, ideal per-channel size to fully utilize bandwidth */
+
 size_t ncclKernMaxLocalSize();
 ncclResult_t ncclEnqueueCheck(struct ncclInfo* info);
 ncclResult_t ncclCpuBarrierIn(struct ncclComm* comm, int* isLast);
@@ -31,39 +34,22 @@ ncclResult_t ncclCudaGraphHostSetup(ncclComm_t comm, cudaGraph_t graph);
 struct ncclQueueElem {
   struct ncclWorkElem work;
   struct ncclProxyArgs proxyArgs;
-  struct ncclQueueElem* next;
 };
 
-// Store enqueue elements in a list
-struct ncclQueueElemList {
-  struct ncclQueueElem* head;
-  struct ncclQueueElem* tail;
-};
+typedef ncclRecyclableList<struct ncclQueueElem> ncclQueueElemList;
 
 // Structure passed to CUDA graph
 struct ncclQueueInfo {
   ncclComm_t comm;
   int maxChannels;    // Dynamic version of gridDim
   ncclResult_t ret;   // Return value of host setup call
-  struct ncclQueueElemList elemList;
+  ncclQueueElemList* elemList;
 };
 
-// Get next element from enqueue list
-static ncclResult_t ncclAddQueueElem(struct ncclQueueInfo* eqInfo, struct ncclQueueElem** elemOut) {
-  if (eqInfo == NULL) return ncclInternalError;
-  struct ncclQueueElemList* list = &eqInfo->elemList;
-  if (list->tail != NULL) {
-    *elemOut = list->tail;
-    memset(*elemOut, 0, sizeof(struct ncclWorkElem) + sizeof(struct ncclProxyArgs));
-  } else {
-    NCCLCHECK(ncclCalloc(&list->tail, 1));
-    *elemOut = list->tail;
-    list->head = list->tail;
-  }
-  if (list->tail->next == NULL) {
-    NCCLCHECK(ncclCalloc(&list->tail->next, 1));
-  }
-  list->tail = list->tail->next;
+static ncclResult_t ncclCreateQueueInfo(struct ncclQueueInfo** eqInfo, ncclComm_t comm) {
+  NCCLCHECK(ncclCalloc(eqInfo, 1));
+  (*eqInfo)->comm = comm;
+  (*eqInfo)->elemList = new ncclQueueElemList();
   return ncclSuccess;
 }
 
@@ -72,7 +58,7 @@ static ncclResult_t ncclResetQueueInfo(struct ncclQueueInfo* eqInfo) {
   if (eqInfo == NULL) return ncclInternalError;
   eqInfo->maxChannels = 0;
   eqInfo->ret = ncclSuccess;
-  eqInfo->elemList.tail = eqInfo->elemList.head;
+  eqInfo->elemList->recycle();
   return ncclSuccess;
 }
 
@@ -81,12 +67,7 @@ static ncclResult_t ncclResetQueueInfo(struct ncclQueueInfo* eqInfo) {
 static void ncclDestroyQueueInfo(void* ptr) {
   if (ptr == NULL) return;
   struct ncclQueueInfo* eqInfo = (struct ncclQueueInfo*)ptr;
-  struct ncclQueueElem* head = eqInfo->elemList.head;
-  while (head != NULL) {
-    struct ncclQueueElem* temp = head;
-    head = head->next;
-    free(temp);
-  }
+  delete eqInfo->elemList;
   free(eqInfo);
 }
 #endif // End include guard

src/include/graph.h

@@ -14,6 +14,7 @@
 #include <stdlib.h>
 #include <ctype.h>
 #include <stdio.h>
+#include <sched.h>
 
 ncclResult_t ncclTopoCudaPath(int cudaDev, char** path);
 
@@ -36,8 +37,8 @@ ncclResult_t ncclTopoCheckGdr(struct ncclTopoSystem* topo, int64_t busId, int ne
 ncclResult_t ncclTopoGetIntraNetDev(struct ncclTopoSystem* system, int rank, struct ncclTopoGraph* graph, int channelId, int type, int* dev);
 ncclResult_t ncclTopoGetLinkType(struct ncclTopoSystem* system, int cudaDev1, int cudaDev2, bool* isXGMI, bool direct_only=false);
 
-// Set CPU affinity
-ncclResult_t ncclTopoSetAffinity(struct ncclTopoSystem* system, int rank);
+// Find CPU affinity
+ncclResult_t ncclTopoGetCpuAffinity(struct ncclTopoSystem* system, int rank, cpu_set_t* affinity);
 
 #define NCCL_TOPO_CPU_ARCH_X86 1
 #define NCCL_TOPO_CPU_ARCH_POWER 2
@@ -107,6 +108,6 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
 
 ncclResult_t ncclTopoTuneModel(struct ncclComm* comm, int minCompCap, int maxCompCap, struct ncclTopoGraph* treeGraph, struct ncclTopoGraph* ringGraph, struct ncclTopoGraph* collNetGraph);
 #include "info.h"
-ncclResult_t ncclTopoGetAlgoTime(struct ncclInfo* info, int algorithm, int protocol, float* time);
+ncclResult_t ncclTopoGetAlgoTime(struct ncclInfo* info, int algorithm, int protocol, int numPipeOps, float* time);
 
 #endif

src/include/p2p.h

@@ -12,32 +12,16 @@
 struct ncclP2Pinfo {
   void* buff;
   ssize_t nbytes;
-  struct ncclP2Pinfo* next;
 };
 
-struct ncclP2Plist {
-  struct ncclP2Pinfo *head;
-  struct ncclP2Pinfo *tail;
-};
+typedef ncclRecyclableList<struct ncclP2Pinfo> ncclP2Plist;
 
-static ncclResult_t enqueueP2pInfo(ncclP2Plist* p2p, void* buff, ssize_t nBytes) {
-  if (p2p == NULL) return ncclInternalError;
+static ncclResult_t ncclSaveP2pInfo(ncclP2Plist* &p2p, void* buff, ssize_t nBytes) {
+  if (p2p == NULL) p2p = new ncclP2Plist();
   struct ncclP2Pinfo* next;
-  NCCLCHECK(ncclCalloc(&next, 1));
+  NCCLCHECK(p2p->getNewElem(&next));
   next->buff = buff;
   next->nbytes = nBytes;
-  if (p2p->tail != NULL) p2p->tail->next = next;
-  p2p->tail = next;
-  if (p2p->head == NULL) p2p->head = next;
-  return ncclSuccess;
-}
-
-static ncclResult_t dequeueP2pInfo(ncclP2Plist* p2p) {
-  if (p2p == NULL) return ncclInternalError;
-  struct ncclP2Pinfo* temp = p2p->head;
-  p2p->head = p2p->head->next;
-  if (p2p->tail == temp) p2p->tail = NULL;
-  free(temp);
   return ncclSuccess;
 }
 #endif

src/include/socket.h

@@ -31,12 +31,13 @@ union socketAddress {
   struct sockaddr_in6 sin6;
 };
 
-/* Format a string representation of a (struct sockaddr *) socket address using getnameinfo()
+/* Format a string representation of a (union socketAddress *) socket address using getnameinfo()
  *
  * Output: "IPv4/IPv6 address<port>"
  */
-static inline const char *socketToString(struct sockaddr *saddr, char *buf) {
-  if (buf == NULL || saddr == NULL) return NULL;
+static inline const char *socketToString(union socketAddress *addr, char *buf) {
+  if (buf == NULL || addr == NULL) return NULL;
+  struct sockaddr *saddr = &addr->sa;
   if (saddr->sa_family != AF_INET && saddr->sa_family != AF_INET6) { buf[0]='\0'; return buf; }
   char host[NI_MAXHOST], service[NI_MAXSERV];
   (void) getnameinfo(saddr, sizeof(union socketAddress), host, NI_MAXHOST, service, NI_MAXSERV, NI_NUMERICHOST|NI_NUMERICSERV);
@@ -44,8 +45,9 @@ static inline const char *socketToString(struct sockaddr *saddr, char *buf) {
   return buf;
 }
 
-static inline uint16_t socketToPort(struct sockaddr *saddr) {
-  return ntohs(saddr->sa_family == AF_INET ? ((struct sockaddr_in*)saddr)->sin_port : ((struct sockaddr_in6*)saddr)->sin6_port);
+static inline uint16_t socketToPort(union socketAddress *addr) {
+  struct sockaddr *saddr = &addr->sa;
+  return ntohs(saddr->sa_family == AF_INET ? addr->sin.sin_port : addr->sin6.sin6_port);
 }
 
 /* Allow the user to force the IPv4/IPv6 interface selection */
@@ -86,7 +88,7 @@ static int findInterfaces(const char* prefixList, char* names, union socketAddre
     if (family != AF_INET && family != AF_INET6)
       continue;
 
-    TRACE(NCCL_INIT|NCCL_NET,"Found interface %s:%s", interface->ifa_name, socketToString(interface->ifa_addr, line));
+    TRACE(NCCL_INIT|NCCL_NET,"Found interface %s:%s", interface->ifa_name, socketToString((union socketAddress *)interface->ifa_addr, line));
 
     /* Allow the caller to force the socket family type */
     if (sock_family != -1 && family != sock_family)
@@ -195,13 +197,13 @@ static int findInterfaceMatchSubnet(char* ifNames, union socketAddress* localAdd
     // Store the interface name
     strncpy(ifNames+found*ifNameMaxSize, interface->ifa_name, ifNameMaxSize);
 
-    TRACE(NCCL_INIT|NCCL_NET,"NET : Found interface %s:%s in the same subnet as remote address %s", interface->ifa_name, socketToString(&(localAddrs[found].sa), line), socketToString(&(remoteAddr->sa), line_a));
+    TRACE(NCCL_INIT|NCCL_NET,"NET : Found interface %s:%s in the same subnet as remote address %s", interface->ifa_name, socketToString(localAddrs+found, line), socketToString(remoteAddr, line_a));
     found++;
     if (found == maxIfs) break;
   }
 
   if (found == 0) {
-    WARN("Net : No interface found in the same subnet as remote address %s", socketToString(&(remoteAddr->sa), line_a));
+    WARN("Net : No interface found in the same subnet as remote address %s", socketToString(remoteAddr, line_a));
   }
   freeifaddrs(interfaces);
   return found;
@@ -339,7 +341,7 @@ static ncclResult_t createListenSocket(int *fd, union socketAddress *localAddr)
   localAddr->sin.sin_port = htons(port++);
 #endif
 
-  if (socketToPort(&localAddr->sa)) {
+  if (socketToPort(localAddr)) {
     // Port is forced by env. Make sure we get the port.
     int opt = 1;
 #if defined(SO_REUSEPORT)
@@ -358,7 +360,7 @@ static ncclResult_t createListenSocket(int *fd, union socketAddress *localAddr)
 
 #ifdef ENABLE_TRACE
   char line[SOCKET_NAME_MAXLEN+1];
-  TRACE(NCCL_INIT|NCCL_NET,"Listening on socket %s", socketToString(&localAddr->sa, line));
+  TRACE(NCCL_INIT|NCCL_NET,"Listening on socket %s", socketToString(localAddr, line));
 #endif
 
   /* Put the socket in listen mode
@@ -370,10 +372,12 @@ static ncclResult_t createListenSocket(int *fd, union socketAddress *localAddr)
 }
 
 static ncclResult_t connectAddress(int* fd, union socketAddress* remoteAddr) {
+  char line[SOCKET_NAME_MAXLEN+1];
   /* IPv4/IPv6 support */
   int family = remoteAddr->sa.sa_family;
   if (family != AF_INET && family != AF_INET6) {
-    WARN("Error : connecting to address with family %d is neither AF_INET(%d) nor AF_INET6(%d)", family, AF_INET, AF_INET6);
+    WARN("Net : connecting to address %s with family %d is neither AF_INET(%d) nor AF_INET6(%d)",
+         socketToString(remoteAddr, line), family, AF_INET, AF_INET6);
     return ncclInternalError;
   }
   int salen = (family == AF_INET) ? sizeof(sockaddr_in) : sizeof(sockaddr_in6);
@@ -392,8 +396,7 @@ static ncclResult_t connectAddress(int* fd, union socketAddress* remoteAddr) {
     SYSCHECK(setsockopt(*fd, SOL_SOCKET, SO_SNDBUF, (char*)&bufsize, sizeof(int)), "setsockopt");
     SYSCHECK(setsockopt(*fd, SOL_SOCKET, SO_RCVBUF, (char*)&bufsize, sizeof(int)), "setsockopt");*/
 
-  char line[SOCKET_NAME_MAXLEN+1];
-  TRACE(NCCL_INIT|NCCL_NET,"Connecting to socket %s", socketToString(&remoteAddr->sa, line));
+  TRACE(NCCL_INIT|NCCL_NET,"Connecting to socket %s", socketToString(remoteAddr, line));
 
   int ret;
   int timedout_retries = 0;
@@ -409,25 +412,26 @@ retry:
       goto retry;
     }
   }
-  WARN("Connect to %s failed : %s", socketToString(&remoteAddr->sa, line), strerror(errno));
+  WARN("Net : Connect to %s failed : %s", socketToString(remoteAddr, line), strerror(errno));
   return ncclSystemError;
 }
 
 #define NCCL_SOCKET_SEND 0
 #define NCCL_SOCKET_RECV 1
-static ncclResult_t socketProgressOpt(int op, int fd, void* ptr, int size, int* offset, int block) {
+static ncclResult_t socketProgressOpt(int op, int fd, union socketAddress *addr, void* ptr, int size, int* offset, int block) {
   int bytes = 0;
   char* data = (char*)ptr;
+  char line[SOCKET_NAME_MAXLEN+1];
   do {
     if (op == NCCL_SOCKET_RECV) bytes = recv(fd, data+(*offset), size-(*offset), block ? 0 : MSG_DONTWAIT);
     if (op == NCCL_SOCKET_SEND) bytes = send(fd, data+(*offset), size-(*offset), block ? 0 : MSG_DONTWAIT);
     if (op == NCCL_SOCKET_RECV && bytes == 0) {
-      WARN("Net : Connection closed by remote peer");
+      WARN("Net : Connection closed by remote peer %s", socketToString(addr, line));
       return ncclSystemError;
     }
     if (bytes == -1) {
       if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN) {
-        WARN("Call to recv failed : %s", strerror(errno));
+        WARN("Net : Call to recv from %s failed : %s", socketToString(addr, line), strerror(errno));
         return ncclSystemError;
       } else {
         bytes = 0;
@@ -438,25 +442,25 @@ static ncclResult_t socketProgressOpt(int op, int fd, void* ptr, int size, int*
   return ncclSuccess;
 }
 
-static ncclResult_t socketProgress(int op, int fd, void* ptr, int size, int* offset) {
-  return socketProgressOpt(op, fd, ptr, size, offset, 0);
+static ncclResult_t socketProgress(int op, int fd, union socketAddress *addr, void* ptr, int size, int* offset) {
+  return socketProgressOpt(op, fd, addr, ptr, size, offset, 0);
 }
 
-static ncclResult_t socketWait(int op, int fd, void* ptr, int size, int* offset) {
+static ncclResult_t socketWait(int op, int fd, union socketAddress *addr, void* ptr, int size, int* offset) {
   while (*offset < size)
-    NCCLCHECK(socketProgressOpt(op, fd, ptr, size, offset, 1));
+    NCCLCHECK(socketProgressOpt(op, fd, addr, ptr, size, offset, 1));
   return ncclSuccess;
 }
 
-static ncclResult_t socketSend(int fd, void* ptr, int size) {
+static ncclResult_t socketSend(int fd, union socketAddress *addr, void* ptr, int size) {
   int offset = 0;
-  NCCLCHECK(socketWait(NCCL_SOCKET_SEND, fd, ptr, size, &offset));
+  NCCLCHECK(socketWait(NCCL_SOCKET_SEND, fd, addr, ptr, size, &offset));
   return ncclSuccess;
 }
 
-static ncclResult_t socketRecv(int fd, void* ptr, int size) {
+static ncclResult_t socketRecv(int fd, union socketAddress *addr, void* ptr, int size) {
   int offset = 0;
-  NCCLCHECK(socketWait(NCCL_SOCKET_RECV, fd, ptr, size, &offset));
+  NCCLCHECK(socketWait(NCCL_SOCKET_RECV, fd, addr, ptr, size, &offset));
   return ncclSuccess;
 }
 

src/include/utils.h

@@ -37,4 +37,76 @@ static long log2i(long n) {
  return l;
 }
 
+// Recyclable list that avoids frequent malloc/free
+template<typename T>
+struct ncclListElem {
+  T data;
+  struct ncclListElem* next;
+};
+
+template<typename T>
+class ncclRecyclableList {
+ private:
+  struct ncclListElem<T>* head;
+  struct ncclListElem<T>* tail;
+  struct ncclListElem<T>* cursor;
+  int n;
+
+ public:
+  ncclRecyclableList() {
+    tail = cursor = head = NULL;
+    n = 0;
+  }
+
+  int count() const { return n; }
+
+  // Get a new element from the list and return pointer
+  ncclResult_t getNewElem(T** dataOut) {
+    if (tail != NULL) {
+      *dataOut = &tail->data;
+      memset(*dataOut, 0, sizeof(T));
+    } else {
+      NCCLCHECK(ncclCalloc(&tail, 1));
+      *dataOut = &tail->data;
+      cursor = head = tail;
+    }
+    if (tail->next == NULL) {
+      NCCLCHECK(ncclCalloc(&tail->next, 1));
+    }
+    tail = tail->next;
+    n += 1;
+    return ncclSuccess;
+  }
+
+  T* begin() {
+    if (head == NULL || head == tail) return NULL;
+    cursor = head->next;
+    return &head->data;
+  }
+
+  // Get next element from the list during an iteration
+  T* getNext() {
+    // tail always points to the next element to be enqueued
+    // hence does not contain valid data
+    if (cursor == NULL || cursor == tail) return NULL;
+    T* rv = &cursor->data;
+    cursor = cursor->next;
+    return rv;
+  }
+
+  // Recycle the list without freeing the space
+  void recycle() {
+    tail = cursor = head;
+    n = 0;
+  }
+
+  ~ncclRecyclableList() {
+    while (head != NULL) {
+      struct ncclListElem<T>* temp = head;
+      head = head->next;
+      free(temp);
+    }
+  }
+};
+
 #endif