easyblock = 'ConfigureMake'

name = 'NVSHMEM'
version = '2.5.0'
versionsuffix = '-CUDA-%(cudaver)s'

homepage = 'https://developer.nvidia.com/nvshmem'
description = """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.
"""

toolchain = {'name': 'gompi', 'version': '2022a'}

download_instructions = """The sources of NVSHMEM can be downloaded at NVIDIA's webpage when you have signed up for
their (free) developer program:
https://developer.nvidia.com/nvshmem-downloads"""

sources = ['%(namelower)s_src_%(version)s-19.txz']
checksums = ['dd800b40f1d296e1d3ed2a9885adcfe745c3e57582bc809860e87bd32abcdc60']

builddependencies = [
    ('Autotools', '20220317'),
    ('pkgconf', '1.8.0'),
]

dependencies = [
    ('CUDA', '11.7.0', '', SYSTEM),
    ('UCX-CUDA', '1.12.1', versionsuffix),
    ('NCCL', '2.12.12', versionsuffix),
]

skipsteps = ['configure']

prebuildopts = 'export %s &&' % ' '.join([
    'NVSHMEM_USE_GDRCOPY=1',
    'GDRCOPY_HOME=${EBROOTGDRCOPY}',

    'MPI_HOME=${EBROOTOPENMPI}',
    'NVSHMEM_MPI_SUPPORT=1',
    'NVSHMEMTEST_USE_MPI_LAUNCHER=1',

    'NCCL_HOME=${EBROOTNCCL}',
    'NVSHMEM_USE_NCCL=1',

    'NVSHMEM_BUILDDIR=%(builddir)s',
    'NVSHMEM_EXAMPLES_BUILDDIR=${NVSHMEM_BUILDDIR}/examples/obj',
    'NVSHMEM_OTHERTEST_BUILDDIR=${NVSHMEM_BUILDDIR}/othertest/obj',
    'NVSHMEM_TEST_BUILDDIR=${NVSHMEM_BUILDDIR}/test/obj',
    'NVSHMEM_PERFTEST_BUILDDIR=${NVSHMEM_BUILDDIR}/perftest/obj',

    'NVSHMEM_PREFIX=%(installdir)s',
    'NVSHMEM_EXAMPLES_INSTALL=${NVSHMEM_PREFIX}/examples',
    'NVSHMEM_OTHERTEST_INSTALL=${NVSHMEM_PREFIX}/othertest',
    'NVSHMEM_PERFTEST_INSTALL=${NVSHMEM_PREFIX}/perftest',
    'NVSHMEM_TEST_INSTALL=${NVSHMEM_PREFIX}/test',
])

preinstallopts = prebuildopts

sanity_check_paths = {
    'files': ['lib/libnvshmem.a', 'lib/nvshmem_bootstrap_mpi.%s' % SHLIB_EXT],
    'dirs': ['include']
}

modextravars = {'NVSHMEM_HOME': '%(installdir)s'}

moduleclass = 'devel'
