Added code to help DiHydrogen find cuDNN and CUB (#19783)

* Added code to help DiHydrogen find cuDNN and CUB

* Cleaning up dependencies on CUB and adding guards for when newer
versions of CUDA include CUB and it should be excluded.

* Changed Hydrogen to disable half support by default.

* Have LBANN force Hydrogen and DiHydrogen to build without half when the variant is disabled.

* Added explicit variants to enusre that if LBANN is build without Cuda,
Aluminum, or Half support, it enforces those constraints for Hydrogen
and DiHydrogen.  Cleaned up the use of Python extend versus append in
LBANN and DiHydrogen recipes.

* Fixed Flake8

4 files changed, 52 insertions, 36 deletions

diff --git a/var/spack/repos/builtin/packages/aluminum/package.py b/var/spack/repos/builtin/packages/aluminum/package.py
index 69098362ae..8997737f91 100644
--- a/var/spack/repos/builtin/packages/aluminum/package.py
+++ b/var/spack/repos/builtin/packages/aluminum/package.py
@@ -42,7 +42,7 @@ class Aluminum(CMakePackage, CudaPackage):
     depends_on('mpi')
     depends_on('nccl', when='+nccl')
     depends_on('hwloc@1.11:')
-    depends_on('cub', when='@:0.1,0.6.0: +cuda')
+    depends_on('cub', when='@:0.1,0.6.0: +cuda ^cuda@:10.99')
 
     generator = 'Ninja'
     depends_on('ninja', type='build')
@@ -64,9 +64,9 @@ class Aluminum(CMakePackage, CudaPackage):
             args.append(
                 '-DALUMINUM_ENABLE_MPI_CUDA:BOOL=%s' % ('+ht' in spec))
 
-        if '@:0.1,0.6.0:':
+        if '@:0.1,0.6.0:' and spec.satisfies('^cuda@:10.99'):
             args.append(
-                '-DCUB_DIR:FILEPATH=%s' % spec['cub'].prefix.include)
+                '-DCUB_DIR:FILEPATH=%s' % spec['cub'].prefix)
 
         # Add support for OS X to find OpenMP (LLVM installed via brew)
         if self.spec.satisfies('%clang platform=darwin'):
diff --git a/var/spack/repos/builtin/packages/dihydrogen/package.py b/var/spack/repos/builtin/packages/dihydrogen/package.py
index b9d8764d38..c25653bcb3 100644
--- a/var/spack/repos/builtin/packages/dihydrogen/package.py
+++ b/var/spack/repos/builtin/packages/dihydrogen/package.py
@@ -150,4 +150,13 @@ class Dihydrogen(CMakePackage, CudaPackage):
                         ' '.join(self.cuda_flags(cuda_arch))
                     ))
 
+        if '+cuda' in spec or '+legacy' in spec:
+            args.append('-DcuDNN_DIR={0}'.format(
+                spec['cudnn'].prefix))
+
+        if spec.satisfies('^cuda@:10.99'):
+            if '+cuda' in spec or '+legacy' in spec:
+                args.append('-DCUB_DIR={0}'.format(
+                    spec['cub'].prefix))
+
         return args
diff --git a/var/spack/repos/builtin/packages/hydrogen/package.py b/var/spack/repos/builtin/packages/hydrogen/package.py
index 3ffba8b812..d6bea1bebd 100644
--- a/var/spack/repos/builtin/packages/hydrogen/package.py
+++ b/var/spack/repos/builtin/packages/hydrogen/package.py
@@ -61,7 +61,7 @@ class Hydrogen(CMakePackage, CudaPackage):
             description='Builds with Aluminum communication library')
     variant('omp_taskloops', default=False,
             description='Use OpenMP taskloops instead of parallel for loops.')
-    variant('half', default=True,
+    variant('half', default=False,
             description='Builds with support for FP16 precision data types')
     depends_on('cmake@3.17.0:', type='build')
     depends_on('mpi')
diff --git a/var/spack/repos/builtin/packages/lbann/package.py b/var/spack/repos/builtin/packages/lbann/package.py
index c4d40d4be0..cdb98731c2 100644
--- a/var/spack/repos/builtin/packages/lbann/package.py
+++ b/var/spack/repos/builtin/packages/lbann/package.py
@@ -69,8 +69,11 @@ class Lbann(CMakePackage, CudaPackage):
 
     # Add Hydrogen variants
     depends_on('hydrogen +openmp_blas +shared +int64')
+    depends_on('hydrogen ~al', when='~al')
     depends_on('hydrogen +al', when='+al')
+    depends_on('hydrogen ~cuda', when='~cuda')
     depends_on('hydrogen +cuda', when='+cuda')
+    depends_on('hydrogen ~half', when='~half')
     depends_on('hydrogen +half', when='+half')
     depends_on('hydrogen build_type=Debug', when='build_type=Debug')
 
@@ -88,9 +91,12 @@ class Lbann(CMakePackage, CudaPackage):
     depends_on('aluminum +cuda +nccl +ht +cuda_rma', when='+al +cuda')
 
     depends_on('dihydrogen +openmp', when='+dihydrogen')
+    depends_on('dihydrogen ~cuda', when='+dihydrogen ~cuda')
     depends_on('dihydrogen +cuda', when='+dihydrogen +cuda')
+    depends_on('dihydrogen ~al', when='+dihydrogen ~al')
     depends_on('dihydrogen +al', when='+dihydrogen +al')
     depends_on('dihydrogen +legacy +cuda', when='+distconv')
+    depends_on('dihydrogen ~half', when='+dihydrogen ~half')
     depends_on('dihydrogen +half', when='+dihydrogen +half')
     depends_on('dihydrogen@0.1', when='@0.101:0.101.99 +dihydrogen')
     depends_on('dihydrogen@:0.0,0.2:', when='@:0.90,0.102: +dihydrogen')
@@ -98,7 +104,7 @@ class Lbann(CMakePackage, CudaPackage):
 
     depends_on('cudnn', when='@0.90:0.100.99 +cuda')
     depends_on('cudnn@8.0.2:', when='@:0.90,0.101: +cuda')
-    depends_on('cub', when='@0.94:0.98.2 +cuda')
+    depends_on('cub', when='@0.94:0.98.2 +cuda ^cuda@:10.99')
     depends_on('mpi')
     depends_on('hwloc@1.11:', when='@:0.90,0.102:')
     depends_on('hwloc@1.11:1.11.99', when='@0.95:0.101.99')
@@ -187,23 +193,23 @@ class Lbann(CMakePackage, CudaPackage):
             '-DProtobuf_DIR={0}'.format(spec['protobuf'].prefix)])
 
         if spec.satisfies('@:0.90') or spec.satisfies('@0.95:'):
