AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allocate transport memory based on numa node (#97)

Browse Source
This commit is contained in:
Wenkai Du
2019-07-15 11:45:38 -07:00
committed by GitHub
parent 7b6332d3d0
commit 43bd6f5fbf
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/init.cu
+9
View File
@@ -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;
}
}