-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsetup.py
More file actions
68 lines (57 loc) · 1.33 KB
/
setup.py
File metadata and controls
68 lines (57 loc) · 1.33 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
import os
import torch
import glob
from setuptools import find_packages, setup
from torch.utils.cpp_extension import (
CUDAExtension,
BuildExtension
)
# ------------------------------------------- #
LIBRARY_NAME = "ddgs"
CSRC_DIR = "csrc"
CSRC_DIR_ABSOLUTE = os.path.join(os.path.abspath(os.path.dirname(__file__)), "csrc")
if torch.__version__ >= "2.6.0": # TODO: what is this ?
limitedAPI = True
else:
limitedAPI = False
# ------------------------------------------- #
def get_extension():
linkArgs = [
"-O3"
]
compileArgs = {
"cxx": [
"-O3",
"-DPy_LIMITED_API=0x03090000",
],
"nvcc": [
"-O3",
],
}
srcDir = os.path.join(CSRC_DIR, "src")
srcs = list((
glob.glob(os.path.join(srcDir, "*.cpp")) +
glob.glob(os.path.join(srcDir, "*.cu"))
))
srcs.append(
os.path.join(CSRC_DIR, "ext.cpp")
)
includeDirs = [
os.path.join(CSRC_DIR_ABSOLUTE, "include"),
os.path.join(CSRC_DIR_ABSOLUTE, "external")
]
return CUDAExtension(
f"{LIBRARY_NAME}._C",
srcs,
include_dirs=includeDirs,
extra_compile_args=compileArgs,
extra_link_args=linkArgs,
py_limited_api=limitedAPI
)
# ------------------------------------------- #
setup(
packages=find_packages(),
ext_modules=[ get_extension() ],
cmdclass={"build_ext": BuildExtension},
options={"bdist_wheel": {"py_limited_api": "cp39"}} if limitedAPI else {},
)