Overhaul PyTorch package (#12420)

* Overhaul PyTorch package

* Flake8

* Distributed requires MPI, NCCL not specified correctly

* Binary was being specified incorrectly

* Run tests at install time

* Disable GLOO in addition to GLOO_IBVERBS

1 files changed, 155 insertions, 30 deletions

diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index 9a75cd956a..ce432ee33b 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -10,8 +10,9 @@ class PyTorch(PythonPackage):
     """Tensors and Dynamic neural networks in Python
     with strong GPU acceleration."""
 
-    homepage = "http://pytorch.org/"
+    homepage = "https://pytorch.org/"
     git      = "https://github.com/pytorch/pytorch.git"
+    install_time_test_callbacks = ['install_test', 'import_module_test']
 
     maintainers = ['adamjstewart']
     import_modules = [
@@ -48,6 +49,8 @@ class PyTorch(PythonPackage):
         'torch.utils.data', 'torch.utils.data._utils'
     ]
 
+    version('master', branch='master', 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)
     version('1.0.0', tag='v1.0.0', submodules=True)
@@ -55,59 +58,181 @@ class PyTorch(PythonPackage):
     version('0.4.0', tag='v0.4.0', submodules=True)
     version('0.3.1', tag='v0.3.1', submodules=True)
 
-    variant('cuda', default='False', description='Add GPU support')
-    variant('cudnn', default='False', description='Add cuDNN support')
-    variant('nccl', default='False', description='Add NCCL support')
-    variant('mkldnn', default='False', description='Add Intel MKL DNN support')
-    variant('magma', default='False', description='Add MAGMA support')
+    variant('cuda', default=True, description='Enables CUDA build')
+    variant('cudnn', default=False, description='Enables the cuDNN build')
+    variant('magma', default=False, description='Enables the MAGMA build')
+    variant('fbgemm', default=False, description='Enables the FBGEMM build')
+    variant('test', default=False, description='Enables the test build')
+    variant('miopen', default=False, description='Enables the MIOpen build')
+    variant('mkldnn', default=False, description='Enables use of MKLDNN')
+    variant('nnpack', default=False, description='Enables NNPACK build')
+    variant('qnnpack', default=False, description='Enables QNNPACK build (quantized 8-bit operators)')
+    variant('distributed', default=False, description='Enables distributed (c10d, gloo, mpi, etc.) build')
+    variant('nccl', default=False, description='Use Spack-installed NCCL')
+    variant('caffe2', default=False, description='Enables Caffe2 operators build')
+    variant('gloo', default=False, description='Enables features related to distributed support')
+    variant('opencv', default=False, description='Enables use of OpenCV for additional operators')
+    variant('openmp', default=True, description='Enables use of OpenMP for parallelization')
+    variant('ffmpeg', default=False, description='Enables use of ffmpeg for additional operators')
+    variant('leveldb', default=False, description='Enables use of LevelDB for storage')
+    variant('lmdb', default=False, description='Enables use of LMDB for storage')
+    variant('binary', default=False, description='Enables the additional binaries/ build')
+    variant('redis', default=False, description='Use Redis for distributed workflows')
+    variant('zstd', default=False, description='Enables use of ZSTD')
+    variant('tbb', default=False, description='Enables TBB support')
 
     conflicts('+cudnn', when='~cuda')
-    conflicts('+nccl', when='~cuda')
     conflicts('+magma', when='~cuda')
-    conflicts('+mkldnn', when='@:0.3.2')
-
+    conflicts('+fbgemm', when='@:0.4')
+    conflicts('+miopen', when='@:0.4')
+    conflicts('+mkldnn', when='@:0.3')
+    conflicts('+qnnpack', when='@:0.4')
+    conflicts('+nccl', when='~cuda')
+    conflicts('+opencv', when='@:0.4')
+    conflicts('+ffmpeg', when='@:0.4')
+    conflicts('+leveldb', when='@:0.4')
+    conflicts('+lmdb', when='@:0.4')
+    conflicts('+binary', when='@:0.4')
+    conflicts('+redis', when='@:1.0')
+    conflicts('+zstd', when='@:1.0')
+    conflicts('+tbb', when='@:1.1')
+
+    # Required dependencies
+    depends_on('cmake@3.5:', type='build')
+    depends_on('python@2.7:2.8,3.5:', type=('build', 'run'))
     depends_on('py-setuptools', type='build')
-    depends_on('py-cffi', type='build')
     depends_on('py-numpy', type=('run', 'build'))
+    depends_on('py-future', when='@1.1: ^python@:2', type='build')
+    depends_on('py-pyyaml', type=('run', 'build'))
+    depends_on('py-typing', when='@0.4: ^python@:3.4', type=('run', 'build'))
     depends_on('blas')
     depends_on('lapack')
-    depends_on('py-pyyaml', type=('run', 'build'))
-    depends_on('py-typing', when='@0.3.2: ^python@:3.4', type=('run', 'build'))
-    depends_on('intel-mkl', when='+mkl')
+
+    # Optional dependencies
     depends_on('cuda@7.5:', when='+cuda', type=('build', 'link', 'run'))
     depends_on('cuda@9:', when='@1.1:+cuda', type=('build', 'link', 'run'))
-    depends_on('cudnn@6:', when='+cuda+cudnn')
-    depends_on('cudnn@7:', when='@1.1:+cuda+cudnn')
-    depends_on('nccl', when='+cuda+nccl')
-    depends_on('magma+shared', when='+cuda+magma')
+    depends_on('cudnn@6:', when='+cudnn')
+    depends_on('cudnn@7:', when='@1.1:+cudnn')
+    depends_on('magma', when='+magma')
+    # TODO: add dependency: https://github.com/pytorch/FBGEMM
+    depends_on('fbgemm', when='+fbgemm')
+    # TODO: add dependency: https://github.com/ROCmSoftwarePlatform/MIOpen
+    depends_on('miopen', when='+miopen')
+    depends_on('mkl', when='+mkldnn')
+    # TODO: add dependency: https://github.com/Maratyszcza/NNPACK
+    depends_on('nnpack', when='+nnpack')
+    # TODO: add dependency: https://github.com/pytorch/QNNPACK
+    depends_on('qnnpack', when='+qnnpack')
+    depends_on('mpi', when='+distributed')
+    depends_on('nccl', when='+nccl')
+    depends_on('gloo', when='+gloo')
+    depends_on('opencv', when='+opencv')
+    depends_on('llvm-openmp', when='%clang platform=darwin +openmp')
+    depends_on('ffmpeg', when='+ffmpeg')
+    depends_on('leveldb', when='+leveldb')
+    depends_on('lmdb', when='+lmdb')
+    depends_on('redis', when='+redis')
+    depends_on('zstd', when='+zstd')
+    depends_on('tbb', when='+tbb')
 
     def setup_environment(self, build_env, run_env):
+        def enable_or_disable(variant, keyword='USE', var=None, newer=False):
+            """Set environment variable to enable or disable support for a
+            particular variant.
+
+            Parameters:
+                variant (str): the variant to check
+                keyword (str): the prefix to use for enabling/disabling
+                var (str): CMake variable to set. Defaults to variant.upper()
+                newer (bool): newer variants that never used NO_*
+            """
+            if var is None:
+                var = variant.upper()
+
+            # Version 1.1.0 switched from NO_* to USE_* or BUILD_*
+            # 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')
+                else:
+                    build_env.set(keyword + '_' + var, 'OFF')
+            else:
+                if '+' + variant in self.spec:
+                    build_env.unset('NO_' + var)
+                else:
+                    build_env.set('NO_' + var, 'ON')
+
         build_env.set('MAX_JOBS', make_jobs)
 
+        enable_or_disable('cuda')
         if '+cuda' in self.spec:
             build_env.set('CUDA_HOME', self.spec['cuda'].prefix)
-        else:
-            build_env.set('NO_CUDA', 'TRUE')
 
+        enable_or_disable('cudnn')
         if '+cudnn' in self.spec:
             build_env.set('CUDNN_LIB_DIR',
-                          self.spec['cudnn'].prefix.lib)
+                          self.spec['cudnn'].libs.directories[0])
             build_env.set('CUDNN_INCLUDE_DIR',
                           self.spec['cudnn'].prefix.include)
-        else:
-            build_env.set('NO_CUDNN', 'TRUE')
+            build_env.set('CUDNN_LIBRARY', self.spec['cudnn'].libs[0])
 
-        if '+nccl' in self.spec:
-            build_env.set('NCCL_ROOT_DIR', self.spec['nccl'].prefix)
-        else:
-            build_env.set('NO_SYSTEM_NCCL', 'TRUE')
+        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])
 
+        enable_or_disable('mkldnn')
         if '+mkldnn' in self.spec:
             build_env.set('MKLDNN_HOME', self.spec['intel-mkl'].prefix)
-        else:
-            build_env.set('NO_MKLDNN', 'TRUE')
 
-        build_env.set('NO_NNPACK', 'TRUE')
+        enable_or_disable('nnpack')
+        enable_or_disable('qnnpack')
+        enable_or_disable('distributed')
 
-        build_env.set('PYTORCH_BUILD_VERSION', str(self.version))
+        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)
+
+        enable_or_disable('caffe2', keyword='BUILD', var='CAFFE2_OPS')
+        enable_or_disable('gloo', newer=True)
+        enable_or_disable('gloo', var='GLOO_IBVERBS', newer=True)
+        enable_or_disable('opencv', newer=True)
+        enable_or_disable('openmp', newer=True)
+        enable_or_disable('ffmpeg', newer=True)
+        enable_or_disable('leveldb', newer=True)
+        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)
+
+        # BLAS to be used by Caffe2. Can be MKL, Eigen, ATLAS, or OpenBLAS.
+        if '^mkl' in self.spec:
+            build_env.set('BLAS', 'MKL')
+        elif '^eigen' in self.spec:
+            build_env.set('BLAS', 'Eigen')
+        elif '^atlas' in self.spec:
+            build_env.set('BLAS', 'ATLAS')
+        elif '^openblas' in self.spec:
+            build_env.set('BLAS', 'OpenBLAS')
+
+        enable_or_disable('redis', newer=True)
+        enable_or_disable('zstd', newer=True)
+        enable_or_disable('tbb', newer=True)
+
+    def test(self):
+        pass
+
+    def install_test(self):
+        with working_dir('test'):
+            python('run_test.py')