From 9ddf0e06498d4b85082f767e616a8ca4d3aaec9b Mon Sep 17 00:00:00 2001
From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com>
Date: Mon, 17 Oct 2022 10:24:59 -0700
Subject: [PATCH] Support P2P with invisible devices (#636)

* Support P2P with invisible devices

* Update copyright year
---
 src/graph/xml.cc            | 29 ++++++++++++-----------------
 src/include/rocm_smi_wrap.h |  4 +++-
 src/misc/rocm_smi_wrap.cc   | 34 ++++++++++++++++++++++++----------
 src/transport/p2p.cc        |  2 +-
 4 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/src/graph/xml.cc b/src/graph/xml.cc
index 9b4a610aed..99dd5ab9f7 100644
--- a/src/graph/xml.cc
+++ b/src/graph/xml.cc
@@ -602,7 +602,7 @@ ncclResult_t ncclTopoGetXmlFromSys(struct ncclXmlNode* pciNode, struct ncclXml*
   return ncclSuccess;
 }
 
-ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvmlDev, struct ncclXml* xml, struct ncclXmlNode** gpuNodeRet) {
+ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev, struct ncclXml* xml, struct ncclXmlNode** gpuNodeRet) {
   struct ncclXmlNode* gpuNode = NULL;
   NCCLCHECK(xmlGetSub(pciNode, "gpu", &gpuNode));
   if (gpuNode == NULL) NCCLCHECK(xmlAddNode(xml, pciNode, "gpu", &gpuNode));
@@ -612,12 +612,12 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvm
   int dev = -1;
   NCCLCHECK(xmlGetAttrIndex(gpuNode, "dev", &index));
   if (index == -1) {
-    if (nvmlDev == NULL) {
+    if (rocmDev == -1) {
       const char* busId;
       NCCLCHECK(xmlGetAttr(pciNode, "busid", &busId));
       if (busId == NULL || hipDeviceGetByPCIBusId(&dev, busId) != hipSuccess) dev = -1;
     } else {
-      NCCLCHECK(ncclNvmlDeviceGetIndex(nvmlDev, (unsigned int*)&dev));
+      dev = rocmDev;
     }
     NCCLCHECK(xmlSetAttrInt(gpuNode, "dev", dev));
   }
@@ -627,13 +627,9 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvm
   NCCLCHECK(xmlGetAttrIndex(gpuNode, "sm", &index));
   if (index == -1) {
     int cudaMajor, cudaMinor;
-    if (nvmlDev == NULL) {
-      hipDeviceProp_t devProp;
-      CUDACHECK(hipGetDeviceProperties(&devProp, dev));
-      cudaMajor = devProp.major; cudaMinor = devProp.minor;
-    } else {
-      NCCLCHECK(ncclNvmlDeviceGetCudaComputeCapability(nvmlDev, &cudaMajor, &cudaMinor));
-    }
+    hipDeviceProp_t devProp;
+    CUDACHECK(hipGetDeviceProperties(&devProp, 0));
+    cudaMajor = devProp.major; cudaMinor = devProp.minor;
     NCCLCHECK(xmlSetAttrInt(gpuNode, "sm", cudaMajor*10+cudaMinor));
   }
   int sm;
@@ -643,7 +639,7 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvm
   NCCLCHECK(xmlGetAttrIndex(gpuNode, "gcn", &index));
   if (index == -1) {
     hipDeviceProp_t devProp;
-    CUDACHECK(hipGetDeviceProperties(&devProp, dev));
+    CUDACHECK(hipGetDeviceProperties(&devProp, 0));
     gcn = devProp.gcnArch;
     NCCLCHECK(xmlSetAttrInt(gpuNode, "gcn", gcn));
   }
@@ -653,7 +649,7 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvm
   NCCLCHECK(xmlGetAttrIndex(gpuNode, "arch", &index));
   if (index == -1) {
     hipDeviceProp_t devProp;
-    CUDACHECK(hipGetDeviceProperties(&devProp, dev));
+    CUDACHECK(hipGetDeviceProperties(&devProp, 0));
     memcpy(&arch.arch, &devProp.arch, sizeof(hipDeviceArch_t));
     NCCLCHECK(xmlSetAttrInt(gpuNode, "arch", arch.value));
   }
@@ -665,9 +661,8 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvm
 #if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__)
     const char* busId;
     NCCLCHECK(xmlGetAttr(pciNode, "busid", &busId));
-    if (busId == NULL || hipDeviceGetByPCIBusId(&dev, busId) != hipSuccess) return ncclInternalError;
-    int deviceCnt;
-    CUDACHECK(hipGetDeviceCount(&deviceCnt));
+    uint32_t deviceCnt;
+    NCCLCHECK(rocm_smi_getNumDevice(&deviceCnt));
     for (int i=0; i<deviceCnt; i++) {
       if (i != dev) {
         RSMI_IO_LINK_TYPE rsmi_type;
@@ -675,7 +670,7 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, nvmlDevice_t nvm
         if (rocm_smi_getLinkInfo(dev, i, &rsmi_type, &hops, &count) == ncclSuccess) {
           if (rsmi_type == RSMI_IOLINK_TYPE_XGMI && hops == 1) {
             char busIdStr[] = "00000000:00:00.0";
-            CUDACHECK(hipDeviceGetPCIBusId(busIdStr, sizeof(busIdStr), i));
+            NCCLCHECK(rocm_smi_getDevicePciBusIdString(i, busIdStr, sizeof(busIdStr)));
             char lowerId[NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE];
             for (int c=0; c<NVML_DEVICE_PCI_BUS_ID_BUFFER_SIZE; c++) {
               lowerId[c] = tolower(busIdStr[c]);
@@ -779,7 +774,7 @@ ncclResult_t ncclTopoFillGpu(struct ncclXml* xml, const char* busId, struct nccl
   if (rocmsmiInit == 1) {
     if (rocm_smi_getDeviceIndexByPciBusId(busId, &devIndex) != ncclSuccess) devIndex = -1;
   }
-  NCCLCHECK(ncclTopoGetXmlFromGpu(node, NULL, xml, gpuNode));
+  NCCLCHECK(ncclTopoGetXmlFromGpu(node, devIndex, xml, gpuNode));
 #else
   nvmlDevice_t nvmlDev = NULL;
   if (ncclNvmlDeviceGetHandleByPciBusId(busId, &nvmlDev) != ncclSuccess) nvmlDev = NULL;
diff --git a/src/include/rocm_smi_wrap.h b/src/include/rocm_smi_wrap.h
index a118d56a26..417e280a83 100644
--- a/src/include/rocm_smi_wrap.h
+++ b/src/include/rocm_smi_wrap.h
@@ -1,5 +1,5 @@
 /*
-Copyright (c) 2021 Advanced Micro Devices, Inc. All rights reserved.
+Copyright (c) 2021-2022 Advanced Micro Devices, Inc. All rights reserved.
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -30,6 +30,8 @@ THE SOFTWARE.
 #include "nccl.h"
 
 ncclResult_t rocm_smi_init();
+ncclResult_t rocm_smi_getNumDevice(uint32_t* num_devs);
+ncclResult_t rocm_smi_getDevicePciBusIdString(uint32_t deviceIndex, char* pciBusId, size_t len);
 ncclResult_t rocm_smi_getDeviceIndexByPciBusId(const char* pciBusId, uint32_t* deviceIndex);
 ncclResult_t rocm_smi_getLinkInfo(int srcDev, int dstDev, RSMI_IO_LINK_TYPE* rsmi_type, int *hops, int *count);
 
diff --git a/src/misc/rocm_smi_wrap.cc b/src/misc/rocm_smi_wrap.cc
index 3e38e7d51c..a233ad3745 100644
--- a/src/misc/rocm_smi_wrap.cc
+++ b/src/misc/rocm_smi_wrap.cc
@@ -1,5 +1,5 @@
 /*
-Copyright (c) 2021 Advanced Micro Devices, Inc. All rights reserved.
+Copyright (c) 2021-2022 Advanced Micro Devices, Inc. All rights reserved.
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -42,6 +42,28 @@ ncclResult_t rocm_smi_init() {
   return ncclSuccess;
 }
 
+ncclResult_t rocm_smi_getNumDevice(uint32_t* num_devs) {
+  ROCMSMICHECK(rsmi_num_monitor_devices(num_devs));
+  return ncclSuccess;
+}
+
+ncclResult_t rocm_smi_getDevicePciBusIdString(uint32_t deviceIndex, char* busId, size_t len) {
+  uint64_t id;
+  ROCMSMICHECK(rsmi_dev_pci_id_get(deviceIndex, &id));
+  /** rocm_smi's bus ID format
+   *  | Name     | Field   |
+   *  ---------- | ------- |
+   *  | Domain   | [64:32] |
+   *  | Reserved | [31:16] |
+   *  | Bus      | [15: 8] |
+   *  | Device   | [ 7: 3] |
+   *  | Function | [ 2: 0] |
+   **/
+  snprintf(busId, len, "%04lx:%02lx:%02lx.%01lx", (id) >> 32, (id & 0xff00) >> 8, (id & 0xf0) >> 4, (id & 0x3));
+  return ncclSuccess;
+}
+
+
 ncclResult_t rocm_smi_getDeviceIndexByPciBusId(const char* pciBusId, uint32_t* deviceIndex) {
   uint32_t i, num_devs = 0;
   int64_t busid;
@@ -74,15 +96,7 @@ ncclResult_t rocm_smi_getDeviceIndexByPciBusId(const char* pciBusId, uint32_t* d
   }
 }
 
-ncclResult_t rocm_smi_getLinkInfo(int srcDev, int dstDev, RSMI_IO_LINK_TYPE* rsmi_type, int *hops, int *count) {
-  char srcStr[] = "00000000:00:00.0", dstStr[] = "00000000:00:00.0";
-  uint32_t srcIndex, dstIndex;
-
-  CUDACHECK(hipDeviceGetPCIBusId(srcStr, sizeof(srcStr), srcDev));
-  CUDACHECK(hipDeviceGetPCIBusId(dstStr, sizeof(dstStr), dstDev));
-  NCCLCHECK(rocm_smi_getDeviceIndexByPciBusId(srcStr, &srcIndex));
-  NCCLCHECK(rocm_smi_getDeviceIndexByPciBusId(dstStr, &dstIndex));
-
+ncclResult_t rocm_smi_getLinkInfo(int srcIndex, int dstIndex, RSMI_IO_LINK_TYPE* rsmi_type, int *hops, int *count) {
   uint64_t rsmi_hops, rsmi_weight;
   ROCMSMICHECK(rsmi_topo_get_link_type(srcIndex, dstIndex, &rsmi_hops, rsmi_type));
   ROCMSMICHECK(rsmi_topo_get_link_weight(srcIndex, dstIndex, &rsmi_weight));
diff --git a/src/transport/p2p.cc b/src/transport/p2p.cc
index d6bfaf29c8..ff092d2a43 100644
--- a/src/transport/p2p.cc
+++ b/src/transport/p2p.cc
@@ -121,7 +121,7 @@ ncclResult_t p2pCanConnect(int* ret, struct ncclTopoSystem* topo, struct ncclTop
   int cudaDev1 = busIdToCudaDev(info1->busId);
   int cudaDev2 = busIdToCudaDev(info2->busId);
   if (cudaDev1 == -1 || cudaDev2 == -1) {
-#if CUDART_VERSION >= 10010
+#if defined(__HIP_PLATFORM_HCC__) || defined(__HCC__) || defined(__HIPCC__) || CUDART_VERSION >= 10010
     // CUDA 10.1 and later can use P2P with invisible devices.
     return ncclSuccess;
 #else