NVSHMEM package (#21279)

var/spack/repos/builtin/packages/dihydrogen/package.py

@@ -66,6 +66,8 @@ class Dihydrogen(CMakePackage, CudaPackage):
             description='CUDA architecture',
             values=spack.variant.auto_or_any_combination_of(*cuda_arch_values))
 
+    conflicts('~cuda', when='+nvshmem')
+
     depends_on('mpi')
     depends_on('catch2', type='test')
 
@@ -120,6 +122,8 @@ class Dihydrogen(CMakePackage, CudaPackage):
 
     depends_on('llvm-openmp', when='%apple-clang +openmp')
 
+    depends_on('nvshmem', when='+nvshmem')
+
     illegal_cuda_arch_values = [
         '10', '11', '12', '13',
         '20', '21',

var/spack/repos/builtin/packages/lbann/package.py

@@ -60,12 +60,14 @@ class Lbann(CMakePackage, CudaPackage):
     variant('opencv', default=True,
             description='Builds with support for image processing with OpenCV')
     variant('vtune', default=False, description='Builds with support for Intel VTune')
+    variant('nvshmem', default=False, description='Support for NVSHMEM')
 
     # Variant Conflicts
     conflicts('@:0.90,0.99:', when='~conduit')
     conflicts('@0.90:0.101.99', when='+fft')
     conflicts('~cuda', when='+nvprof')
     conflicts('~hwloc', when='+al')
+    conflicts('~cuda', when='+nvshmem')
 
     depends_on('cmake@3.17.0:', type='build')
 
@@ -105,6 +107,8 @@ class Lbann(CMakePackage, CudaPackage):
     depends_on('dihydrogen +distconv +cuda', when='+distconv')
     depends_on('dihydrogen ~half', when='+dihydrogen ~half')
     depends_on('dihydrogen +half', when='+dihydrogen +half')
+    depends_on('dihydrogen ~nvshmem', when='+dihydrogen ~nvshmem')
+    depends_on('dihydrogen +nvshmem', when='+dihydrogen +nvshmem')
     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')
     conflicts('~dihydrogen', when='+distconv')
@@ -168,6 +172,8 @@ class Lbann(CMakePackage, CudaPackage):
 
     depends_on('llvm-openmp', when='%apple-clang')
 
+    depends_on('nvshmem', when='+nvshmem')
+
     generator = 'Ninja'
     depends_on('ninja', type='build')
 
@@ -210,6 +216,7 @@ def cmake_args(self):
             '-DLBANN_WITH_CONDUIT:BOOL=%s' % ('+conduit' in spec),
             '-DLBANN_WITH_CUDA:BOOL=%s' % ('+cuda' in spec),
             '-DLBANN_WITH_CUDNN:BOOL=%s' % ('+cuda' in spec),
+            '-DLBANN_WITH_NVSHMEM:BOOL=%s' % ('+nvshmem' in spec),
             '-DLBANN_WITH_FFT:BOOL=%s' % ('+fft' in spec),
             '-DLBANN_WITH_TBINF=OFF',
             '-DLBANN_WITH_UNIT_TESTING:BOOL=%s' % (self.run_tests),

var/spack/repos/builtin/packages/nvshmem/package.py

@@ -0,0 +1,53 @@
+# Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
+# Spack Project Developers. See the top-level COPYRIGHT file for details.
+#
+# SPDX-License-Identifier: (Apache-2.0 OR MIT)
+
+from spack import *
+import re
+
+
+class Nvshmem(MakefilePackage, CudaPackage):
+    """NVSHMEM is a parallel programming interface based on OpenSHMEM that
+    provides efficient and scalable communication for NVIDIA GPU
+    clusters. NVSHMEM creates a global address space for data that spans
+    the memory of multiple GPUs and can be accessed with fine-grained
+    GPU-initiated operations, CPU-initiated operations, and operations on
+    CUDA streams."""
+
+    homepage = "https://developer.nvidia.com/nvshmem"
+    url      = "https://developer.nvidia.com/nvshmem-src-203-0"
+
+    maintainers = ['bvanessen']
+
+    version('2.0.3-0', sha256='20da93e8508511e21aaab1863cb4c372a3bec02307b932144a7d757ea5a1bad2', extension='txz')
+
+    conflicts('~cuda')
+
+    def url_for_version(self, version):
+        ver_str = '{0}'.format(version)
+        ver = re.sub('[.]', '', ver_str)
+        url_fmt = "https://developer.nvidia.com/nvshmem-src-{0}"
+        return url_fmt.format(ver)
+
+    depends_on('mpi')
+    depends_on('gdrcopy')
+
+    def setup_build_environment(self, env):
+        env.append_flags(
+            'NVSHMEM_PREFIX', self.prefix)
+        env.append_flags(
+            'NVSHMEM_MPI_SUPPORT', '1')
+        env.append_flags(
+            'NVSHMEM_USE_GDRCOPY', '1')
+
+        if self.spec.satisfies('^spectrum-mpi') or self.spec.satisfies('^openmpi'):
+            env.append_flags(
+                'NVSHMEM_MPI_IS_OMPI', '1')
+            env.append_flags(
+                'NVSHMEM_SHMEM_SUPPORT', '1')
+        else:
+            env.append_flags(
+                'NVSHMEM_MPI_IS_OMPI', '0')
+            env.append_flags(
+                'NVSHMEM_SHMEM_SUPPORT', '0')