Support WSL2 (#1173)

2024-05-10 07:31:12 -07:00
@@ -75,7 +75,8 @@ ncclResult_t ncclInitKernelsForDevice(int cudaArch, size_t* maxStackSize) {

    if (maxStackSize) {
      cudaFuncAttributes attr = {0};
-      CUDACHECKGOTO(cudaFuncGetAttributes(&attr, fn), result, ignore0);
+      if (cudaFuncGetAttributes(&attr, fn) != cudaSuccess)
+        WARN("Failed to get kernel attributes");
      if (attr.localSizeBytes > *maxStackSize) *maxStackSize = attr.localSizeBytes;
    ignore0:;
    }
@@ -100,7 +101,6 @@ ncclResult_t ncclInitKernelsForDevice(int cudaArch, size_t* maxStackSize) {
 /*****************************************************************************/
 /*       Launch system : synchronization and CUDA kernel launch              */
 /*****************************************************************************/
-
 static void appendWorkElemColl(
    struct ncclComm* comm, struct ncclKernelPlan* plan, int channelId,
    int funcIndex, struct ncclWorkElem const *elem) {
@@ -2189,4 +2189,4 @@ ncclResult_t ncclRedOpDestroy(ncclRedOp_t op, ncclComm_t comm) {
  comm->userRedOpFreeHead = ix;
  TRACE_CALL("ncclRedOpDestroy(%d,%p)", op, comm);
  return ncclSuccess;
-}
+}
@@ -23,6 +23,8 @@ THE SOFTWARE.
 #include "core.h"
 #include "utils.h"

+static int is_wsl2 = -1;
+
 #define ROCMSMICHECK(cmd) do {               \
  rsmi_status_t ret = cmd;                   \
  if( ret != RSMI_STATUS_SUCCESS ) {         \
@@ -35,6 +37,12 @@ THE SOFTWARE.


 ncclResult_t rocm_smi_init() {
+  if (is_wsl2 == -1)
+    is_wsl2 = (access("/dev/dxg", F_OK) == -1) ? 0 : 1;
+  if (is_wsl2) {
+    INFO(NCCL_INIT, "Not using rocm_smi_lib due to WSL2 environment detected.");
+    return ncclSuccess;
+  }
 #ifdef USE_ROCM_SMI_THREAD_ONLY_MUTEX
  ROCMSMICHECK(rsmi_init(RSMI_INIT_FLAG_THRAD_ONLY_MUTEX));
 #else
@@ -47,76 +55,94 @@ ncclResult_t rocm_smi_init() {
 }

 ncclResult_t rocm_smi_getNumDevice(uint32_t* num_devs) {
-  ROCMSMICHECK(rsmi_num_monitor_devices(num_devs));
+  if (is_wsl2)
+    CUDACHECK(cudaGetDeviceCount((int *)num_devs));
+  else
+    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 & 0xf8) >> 3, (id & 0x7));
+  if (is_wsl2) {
+    CUDACHECK(cudaDeviceGetPCIBusId(busId, len, deviceIndex));
+  } else {
+    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 & 0xf8) >> 3, (id & 0x7));
+  }
  return ncclSuccess;
 }


 ncclResult_t rocm_smi_getDeviceIndexByPciBusId(const char* pciBusId, uint32_t* deviceIndex) {
-  uint32_t i, num_devs = 0;
-  int64_t busid;
-
-  busIdToInt64(pciBusId, &busid);
-  /** convert to rocm_smi's bus ID format
-   *  | Name     | Field   |
-   *  ---------- | ------- |
-   *  | Domain   | [64:32] |
-   *  | Reserved | [31:16] |
-   *  | Bus      | [15: 8] |
-   *  | Device   | [ 7: 3] |
-   *  | Function | [ 2: 0] |
-   **/
-  busid = ((busid&0xffff00000L)<<12)+((busid&0xff000L)>>4)+((busid&0xff0L)>>1)+(busid&0x7L);
-  ROCMSMICHECK(rsmi_num_monitor_devices(&num_devs));
-  for (i = 0; i < num_devs; i++) {
-    uint64_t bdfid;
-    ROCMSMICHECK(rsmi_dev_pci_id_get(i, &bdfid));
-    if (bdfid == busid) break;
-  }
-
-  if (i < num_devs) {
-    *deviceIndex = i;
+  if (is_wsl2) {
+    CUDACHECK(hipDeviceGetByPCIBusId((int *)deviceIndex, pciBusId));
    return ncclSuccess;
-  }
-  else {
-    WARN("rocm_smi_lib: %s device index not found", pciBusId);
-    return ncclInternalError;
+  } else {
+    uint32_t i, num_devs = 0;
+    int64_t busid;
+
+    busIdToInt64(pciBusId, &busid);
+    /** convert to rocm_smi's bus ID format
+     *  | Name     | Field   |
+     *  ---------- | ------- |
+     *  | Domain   | [64:32] |
+     *  | Reserved | [31:16] |
+     *  | Bus      | [15: 8] |
+     *  | Device   | [ 7: 3] |
+     *  | Function | [ 2: 0] |
+     **/
+    busid = ((busid&0xffff00000L)<<12)+((busid&0xff000L)>>4)+((busid&0xff0L)>>1)+(busid&0x7L);
+    ROCMSMICHECK(rsmi_num_monitor_devices(&num_devs));
+    for (i = 0; i < num_devs; i++) {
+      uint64_t bdfid;
+      ROCMSMICHECK(rsmi_dev_pci_id_get(i, &bdfid));
+      if (bdfid == busid) break;
+    }
+
+    if (i < num_devs) {
+      *deviceIndex = i;
+      return ncclSuccess;
+    }
+    else {
+      WARN("rocm_smi_lib: %s device index not found", pciBusId);
+      return ncclInternalError;
+    }
  }
 }

 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));
