Bugfixes for lbann sw stack (#19903)

* Added guard for setting CUB_DIR to only when cuda variant is true

* Added support for OpenMP on OSX platforms

* Updated the way that LBANN, Hydrogen, and DiHydrogen handle
apple-clang with OpenMP and Clang installed on OS X via brew.

* Fixed bug in spec resolution

* Fixed merge conflict

* Fixed typo

* Fixed flake8

4 files changed, 80 insertions, 30 deletions

diff --git a/var/spack/repos/builtin/packages/aluminum/package.py b/var/spack/repos/builtin/packages/aluminum/package.py
index 8997737f91..8b3453dd54 100644
--- a/var/spack/repos/builtin/packages/aluminum/package.py
+++ b/var/spack/repos/builtin/packages/aluminum/package.py
@@ -64,7 +64,7 @@ class Aluminum(CMakePackage, CudaPackage):
             args.append(
                 '-DALUMINUM_ENABLE_MPI_CUDA:BOOL=%s' % ('+ht' in spec))
 
-        if '@:0.1,0.6.0:' and spec.satisfies('^cuda@:10.99'):
+        if spec.satisfies('@:0.1,0.6.0: +cuda ^cuda@:10.99'):
             args.append(
                 '-DCUB_DIR:FILEPATH=%s' % spec['cub'].prefix)
 
diff --git a/var/spack/repos/builtin/packages/dihydrogen/package.py b/var/spack/repos/builtin/packages/dihydrogen/package.py
index c25653bcb3..03a3f8f671 100644
--- a/var/spack/repos/builtin/packages/dihydrogen/package.py
+++ b/var/spack/repos/builtin/packages/dihydrogen/package.py
@@ -3,6 +3,7 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
+import os
 from spack import *
 
 
@@ -110,6 +111,8 @@ class Dihydrogen(CMakePackage, CudaPackage):
     depends_on('py-breathe', type='build', when='+docs')
     depends_on('doxygen', type='build', when='+docs')
 
+    depends_on('llvm-openmp', when='%apple-clang +openmp')
+
     illegal_cuda_arch_values = [
         '10', '11', '12', '13',
         '20', '21',
@@ -159,4 +162,26 @@ class Dihydrogen(CMakePackage, CudaPackage):
                 args.append('-DCUB_DIR={0}'.format(
                     spec['cub'].prefix))
 
+        # Add support for OpenMP with external (Brew) clang
+        if spec.satisfies('%clang +openmp platform=darwin'):
+            clang = self.compiler.cc
+            clang_bin = os.path.dirname(clang)
+            clang_root = os.path.dirname(clang_bin)
+            args.extend([
+                '-DOpenMP_CXX_FLAGS=-fopenmp=libomp',
+                '-DOpenMP_CXX_LIB_NAMES=libomp',
+                '-DOpenMP_libomp_LIBRARY={0}/lib/libomp.dylib'.format(
+                    clang_root)])
+
         return args
+
+    def setup_build_environment(self, env):
+        if self.spec.satisfies('%apple-clang +openmp'):
+            env.append_flags(
+                'CPPFLAGS', self.compiler.openmp_flag)
+            env.append_flags(
+                'CFLAGS', self.spec['llvm-openmp'].headers.include_flags)
+            env.append_flags(
+                'CXXFLAGS', self.spec['llvm-openmp'].headers.include_flags)
+            env.append_flags(
+                'LDFLAGS', self.spec['llvm-openmp'].libs.ld_flags)
diff --git a/var/spack/repos/builtin/packages/hydrogen/package.py b/var/spack/repos/builtin/packages/hydrogen/package.py
index d6bea1bebd..082d8316d7 100644
--- a/var/spack/repos/builtin/packages/hydrogen/package.py
+++ b/var/spack/repos/builtin/packages/hydrogen/package.py
@@ -33,8 +33,8 @@ class Hydrogen(CMakePackage, CudaPackage):
 
     variant('shared', default=True,
             description='Enables the build of shared libraries')
-    variant('hybrid', default=True,
-            description='Make use of OpenMP within MPI packing/unpacking')
+    variant('openmp', default=True,
+            description='Make use of OpenMP within CPU-kernels')
     variant('openmp_blas', default=False,
             description='Use OpenMP for threading in the BLAS library')
     variant('quad', default=False,
@@ -63,28 +63,29 @@ class Hydrogen(CMakePackage, CudaPackage):
             description='Use OpenMP taskloops instead of parallel for loops.')
     variant('half', default=False,
             description='Builds with support for FP16 precision data types')
+
+    conflicts('~openmp', when='+omp_taskloops')
+
     depends_on('cmake@3.17.0:', type='build')
     depends_on('mpi')
     depends_on('hwloc@1.11:')
 
     # Note that #1712 forces us to enumerate the different blas variants
-    depends_on('openblas', when='blas=openblas ~openmp_blas ~int64_blas')
-    depends_on('openblas +ilp64', when='blas=openblas ~openmp_blas +int64_blas')
-    depends_on('openblas threads=openmp', when='blas=openblas +openmp_blas ~int64_blas')
-    depends_on('openblas threads=openmp +lip64', when='blas=openblas +openmp_blas +int64_blas')
+    depends_on('openblas', when='blas=openblas')
+    depends_on('openblas +ilp64', when='blas=openblas +int64_blas')
+    depends_on('openblas threads=openmp', when='blas=openblas +openmp_blas')
 
-    depends_on('intel-mkl', when="blas=mkl ~openmp_blas ~int64_blas")
-    depends_on('intel-mkl +ilp64', when="blas=mkl ~openmp_blas +int64_blas")
-    depends_on('intel-mkl threads=openmp', when='blas=mkl +openmp_blas ~int64_blas')
-    depends_on('intel-mkl@2017.1 +openmp +ilp64', when='blas=mkl +openmp_blas +int64_blas')
+    depends_on('intel-mkl', when="blas=mkl")
+    depends_on('intel-mkl +ilp64', when="blas=mkl +int64_blas")
+    depends_on('intel-mkl threads=openmp', when='blas=mkl +openmp_blas')
 
     depends_on('veclibfort', when='blas=accelerate')
     conflicts('blas=accelerate +openmp_blas')
 
-    depends_on('essl -cuda', when='blas=essl -openmp_blas ~int64_blas')
-    depends_on('essl -cuda +ilp64', when='blas=essl -openmp_blas +int64_blas')
-    depends_on('essl threads=openmp', when='blas=essl +openmp_blas ~int64_blas')
-    depends_on('essl threads=openmp +ilp64', when='blas=essl +openmp_blas +int64_blas')
+    depends_on('essl', when='blas=essl')
+    depends_on('essl -cuda', when='blas=essl -openmp_blas')
+    depends_on('essl +ilp64', when='blas=essl +int64_blas')
+    depends_on('essl threads=openmp', when='blas=essl +openmp_blas')
     depends_on('netlib-lapack +external-blas', when='blas=essl')
 
     # Specify the correct version of Aluminum
@@ -107,6 +108,8 @@ class Hydrogen(CMakePackage, CudaPackage):
     depends_on('cub', when='^cuda@:10.99')
     depends_on('half', when='+half')
 
+    depends_on('llvm-openmp', when='%apple-clang +openmp')
+
     conflicts('@0:0.98', msg="Hydrogen did not exist before v0.99. " +
               "Did you mean to use Elemental instead?")
 
@@ -126,7 +129,7 @@ class Hydrogen(CMakePackage, CudaPackage):
         args = [
             '-DCMAKE_INSTALL_MESSAGE:STRING=LAZY',
             '-DBUILD_SHARED_LIBS:BOOL=%s'      % ('+shared' in spec),
-            '-DHydrogen_ENABLE_OPENMP:BOOL=%s'       % ('+hybrid' in spec),
+            '-DHydrogen_ENABLE_OPENMP:BOOL=%s'       % ('+openmp' in spec),
             '-DHydrogen_ENABLE_QUADMATH:BOOL=%s'     % ('+quad' in spec),
             '-DHydrogen_USE_64BIT_INTS:BOOL=%s'      % ('+int64' in spec),
             '-DHydrogen_USE_64BIT_BLAS_INTS:BOOL=%s' % ('+int64_blas' in spec),
@@ -140,7 +143,7 @@ class Hydrogen(CMakePackage, CudaPackage):
         ]
 
         # Add support for OS X to find OpenMP (LLVM installed via brew)
-        if self.spec.satisfies('%clang platform=darwin'):
+        if self.spec.satisfies('%clang +openmp platform=darwin'):
             clang = self.compiler.cc
             clang_bin = os.path.dirname(clang)
             clang_root = os.path.dirname(clang_bin)
@@ -173,3 +176,14 @@ class Hydrogen(CMakePackage, CudaPackage):
                     spec['aluminum'].prefix)])
 
         return args
+
+    def setup_build_environment(self, env):
+        if self.spec.satisfies('%apple-clang +openmp'):
+            env.append_flags(
+                'CPPFLAGS', self.compiler.openmp_flag)
+            env.append_flags(
+                'CFLAGS', self.spec['llvm-openmp'].headers.include_flags)
+            env.append_flags(
+                'CXXFLAGS', self.spec['llvm-openmp'].headers.include_flags)
+            env.append_flags(
+                'LDFLAGS', self.spec['llvm-openmp'].libs.ld_flags)
diff --git a/var/spack/repos/builtin/packages/lbann/package.py b/var/spack/repos/builtin/packages/lbann/package.py
index cdb98731c2..11c4ed2f1a 100644
--- a/var/spack/repos/builtin/packages/lbann/package.py
+++ b/var/spack/repos/builtin/packages/lbann/package.py
@@ -4,7 +4,6 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 import os
-import sys
 from spack import *
 
 
@@ -68,7 +67,7 @@ class Lbann(CMakePackage, CudaPackage):
     depends_on('hydrogen@1.5.0:', when='@:0.90,0.102:')
 
     # Add Hydrogen variants
-    depends_on('hydrogen +openmp_blas +shared +int64')
+    depends_on('hydrogen +openmp +openmp_blas +shared +int64')
     depends_on('hydrogen ~al', when='~al')
     depends_on('hydrogen +al', when='+al')
     depends_on('hydrogen ~cuda', when='~cuda')
