Use posix_memalign for network buffer allocation on host memory (#221)

* Use posix_memalign for network buffer allocation on host memory

* ib-test: add ability to specify run iterations

* ib-test: define iterations as multiple of default cycles

* Add checking to posix_memalign return value

src/graph/topo.cc

@@ -20,6 +20,7 @@
 #endif
 #include "xml.h"
 #include "cpuset.h"
+#include <numa.h>
 
 #define BUSID_SIZE (sizeof("0000:00:00.0"))
 #define BUSID_REDUCED_SIZE (sizeof("0000:00"))
@@ -733,6 +734,11 @@ ncclResult_t ncclTopoSetAffinity(struct ncclTopoSystem* system, int rank) {
     NCCLCHECK(ncclCpusetToStr(&finalMask, affinityStr));
     INFO(NCCL_INIT, "Setting affinity for GPU %d to %s", gpu->gpu.dev, affinityStr);
     SYSCHECK(sched_setaffinity(0, sizeof(cpu_set_t), &finalMask), "sched_setaffinity");
+    int ret = numa_run_on_node(cpu->id);
+    if (ret != 0)
+      INFO(NCCL_INIT, "Failed to run on numa node %ld", cpu->id);
+    else
+      INFO(NCCL_INIT, "Thread is set to run on numa node %ld", cpu->id);
   }
   return ncclSuccess;
 }

src/include/devcomm.h

@@ -92,6 +92,9 @@ static_assert(NCCL_LL_CLEAN_MASK % NCCL_STEPS == 0, "Invalid NCCL_LL_CLEAN_MASK
 struct ncclConnInfo {
   // Regular comm mechanism
   char *buffs[NCCL_NUM_PROTOCOLS]; // Local for recv, remote for send
+#ifdef USE_MEMALIGN
+  char *hostBuffs[NCCL_NUM_PROTOCOLS]; // Local for recv, remote for send
+#endif
   uint64_t *tail;     // Local for recv, remote for send
   uint64_t *head;     // Local for send, remote for recv
   uint64_t *opCountLoc; // opCount of local rank

src/transport/net.cc

@@ -26,6 +26,9 @@ struct netSendResources {
   int netDev;
   int useGdr;
   char* buffers[LOC_COUNT];
+#ifdef USE_MEMALIGN
+  char* hostBuffer;
+#endif
   int buffSizes[LOC_COUNT];
   void* mhandles[LOC_COUNT];
   void** mhandlesProto[NCCL_NUM_PROTOCOLS];
@@ -93,10 +96,24 @@ ncclResult_t netSendSetup(struct ncclTopoSystem* topo, struct ncclTopoGraph* gra
   }
   char line[16];
   if (resources->buffSizes[LOC_HOSTMEM]) {
+#ifdef USE_MEMALIGN
+    int page_size = getpagesize();
+    if (posix_memalign((void **)&resources->hostBuffer, page_size, resources->buffSizes[LOC_HOSTMEM])) {
+      WARN("Failed to posix_memalign %d bytes", resources->buffSizes[LOC_HOSTMEM]);
+      return ncclSystemError;
+    }
+    CUDACHECK(hipHostRegister(resources->hostBuffer, resources->buffSizes[LOC_HOSTMEM], hipHostRegisterMapped));
+    CUDACHECK(hipHostGetDevicePointer((void **)resources->buffers+LOC_HOSTMEM, resources->hostBuffer, 0));
+#else
     NCCLCHECK(ncclCudaHostCalloc(resources->buffers+LOC_HOSTMEM, resources->buffSizes[LOC_HOSTMEM]));
+#endif
     int status[1] = {-1};
     line[0]= 0;
+#ifdef USE_MEMALIGN
+    if (!move_pages(0, 1, (void **)&resources->hostBuffer, NULL, status, 0))
+#else
     if (!move_pages(0, 1, (void **)resources->buffers+LOC_HOSTMEM, NULL, status, 0))
+#endif
       sprintf(line, "/MEM%d", status[0]);
   }
 
@@ -105,6 +122,12 @@ ncclResult_t netSendSetup(struct ncclTopoSystem* topo, struct ncclTopoGraph* gra
   for (int p=0; p<NCCL_NUM_PROTOCOLS; p++) {
     resources->mhandlesProto[p] = resources->mhandles+protoLoc[p];
     send->conn.buffs[p] = resources->buffers[protoLoc[p]] + offsets[protoLoc[p]];
+#ifdef USE_MEMALIGN
+    if (protoLoc[p] == LOC_HOSTMEM)
+      send->conn.hostBuffs[p] = resources->hostBuffer + offsets[protoLoc[p]];
+    else
+      send->conn.hostBuffs[p] = send->conn.buffs[p];
+#endif
     offsets[protoLoc[p]] += buffSizes[p];
   }
 
@@ -182,7 +205,11 @@ ncclResult_t netSendConnect(struct ncclConnect* connectInfo, int nranks, int ran
     NCCLCHECK(ncclNetRegMr(resources->netSendComm, resources->buffers[LOC_DEVMEM], resources->buffSizes[LOC_DEVMEM], NCCL_PTR_CUDA, &resources->mhandles[LOC_DEVMEM]));
   }
   if (resources->buffSizes[LOC_HOSTMEM]) {
+#ifdef USE_MEMALIGN
+    NCCLCHECK(ncclNetRegMr(resources->netSendComm, resources->hostBuffer, resources->buffSizes[LOC_HOSTMEM], NCCL_PTR_HOST, &resources->mhandles[LOC_HOSTMEM]));
+#else
     NCCLCHECK(ncclNetRegMr(resources->netSendComm, resources->buffers[LOC_HOSTMEM], resources->buffSizes[LOC_HOSTMEM], NCCL_PTR_HOST, &resources->mhandles[LOC_HOSTMEM]));
+#endif
   }
   return ncclSuccess;
 }
@@ -213,7 +240,12 @@ ncclResult_t netSendFree(void* transportResources) {
     if (resources->buffers[l])
       NCCLCHECK(ncclNetDeregMr(resources->netSendComm, resources->mhandles[l]));
   }
+#ifdef USE_MEMALIGN
+  CUDACHECK(hipHostUnregister(resources->hostBuffer));
+  free(resources->hostBuffer);
+#else
   NCCLCHECK(ncclCudaHostFree(resources->buffers[LOC_HOSTMEM]));
+#endif
   CUDACHECK(hipFree(resources->buffers[LOC_DEVMEM]));
   NCCLCHECK(ncclNetCloseSend(resources->netSendComm));
   free(resources);
@@ -251,7 +283,11 @@ ncclResult_t netSendProxy(struct ncclProxyArgs* args) {
   if (args->state == ncclProxyOpProgress) {
     int p = args->protocol;
     int stepSize = args->connector->comm->buffSizes[p] / NCCL_STEPS;
+#ifdef USE_MEMALIGN
+    char* localBuff = args->connector->conn.hostBuffs[p];
+#else
     char* localBuff = args->connector->conn.buffs[p];
+#endif
     void* mhandle = *(resources->mhandlesProto[p]);
     args->idle = 1;
     if (args->head < args->end) {

src/transport/net_ib.cc

@@ -291,7 +291,7 @@ struct ncclIbListenComm {
   int fd;
 };
 
-struct ncclIbSendFifo {
+struct alignas(64) ncclIbSendFifo {
   uint64_t addr;
   int      size;
   uint32_t seq;