diff --git a/var/spack/repos/builtin/packages/dihydrogen/package.py b/var/spack/repos/builtin/packages/dihydrogen/package.py index c2395e9d33..befa4df243 100644 --- a/var/spack/repos/builtin/packages/dihydrogen/package.py +++ b/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', diff --git a/var/spack/repos/builtin/packages/lbann/package.py b/var/spack/repos/builtin/packages/lbann/package.py index 41d25a7fb1..ede920e72f 100644 --- a/var/spack/repos/builtin/packages/lbann/package.py +++ b/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), diff --git a/var/spack/repos/builtin/packages/nvshmem/package.py b/var/spack/repos/builtin/packages/nvshmem/package.py new file mode 100644 index 0000000000..6514ef6595 --- /dev/null +++ b/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')