Skip to content

Commit ce15fbb

Browse files
albestroalbestro
authored
nvshmem: introduce cmake, add new versions 3 and libfabric variant (#504)
* introduce cmake build system * introduce libfabric variant * WIP: introduce new version with more complex filename * confirm language dependency * ignore PMIx support (default is disabled) * minor changes to cmake packaging * minor fix * address comment about version constraints * minor change
1 parent 46eb614 commit ce15fbb

File tree

1 file changed

+80
-6
lines changed
  • repos/spack_repo/builtin/packages/nvshmem

1 file changed

+80
-6
lines changed

‎repos/spack_repo/builtin/packages/nvshmem/package.py‎

Lines changed: 80 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,15 @@
22
#
33
# SPDX-License-Identifier: (Apache-2.0 OR MIT)
44

5-
from spack_repo.builtin.build_systems.cuda import CudaPackage
5+
from spack_repo.builtin.build_systems import cmake, makefile
6+
from spack_repo.builtin.build_systems.cmake import CMakePackage, generator
7+
from spack_repo.builtin.build_systems.cuda import CudaPackage, conflicts
68
from spack_repo.builtin.build_systems.makefile import MakefilePackage
79

810
from spack.package import *
911

1012

11-
class Nvshmem(MakefilePackage, CudaPackage):
13+
class Nvshmem(MakefilePackage, CMakePackage, CudaPackage):
1214
"""NVSHMEM is a parallel programming interface based on OpenSHMEM that
1315
provides efficient and scalable communication for NVIDIA GPU
1416
clusters. NVSHMEM creates a global address space for data that spans
@@ -22,6 +24,8 @@ class Nvshmem(MakefilePackage, CudaPackage):
2224

2325
license("BSD-3-Clause-Open-MPI")
2426

27+
version("3.3.9", sha256="ba41e9ad6650cf99c1a60a3e47c19d1d97d814add7d35ea72337520ae13eeb59")
28+
version("3.2.5-1", sha256="eb2c8fb3b7084c2db86bd9fd905387909f1dfd483e7b45f7b3c3d5fcf5374b5a")
2529
version("2.7.0-6", sha256="23ed9b0187104dc87d5d2bc1394b6f5ff29e8c19138dc019d940b109ede699df")
2630
version("2.6.0-1", sha256="fc0e8de61b034f3a079dc231b1d0955e665a9f57b5013ee98b6743647bd60417")
2731
version("2.5.0-19", sha256="dd800b40f1d296e1d3ed2a9885adcfe745c3e57582bc809860e87bd32abcdc60")
@@ -30,6 +34,12 @@ class Nvshmem(MakefilePackage, CudaPackage):
3034
version("2.1.2-0", sha256="367211808df99b4575fb901977d9f4347065c61a26642d65887f24d60342a4ec")
3135
version("2.0.3-0", sha256="20da93e8508511e21aaab1863cb4c372a3bec02307b932144a7d757ea5a1bad2")
3236

37+
build_system(
38+
conditional("cmake", when="@2.9.0:"),
39+
conditional("makefile", when="@:2.11"),
40+
default="cmake",
41+
)
42+
3343
variant("cuda", default=True, description="Build with CUDA")
3444
variant("ucx", default=True, description="Build with UCX support")
3545
variant("nccl", default=True, description="Build with NCCL support")
@@ -42,22 +52,86 @@ class Nvshmem(MakefilePackage, CudaPackage):
4252
when="@2.6:",
4353
description="Build with support for GPU initiated communication",
4454
)
45-
conflicts("~cuda")
55+
variant("libfabric", default=False, description="Build with Libfabric support")
56+
57+
generator("ninja")
58+
59+
conflicts("~cuda", msg="NVSHMEM requires CUDA")
4660

4761
def url_for_version(self, version):
4862
ver_str = "{0}".format(version)
4963
directory = ver_str.split("-")[0]
50-
url_fmt = "https://developer.download.nvidia.com/compute/redist/nvshmem/{0}/source/nvshmem_src_{1}.txz"
64+
if version < Version("3.3.9"):
65+
url_fmt = "https://developer.download.nvidia.com/compute/redist/nvshmem/{0}/source/nvshmem_src_{1}.txz"
66+
else:
67+
url_fmt = "https://developer.download.nvidia.com/compute/redist/nvshmem/{0}/source/nvshmem_src_cuda12-all-all-{0}.tar.gz"
5168
return url_fmt.format(directory, version)
5269

53-
depends_on("c", type="build") # generated
54-
depends_on("cxx", type="build") # generated
70+
depends_on("c", type="build")
71+
depends_on("cxx", type="build")
72+
73+
depends_on("cuda@11:", when="@3.2.5:")
74+
75+
with default_args(when="build_system=cmake", type="build"):
76+
depends_on("[email protected]:")
77+
depends_on("ninja")
5578

5679
depends_on("mpi", when="+mpi")
80+
5781
depends_on("ucx", when="+ucx")
82+
depends_on("[email protected]:", when="@3: +ucx")
83+
5884
depends_on("gdrcopy", when="+gdrcopy")
85+
conflicts("~gdrcopy", when="~ucx")
86+
depends_on("gdrcopy@2:", when="@3: +gdrcopy")
87+
5988
depends_on("nccl", when="+nccl")
89+
depends_on("nccl@2:", when="@3: +nccl")
90+
91+
depends_on("libfabric", when="+libfabric")
92+
depends_on("[email protected]:", when="@3: +libfabric")
93+
94+
95+
class CMakeBuilder(cmake.CMakeBuilder):
96+
def cmake_args(self):
97+
config = [
98+
self.define("CMAKE_CUDA_ARCHITECTURES", self.spec.variants["cuda_arch"].values),
99+
self.define_from_variant("NVSHMEM_MPI_SUPPORT", "mpi"),
100+
self.define_from_variant("NVSHMEM_LIBFABRIC_SUPPORT", "libfabric"),
101+
self.define_from_variant("NVSHMEM_UCX_SUPPORT", "ucx"),
102+
self.define_from_variant("NVSHMEM_USE_NCCL", "nccl"),
103+
self.define_from_variant("NVSHMEM_USE_GDRCOPY", "gdrcopy"),
104+
self.define_from_variant("NVSHMEM_SHMEM_SUPPORT", "shmem"),
105+
self.define("NVSHMEM_IBRC_SUPPORT", False),
106+
self.define("NVSHMEM_BUILD_PYTHON_LIB", False),
107+
self.define("NVSHMEM_BUILD_EXAMPLES", False),
108+
self.define("NVSHMEM_BUILD_HYDRA_LAUNCHER", False),
109+
self.define("NVSHMEM_BUILD_TESTS", False),
110+
self.define("NVSHMEM_BUILD_TXZ_PACKAGE", False),
111+
]
112+
113+
if "+mpi" in self.spec:
114+
config.append(self.define("MPI_HOME", self.spec["mpi"].prefix))
115+
116+
if "+libfabric" in self.spec:
117+
config.append(self.define("LIBFABRIC_HOME", self.spec["libfabric"].prefix))
118+
119+
if "+ucx" in self.spec:
120+
config.append(self.define("UCX_HOME", self.spec["ucx"].prefix))
121+
122+
if "+nccl" in self.spec:
123+
config.append(self.define("NCCL_HOME", self.spec["nccl"].prefix))
124+
125+
if "+gdrcopy" in self.spec:
126+
config.append(self.define("GDRCOPY_HOME", self.spec["gdrcopy"].prefix))
127+
128+
if "+shmem" in self.spec:
129+
config.append(self.define("SHMEM_HOME", self.spec["shmem"].prefix))
130+
131+
return config
132+
60133

134+
class MakeBuilder(makefile.MakefileBuilder):
61135
def setup_build_environment(self, env: EnvironmentModifications) -> None:
62136
env.set("CUDA_HOME", self.spec["cuda"].prefix)
63137
env.set("NVSHMEM_PREFIX", self.prefix)

0 commit comments

Comments
 (0)