Commit e8d427f3 authored by YdrMaster's avatar YdrMaster
Browse files

issue/291/fix: 动态加载 cuda 路径,避免依赖环境变量


Signed-off-by: default avatarYdrMaster <ydrml@hotmail.com>
parent f72b8653
local CUDA_ROOT = os.getenv("CUDA_ROOT") or os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH")
local CUDNN_ROOT = os.getenv("CUDNN_ROOT") or os.getenv("CUDNN_HOME") or os.getenv("CUDNN_PATH") local CUDNN_ROOT = os.getenv("CUDNN_ROOT") or os.getenv("CUDNN_HOME") or os.getenv("CUDNN_PATH")
if CUDA_ROOT ~= nil then
add_includedirs(CUDA_ROOT .. "/include")
end
if CUDNN_ROOT ~= nil then if CUDNN_ROOT ~= nil then
add_includedirs(CUDNN_ROOT .. "/include") add_includedirs(CUDNN_ROOT .. "/include")
end end
...@@ -21,6 +16,14 @@ target("infiniop-cuda") ...@@ -21,6 +16,14 @@ target("infiniop-cuda")
end end
add_cugencodes("native") add_cugencodes("native")
on_load(function (target)
import("lib.detect.find_tool")
local nvcc = find_tool("nvcc")
if nvcc ~= nil then
target:add("linkdirs", path.directory(path.directory(nvcc.program)) .. "/lib64/stubs")
end
end)
if is_plat("windows") then if is_plat("windows") then
add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler") add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler")
add_cuflags("-Xcompiler=/W3", "-Xcompiler=/WX") add_cuflags("-Xcompiler=/W3", "-Xcompiler=/WX")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment