From c45183d45c98c151ae4227dd415859249afae766 Mon Sep 17 00:00:00 2001
From: Nick Hagerty <hagertynl@ornl.gov>
Date: Thu, 14 Dec 2023 15:38:45 -0500
Subject: [PATCH] Updated CMake build system to allow for FFT_KOKKOS parameter.
 Updated CMakeLists.txt to print the correct value when summarizing.

---
 cmake/CMakeLists.txt                | 16 +---------------
 cmake/Modules/Packages/KOKKOS.cmake | 23 +++++++++++++++++++----
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt
index 28e02bbee7..aacaca4e6c 100644
--- a/cmake/CMakeLists.txt
+++ b/cmake/CMakeLists.txt
@@ -984,21 +984,7 @@ if(PKG_KSPACE)
     message(STATUS "Using non-threaded FFTs")
   endif()
   if(PKG_KOKKOS)
-    if(Kokkos_ENABLE_CUDA)
-      if(FFT STREQUAL "KISS")
-        message(STATUS "Kokkos FFT: KISS")
-      else()
-        message(STATUS "Kokkos FFT: cuFFT")
-      endif()
-    elseif(Kokkos_ENABLE_HIP)
-      if(FFT STREQUAL "KISS")
-        message(STATUS "Kokkos FFT: KISS")
-      else()
-        message(STATUS "Kokkos FFT: hipFFT")
-      endif()
-    else()
-      message(STATUS "Kokkos FFT: ${FFT}")
-    endif()
+    message(STATUS "Kokkos FFT: ${FFT_KOKKOS}")
   endif()
 endif()
 if(BUILD_DOC)
diff --git a/cmake/Modules/Packages/KOKKOS.cmake b/cmake/Modules/Packages/KOKKOS.cmake
index 0edd9a3baa..eb20f93956 100644
--- a/cmake/Modules/Packages/KOKKOS.cmake
+++ b/cmake/Modules/Packages/KOKKOS.cmake
@@ -131,16 +131,31 @@ if(PKG_KSPACE)
   list(APPEND KOKKOS_PKG_SOURCES ${KOKKOS_PKG_SOURCES_DIR}/fft3d_kokkos.cpp
                                  ${KOKKOS_PKG_SOURCES_DIR}/grid3d_kokkos.cpp
                                  ${KOKKOS_PKG_SOURCES_DIR}/remap_kokkos.cpp)
+  set(FFT_KOKKOS_VALUES KISS FFTW3 MKL HIPFFT CUFFT)
+  set_property(CACHE FFT_KOKKOS PROPERTY STRINGS ${FFT_KOKKOS_VALUES})
+  validate_option(FFT_KOKKOS FFT_KOKKOS_VALUES)
+  string(TOUPPER ${FFT_KOKKOS} FFT_KOKKOS)
+
   if(Kokkos_ENABLE_CUDA)
-    if(NOT (FFT STREQUAL "KISS"))
-      target_compile_definitions(lammps PRIVATE -DFFT_CUFFT)
+    if(NOT ((FFT_KOKKOS STREQUAL "KISS") OR (FFT_KOKKOS STREQUAL "CUFFT")))
+      message(FATAL_ERROR "The CUDA backend of Kokkos requires either KISS FFT or CUFFT.")
+    elseif(FFT_KOKKOS STREQUAL "KISS")
+      message(WARNING "Using KISS FFT with the CUDA backend of Kokkos may be sub-optimal.")
+      target_compile_definitions(lammps PRIVATE -DFFT_KOKKOS_KISS)
+    elseif(FFT_KOKKOS STREQUAL "CUFFT")
+      target_compile_definitions(lammps PRIVATE -DFFT_KOKKOS_CUFFT)
       target_link_libraries(lammps PRIVATE cufft)
     endif()
   elseif(Kokkos_ENABLE_HIP)
-    if(NOT (FFT STREQUAL "KISS"))
+    if(NOT ((FFT_KOKKOS STREQUAL "KISS") OR (FFT_KOKKOS STREQUAL "HIPFFT")))
+      message(FATAL_ERROR "The HIP backend of Kokkos requires either KISS FFT or HIPFFT.")
+    elseif(FFT_KOKKOS STREQUAL "KISS")
+      message(WARNING "Using KISS FFT with the HIP backend of Kokkos may be sub-optimal.")
+      target_compile_definitions(lammps PRIVATE -DFFT_KOKKOS_KISS)
+    elseif(FFT_KOKKOS STREQUAL "HIPFFT")
       include(DetectHIPInstallation)
       find_package(hipfft REQUIRED)
-      target_compile_definitions(lammps PRIVATE -DFFT_HIPFFT)
+      target_compile_definitions(lammps PRIVATE -DFFT_KOKKOS_HIPFFT)
       target_link_libraries(lammps PRIVATE hip::hipfft)
     endif()
   endif()