|
from setuptools import setup |
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension |
|
import os |
|
|
|
cuda_home = os.getenv('CUDA_HOME', '/usr/local/cuda') |
|
|
|
setup( |
|
name='pc_util', |
|
version='1.0', |
|
ext_modules=[ |
|
CUDAExtension('pc_util', [ |
|
'src/pointnet2_api.cpp', |
|
'src/ball_query.cpp', |
|
'src/ball_query_gpu.cu', |
|
'src/group_points.cpp', |
|
'src/group_points_gpu.cu', |
|
'src/interpolate.cpp', |
|
'src/interpolate_gpu.cu', |
|
'src/sampling.cpp', |
|
'src/sampling_gpu.cu', |
|
'src/cluster.cpp', |
|
'src/cluster_gpu.cu', |
|
], |
|
include_dirs=[os.path.join(cuda_home, 'include')], |
|
extra_compile_args={ |
|
'cxx': ['-g', '-O2'], |
|
'nvcc': ['-O2', '--compiler-options', "'-fPIC'"]}) |
|
], |
|
cmdclass={'build_ext': BuildExtension} |
|
) |
|
|
