Update nvshmem package to use public URL (#23933)
This commit is contained in:
parent
ef56462867
commit
f31bc986a8
@ -4,7 +4,6 @@
|
||||
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
|
||||
|
||||
from spack import *
|
||||
import re
|
||||
|
||||
|
||||
class Nvshmem(MakefilePackage, CudaPackage):
|
||||
@ -16,39 +15,70 @@ class Nvshmem(MakefilePackage, CudaPackage):
|
||||
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')
|
||||
version('2.1.2-0', sha256='367211808df99b4575fb901977d9f4347065c61a26642d65887f24d60342a4ec')
|
||||
version('2.0.3-0', sha256='20da93e8508511e21aaab1863cb4c372a3bec02307b932144a7d757ea5a1bad2')
|
||||
|
||||
variant('cuda', default=True, description='Build with CUDA')
|
||||
variant('ucx', default=True, description='Build with UCX support')
|
||||
variant('nccl', default=True, description='Build with NCCL support')
|
||||
variant('gdrcopy', default=True, description='Build with gdrcopy support')
|
||||
variant('mpi', default=True, description='Build with MPI support')
|
||||
variant('shmem', default=False, description='Build with shmem support')
|
||||
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)
|
||||
directory = ver_str.split('-')[0]
|
||||
url_fmt = "https://developer.download.nvidia.com/compute/redist/nvshmem/{0}/source/nvshmem_src_{1}.txz"
|
||||
return url_fmt.format(directory, version)
|
||||
|
||||
depends_on('mpi')
|
||||
depends_on('gdrcopy')
|
||||
depends_on('mpi', when='+mpi')
|
||||
depends_on('ucx', when='+ucx')
|
||||
depends_on('gdrcopy', when='+gdrcopy')
|
||||
depends_on('nccl', when='+nccl')
|
||||
|
||||
def setup_build_environment(self, env):
|
||||
env.append_flags(
|
||||
'CUDA_HOME', self.spec['cuda'].prefix)
|
||||
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'):
|
||||
if '+ucx' in self.spec:
|
||||
env.append_flags(
|
||||
'NVSHMEM_MPI_IS_OMPI', '1')
|
||||
'NVSHMEM_UCX_SUPPORT', '1')
|
||||
env.append_flags(
|
||||
'UCX_HOME', self.spec['ucx'].prefix)
|
||||
|
||||
if '+gdrcopy' in self.spec:
|
||||
env.append_flags(
|
||||
'NVSHMEM_USE_GDRCOPY', '1')
|
||||
env.append_flags(
|
||||
'GDRCOPY_HOME', self.spec['gdrcopy'].prefix)
|
||||
|
||||
if '+nccl' in self.spec:
|
||||
env.append_flags(
|
||||
'NVSHMEM_USE_NCCL', '1')
|
||||
env.append_flags(
|
||||
'NCCL_HOME', self.spec['nccl'].prefix)
|
||||
|
||||
if '+mpi' in self.spec:
|
||||
env.append_flags(
|
||||
'NVSHMEM_MPI_SUPPORT', '1')
|
||||
env.append_flags(
|
||||
'MPI_HOME', self.spec['mpi'].prefix)
|
||||
|
||||
if self.spec.satisfies('^spectrum-mpi') or self.spec.satisfies('^openmpi'):
|
||||
env.append_flags(
|
||||
'NVSHMEM_MPI_IS_OMPI', '1')
|
||||
else:
|
||||
env.append_flags(
|
||||
'NVSHMEM_MPI_IS_OMPI', '0')
|
||||
|
||||
if '+shmem' in self.spec:
|
||||
env.append_flags(
|
||||
'NVSHMEM_SHMEM_SUPPORT', '1')
|
||||
else:
|
||||
env.append_flags(
|
||||
'NVSHMEM_MPI_IS_OMPI', '0')
|
||||
env.append_flags(
|
||||
'NVSHMEM_SHMEM_SUPPORT', '0')
|
||||
'SHMEM_HOME', self.spec['mpi'].prefix)
|
||||
|
Loading…
Reference in New Issue
Block a user