diff options
| author | ruki <[email protected]> | 2022-01-25 00:27:06 +0800 |
|---|---|---|
| committer | ruki <[email protected]> | 2022-01-25 00:27:06 +0800 |
| commit | dee5ebd8220b3695a0ebf43d88fee06d910766ee (patch) | |
| tree | cdb1641b18baced30fbad284fa29455395454f7f | |
| parent | a1ecc4de70ff0e108bcf89fbccfd25189239071c (diff) | |
improve find cuda devices
| -rw-r--r-- | xmake/modules/lib/detect/find_cudadevices.lua | 6 | ||||
| -rw-r--r-- | xmake/rules/cuda/gencodes/xmake.lua | 9 |
2 files changed, 11 insertions, 4 deletions
diff --git a/xmake/modules/lib/detect/find_cudadevices.lua b/xmake/modules/lib/detect/find_cudadevices.lua index c55cc26f8..d2c83c663 100644 --- a/xmake/modules/lib/detect/find_cudadevices.lua +++ b/xmake/modules/lib/detect/find_cudadevices.lua @@ -113,7 +113,7 @@ function _parse_result(lines, verbose) end -- find devices -function _find_devices(verbose) +function _find_devices(verbose, envs) -- find nvcc local nvcc = assert(find_tool("nvcc"), "nvcc not found") @@ -131,7 +131,7 @@ function _find_devices(verbose) { function () local args = { sourcefile, "-run", "-o", outfile , '-DPRINT_SUFFIX="' .. _PRINT_SUFFIX .. '"' } - return os.iorunv(nvcc.program, args) + return os.iorunv(nvcc.program, args, {envs = envs}) end, catch { @@ -187,7 +187,7 @@ function _get_devices(opt) end local verbose = opt.verbose or option.get("verbose") or option.get("diagnosis") - local devices = _find_devices(verbose) + local devices = _find_devices(verbose, opt.envs) if devices then cachedata = { succeed = true, data = devices } else diff --git a/xmake/rules/cuda/gencodes/xmake.lua b/xmake/rules/cuda/gencodes/xmake.lua index c7ee45050..f83ce4740 100644 --- a/xmake/rules/cuda/gencodes/xmake.lua +++ b/xmake/rules/cuda/gencodes/xmake.lua @@ -52,7 +52,14 @@ rule("cuda.gencodes") end archs = archs:trim():lower() if archs == "native" then - local device = find_cudadevices({ skip_compute_mode_prohibited = true, order_by_flops = true })[1] + local cuda_envs + for _, toolchain_inst in ipairs(target:toolchains()) do + if toolchain_inst:name() == "cuda" then + cuda_envs = toolchain_inst:runenvs() + break + end + end + local device = find_cudadevices({skip_compute_mode_prohibited = true, order_by_flops = true, envs = cuda_envs})[1] if device then return nf_cugencode("sm_" .. device.major .. device.minor) end |