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

1 files changed, 48 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/gromacs/package.py b/var/spack/repos/builtin/packages/gromacs/package.py
index a8ede744d7..6f95180701 100644
--- a/var/spack/repos/builtin/packages/gromacs/package.py
+++ b/var/spack/repos/builtin/packages/gromacs/package.py
@@ -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 @@ class Gromacs(CMakePackage, CudaPackage):
                     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 @@ class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
 
         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 @@ class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
             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 @@ class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
         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",
+                ),
+            )