nvidia.lua 3.82 KB
Newer Older
PanZezhongQY's avatar
PanZezhongQY committed
1
2
3
4
5
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

6
local CUTLASS_ROOT = os.getenv("CUTLASS_ROOT") or os.getenv("CUTLASS_HOME") or os.getenv("CUTLASS_PATH")
7
local CUTE_ROOT = os.getenv("CUTE_ROOT") or os.getenv("CUTE_HOME") or os.getenv("CUTE_PATH")
8
9
10

if CUTLASS_ROOT ~= nil then
    add_includedirs(CUTLASS_ROOT)
11
    add_includedirs(CUTE_ROOT)
12
13
end

14
target("infiniop-nvidia")
PanZezhongQY's avatar
PanZezhongQY committed
15
    set_kind("static")
PanZezhong's avatar
PanZezhong committed
16
    add_deps("infini-utils")
PanZezhongQY's avatar
PanZezhongQY committed
17
18
    on_install(function (target) end)

19
    set_policy("build.cuda.devlink", true)
PanZezhongQY's avatar
PanZezhongQY committed
20
    set_toolchains("cuda")
21
    add_links("cudart", "cublas")
22
23
24
    if has_config("cudnn") then
        add_links("cudnn")
    end
25
    add_cugencodes("native")
PanZezhongQY's avatar
PanZezhongQY committed
26

27
28
29
30
    on_load(function (target)
        import("lib.detect.find_tool")
        local nvcc = find_tool("nvcc")
        if nvcc ~= nil then
pengcheng888's avatar
pengcheng888 committed
31
32
33
34
35
36
37
            if is_plat("windows") then
                nvcc_path = os.iorun("where nvcc"):match("(.-)\r?\n")
            else
                nvcc_path = nvcc.program
            end

            target:add("linkdirs", path.directory(path.directory(nvcc_path)) .. "/lib64/stubs")
38
            target:add("links", "cuda")
39
40
41
42
43

            local cuda_arch = get_config("cuda_arch")
            if cuda_arch ~= nil then
                target:add("cu-cxxflags", "-arch=", cuda_arch)
            end
44
45
46
        end
    end)

PanZezhongQY's avatar
PanZezhongQY committed
47
48
    if is_plat("windows") then
        add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler")
49
        add_cuflags("-Xcompiler=/W3", "-Xcompiler=/WX")
50
        add_cxxflags("/FS")
PanZezhongQY's avatar
PanZezhongQY committed
51
52
53
54
        if CUDNN_ROOT ~= nil then
            add_linkdirs(CUDNN_ROOT .. "\\lib\\x64")
        end
    else
55
        add_cuflags("-Xcompiler=-Wall", "-Xcompiler=-Werror")
PanZezhongQY's avatar
PanZezhongQY committed
56
        add_cuflags("-Xcompiler=-fPIC")
57
        add_cuflags("--extended-lambda")
PanZezhongQY's avatar
PanZezhongQY committed
58
59
        add_culdflags("-Xcompiler=-fPIC")
        add_cxxflags("-fPIC")
PanZezhong's avatar
PanZezhong committed
60
61
62
63
        add_cuflags("--expt-relaxed-constexpr")
        if CUDNN_ROOT ~= nil then
            add_linkdirs(CUDNN_ROOT .. "/lib")
        end
PanZezhongQY's avatar
PanZezhongQY committed
64
65
    end

66
67
    add_cuflags("-Xcompiler=-Wno-error=deprecated-declarations")

PanZezhongQY's avatar
PanZezhongQY committed
68
    set_languages("cxx17")
69
70
71
72
73
    add_files("../src/infiniop/devices/nvidia/*.cu", "../src/infiniop/ops/*/nvidia/*.cu")

    if has_config("ninetoothed") then
        add_files("../build/ninetoothed/*.c")
    end
PanZezhongQY's avatar
PanZezhongQY committed
74
target_end()
75

76
target("infinirt-nvidia")
77
    set_kind("static")
PanZezhong's avatar
PanZezhong committed
78
    add_deps("infini-utils")
79
    on_install(function (target) end)
80

81
    set_policy("build.cuda.devlink", true)
82
83
    set_toolchains("cuda")
    add_links("cudart")
84
85
86

    if is_plat("windows") then
        add_cuflags("-Xcompiler=/utf-8", "--expt-relaxed-constexpr", "--allow-unsupported-compiler")
87
        add_cxxflags("/FS")
88
89
90
91
92
93
94
95
96
    else
        add_cuflags("-Xcompiler=-fPIC")
        add_culdflags("-Xcompiler=-fPIC")
        add_cxflags("-fPIC")
    end

    set_languages("cxx17")
    add_files("../src/infinirt/cuda/*.cu")
target_end()
97

98
target("infiniccl-nvidia")
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
    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")
126

127
target_end()