gromacs: add env mods for cufftmp w/ gcc (#45887)

This commit is contained in:
snehring 2024-08-23 11:37:38 -05:00 committed by GitHub
parent fff8165f2f
commit 86050decb9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -158,6 +158,12 @@ class Gromacs(CMakePackage, CudaPackage):
conflicts(
"+mdrun_only", when="@2021:", msg="mdrun-only build option was removed for GROMACS 2021."
)
variant(
"nvshmem",
default=False,
description="Enable nvshmem support for nvidia gpus",
when="+cuda+mpi",
)
variant("openmp", default=True, description="Enables OpenMP at configure time")
variant("openmp_max_threads", default="none", description="Max number of OpenMP threads")
conflicts(
@ -302,6 +308,7 @@ class Gromacs(CMakePackage, CudaPackage):
depends_on("cp2k@8.1:", when="+cp2k")
depends_on("nvhpc", when="+cufftmp")
depends_on("nvhpc", when="+nvshmem")
depends_on("heffte", when="+heffte")
requires(
@ -410,6 +417,20 @@ def patch(self):
r"-gencode;arch=compute_20,code=sm_21;?", "", "cmake/gmxManageNvccConfig.cmake"
)
def setup_run_environment(self, env):
if self.spec.satisfies("+cufftmp"):
env.append_path(
"LD_LIBRARY_PATH",
join_path(
self.spec["nvhpc"].prefix,
f"Linux_{self.spec.target.family}",
self.spec["nvhpc"].version,
"comm_libs",
"nvshmem",
"lib",
),
)
class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
@run_after("build")
@ -545,6 +566,9 @@ def cmake_args(self):
if "+cuda" in self.spec:
options.append("-DCUDA_TOOLKIT_ROOT_DIR:STRING=" + self.spec["cuda"].prefix)
if not self.spec.satisfies("cuda_arch=none"):
cuda_arch = self.spec.variants["cuda_arch"].value
options.append(f"-DGMX_CUDA_TARGET_SM:STRING={';'.join(cuda_arch)}")
options.append("-DGMX_EXTERNAL_LAPACK:BOOL=ON")
if self.spec["lapack"].libs:
@ -660,6 +684,16 @@ def cmake_args(self):
options.append(
"-DGMX_OPENMP_MAX_THREADS=%s" % self.spec.variants["openmp_max_threads"].value
)
if self.spec.satisfies("+nvshmem"):
options.append("-DGMX_NVSHMEM:BOOL=ON")
nvshmem_root = join_path(
self.spec["nvhpc"].prefix,
f"Linux_{self.spec.target.family}",
self.spec["nvhpc"].version,
"comm_libs",
"nvshmem",
)
options.append(f"-DNVSHMEM_ROOT={nvshmem_root}")
if self.spec["lapack"].name in INTEL_MATH_LIBRARIES:
# fftw-api@3 is provided by intel-mkl or intel-parllel-studio
@ -703,3 +737,17 @@ def cmake_args(self):
else:
options.append("-DGMX_VERSION_STRING_OF_FORK=spack")
return options
def setup_build_environment(self, env):
if self.spec.satisfies("+cufftmp"):
env.append_path(
"LD_LIBRARY_PATH",
join_path(
self.spec["nvhpc"].prefix,
f"Linux_{self.spec.target.family}",
self.spec["nvhpc"].version,
"comm_libs",
"nvshmem",
"lib",
),
)