Determine fine grained memory availability at RCCL bootstrapping (#471)

2021-11-19 08:12:53 -08:00
parent 7b8b54955b
commit e9bf01fb7e
5 changed files with 14 additions and 24 deletions
@@ -613,7 +613,18 @@ static ncclResult_t fillInfo(struct ncclComm* comm, struct ncclPeerInfo* info, u

  info->busId = comm->busId;

-  NCCLCHECK(ncclGpuGdrSupport(&info->gdrSupport));
+  // detect if fine grained memory is available on this GPU
+  int *ptr;
+  if (hipExtMallocWithFlags((void**)&ptr, sizeof(int), hipDeviceMallocFinegrained) == hipSuccess) {
+    CUDACHECK(hipFree(ptr));
+    info->hasFineGrain = true;
+    NCCLCHECK(ncclGpuGdrSupport(&info->gdrSupport));
+  }
+  else {
+    info->hasFineGrain = false;
+    info->gdrSupport = 0;
+  }
+
  return ncclSuccess;
 }