-            args.extend([
+            args.append(
                 '-DHydrogen_DIR={0}/CMake/hydrogen'.format(
-                    spec['hydrogen'].prefix)])
+                    spec['hydrogen'].prefix))
         elif spec.satisfies('@0.94'):
-            args.extend([
+            args.append(
                 '-DElemental_DIR={0}/CMake/elemental'.format(
-                    spec['elemental'].prefix)])
+                    spec['elemental'].prefix))
 
         if spec.satisfies('@0.94:0.98.2'):
-            args.extend(['-DLBANN_WITH_NCCL:BOOL=%s' %
-                         ('+cuda +nccl' in spec)])
+            args.append('-DLBANN_WITH_NCCL:BOOL=%s' %
+                        ('+cuda +nccl' in spec))
 
         if '+vtune' in spec:
-            args.extend(['-DVTUNE_DIR={0}'.format(spec['vtune'].prefix)])
+            args.append('-DVTUNE_DIR={0}'.format(spec['vtune'].prefix))
 
         if '+al' in spec:
-            args.extend(['-DAluminum_DIR={0}'.format(spec['aluminum'].prefix)])
+            args.append('-DAluminum_DIR={0}'.format(spec['aluminum'].prefix))
 
         if '+conduit' in spec:
             args.extend([
@@ -223,31 +229,32 @@ class Lbann(CMakePackage, CudaPackage):
                         clang_root)])
 
         if '+opencv' in spec:
-            args.extend(['-DOpenCV_DIR:STRING={0}'.format(
-                spec['opencv'].prefix)])
+            args.append('-DOpenCV_DIR:STRING={0}'.format(
+                spec['opencv'].prefix))
 
         if '+cuda' in spec:
-            args.extend([
+            args.append(
                 '-DCUDA_TOOLKIT_ROOT_DIR={0}'.format(
-                    spec['cuda'].prefix)])
-            args.extend([
+                    spec['cuda'].prefix))
+            args.append(
                 '-DcuDNN_DIR={0}'.format(
-                    spec['cudnn'].prefix)])
+                    spec['cudnn'].prefix))
             if spec.satisfies('@0.94:0.98.2'):
-                args.extend(['-DCUB_DIR={0}'.format(
-                    spec['cub'].prefix)])
+                if spec.satisfies('^cuda@:10.99'):
+                    args.append('-DCUB_DIR={0}'.format(
+                        spec['cub'].prefix))
                 if '+nccl' in spec:
-                    args.extend([
+                    args.append(
                         '-DNCCL_DIR={0}'.format(
-                            spec['nccl'].prefix)])
+                            spec['nccl'].prefix))
 
         if spec.satisfies('@:0.90') or spec.satisfies('@0.100:'):
-            args.extend([
-                '-DLBANN_WITH_DIHYDROGEN:BOOL=%s' % ('+dihydrogen' in spec)])
+            args.append(
+                '-DLBANN_WITH_DIHYDROGEN:BOOL=%s' % ('+dihydrogen' in spec))
 
         if spec.satisfies('@:0.90') or spec.satisfies('@0.101:'):
-            args.extend([
-                '-DLBANN_WITH_DISTCONV:BOOL=%s' % ('+distconv' in spec)])
+            args.append(
+                '-DLBANN_WITH_DISTCONV:BOOL=%s' % ('+distconv' in spec))
 
         return args
 
@@ -270,20 +277,20 @@ class Lbann(CMakePackage, CudaPackage):
             '-DLBANN_HOME=.'])
 
         if spec.variants['dtype'].value == 'float':
-            args.extend(['-DDATATYPE=4'])
+            args.append('-DDATATYPE=4')
         elif spec.variants['dtype'].value == 'double':
-            args.extend(['-DDATATYPE=8'])
+            args.append('-DDATATYPE=8')
 
         if '+opencv' in spec:
-            args.extend(['-DOpenCV_DIR:STRING={0}'.format(
-                spec['opencv'].prefix)])
+            args.append('-DOpenCV_DIR:STRING={0}'.format(
+                spec['opencv'].prefix))
 
         if '+cudnn' in spec:
-            args.extend(['-DcuDNN_DIR={0}'.format(
-                spec['cudnn'].prefix)])
+            args.append('-DcuDNN_DIR={0}'.format(
+                spec['cudnn'].prefix))
 
-        if '+cub' in spec:
-            args.extend(['-DCUB_DIR={0}'.format(
-                spec['cub'].prefix)])
+        if '+cub' in spec and spec.satisfies('^cuda@:10.99'):
+            args.append('-DCUB_DIR={0}'.format(
+                spec['cub'].prefix))
 
         return args