@@ -153,6 +152,8 @@ class Lbann(CMakePackage, CudaPackage):
     depends_on('catch2', type='test')
     depends_on('clara')
 
+    depends_on('llvm-openmp', when='%apple-clang')
+
     generator = 'Ninja'
     depends_on('ninja', type='build')
 
@@ -169,6 +170,17 @@ class Lbann(CMakePackage, CudaPackage):
             '-DCNPY_DIR={0}'.format(spec['cnpy'].prefix),
         ]
 
+    def setup_build_environment(self, env):
+        if self.spec.satisfies('%apple-clang'):
+            env.append_flags(
+                'CPPFLAGS', self.compiler.openmp_flag)
+            env.append_flags(
+                'CFLAGS', self.spec['llvm-openmp'].headers.include_flags)
+            env.append_flags(
+                'CXXFLAGS', self.spec['llvm-openmp'].headers.include_flags)
+            env.append_flags(
+                'LDFLAGS', self.spec['llvm-openmp'].libs.ld_flags)
+
     # Get any recent versions or non-numeric version
     # Note that develop > numeric and non-develop < numeric
     @when('@:0.90,0.94:')
@@ -216,17 +228,16 @@ class Lbann(CMakePackage, CudaPackage):
                 '-DLBANN_CONDUIT_DIR={0}'.format(spec['conduit'].prefix),
                 '-DConduit_DIR={0}'.format(spec['conduit'].prefix)])
 
-        # Add support for OpenMP
-        if spec.satisfies('%clang') or spec.satisfies('%apple-clang'):
-            if sys.platform == 'darwin':
-                clang = self.compiler.cc
-                clang_bin = os.path.dirname(clang)
-                clang_root = os.path.dirname(clang_bin)
-                args.extend([
-                    '-DOpenMP_CXX_FLAGS=-fopenmp=libomp',
-                    '-DOpenMP_CXX_LIB_NAMES=libomp',
-                    '-DOpenMP_libomp_LIBRARY={0}/lib/libomp.dylib'.format(
-                        clang_root)])
+        # Add support for OpenMP with external (Brew) clang
+        if spec.satisfies('%clang platform=darwin'):
+            clang = self.compiler.cc
+            clang_bin = os.path.dirname(clang)
+            clang_root = os.path.dirname(clang_bin)
+            args.extend([
+                '-DOpenMP_CXX_FLAGS=-fopenmp=libomp',
+                '-DOpenMP_CXX_LIB_NAMES=libomp',
+                '-DOpenMP_libomp_LIBRARY={0}/lib/libomp.dylib'.format(
+                    clang_root)])
 
         if '+opencv' in spec:
             args.append('-DOpenCV_DIR:STRING={0}'.format(