-
Notifications
You must be signed in to change notification settings - Fork 1
/
setup.py
102 lines (92 loc) · 4.64 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
from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import torch
import os
# setup(
# name='c3d',
# description='Package for Continuous 3D Loss',
# author='Minghan Zhu',
# author_email='[email protected]',
# url='https://github.com/minghanz/c3d',
# packages=find_packages(),
# )
torch_vs = (torch.__version__).split('.')
torch_version = float(torch_vs[0]) + 0.1 * float(torch_vs[1])
if torch_version > 1.2:
######### for pytorch 1.6
cxx_args = ['-std=c++14'] # for pytorch 1.6. For pytorch 1.2 use c++11 or do not use extra_compile_args
nvcc_args = [
'-gencode', 'arch=compute_50,code=sm_50',
'-gencode', 'arch=compute_52,code=sm_52',
'-gencode', 'arch=compute_60,code=sm_60',
'-gencode', 'arch=compute_61,code=sm_61',
'-gencode', 'arch=compute_61,code=compute_61'
]
extra_compile_args={'cxx': cxx_args, 'nvcc': nvcc_args}
else:
######### for pytorch 1.2
extra_compile_args = dict()
# ### only needed for cluster if the compiler is too old (when nvcc found anaconda3/envs/tp36dup/gcc/libexec/gcc/x86_64-unknown-linux-gnu/5.2.0/cc1plus which cannot find libisl.so.10 which is needed)
# ### https://stackoverflow.com/questions/6622454/cuda-incompatible-with-my-gcc-version/6622751#6622751
# extra_compile_args['cxx'] = ['-std=c++11']
# extra_compile_args['nvcc'] = []
# extra_compile_args['nvcc'] += ['--compiler-bindir', '/sw/arcts/centos7/gcc/8.2.0/bin/g++']
# # extra_compile_args = dict()
# os.environ["CC"] = "/sw/arcts/centos7/gcc/8.2.0/bin/gcc"
# os.environ["CXX"] = "/sw/arcts/centos7/gcc/8.2.0/bin/g++"
setup(
name='c3d',
description='Package for Continuous 3D Loss',
author='Minghan Zhu',
author_email='[email protected]',
url='https://github.com/minghanz/c3d',
ext_modules=[
CUDAExtension('c3d.cvo_ops.sub_cuda', [
'c3d/cvo_ops/custom_ori/sub_cuda.cpp',
'c3d/cvo_ops/custom_ori/sub_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.sub_norm_cuda_half_paral', [
'c3d/cvo_ops/custom_norm/sub_norm_cuda.cpp',
'c3d/cvo_ops/custom_norm/sub_norm_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_with_normal_output', [
'c3d/cvo_ops/custom_dense_with_normal/cvo_dense_with_normal_cuda.cpp',
'c3d/cvo_ops/custom_dense_with_normal/cvo_dense_with_normal_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_normal', [
'c3d/cvo_ops/custom_dense_normal/cvo_dense_normal_cuda.cpp',
'c3d/cvo_ops/custom_dense_normal/cvo_dense_normal_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_dist', [
'c3d/cvo_ops/custom_dense_dist/cvo_dense_samp_cuda.cpp',
'c3d/cvo_ops/custom_dense_dist/cvo_dense_samp_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_angle', [
'c3d/cvo_ops/custom_dense_angle/cvo_dense_angle_cuda.cpp',
'c3d/cvo_ops/custom_dense_angle/cvo_dense_angle_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_samp', [
'c3d/cvo_ops/custom_dense/cvo_dense_samp_cuda.cpp',
'c3d/cvo_ops/custom_dense/cvo_dense_samp_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cross_subtract_cuda', [
'c3d/cvo_ops/custom_cross_subtract/cross_subtract_cuda.cpp',
'c3d/cvo_ops/custom_cross_subtract/cross_subtract_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cross_prod_cuda', [
'c3d/cvo_ops/custom_cross_prod/cross_prod_cuda.cpp',
'c3d/cvo_ops/custom_cross_prod/cross_prod_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_Sigma', [
'c3d/cvo_ops/custom_dense_Sigma/cvo_dense_Sigma_cuda.cpp',
'c3d/cvo_ops/custom_dense_Sigma/cvo_dense_Sigma_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
CUDAExtension('c3d.cvo_ops.cvo_dense_Sigma_grid', [
'c3d/cvo_ops/custom_dense_Sigma/cvo_dense_Sigma_grid_cuda.cpp',
'c3d/cvo_ops/custom_dense_Sigma/cvo_dense_Sigma_grid_cuda_kernel.cu',
], extra_compile_args=extra_compile_args),
],
packages=find_packages(),
cmdclass={
'build_ext': BuildExtension
})