1818# Monkey-patch setuptools to compile in parallel
1919################################################################################
2020
21- def parallelCCompile (self , sources , output_dir = None , macros = None , include_dirs = None , debug = 0 , extra_preargs = None , extra_postargs = None , depends = None ):
21+
22+ def parallelCCompile (self , sources , output_dir = None , macros = None ,
23+ include_dirs = None , debug = 0 , extra_preargs = None ,
24+ extra_postargs = None , depends = None ):
2225 # those lines are copied from distutils.ccompiler.CCompiler directly
23- macros , objects , extra_postargs , pp_opts , build = self ._setup_compile (output_dir , macros , include_dirs , sources , depends , extra_postargs )
26+ macros , objects , extra_postargs , pp_opts , build = self ._setup_compile (
27+ output_dir , macros , include_dirs , sources , depends , extra_postargs )
2428 cc_args = self ._get_cc_args (pp_opts , debug , extra_preargs )
2529
2630 # compile using a thread pool
2731 import multiprocessing .pool
32+
2833 def _single_compile (obj ):
2934 src , ext = build [obj ]
3035 self ._compile (obj , src , ext , cc_args , extra_postargs , pp_opts )
@@ -39,6 +44,7 @@ def _single_compile(obj):
3944# Custom build commands
4045################################################################################
4146
47+
4248class build_deps (Command ):
4349 user_options = []
4450
@@ -73,6 +79,7 @@ def run(self):
7379
7480
7581class build_ext (setuptools .command .build_ext .build_ext ):
82+
7683 def run (self ):
7784 # Print build options
7885 if WITH_NUMPY :
@@ -116,13 +123,15 @@ class build(distutils.command.build.build):
116123
117124
118125class install (setuptools .command .install .install ):
126+
119127 def run (self ):
120128 if not self .skip_build :
121129 self .run_command ('build_deps' )
122130 setuptools .command .install .install .run (self )
123131
124132
125133class clean (distutils .command .clean .clean ):
134+
126135 def run (self ):
127136 import glob
128137 with open ('.gitignore' , 'r' ) as f :
@@ -138,7 +147,6 @@ def run(self):
138147 distutils .command .clean .clean .run (self )
139148
140149
141-
142150################################################################################
143151# Configure compile flags
144152################################################################################
@@ -166,21 +174,21 @@ def run(self):
166174extra_link_args .append ('-L' + lib_path )
167175
168176# we specify exact lib names to avoid conflict with lua-torch installs
169- TH_LIB = os .path .join (lib_path , 'libTH.so.1' )
170- THS_LIB = os .path .join (lib_path , 'libTHS.so.1' )
171- THC_LIB = os .path .join (lib_path , 'libTHC.so.1' )
172- THCS_LIB = os .path .join (lib_path , 'libTHCS.so.1' )
173- THNN_LIB = os .path .join (lib_path , 'libTHNN.so.1' )
177+ TH_LIB = os .path .join (lib_path , 'libTH.so.1' )
178+ THS_LIB = os .path .join (lib_path , 'libTHS.so.1' )
179+ THC_LIB = os .path .join (lib_path , 'libTHC.so.1' )
180+ THCS_LIB = os .path .join (lib_path , 'libTHCS.so.1' )
181+ THNN_LIB = os .path .join (lib_path , 'libTHNN.so.1' )
174182THCUNN_LIB = os .path .join (lib_path , 'libTHCUNN.so.1' )
175- THPP_LIB = os .path .join (lib_path , 'libTHPP.so.1' )
183+ THPP_LIB = os .path .join (lib_path , 'libTHPP.so.1' )
176184if platform .system () == 'Darwin' :
177- TH_LIB = os .path .join (lib_path , 'libTH.1.dylib' )
178- THS_LIB = os .path .join (lib_path , 'libTHS.1.dylib' )
179- THC_LIB = os .path .join (lib_path , 'libTHC.1.dylib' )
180- THCS_LIB = os .path .join (lib_path , 'libTHCS.1.dylib' )
181- THNN_LIB = os .path .join (lib_path , 'libTHNN.1.dylib' )
185+ TH_LIB = os .path .join (lib_path , 'libTH.1.dylib' )
186+ THS_LIB = os .path .join (lib_path , 'libTHS.1.dylib' )
187+ THC_LIB = os .path .join (lib_path , 'libTHC.1.dylib' )
188+ THCS_LIB = os .path .join (lib_path , 'libTHCS.1.dylib' )
189+ THNN_LIB = os .path .join (lib_path , 'libTHNN.1.dylib' )
182190 THCUNN_LIB = os .path .join (lib_path , 'libTHCUNN.1.dylib' )
183- THPP_LIB = os .path .join (lib_path , 'libTHPP.1.dylib' )
191+ THPP_LIB = os .path .join (lib_path , 'libTHPP.1.dylib' )
184192
185193main_compile_args = ['-D_THP_CORE' ]
186194main_libraries = ['shm' ]
@@ -267,70 +275,70 @@ def make_relative_rpath(path):
267275packages = find_packages (exclude = ('tools.*' ,))
268276
269277C = Extension ("torch._C" ,
270- libraries = main_libraries ,
271- sources = main_sources ,
272- language = 'c++' ,
273- extra_compile_args = main_compile_args + extra_compile_args ,
274- include_dirs = include_dirs ,
275- extra_link_args = extra_link_args + main_link_args + [make_relative_rpath ('lib' )],
276- )
278+ libraries = main_libraries ,
279+ sources = main_sources ,
280+ language = 'c++' ,
281+ extra_compile_args = main_compile_args + extra_compile_args ,
282+ include_dirs = include_dirs ,
283+ extra_link_args = extra_link_args + main_link_args + [make_relative_rpath ('lib' )],
284+ )
277285extensions .append (C )
278286
279287DL = Extension ("torch._dl" ,
280- sources = ["torch/csrc/dl.c" ],
281- language = 'c' ,
282- )
288+ sources = ["torch/csrc/dl.c" ],
289+ language = 'c' ,
290+ )
283291extensions .append (DL )
284292
285293THNN = Extension ("torch._thnn._THNN" ,
286- sources = ['torch/csrc/nn/THNN.cpp' ],
287- language = 'c++' ,
288- extra_compile_args = extra_compile_args ,
289- include_dirs = include_dirs ,
290- extra_link_args = extra_link_args + [
291- TH_LIB ,
292- THNN_LIB ,
293- make_relative_rpath ('../lib' ),
294- ]
295- )
294+ sources = ['torch/csrc/nn/THNN.cpp' ],
295+ language = 'c++' ,
296+ extra_compile_args = extra_compile_args ,
297+ include_dirs = include_dirs ,
298+ extra_link_args = extra_link_args + [
299+ TH_LIB ,
300+ THNN_LIB ,
301+ make_relative_rpath ('../lib' ),
302+ ]
303+ )
296304extensions .append (THNN )
297305
298306if WITH_CUDA :
299307 THCUNN = Extension ("torch._thnn._THCUNN" ,
300- sources = ['torch/csrc/nn/THCUNN.cpp' ],
301- language = 'c++' ,
302- extra_compile_args = extra_compile_args ,
303- include_dirs = include_dirs ,
304- extra_link_args = extra_link_args + [
305- TH_LIB ,
306- THC_LIB ,
307- THCUNN_LIB ,
308- make_relative_rpath ('../lib' ),
309- ]
310- )
308+ sources = ['torch/csrc/nn/THCUNN.cpp' ],
309+ language = 'c++' ,
310+ extra_compile_args = extra_compile_args ,
311+ include_dirs = include_dirs ,
312+ extra_link_args = extra_link_args + [
313+ TH_LIB ,
314+ THC_LIB ,
315+ THCUNN_LIB ,
316+ make_relative_rpath ('../lib' ),
317+ ]
318+ )
311319 extensions .append (THCUNN )
312320
313- version = "0.1"
321+ version = "0.1"
314322if os .getenv ('PYTORCH_BUILD_VERSION' ):
315323 version = os .getenv ('PYTORCH_BUILD_VERSION' ) \
316- + '_' + os .getenv ('PYTORCH_BUILD_NUMBER' )
324+ + '_' + os .getenv ('PYTORCH_BUILD_NUMBER' )
317325
318326setup (name = "torch" , version = version ,
319- ext_modules = extensions ,
320- cmdclass = {
321- 'build' : build ,
322- 'build_ext' : build_ext ,
323- 'build_deps' : build_deps ,
324- 'build_module' : build_module ,
325- 'install' : install ,
326- 'clean' : clean ,
327- },
328- packages = packages ,
329- package_data = {'torch' : [
330- 'lib/*.so*' , 'lib/*.dylib*' ,
331- 'lib/torch_shm_manager' ,
332- 'lib/*.h' ,
333- 'lib/include/TH/*.h' , 'lib/include/TH/generic/*.h' ,
334- 'lib/include/THC/*.h' , 'lib/include/THC/generic/*.h' ]},
335- install_requires = ['pyyaml' ],
336- )
327+ ext_modules = extensions ,
328+ cmdclass = {
329+ 'build' : build ,
330+ 'build_ext' : build_ext ,
331+ 'build_deps' : build_deps ,
332+ 'build_module' : build_module ,
333+ 'install' : install ,
334+ 'clean' : clean ,
335+ },
336+ packages = packages ,
337+ package_data = {'torch' : [
338+ 'lib/*.so*' , 'lib/*.dylib*' ,
339+ 'lib/torch_shm_manager' ,
340+ 'lib/*.h' ,
341+ 'lib/include/TH/*.h' , 'lib/include/TH/generic/*.h' ,
342+ 'lib/include/THC/*.h' , 'lib/include/THC/generic/*.h' ]},
343+ install_requires = ['pyyaml' ],
344+ )
0 commit comments