gromacs: add cufftmp variant to enable distributed PME when using CUDA (#36552)

Co-authored-by: Massimiliano Culpo <massimiliano.culpo@gmail.com>

1 files changed, 10 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/gromacs/package.py b/var/spack/repos/builtin/packages/gromacs/package.py
index e49dad2ef9..68fcd7c693 100644
--- a/var/spack/repos/builtin/packages/gromacs/package.py
+++ b/var/spack/repos/builtin/packages/gromacs/package.py
@@ -87,6 +87,7 @@ class Gromacs(CMakePackage, CudaPackage):
         description="Produces a double precision version of the executables",
     )
     variant("plumed", default=False, description="Enable PLUMED support")
+    variant("cufftmp", default=False, when="+cuda+mpi", description="Enable Multi GPU FFT support")
     variant("opencl", default=False, description="Enable OpenCL support")
     variant("sycl", default=False, description="Enable SYCL support")
     variant("nosuffix", default=False, description="Disable default suffixes")
@@ -249,6 +250,8 @@ class Gromacs(CMakePackage, CudaPackage):
     depends_on("cp2k@8.1:", when="+cp2k")
     depends_on("dbcsr", when="+cp2k")
 
+    depends_on("nvhpc", when="+cufftmp")
+
     patch("gmxDetectCpu-cmake-3.14.patch", when="@2018:2019.3^cmake@3.14.0:")
     patch("gmxDetectSimd-cmake-3.14.patch", when="@5.0:2017^cmake@3.14.0:")
 
@@ -478,6 +481,13 @@ class CMakeBuilder(spack.build_systems.cmake.CMakeBuilder):
             options.append("-DGMX_CP2K:BOOL=ON")
             options.append("-DCP2K_DIR:STRING={0}".format(self.spec["cp2k"].prefix))
 
+        if "+cufftmp" in self.spec:
+            options.append("-DGMX_USE_CUFFTMP=ON")
+            options.append(
+                f'-DcuFFTMp_ROOT={self.spec["nvhpc"].prefix}/Linux_{self.spec.target.family}'
+                + f'/{self.spec["nvhpc"].version}/math_libs'
+            )
+
         # Activate SIMD based on properties of the target
         target = self.spec.target
         if target >= "zen4":