diff options
| author | Adam J. Stewart <ajstewart426@gmail.com> | 2019-11-07 10:05:50 -0600 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2019-11-07 10:05:50 -0600 |
| commit | 344e4e0d1c139d49caa058cbd672da3e49c2715a (patch) | |
| tree | c3afcd75f64a773fa59dcee09c56053478740777 | |
| parent | be56fe06777451f89ff30e126765945da06247ad (diff) | |
| download | spack-344e4e0d1c139d49caa058cbd672da3e49c2715a.tar.gz spack-344e4e0d1c139d49caa058cbd672da3e49c2715a.tar.bz2 spack-344e4e0d1c139d49caa058cbd672da3e49c2715a.tar.xz spack-344e4e0d1c139d49caa058cbd672da3e49c2715a.zip | |
Speed up PyTorch build (#13621)
| -rw-r--r-- | var/spack/repos/builtin/packages/py-torch/package.py | 24 |
1 files changed, 18 insertions, 6 deletions
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py index 2b609047b5..83000acb32 100644 --- a/var/spack/repos/builtin/packages/py-torch/package.py +++ b/var/spack/repos/builtin/packages/py-torch/package.py @@ -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 @@ class PyTorch(PythonPackage): 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 @@ class PyTorch(PythonPackage): 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 @@ class PyTorch(PythonPackage): 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 @@ class PyTorch(PythonPackage): 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) |