from setuptools import setup, Extension, distutils, Command, find_packages import setuptools.command.build_ext import setuptools.command.install import setuptools.command.develop import setuptools.command.build_py import distutils.unixccompiler import distutils.command.build import distutils.command.clean import platform import subprocess import shutil import sys import os from tools.setup_helpers.env import check_env_flag from tools.setup_helpers.cuda import WITH_CUDA, CUDA_HOME from tools.setup_helpers.cudnn import WITH_CUDNN, CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR from tools.setup_helpers.split_types import split_types DEBUG = check_env_flag('DEBUG') WITH_DISTRIBUTED = not check_env_flag('NO_DISTRIBUTED') WITH_DISTRIBUTED_MW = WITH_DISTRIBUTED and check_env_flag('WITH_DISTRIBUTED_MW') WITH_NCCL = WITH_CUDA and platform.system() != 'Darwin' SYSTEM_NCCL = False ################################################################################ # Workaround setuptools -Wstrict-prototypes warnings # I lifted this code from https://stackoverflow.com/a/29634231/23845 ################################################################################ import distutils.sysconfig cfg_vars = distutils.sysconfig.get_config_vars() for key, value in cfg_vars.items(): if type(value) == str: cfg_vars[key] = value.replace("-Wstrict-prototypes", "") ################################################################################ # Monkey-patch setuptools to compile in parallel ################################################################################ original_link = distutils.unixccompiler.UnixCCompiler.link def parallelCCompile(self, sources, output_dir=None, macros=None, include_dirs=None, debug=0, extra_preargs=None, extra_postargs=None, depends=None): # those lines are copied from distutils.ccompiler.CCompiler directly macros, objects, extra_postargs, pp_opts, build = self._setup_compile( output_dir, macros, include_dirs, sources, depends, extra_postargs) cc_args = self._get_cc_args(pp_opts, debug, extra_preargs) # compile using a thread pool import multiprocessing.pool def _single_compile(obj): src, ext = build[obj] self._compile(obj, src, ext, cc_args, extra_postargs, pp_opts) num_jobs = multiprocessing.cpu_count() max_jobs = os.getenv("MAX_JOBS") if max_jobs is not None: num_jobs = min(num_jobs, int(max_jobs)) multiprocessing.pool.ThreadPool(num_jobs).map(_single_compile, objects) return objects def patched_link(self, *args, **kwargs): _cxx = self.compiler_cxx self.compiler_cxx = None result = original_link(self, *args, **kwargs) self.compiler_cxx = _cxx return result distutils.ccompiler.CCompiler.compile = parallelCCompile distutils.unixccompiler.UnixCCompiler.link = patched_link ################################################################################ # Custom build commands ################################################################################ dep_libs = [ 'TH', 'THS', 'THNN', 'THC', 'THCS', 'THCUNN', 'nccl', 'THPP', 'libshm', 'ATen', 'gloo', 'THD', 'nanopb', ] def build_libs(libs): for lib in libs: assert lib in dep_libs, 'invalid lib: {}'.format(lib) build_libs_cmd = ['bash', 'torch/lib/build_libs.sh'] my_env = os.environ.copy() my_env["PYTORCH_PYTHON"] = sys.executable if WITH_CUDA: build_libs_cmd += ['--with-cuda'] if subprocess.call(build_libs_cmd + libs, env=my_env) != 0: sys.exit(1) class build_deps(Command): user_options = [] def initialize_options(self): pass def finalize_options(self): pass def run(self): libs = ['TH', 'THS', 'THNN'] if WITH_CUDA: libs += ['THC', 'THCS', 'THCUNN'] if WITH_NCCL and not SYSTEM_NCCL: libs += ['nccl'] libs += ['THPP', 'libshm', 'ATen', 'nanopb'] if WITH_DISTRIBUTED: if sys.platform.startswith('linux'): libs += ['gloo'] libs += ['THD'] build_libs(libs) from tools.nnwrap import generate_wrappers as generate_nn_wrappers generate_nn_wrappers() build_dep_cmds = {} for lib in dep_libs: # wrap in function to capture lib class build_dep(build_deps): description = 'Build {} external library'.format(lib) def run(self): build_libs([self.lib]) build_dep.lib = lib build_dep_cmds['build_' + lib.lower()] = build_dep class build_module(Command): user_options = [] def initialize_options(self): pass def finalize_options(self): pass def run(self): self.run_command('build_py') self.run_command('build_ext') class build_py(setuptools.command.build_py.build_py): def run(self): self.create_version_file() setuptools.command.build_py.build_py.run(self) @staticmethod def create_version_file(): global version, cwd print('-- Building version ' + version) version_path = os.path.join(cwd, 'torch', 'version.py') with open(version_path, 'w') as f: f.write("__version__ = '{}'\n".format(version)) class develop(setuptools.command.develop.develop): def run(self): build_py.create_version_file() setuptools.command.develop.develop.run(self) def monkey_patch_THD_link_flags(): ''' THD's dynamic link deps are not determined until after build_deps is run So, we need to monkey-patch them in later ''' # read tmp_install_path/THD_deps.txt for THD's dynamic linkage deps with open(tmp_install_path + '/THD_deps.txt', 'r') as f: thd_deps_ = f.read() thd_deps = [] # remove empty lines for l in thd_deps_.split(';'): if l != '': thd_deps.append(l) C.extra_link_args += thd_deps class build_ext(setuptools.command.build_ext.build_ext): def run(self): # Print build options if WITH_NUMPY: print('-- Building with NumPy bindings') else: print('-- NumPy not found') if WITH_CUDNN: print('-- Detected cuDNN at ' + CUDNN_LIB_DIR + ', ' + CUDNN_INCLUDE_DIR) else: print('-- Not using cuDNN') if WITH_CUDA: print('-- Detected CUDA at ' + CUDA_HOME) else: print('-- Not using CUDA') if WITH_NCCL and SYSTEM_NCCL: print('-- Using system provided NCCL library') elif WITH_NCCL: print('-- Building NCCL library') else: print('-- Not using NCCL') if WITH_DISTRIBUTED: print('-- Building with distributed package ') monkey_patch_THD_link_flags() else: print('-- Building without distributed package') # cwrap depends on pyyaml, so we can't import it earlier from tools.cwrap import cwrap from tools.cwrap.plugins.THPPlugin import THPPlugin from tools.cwrap.plugins.ArgcountSortPlugin import ArgcountSortPlugin from tools.cwrap.plugins.AutoGPU import AutoGPU from tools.cwrap.plugins.BoolOption import BoolOption from tools.cwrap.plugins.KwargsPlugin import KwargsPlugin from tools.cwrap.plugins.NullableArguments import NullableArguments from tools.cwrap.plugins.CuDNNPlugin import CuDNNPlugin from tools.cwrap.plugins.WrapDim import WrapDim from tools.cwrap.plugins.AssertNDim import AssertNDim from tools.cwrap.plugins.Broadcast import Broadcast from tools.cwrap.plugins.ProcessorSpecificPlugin import ProcessorSpecificPlugin from tools.autograd.gen_variable_type import gen_variable_type thp_plugin = THPPlugin() cwrap('torch/csrc/generic/TensorMethods.cwrap', plugins=[ ProcessorSpecificPlugin(), BoolOption(), thp_plugin, AutoGPU(condition='IS_CUDA'), ArgcountSortPlugin(), KwargsPlugin(), AssertNDim(), WrapDim(), Broadcast() ]) cwrap('torch/csrc/cudnn/cuDNN.cwrap', plugins=[ CuDNNPlugin(), NullableArguments() ]) # Build ATen based Variable classes autograd_gen_dir = 'torch/csrc/autograd/generated' if not os.path.exists(autograd_gen_dir): os.mkdir(autograd_gen_dir) gen_variable_type( 'torch/lib/build/ATen/ATen/Declarations.yaml', autograd_gen_dir) # It's an old-style class in Python 2.7... setuptools.command.build_ext.build_ext.run(self) class build(distutils.command.build.build): sub_commands = [ ('build_deps', lambda self: True), ] + distutils.command.build.build.sub_commands class install(setuptools.command.install.install): def run(self): if not self.skip_build: self.run_command('build_deps') setuptools.command.install.install.run(self) class clean(distutils.command.clean.clean): def run(self): import glob with open('.gitignore', 'r') as f: ignores = f.read() for wildcard in filter(bool, ignores.split('\n')): for filename in glob.glob(wildcard): try: os.remove(filename) except OSError: shutil.rmtree(filename, ignore_errors=True) # It's an old-style class in Python 2.7... distutils.command.clean.clean.run(self) ################################################################################ # Configure compile flags ################################################################################ include_dirs = [] library_dirs = [] extra_link_args = [] extra_compile_args = ['-std=c++11', '-Wno-write-strings', # Python 2.6 requires -fno-strict-aliasing, see # http://legacy.python.org/dev/peps/pep-3123/ '-fno-strict-aliasing'] cwd = os.path.dirname(os.path.abspath(__file__)) lib_path = os.path.join(cwd, "torch", "lib") # Check if you remembered to check out submodules def check_file(f): if not os.path.exists(f): print("Could not find {}".format(f)) print("Did you run 'git submodule update --init'?") sys.exit(1) check_file(os.path.join(lib_path, "gloo", "CMakeLists.txt")) check_file(os.path.join(lib_path, "nanopb", "CMakeLists.txt")) check_file(os.path.join(lib_path, "pybind11", "CMakeLists.txt")) tmp_install_path = lib_path + "/tmp_install" include_dirs += [ cwd, os.path.join(cwd, "torch", "csrc"), lib_path + "/pybind11/include", tmp_install_path + "/include", tmp_install_path + "/include/TH", tmp_install_path + "/include/THPP", tmp_install_path + "/include/THNN", tmp_install_path + "/include/ATen", ] library_dirs.append(lib_path) # we specify exact lib names to avoid conflict with lua-torch installs TH_LIB = os.path.join(lib_path, 'libTH.so.1') THS_LIB = os.path.join(lib_path, 'libTHS.so.1') THC_LIB = os.path.join(lib_path, 'libTHC.so.1') THCS_LIB = os.path.join(lib_path, 'libTHCS.so.1') THNN_LIB = os.path.join(lib_path, 'libTHNN.so.1') THCUNN_LIB = os.path.join(lib_path, 'libTHCUNN.so.1') ATEN_LIB = os.path.join(lib_path, 'libATen.so.1') THD_LIB = os.path.join(lib_path, 'libTHD.a') NCCL_LIB = os.path.join(lib_path, 'libnccl.so.1') if platform.system() == 'Darwin': TH_LIB = os.path.join(lib_path, 'libTH.1.dylib') THS_LIB = os.path.join(lib_path, 'libTHS.1.dylib') THC_LIB = os.path.join(lib_path, 'libTHC.1.dylib') THCS_LIB = os.path.join(lib_path, 'libTHCS.1.dylib') THNN_LIB = os.path.join(lib_path, 'libTHNN.1.dylib') THCUNN_LIB = os.path.join(lib_path, 'libTHCUNN.1.dylib') ATEN_LIB = os.path.join(lib_path, 'libATen.1.dylib') NCCL_LIB = os.path.join(lib_path, 'libnccl.1.dylib') # static library only NANOPB_STATIC_LIB = os.path.join(lib_path, 'libprotobuf-nanopb.a') if WITH_NCCL and (subprocess.call('ldconfig -p | grep libnccl >/dev/null', shell=True) == 0 or subprocess.call('/sbin/ldconfig -p | grep libnccl >/dev/null', shell=True) == 0): SYSTEM_NCCL = True main_compile_args = ['-D_THP_CORE'] main_libraries = ['shm'] main_link_args = [TH_LIB, THS_LIB, THNN_LIB, ATEN_LIB, NANOPB_STATIC_LIB] main_sources = [ "torch/csrc/onnx.pb.cpp", "torch/csrc/onnx.cpp", "torch/csrc/PtrWrapper.cpp", "torch/csrc/Module.cpp", "torch/csrc/Generator.cpp", "torch/csrc/Size.cpp", "torch/csrc/Exceptions.cpp", "torch/csrc/Storage.cpp", "torch/csrc/DynamicTypes.cpp", "torch/csrc/byte_order.cpp", "torch/csrc/utils.cpp", "torch/csrc/expand_utils.cpp", "torch/csrc/utils/invalid_arguments.cpp", "torch/csrc/utils/object_ptr.cpp", "torch/csrc/utils/tuple_parser.cpp", "torch/csrc/allocators.cpp", "torch/csrc/serialization.cpp", "torch/csrc/jit/assert.cpp", "torch/csrc/jit/init.cpp", "torch/csrc/jit/ir.cpp", "torch/csrc/jit/python_ir.cpp", "torch/csrc/jit/test_jit.cpp", "torch/csrc/jit/tracer.cpp", "torch/csrc/jit/python_tracer.cpp", "torch/csrc/jit/interned_strings.cpp", "torch/csrc/jit/passes/graph_fuser.cpp", "torch/csrc/jit/passes/init_pass.cpp", "torch/csrc/jit/passes/dead_code_elimination.cpp", "torch/csrc/autograd/init.cpp", "torch/csrc/autograd/engine.cpp", "torch/csrc/autograd/function.cpp", "torch/csrc/autograd/variable.cpp", "torch/csrc/autograd/saved_variable.cpp", "torch/csrc/autograd/input_buffer.cpp", "torch/csrc/autograd/python_function.cpp", "torch/csrc/autograd/python_cpp_function.cpp", "torch/csrc/autograd/python_variable.cpp", "torch/csrc/autograd/python_engine.cpp", "torch/csrc/autograd/python_hook.cpp", "torch/csrc/autograd/functions/jit_closure.cpp", "torch/csrc/autograd/generated/VariableType.cpp", "torch/csrc/autograd/generated/Functions.cpp", "torch/csrc/autograd/generated/python_variable_methods.cpp", "torch/csrc/autograd/functions/batch_normalization.cpp", "torch/csrc/autograd/functions/convolution.cpp", "torch/csrc/autograd/functions/basic_ops.cpp", "torch/csrc/autograd/functions/tensor.cpp", "torch/csrc/autograd/functions/accumulate_grad.cpp", "torch/csrc/autograd/functions/special.cpp", "torch/csrc/autograd/functions/utils.cpp", "torch/csrc/autograd/functions/init.cpp", "torch/csrc/autograd/functions/onnx/convolution.cpp", "torch/csrc/autograd/functions/onnx/batch_normalization.cpp", "torch/csrc/onnx/export.cpp", ] main_sources += split_types("torch/csrc/Tensor.cpp") try: import numpy as np include_dirs += [np.get_include()] extra_compile_args += ['-DWITH_NUMPY'] WITH_NUMPY = True except ImportError: WITH_NUMPY = False if WITH_DISTRIBUTED: extra_compile_args += ['-DWITH_DISTRIBUTED'] main_sources += [ "torch/csrc/distributed/Module.cpp", "torch/csrc/distributed/utils.cpp", ] if WITH_DISTRIBUTED_MW: main_sources += [ "torch/csrc/distributed/Tensor.cpp", "torch/csrc/distributed/Storage.cpp", ] extra_compile_args += ['-DWITH_DISTRIBUTED_MW'] include_dirs += [tmp_install_path + "/include/THD"] main_link_args += [THD_LIB] if WITH_CUDA: cuda_lib_dirs = ['lib64', 'lib'] cuda_include_path = os.path.join(CUDA_HOME, 'include') for lib_dir in cuda_lib_dirs: cuda_lib_path = os.path.join(CUDA_HOME, lib_dir) if os.path.exists(cuda_lib_path): break include_dirs.append(cuda_include_path) include_dirs.append(tmp_install_path + "/include/THCUNN") library_dirs.append(cuda_lib_path) extra_link_args.append('-Wl,-rpath,' + cuda_lib_path) extra_compile_args += ['-DWITH_CUDA'] extra_compile_args += ['-DCUDA_LIB_PATH=' + cuda_lib_path] main_libraries += ['cudart', 'nvToolsExt', 'nvrtc', 'cuda'] main_link_args += [THC_LIB, THCS_LIB, THCUNN_LIB] main_sources += [ "torch/csrc/cuda/Module.cpp", "torch/csrc/cuda/Storage.cpp", "torch/csrc/cuda/Stream.cpp", "torch/csrc/cuda/AutoGPU.cpp", "torch/csrc/cuda/utils.cpp", "torch/csrc/cuda/expand_utils.cpp", "torch/csrc/cuda/serialization.cpp", "torch/csrc/jit/fusion_compiler.cpp", ] main_sources += split_types("torch/csrc/cuda/Tensor.cpp") if WITH_NCCL: if SYSTEM_NCCL: main_libraries += ['nccl'] else: main_link_args += [NCCL_LIB] extra_compile_args += ['-DWITH_NCCL'] if WITH_CUDNN: main_libraries += ['cudnn'] include_dirs.append(CUDNN_INCLUDE_DIR) library_dirs.append(CUDNN_LIB_DIR) main_sources += [ "torch/csrc/cudnn/BatchNorm.cpp", "torch/csrc/cudnn/Conv.cpp", "torch/csrc/cudnn/cuDNN.cpp", "torch/csrc/cudnn/GridSampler.cpp", "torch/csrc/cudnn/AffineGridGenerator.cpp", "torch/csrc/cudnn/Types.cpp", "torch/csrc/cudnn/Handles.cpp", ] extra_compile_args += ['-DWITH_CUDNN'] if DEBUG: extra_compile_args += ['-O0', '-g'] extra_link_args += ['-O0', '-g'] if os.getenv('PYTORCH_BINARY_BUILD') and platform.system() == 'Linux': print('PYTORCH_BINARY_BUILD found. Static linking libstdc++ on Linux') # get path of libstdc++ and link manually. # for reasons unknown, -static-libstdc++ doesn't fully link some symbols CXXNAME = os.getenv('CXX', 'g++') STDCPP_LIB = subprocess.check_output([CXXNAME, '-print-file-name=libstdc++.a']) STDCPP_LIB = STDCPP_LIB[:-1] if type(STDCPP_LIB) != str: # python 3 STDCPP_LIB = STDCPP_LIB.decode(sys.stdout.encoding) main_link_args += [STDCPP_LIB] version_script = os.path.abspath("tools/pytorch.version") extra_link_args += ['-Wl,--version-script=' + version_script] def make_relative_rpath(path): if platform.system() == 'Darwin': return '-Wl,-rpath,@loader_path/' + path else: return '-Wl,-rpath,$ORIGIN/' + path ################################################################################ # Declare extensions and package ################################################################################ extensions = [] packages = find_packages(exclude=('tools', 'tools.*',)) C = Extension("torch._C", libraries=main_libraries, sources=main_sources, language='c++', extra_compile_args=main_compile_args + extra_compile_args, include_dirs=include_dirs, library_dirs=library_dirs, extra_link_args=extra_link_args + main_link_args + [make_relative_rpath('lib')], ) extensions.append(C) DL = Extension("torch._dl", sources=["torch/csrc/dl.c"], language='c', ) extensions.append(DL) THNN = Extension("torch._thnn._THNN", sources=['torch/csrc/nn/THNN.cpp'], language='c++', extra_compile_args=extra_compile_args, include_dirs=include_dirs, extra_link_args=extra_link_args + [ TH_LIB, THNN_LIB, make_relative_rpath('../lib'), ] ) extensions.append(THNN) if WITH_CUDA: THCUNN = Extension("torch._thnn._THCUNN", sources=['torch/csrc/nn/THCUNN.cpp'], language='c++', extra_compile_args=extra_compile_args, include_dirs=include_dirs, extra_link_args=extra_link_args + [ TH_LIB, THC_LIB, THCUNN_LIB, make_relative_rpath('../lib'), ] ) extensions.append(THCUNN) version = '0.2.0' if os.getenv('PYTORCH_BUILD_VERSION'): assert os.getenv('PYTORCH_BUILD_NUMBER') is not None version = os.getenv('PYTORCH_BUILD_VERSION') \ + '_' + os.getenv('PYTORCH_BUILD_NUMBER') else: try: sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=cwd).decode('ascii').strip() version += '+' + sha[:7] except subprocess.CalledProcessError: pass cmdclass = { 'build': build, 'build_py': build_py, 'build_ext': build_ext, 'build_deps': build_deps, 'build_module': build_module, 'develop': develop, 'install': install, 'clean': clean, } cmdclass.update(build_dep_cmds) setup(name="torch", version=version, description="Tensors and Dynamic neural networks in Python with strong GPU acceleration", ext_modules=extensions, cmdclass=cmdclass, packages=packages, package_data={'torch': [ 'lib/*.so*', 'lib/*.dylib*', 'lib/torch_shm_manager', 'lib/*.h', 'lib/include/TH/*.h', 'lib/include/TH/generic/*.h', 'lib/include/THC/*.h', 'lib/include/THC/generic/*.h', 'lib/include/ATen/*.h', ]}, install_requires=['pyyaml', 'numpy'], )