1 # --------------------------------------------------------
3 # Copyright (c) 2015 Microsoft
4 # Licensed under The MIT License [see LICENSE for details]
5 # Written by Ross Girshick
6 # --------------------------------------------------------
9 from os.path import join as pjoin
10 from setuptools import setup
11 from distutils.extension import Extension
12 from Cython.Distutils import build_ext
17 def find_in_path(name, path):
18 "Find a file in a search path"
20 # http://code.activestate.com/recipes/52224-find-a-file-given-a-search-path/
21 for dir in path.split(os.pathsep):
22 binpath = pjoin(dir, name)
23 if os.path.exists(binpath):
24 return os.path.abspath(binpath)
29 """Locate the CUDA environment on the system
31 Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
32 and values giving the absolute path to each directory.
34 Starts by looking for the CUDAHOME env variable. If not found, everything
35 is based on finding 'nvcc' in the PATH.
38 # first check if the CUDAHOME env variable is in use
39 if 'CUDAHOME' in os.environ:
40 home = os.environ['CUDAHOME']
41 nvcc = pjoin(home, 'bin', 'nvcc')
43 # otherwise, search the PATH for NVCC
44 default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin')
45 nvcc = find_in_path('nvcc', os.environ['PATH'] + os.pathsep + default_path)
47 raise EnvironmentError('The nvcc binary could not be '
48 'located in your $PATH. Either add it to your path, or set $CUDAHOME')
49 home = os.path.dirname(os.path.dirname(nvcc))
51 cudaconfig = {'home':home, 'nvcc':nvcc,
52 'include': pjoin(home, 'include'),
53 'lib64': pjoin(home, 'lib64')}
54 for k, v in cudaconfig.iteritems():
55 if not os.path.exists(v):
56 raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v))
62 # Obtain the numpy include directory. This logic works across numpy versions.
64 numpy_include = np.get_include()
65 except AttributeError:
66 numpy_include = np.get_numpy_include()
69 def customize_compiler_for_nvcc(self):
70 """inject deep into distutils to customize how the dispatch
73 If you subclass UnixCCompiler, it's not trivial to get your subclass
74 injected in, and still have the right customizations (i.e.
75 distutils.sysconfig.customize_compiler) run on it. So instead of going
76 the OO route, I have this. Note, it's kindof like a wierd functional
77 subclassing going on."""
79 # tell the compiler it can processes .cu
80 self.src_extensions.append('.cu')
82 # save references to the default compiler_so and _comple methods
83 default_compiler_so = self.compiler_so
86 # now redefine the _compile method. This gets executed for each
87 # object but distutils doesn't have the ability to change compilers
88 # based on source extension: we add it.
89 def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
90 if os.path.splitext(src)[1] == '.cu':
91 # use the cuda for .cu files
92 self.set_executable('compiler_so', CUDA['nvcc'])
93 # use only a subset of the extra_postargs, which are 1-1 translated
94 # from the extra_compile_args in the Extension class
95 postargs = extra_postargs['nvcc']
97 postargs = extra_postargs['gcc']
99 super(obj, src, ext, cc_args, postargs, pp_opts)
100 # reset the default compiler_so, which we might have changed for cuda
101 self.compiler_so = default_compiler_so
103 # inject our redefined _compile method into the class
104 self._compile = _compile
107 # run the customize_compiler
108 class custom_build_ext(build_ext):
109 def build_extensions(self):
110 customize_compiler_for_nvcc(self.compiler)
111 build_ext.build_extensions(self)
116 # "utils.cython_bbox",
117 # ["utils/bbox.pyx"],
118 # extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]},
119 # include_dirs = [numpy_include]
123 # ["nms/cpu_nms.pyx"],
124 # extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]},
125 # include_dirs = [numpy_include]
129 # "rotation.rotate_cython_nms",
130 # ["rotation/rotate_cython_nms.pyx"],
131 # extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]},
132 # include_dirs = [numpy_include]
136 # "rotation.rotate_circle_nms",
137 # ["rotation/rotate_circle_nms.pyx"],
138 # extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]},
139 # include_dirs = [numpy_include]
142 # Extension('nms.gpu_nms',
143 # ['nms/nms_kernel.cu', 'nms/gpu_nms.pyx'],
144 # library_dirs=[CUDA['lib64']],
145 # libraries=['cudart'],
147 # runtime_library_dirs=[CUDA['lib64']],
148 # # this syntax is specific to this build system
149 # # we're only going to use certain compiler args with nvcc and not with
150 # # gcc the implementation of this trick is in customize_compiler() below
151 # extra_compile_args={'gcc': ["-Wno-unused-function"],
152 # 'nvcc': ['-arch=sm_35',
153 # '--ptxas-options=-v',
155 # '--compiler-options',
157 # include_dirs = [numpy_include, CUDA['include']]
159 # Extension('rotation.rotate_gpu_nms',
160 # ['rotation/rotate_nms_kernel.cu', 'rotation/rotate_gpu_nms.pyx'],
161 # library_dirs=[CUDA['lib64']],
162 # libraries=['cudart'],
164 # runtime_library_dirs=[CUDA['lib64']],
165 # # this syntax is specific to this build system
166 # # we're only going to use certain compiler args with nvcc anrbd not with
167 # # gcc the implementation of this trick is in customize_compiler() below
168 # extra_compile_args={'gcc': ["-Wno-unused-function"],
169 # 'nvcc': ['-arch=sm_35',
170 # '--ptxas-options=-v',
172 # '--compiler-options',
174 # include_dirs = [numpy_include, CUDA['include']]
176 Extension('rotation.rbbox_overlaps',
177 ['rotation/rbbox_overlaps_kernel.cu', 'rotation/rbbox_overlaps.pyx'],
178 library_dirs=[CUDA['lib64']],
179 libraries=['cudart'],
181 runtime_library_dirs=[CUDA['lib64']],
182 # this syntax is specific to this build system
183 # we're only going to use certain compiler args with nvcc and not with
184 # gcc the implementation of this trick is in customize_compiler() below
185 extra_compile_args={'gcc': ["-Wno-unused-function"],
186 'nvcc': ['-arch=sm_35',
187 '--ptxas-options=-v',
189 '--compiler-options',
191 include_dirs = [numpy_include, CUDA['include']]
193 # Extension('rotation.rotate_polygon_nms',
194 # ['rotation/rotate_polygon_nms_kernel.cu', 'rotation/rotate_polygon_nms.pyx'],
195 # library_dirs=[CUDA['lib64']],
196 # libraries=['cudart'],
198 # runtime_library_dirs=[CUDA['lib64']],
199 # # this syntax is specific to this build system
200 # # we're only going to use certain compiler args with nvcc and not with
201 # # gcc the implementation of this trick is in customize_compiler() below
202 # extra_compile_args={'gcc': ["-Wno-unused-function"],
203 # 'nvcc': ['-arch=sm_35',
204 # '--ptxas-options=-v',
206 # '--compiler-options',
208 # include_dirs = [numpy_include, CUDA['include']]
212 # 'pycocotools._mask',
213 # sources=['pycocotools/maskApi.c', 'pycocotools/_mask.pyx'],
214 # include_dirs = [numpy_include, 'pycocotools'],
215 # extra_compile_args={
216 # 'gcc': ['-Wno-cpp', '-Wno-unused-function', '-std=c99']},
222 ext_modules=ext_modules,
223 # inject our custom trigger
224 cmdclass={'build_ext': custom_build_ext},