From ecafc1969c6e9d03f11978da48bf936e09e99b14 Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Fri, 10 May 2024 07:31:12 -0700 Subject: [PATCH] Support WSL2 (#1173) --- src/enqueue.cc | 6 +- src/misc/rocm_smi_wrap.cc | 130 +++++++++++++++++++++++--------------- src/transport/net.cc | 7 +- 3 files changed, 86 insertions(+), 57 deletions(-) diff --git a/src/enqueue.cc b/src/enqueue.cc index 1714aa1158..a77f5b1e01 100644 --- a/src/enqueue.cc +++ b/src/enqueue.cc @@ -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; -} \ No newline at end of file +} diff --git a/src/misc/rocm_smi_wrap.cc b/src/misc/rocm_smi_wrap.cc index 4c8987713f..39515172e3 100644 --- a/src/misc/rocm_smi_wrap.cc +++ b/src/misc/rocm_smi_wrap.cc @@ -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; } diff --git a/src/transport/net.cc b/src/transport/net.cc index 2643849522..d0a5943090 100644 --- a/src/transport/net.cc +++ b/src/transport/net.cc @@ -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 } -}; \ No newline at end of file +};