cuda.lua 3.07 KB
Newer Older
PanZezhongQY's avatar
PanZezhongQY committed
1
2
3
4
5
6
7
local CUDNN_ROOT = os.getenv("CUDNN_ROOT") or os.getenv("CUDNN_HOME") or os.getenv("CUDNN_PATH")
if CUDNN_ROOT ~= nil then
    add_includedirs(CUDNN_ROOT .. "/include")
end

target("infiniop-cuda")
    set_kind("static")
PanZezhong's avatar
PanZezhong committed
8
    add_deps("infini-utils")
PanZezhongQY's avatar
PanZezhongQY committed
9
10
    on_install(function (target) end)

11
    set_policy("build.cuda.devlink", true)
PanZezhongQY's avatar
PanZezhongQY committed
12
    set_toolchains("cuda")
13
    add_links("cudart", "cublas")
14
15
16
    if has_config("cudnn") then
        add_links("cudnn")
    end
PanZezhongQY's avatar
PanZezhongQY committed
17
18
    add_cugencodes("native")

19
20
21
22
23
24
25
26
    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)

PanZezhongQY's avatar
PanZezhongQY committed
27
28
    if is_plat("windows") then
        add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler")
29
        add_cuflags("-Xcompiler=/W3", "-Xcompiler=/WX")
30
        add_cxxflags("/FS")
PanZezhongQY's avatar
PanZezhongQY committed
31
32
33
34
        if CUDNN_ROOT ~= nil then
            add_linkdirs(CUDNN_ROOT .. "\\lib\\x64")
        end
    else
35
        add_cuflags("-Xcompiler=-Wall", "-Xcompiler=-Werror")
PanZezhongQY's avatar
PanZezhongQY committed
36
        add_cuflags("-Xcompiler=-fPIC")
37
        add_cuflags("--extended-lambda")
PanZezhongQY's avatar
PanZezhongQY committed
38
39
        add_culdflags("-Xcompiler=-fPIC")
        add_cxxflags("-fPIC")
PanZezhong's avatar
PanZezhong committed
40
41
42
43
        add_cuflags("--expt-relaxed-constexpr")
        if CUDNN_ROOT ~= nil then
            add_linkdirs(CUDNN_ROOT .. "/lib")
        end
PanZezhongQY's avatar
PanZezhongQY committed
44
45
    end

46
47
    add_cuflags("-Xcompiler=-Wno-error=deprecated-declarations")

PanZezhongQY's avatar
PanZezhongQY committed
48
    set_languages("cxx17")
Jiacheng Huang's avatar
Jiacheng Huang committed
49
    add_files("../src/infiniop/devices/cuda/*.cu", "../src/infiniop/ops/*/cuda/*.cu", "../build/ninetoothed/*.c")
PanZezhongQY's avatar
PanZezhongQY committed
50
target_end()
51
52
53

target("infinirt-cuda")
    set_kind("static")
PanZezhong's avatar
PanZezhong committed
54
    add_deps("infini-utils")
55
    on_install(function (target) end)
56

57
    set_policy("build.cuda.devlink", true)
58
59
    set_toolchains("cuda")
    add_links("cudart")
60
61
62

    if is_plat("windows") then
        add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler")
63
        add_cxxflags("/FS")
64
65
66
67
68
69
70
71
72
    else
        add_cuflags("-Xcompiler=-fPIC")
        add_culdflags("-Xcompiler=-fPIC")
        add_cxflags("-fPIC")
    end

    set_languages("cxx17")
    add_files("../src/infinirt/cuda/*.cu")
target_end()
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

target("infiniccl-cuda")
    set_kind("static")
    add_deps("infinirt")
    on_install(function (target) end)
    if has_config("ccl") then
        set_policy("build.cuda.devlink", true)
        set_toolchains("cuda")
        add_links("cudart")

        if not is_plat("windows") then
            add_cuflags("-Xcompiler=-fPIC")
            add_culdflags("-Xcompiler=-fPIC")
            add_cxflags("-fPIC")

            local nccl_root = os.getenv("NCCL_ROOT")
            if nccl_root then
                add_includedirs(nccl_root .. "/include")
                add_links(nccl_root .. "/lib/libnccl.so")
            else
                add_links("nccl") -- Fall back to default nccl linking
            end

            add_files("../src/infiniccl/cuda/*.cu")
        else
            print("[Warning] NCCL is not supported on Windows")
        end
    end
    set_languages("cxx17")
102

103
target_end()