From 680d1f2e585732c49b8ff3d3dd2f43080b8d5a34 Mon Sep 17 00:00:00 2001 From: Nick Hagerty <85945544+hagertnl@users.noreply.github.com> Date: Tue, 27 Aug 2024 01:15:12 -0400 Subject: lammps: improve FFT selection and add fft_kokkos variant (#45969) --- var/spack/repos/builtin/packages/lammps/package.py | 41 +++++++++++++++------- 1 file changed, 28 insertions(+), 13 deletions(-) (limited to 'var') diff --git a/var/spack/repos/builtin/packages/lammps/package.py b/var/spack/repos/builtin/packages/lammps/package.py index c3d9e76109..69163fb93c 100644 --- a/var/spack/repos/builtin/packages/lammps/package.py +++ b/var/spack/repos/builtin/packages/lammps/package.py @@ -600,6 +600,22 @@ class Lammps(CMakePackage, CudaPackage, ROCmPackage, PythonExtension): 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 @@ class Lammps(CMakePackage, CudaPackage, ROCmPackage, PythonExtension): 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("hipfft", when="+kspace+kokkos+rocm") + depends_on("fftw-api@3", when="+kspace fft=fftw3") + 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 @@ class Lammps(CMakePackage, CudaPackage, ROCmPackage, PythonExtension): 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 @@ class Lammps(CMakePackage, CudaPackage, ROCmPackage, PythonExtension): 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 FFTW3 is selected, then CMake will try to detect, if threaded - # FFTW libraries are available and enable them by default. - 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")) + if spec.satisfies("+kspace"): + args.append(self.define_from_variant("FFT", "fft")) + if spec.satisfies("fft=fftw3 ^armpl-gcc") or spec.satisfies("fft=fftw3 ^acfl"): 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 @@ class Lammps(CMakePackage, CudaPackage, ROCmPackage, PythonExtension): # 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: -- cgit v1.2.3-70-g09d2