forked from BUPT-GAMMA/GammaGL
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsetup.py
140 lines (111 loc) · 4.61 KB
/
setup.py
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
69
70
71
72
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import os.path as osp
import platform
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext as _build_ext
# from pybind11.setup_helpers import Pybind11Extension, build_ext
from gammagl.utils.ggl_build_extension import BuildExtension, PyCudaExtension, PyCPUExtension
import subprocess
# TODO will depend on different host
WITH_CUDA = False
ext_root = osp.join("gammagl", "ops")
# multi cpu modules for development
def CPUExtension(ext_lib, module, dep_src_list=[], **kwargs):
ext_dir = osp.join(ext_root, ext_lib)
if len(dep_src_list) != 0 and not (dep_src_list[0].endswith('.cpp') or dep_src_list[0].endswith('.cc')):
raise NameError("Not a valid source name")
for i in range(len(dep_src_list)):
dep_src_list[i] = osp.join(ext_dir, "cpu", dep_src_list[i])
include_dirs = kwargs.get("include_dirs", [])
include_dirs += [ext_root, ext_dir]
kwargs["include_dirs"] = include_dirs
return PyCPUExtension(osp.join(ext_dir, f"_{module}").replace(osp.sep, "."),
sources=[osp.join(ext_dir, "cpu", f"{module}.cpp")] + dep_src_list,
**kwargs
)
# multi cuda modules for development
def CUDAExtension(ext_lib, module, dep_src_list=[], **kwargs):
ext_dir = osp.join(ext_root, ext_lib)
if len(dep_src_list) != 0 \
and (not dep_src_list[0].endswith('.cu')
or not dep_src_list[0].endswith('.cpp')
or not dep_src_list[0].endswith('.cc')):
raise NameError("Not a valid source name")
for i in range(len(dep_src_list)):
dep_src_list[i] = osp.join(ext_dir, "cuda", dep_src_list[i])
include_dirs = kwargs.get("include_dirs", [])
include_dirs += [ext_root, ext_dir]
kwargs["include_dirs"] = include_dirs
return PyCudaExtension(osp.join(ext_dir, f"_{module}_cuda").replace(osp.sep, "."),
sources=[osp.join(ext_dir, "cuda", f"{module}.cu")] + dep_src_list,
**kwargs
)
link_extra_args = []
def is_src_file(filename:str):
return filename.endswith("cpp") \
or filename.endswith("cu")
# Start to include cuda ops, if no cuda found, will only compile cpu ops
# single-library
def load_extensions():
ops_list = ["sparse", "segment", "tensor"]
ops_third_party_deps = [['parallel_hashmap'], [], []]
ops_root = osp.join('gammagl', 'ops')
extensions = []
for i in range(len(ops_list)):
ops_prefix = ops_list[i]
if WITH_CUDA:
ops_types = ["cpu", "cuda"]
else:
ops_types = ["cpu"]
ops_dir = osp.join(ops_root, ops_prefix)
for ops_type in ops_types:
is_cuda_ext = ops_type == "cuda"
src_dir = osp.join(ops_dir, ops_type)
src_files = filter(is_src_file, os.listdir(src_dir))
if not src_files:
continue
if not is_cuda_ext:
extensions.append(PyCPUExtension(
osp.join(ops_dir, f'_{ops_prefix}').replace(osp.sep, "."),
sources=[osp.join(src_dir, f) for f in src_files],
include_dirs=[osp.join('third_party', d) for d in ops_third_party_deps[i]],
))
else:
extensions.append(PyCudaExtension(
osp.join(ops_dir, f'_{ops_prefix}_cuda').replace(osp.sep, "."),
sources=[osp.join(src_dir, f) for f in src_files],
include_dirs=[osp.join('third_party', d) for d in ops_third_party_deps[i]],
))
return extensions
install_requires = ['numpy', 'scipy', 'pytest', 'tensorlayerx', 'rich', 'tqdm', 'pybind11']
classifiers = [
'Development Status :: 3 - Alpha',
'License :: OSI Approved :: Apache Software License',
]
setup(
name="gammagl",
version="0.2.0",
author="BUPT-GAMMA LAB",
maintainer="Tianyu Zhao",
license="Apache-2.0 License",
cmdclass={'build_ext': BuildExtension},
ext_modules=load_extensions(),
description=" ",
url="https://github.com/BUPT-GAMMA/GammaGL",
download_url="https://github.com/BUPT-GAMMA/GammaGL",
python_requires='>=3.7',
packages=find_packages(),
install_requires=install_requires,
classifiers=classifiers,
include_package_data=True,
package_data={
"gammagl": ["*.json"]
}
)
os.chdir('gammagl/mpops/torch_ext')
subprocess.call('python setup.py install', shell=True)
# python setup.py build_ext --inplace
# python setup.py install