zhangyiss/spack
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Speed up PyTorch build (#13621)

Browse Source
This commit is contained in:
Adam J. Stewart
2019-11-07 10:05:50 -06:00
committed by GitHub
parent be56fe0677
commit 344e4e0d1c
1 changed files with 18 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
var/spack/repos/builtin/packages/py-torch/package.py
View File

@@ -101,6 +101,9 @@ class PyTorch(PythonPackage):
# Required dependencies
depends_on('cmake@3.5:', type='build')
# Use Ninja generator to speed up build times
# Automatically used if found
depends_on('ninja@1.5:', type='build')
depends_on('python@2.7:2.8,3.5:', type=('build', 'run'))
depends_on('py-setuptools', type='build')
depends_on('py-numpy', type=('build', 'run'))
@@ -141,11 +144,14 @@ class PyTorch(PythonPackage):
depends_on('tbb', when='+tbb')
# Test dependencies
depends_on('ninja', type='test')
depends_on('py-hypothesis', type='test')
depends_on('py-six', type='test')
depends_on('py-psutil', type='test')
# Both build and install run cmake/make/make install
# Only run once to speed up build times
phases = ['install']
def setup_build_environment(self, env):
def enable_or_disable(variant, keyword='USE', var=None, newer=False):
"""Set environment variable to enable or disable support for a
@@ -178,8 +184,12 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
if 'mkl' in self.spec:
env.prepend_path('CMAKE_PREFIX_PATH', self.spec['mkl'].prefix.mkl)
# Build in parallel to speed up build times
env.set('MAX_JOBS', make_jobs)
# Spack logs have trouble handling colored output
env.set('COLORIZE_OUTPUT', 'OFF')
# Don't use vendored third-party libraries
env.set('BUILD_CUSTOM_PROTOBUF', 'OFF')
env.set('USE_PYTORCH_QNNPACK', 'OFF')
@@ -201,7 +211,6 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
enable_or_disable('fbgemm')
enable_or_disable('test', keyword='BUILD')
enable_or_disable('miopen')
if '+miopen' in self.spec:
env.set('MIOPEN_LIB_DIR', self.spec['miopen'].libs.directories[0])
env.set('MIOPEN_INCLUDE_DIR', self.spec['miopen'].prefix.include)
@@ -224,7 +233,6 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
enable_or_disable('caffe2', keyword='BUILD', var='CAFFE2_OPS')
enable_or_disable('gloo', newer=True)
enable_or_disable('gloo', var='IBVERBS', newer=True)
enable_or_disable('opencv', newer=True)
enable_or_disable('openmp', newer=True)
enable_or_disable('ffmpeg', newer=True)
@@ -235,15 +243,19 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
env.set('PYTORCH_BUILD_VERSION', self.version)
env.set('PYTORCH_BUILD_NUMBER', 0)
# BLAS to be used by Caffe2. Can be MKL, Eigen, ATLAS, or OpenBLAS.
# BLAS to be used by Caffe2
if '^mkl' in self.spec:
env.set('BLAS', 'MKL')
elif '^eigen' in self.spec:
env.set('BLAS', 'Eigen')
elif '^atlas' in self.spec:
env.set('BLAS', 'ATLAS')
elif '^openblas' in self.spec:
env.set('BLAS', 'OpenBLAS')
elif '^veclibfort' in self.spec:
env.set('BLAS', 'vecLib')
elif '^libflame' in self.spec:
env.set('BLAS', 'FLAME')
elif '^eigen' in self.spec:
env.set('BLAS', 'Eigen')
enable_or_disable('redis', newer=True)
enable_or_disable('zstd', newer=True)