From 43bd6f5fbfa0fbee0080956780e5959d5fa49e03 Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Mon, 15 Jul 2019 11:45:38 -0700
Subject: [PATCH] Allocate transport memory based on numa node (#97)

---
 src/init.cu | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/init.cu b/src/init.cu
index 3e9f01d397..4059967477 100644
--- a/src/init.cu
+++ b/src/init.cu
@@ -299,6 +299,8 @@ static ncclResult_t fillInfo(struct ncclInfo* info, int rank, uint64_t commHash)
   return ncclSuccess;
 }
 
+bool SetCpuAffinity(int cudaDev, nvmlDevice_t* nvmlDevice);
+
 template <int type>
 static ncclResult_t selectTransport(struct ncclInfo* myInfo, struct ncclInfo* peerInfo, struct ncclConnect* connect, struct ncclTransport** transportRet, struct ncclRing* ring) {
   for (int t=0; t<NTRANSPORTS; t++) {
@@ -307,8 +309,15 @@ static ncclResult_t selectTransport(struct ncclInfo* myInfo, struct ncclInfo* pe
     ncclTvalue_t ret = 0;
     NCCLCHECK(transport->canConnect(&ret, myInfo->tinfo+t, peerInfo->tinfo+t));
     if (ret > 0) {
+      cpu_set_t affinitySave;
+      nvmlDevice_t nvmlDevice;
+      int cudaDev;
+      CUDACHECK(hipGetDevice(&cudaDev));
+      sched_getaffinity(0, sizeof(cpu_set_t), &affinitySave);
+      SetCpuAffinity(cudaDev, &nvmlDevice);
       NCCLCHECK(transportComm->setup(myInfo->tinfo+t, peerInfo->tinfo+t, connect, ring));
       *transportRet = transport;
+      sched_setaffinity(0, sizeof(cpu_set_t), &affinitySave);
       return ncclSuccess;
     }
   }