nvshmem: set env variables instead of appending (#25095)

This commit is contained in:
Tim Moon 2021-07-28 09:07:24 -07:00 committed by GitHub
parent 68f696af64
commit c9327649c0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -41,44 +41,30 @@ def url_for_version(self, version):
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.set('CUDA_HOME', self.spec['cuda'].prefix)
env.set('NVSHMEM_PREFIX', self.prefix)
if '+ucx' in self.spec:
env.append_flags(
'NVSHMEM_UCX_SUPPORT', '1')
env.append_flags(
'UCX_HOME', self.spec['ucx'].prefix)
env.set('NVSHMEM_UCX_SUPPORT', '1')
env.set('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)
env.set('NVSHMEM_USE_GDRCOPY', '1')
env.set('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)
env.set('NVSHMEM_USE_NCCL', '1')
env.set('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)
env.set('NVSHMEM_MPI_SUPPORT', '1')
env.set('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')
env.set('NVSHMEM_MPI_IS_OMPI', '1')
else:
env.append_flags(
'NVSHMEM_MPI_IS_OMPI', '0')
env.set('NVSHMEM_MPI_IS_OMPI', '0')
if '+shmem' in self.spec:
env.append_flags(
'NVSHMEM_SHMEM_SUPPORT', '1')
env.append_flags(
'SHMEM_HOME', self.spec['mpi'].prefix)
env.set('NVSHMEM_SHMEM_SUPPORT', '1')
env.set('SHMEM_HOME', self.spec['mpi'].prefix)