-  *hops = 2;
-  *count = 1;
-  if (*rsmi_type == RSMI_IOLINK_TYPE_XGMI && rsmi_weight == 15) {
+  if (is_wsl2) {
+    *rsmi_type = RSMI_IOLINK_TYPE_PCIEXPRESS;
    *hops = 1;
+    *count = 1;
+  } else {
+    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));
+    *hops = 2;
+    *count = 1;
+    if (*rsmi_type == RSMI_IOLINK_TYPE_XGMI && rsmi_weight == 15) {
+      *hops = 1;
 #if defined USE_ROCM_SMI64CONFIG && rocm_smi_VERSION_MAJOR >= 5
-    uint64_t min_bw = 0, max_bw = 0;
-    rsmi_version_t version;
-    ROCMSMICHECK(rsmi_version_get(&version));
-    if (version.major >= 5)
-      ROCMSMICHECK(rsmi_minmax_bandwidth_get(srcIndex, dstIndex, &min_bw, &max_bw));
-    if (max_bw && min_bw)
-      *count = max_bw/min_bw;
+      uint64_t min_bw = 0, max_bw = 0;
+      rsmi_version_t version;
+      ROCMSMICHECK(rsmi_version_get(&version));
+      if (version.major >= 5)
+        ROCMSMICHECK(rsmi_minmax_bandwidth_get(srcIndex, dstIndex, &min_bw, &max_bw));
+      if (max_bw && min_bw)
+        *count = max_bw/min_bw;
 #endif
+    }
  }
  return ncclSuccess;
 }
@@ -591,7 +591,10 @@ static ncclResult_t sharedNetBuffersDestroy(struct ncclProxyState* proxyState, i
 }

 static ncclResult_t proxySharedInit(struct ncclProxyConnection* connection, struct ncclProxyState* proxyState, int nChannels) {
-  NCCLCHECK(sharedNetBuffersInit(proxyState, 1, connection->tpLocalRank, 0, connection->sameProcess, nChannels, NULL, NULL, NULL, NULL));
+  static int is_wsl2 = -1;
+  if (is_wsl2 == -1)
+    is_wsl2 = (access("/dev/dxg", F_OK) == -1) ? 0 : 1;
+  NCCLCHECK(sharedNetBuffersInit(proxyState, is_wsl2 == 0 ? 1 : 0, connection->tpLocalRank, 0, connection->sameProcess, nChannels, NULL, NULL, NULL, NULL));
  return ncclSuccess;
 }

@@ -1644,4 +1647,4 @@ struct ncclTransport netTransport = {
  canConnect,
  { sendSetup, sendConnect, sendFree, proxySharedInit, sendProxySetup, sendProxyConnect, sendProxyFree, sendProxyProgress },
  { recvSetup, recvConnect, recvFree, proxySharedInit, recvProxySetup, recvProxyConnect, recvProxyFree, recvProxyProgress }
-};
+};