hlrs-software-stack/spack
7
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add latest version of py-torch (#13290)

Browse source 
* Add latest version of py-torch

* setup_environment -> setup_build_environment
This commit is contained in:
Adam J. Stewart 2019-10-21 11:17:15 -05:00 committed by GitHub
parent f2d91f7e3c
commit 5090071693
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 30 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
var/spack/repos/builtin/packages/py-torch/package.py
View file

@ -50,6 +50,7 @@ class PyTorch(PythonPackage):
]
version('master', branch='master', submodules=True)
version('1.3.0', tag='v1.3.0', submodules=True)
version('1.2.0', tag='v1.2.0', submodules=True)
version('1.1.0', tag='v1.1.0', submodules=True)
version('1.0.1', tag='v1.0.1', submodules=True)
@ -135,7 +136,13 @@ class PyTorch(PythonPackage):
depends_on('zstd', when='+zstd')
depends_on('tbb', when='+tbb')
def setup_environment(self, build_env, run_env):
# 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')
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
particular variant.
@ -153,43 +160,39 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
# But some newer variants have always used USE_* or BUILD_*
if self.spec.satisfies('@1.1:') or newer:
if '+' + variant in self.spec:
build_env.set(keyword + '_' + var, 'ON')
env.set(keyword + '_' + var, 'ON')
else:
build_env.set(keyword + '_' + var, 'OFF')
env.set(keyword + '_' + var, 'OFF')
else:
if '+' + variant in self.spec:
build_env.unset('NO_' + var)
env.unset('NO_' + var)
else:
build_env.set('NO_' + var, 'ON')
env.set('NO_' + var, 'ON')
build_env.set('MAX_JOBS', make_jobs)
env.set('MAX_JOBS', make_jobs)
enable_or_disable('cuda')
if '+cuda' in self.spec:
build_env.set('CUDA_HOME', self.spec['cuda'].prefix)
env.set('CUDA_HOME', self.spec['cuda'].prefix)
enable_or_disable('cudnn')
if '+cudnn' in self.spec:
build_env.set('CUDNN_LIB_DIR',
self.spec['cudnn'].libs.directories[0])
build_env.set('CUDNN_INCLUDE_DIR',
self.spec['cudnn'].prefix.include)
build_env.set('CUDNN_LIBRARY', self.spec['cudnn'].libs[0])
env.set('CUDNN_LIB_DIR', self.spec['cudnn'].libs.directories[0])
env.set('CUDNN_INCLUDE_DIR', self.spec['cudnn'].prefix.include)
env.set('CUDNN_LIBRARY', self.spec['cudnn'].libs[0])
enable_or_disable('fbgemm')
enable_or_disable('test', keyword='BUILD')
enable_or_disable('miopen')
if '+miopen' in self.spec:
build_env.set('MIOPEN_LIB_DIR',
self.spec['miopen'].libs.directories[0])
build_env.set('MIOPEN_INCLUDE_DIR',
self.spec['miopen'].prefix.include)
build_env.set('MIOPEN_LIBRARY', self.spec['miopen'].libs[0])
env.set('MIOPEN_LIB_DIR', self.spec['miopen'].libs.directories[0])
env.set('MIOPEN_INCLUDE_DIR', self.spec['miopen'].prefix.include)
env.set('MIOPEN_LIBRARY', self.spec['miopen'].libs[0])
enable_or_disable('mkldnn')
if '+mkldnn' in self.spec:
build_env.set('MKLDNN_HOME', self.spec['intel-mkl'].prefix)
env.set('MKLDNN_HOME', self.spec['intel-mkl'].prefix)
enable_or_disable('nnpack')
enable_or_disable('qnnpack')
@ -198,10 +201,9 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
enable_or_disable('nccl')
enable_or_disable('nccl', var='SYSTEM_NCCL')
if '+nccl' in self.spec:
build_env.set('NCCL_ROOT', self.spec['nccl'].prefix)
build_env.set('NCCL_LIB_DIR',
self.spec['nccl'].libs.directories[0])
build_env.set('NCCL_INCLUDE_DIR', self.spec['nccl'].prefix.include)
env.set('NCCL_ROOT', self.spec['nccl'].prefix)
env.set('NCCL_LIB_DIR', self.spec['nccl'].libs.directories[0])
env.set('NCCL_INCLUDE_DIR', self.spec['nccl'].prefix.include)
enable_or_disable('caffe2', keyword='BUILD', var='CAFFE2_OPS')
enable_or_disable('gloo', newer=True)
@ -213,18 +215,18 @@ def enable_or_disable(variant, keyword='USE', var=None, newer=False):
enable_or_disable('lmdb', newer=True)
enable_or_disable('binary', keyword='BUILD', newer=True)
build_env.set('PYTORCH_BUILD_VERSION', self.version)
build_env.set('PYTORCH_BUILD_NUMBER', 0)
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.
if '^mkl' in self.spec:
build_env.set('BLAS', 'MKL')
env.set('BLAS', 'MKL')
elif '^eigen' in self.spec:
build_env.set('BLAS', 'Eigen')
env.set('BLAS', 'Eigen')
elif '^atlas' in self.spec:
build_env.set('BLAS', 'ATLAS')
env.set('BLAS', 'ATLAS')
elif '^openblas' in self.spec:
build_env.set('BLAS', 'OpenBLAS')
env.set('BLAS', 'OpenBLAS')
enable_or_disable('redis', newer=True)
enable_or_disable('zstd', newer=True)