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( conflicts(
"+mdrun_only", when="@2021:", msg="mdrun-only build option was removed for GROMACS 2021." "+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", default=True, description="Enables OpenMP at configure time")
variant("openmp_max_threads", default="none", description="Max number of OpenMP threads") variant("openmp_max_threads", default="none", description="Max number of OpenMP threads")
conflicts( conflicts(
@ -302,6 +308,7 @@ class Gromacs(CMakePackage, CudaPackage):
depends_on("cp2k@8.1:", when="+cp2k") depends_on("cp2k@8.1:", when="+cp2k")
depends_on("nvhpc", when="+cufftmp") depends_on("nvhpc", when="+cufftmp")
depends_on("nvhpc", when="+nvshmem")
depends_on("heffte", when="+heffte") depends_on("heffte", when="+heffte")
requires( requires(
@ -410,6 +417,20 @@ def patch(self):
r"-gencode;arch=compute_20,code=sm_21;?", "", "cmake/gmxManageNvccConfig.cmake" 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): class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
@run_after("build") @run_after("build")
@ -545,6 +566,9 @@ def cmake_args(self):
if "+cuda" in self.spec: if "+cuda" in self.spec:
options.append("-DCUDA_TOOLKIT_ROOT_DIR:STRING=" + self.spec["cuda"].prefix) 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") options.append("-DGMX_EXTERNAL_LAPACK:BOOL=ON")
if self.spec["lapack"].libs: if self.spec["lapack"].libs:
@ -660,6 +684,16 @@ def cmake_args(self):
options.append( options.append(
"-DGMX_OPENMP_MAX_THREADS=%s" % self.spec.variants["openmp_max_threads"].value "-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: if self.spec["lapack"].name in INTEL_MATH_LIBRARIES:
# fftw-api@3 is provided by intel-mkl or intel-parllel-studio # fftw-api@3 is provided by intel-mkl or intel-parllel-studio
@ -703,3 +737,17 @@ def cmake_args(self):
else: else:
options.append("-DGMX_VERSION_STRING_OF_FORK=spack") options.append("-DGMX_VERSION_STRING_OF_FORK=spack")
return options 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",
),
)