lammps: improve FFT selection and add fft_kokkos variant (#45969)

share/spack/gitlab/cloud_pipelines/stacks/aws-pcluster-x86_64_v4/packages.yaml

@@ -26,8 +26,8 @@ packages:
   lammps:
     require:
       - one_of:
-          - "lammps_sizes=bigbig +molecule +kspace +rigid +asphere +opt +openmp +openmp-package +intel ^intel-oneapi-mkl target=x86_64_v4"
+          - "lammps_sizes=bigbig +molecule +kspace +rigid +asphere +opt +openmp +openmp-package +intel fft=mkl ^intel-oneapi-mkl target=x86_64_v4"
-          - "lammps_sizes=bigbig +molecule +kspace +rigid +asphere +opt +openmp +openmp-package ^intel-oneapi-mkl target=x86_64_v3"
+          - "lammps_sizes=bigbig +molecule +kspace +rigid +asphere +opt +openmp +openmp-package fft=mkl ^intel-oneapi-mkl target=x86_64_v3"
         when: "%intel"
   libidn2:
     require:

var/spack/repos/builtin/packages/lammps/package.py

@@ -600,6 +600,22 @@ def url_for_version(self, version):
         values=("single", "double"),
         multi=False,
     )
+    variant(
+        "fft",
+        default="fftw3",
+        when="+kspace",
+        description="FFT library for KSPACE package",
+        values=("kiss", "fftw3", "mkl"),
+        multi=False,
+    )
+    variant(
+        "fft_kokkos",
+        default="fftw3",
+        when="@20240417: +kspace+kokkos",
+        description="FFT library for Kokkos-enabled KSPACE package",
+        values=("kiss", "fftw3", "mkl", "hipfft", "cufft"),
+        multi=False,
+    )
     variant(
         "gpu_precision",
         default="mixed",
@@ -613,8 +629,11 @@ def url_for_version(self, version):
     depends_on("cmake@3.16:", when="@20231121:")
     depends_on("mpi", when="+mpi")
     depends_on("mpi", when="+mpiio")
-    depends_on("fftw-api@3", when="+kspace")
+    depends_on("fftw-api@3", when="+kspace fft=fftw3")
-    depends_on("hipfft", when="+kspace+kokkos+rocm")
+    depends_on("mkl", when="+kspace fft=mkl")
+    depends_on("hipfft", when="+kokkos+kspace+rocm fft_kokkos=hipfft")
+    depends_on("fftw-api@3", when="+kokkos+kspace fft_kokkos=fftw3")
+    depends_on("mkl", when="+kokkos+kspace fft_kokkos=mkl")
     depends_on("voropp+pic", when="+voronoi")
     depends_on("netcdf-c+mpi", when="+user-netcdf")
     depends_on("netcdf-c+mpi", when="+netcdf")
@@ -817,6 +836,10 @@ def cmake_args(self):
                 args.append(self.define_from_variant("HIP_ARCH", "amdgpu_target"))
             else:
                 args.append(self.define("PKG_GPU", False))
+        else:
+            args.append(self.define("EXTERNAL_KOKKOS", True))
+            if spec.satisfies("@20240207: +kokkos+kspace"):
+                args.append(self.define_from_variant("FFT_KOKKOS", "fft_kokkos"))
 
         if spec.satisfies("@20180629:+lib"):
             args.append(self.define("BUILD_LIB", True))
@@ -856,15 +879,9 @@ def cmake_args(self):
                 opt = "{0}_{1}".format(pkg_prefix, pkg.replace("-package", "").upper())
                 args.append(self.define(opt, "+{0}".format(pkg) in spec))
 
-        if "+kspace" in spec:
+        if spec.satisfies("+kspace"):
-            # If FFTW3 is selected, then CMake will try to detect, if threaded
+            args.append(self.define_from_variant("FFT", "fft"))
-            # FFTW libraries are available and enable them by default.
+            if spec.satisfies("fft=fftw3 ^armpl-gcc") or spec.satisfies("fft=fftw3 ^acfl"):
-            if "^fftw" in spec or "^cray-fftw" in spec or "^amdfftw" in spec:
-                args.append(self.define("FFT", "FFTW3"))
-            elif spec["fftw-api"].name in INTEL_MATH_LIBRARIES:
-                args.append(self.define("FFT", "MKL"))
-            elif "^armpl-gcc" in spec or "^acfl" in spec:
-                args.append(self.define("FFT", "FFTW3"))
                 args.append(self.define("FFTW3_LIBRARY", self.spec["fftw-api"].libs[0]))
                 args.append(
                     self.define("FFTW3_INCLUDE_DIR", self.spec["fftw-api"].headers.directories[0])
@@ -874,8 +891,6 @@ def cmake_args(self):
             # for transposing 3d FFT data.
             args.append(self.define("FFT_SINGLE", spec.satisfies("fftw_precision=single")))
 
-        if "+kokkos" in spec:
-            args.append(self.define("EXTERNAL_KOKKOS", True))
         if "+user-adios" in spec or "+adios" in spec:
             args.append(self.define("ADIOS2_DIR", self.spec["adios2"].prefix))
         if "+user-plumed" in spec or "+plumed" in spec: