-
Notifications
You must be signed in to change notification settings - Fork 14
/
setup.py
109 lines (100 loc) · 3.34 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
import argparse
import os
import sys
from setuptools import setup
import torch
from torch.utils.cpp_extension import CppExtension, BuildExtension, CUDAExtension
# Parse command line args.
parser = argparse.ArgumentParser()
parser.add_argument('--with_cuda', action="store_true", default=None)
parser.add_argument('--without_cuda', action="store_true", default=None)
args, unknown = parser.parse_known_args()
sys.argv = sys.argv[:2] + unknown
if args.with_cuda is None:
if args.without_cuda is not None:
args.with_cuda = not args.without_cuda
else:
print("--with_cuda or --without_cuda not specified, using PyTorch to decide...")
args.with_cuda = torch.cuda.is_available()
if args.with_cuda:
print("torch.cuda.is_available says True, proceeding to build with cuda.")
else:
print("torch.cuda.is_available says False, proceeding to build without cuda.")
# Setup global variables.
root_dir = os.path.dirname(os.path.abspath(__file__))
test_dir = os.path.join(root_dir, "tests")
py_dir = os.path.join(root_dir, "python", "SmoothParticleNets")
src_dir = os.path.join(root_dir, "src")
# Create pytest args.
pytest_args = {
'with_cuda': args.with_cuda,
}
fp = open(os.path.join(test_dir, "pytest_args.py"), "w")
for k, v in pytest_args.items():
if isinstance(v, str):
v = "'" + v + "'"
fp.write("%s = %s\n" % (k, str(v)))
fp.close()
# Build kernel_constants.h
# Add path to python source to path.
sys.path.append(py_dir)
from kernels import KERNELS, KERNEL_NAMES, DKERNELS
fp = open(os.path.join(src_dir, "kernel_constants.h"), "w")
fp.write("// THIS FILE IS AUTOGENERATED. DO NOT ALTER.\n")
fp.write("#ifndef __kernel_constants_h__\n")
fp.write("#define __kernel_constants_h__\n")
fp.write("#ifdef __cplusplus\n")
fp.write("extern \"C\" {\n")
fp.write("#endif\n")
fp.write("\n")
fp.write("#include <math.h>\n")
fp.write("#include <stdio.h>\n")
fp.write("\n")
fp.write("#ifdef CUDA\n")
fp.write("__host__ __device__\n")
fp.write("#endif\n")
fp.write("inline\n")
fp.write("float KERNEL_W(float d, float H, int fn) {\n")
fp.write(" float ret = 0.0f;\n")
for i, k in enumerate(KERNEL_NAMES):
fp.write(" if(fn == %d) { ret = (%s); }\n" % (i, KERNELS[k]))
fp.write(" return ret;\n")
fp.write("}\n\n")
fp.write("#ifdef CUDA\n")
fp.write("__host__ __device__\n")
fp.write("#endif\n")
fp.write("inline\n")
fp.write("float KERNEL_DW(float d, float H, int fn) {\n")
fp.write(" float ret = 0.0f;\n")
for i, k in enumerate(KERNEL_NAMES):
fp.write(" if(fn == %d) { ret = (%s); }\n" % (i, DKERNELS[k]))
fp.write(" return ret;\n")
fp.write("}\n\n")
fp.write("#define VALIDATE_KERNEL_ID(fn) (fn >= 0 && fn < %d)" % len(KERNELS))
fp.write("\n")
fp.write("#ifdef __cplusplus\n")
fp.write("}\n")
fp.write("#endif\n")
fp.write("#endif\n")
fp.flush()
fp.close()
# Define extensions.
ext_modules = [
CppExtension('SmoothParticleNets._ext', [
os.path.join(src_dir, 'cpu_layer_funcs.cpp'),
]),
]
if args.with_cuda:
ext_modules.append(CUDAExtension('SmoothParticleNets._extc', [
os.path.join(src_dir, 'cuda_layer_funcs.cpp'),
os.path.join(src_dir, 'gpu_kernels.cu'),
]))
# The main setup call.
setup(
name='SmoothParticleNets',
package_dir={'': 'python'},
packages=['SmoothParticleNets'],
ext_modules=ext_modules,
cmdclass={
'build_ext': BuildExtension
})