use static lambda for initialization rather than KokkosFFT::initialize

yasahi-hpc · Jan 15, 2025 · 7986dd2 · 7986dd2
1 parent 2b59d12
commit 7986dd2
Show file tree

Hide file tree

Showing 17 changed files with 135 additions and 243 deletions.
diff --git a/fft/src/CMakeLists.txt b/fft/src/CMakeLists.txt
@@ -2,15 +2,15 @@
 #
 # SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception
 
-add_library(fft STATIC KokkosFFT_Core.cpp)
+add_library(fft INTERFACE)
 
 target_link_libraries(fft
-    PUBLIC
+INTERFACE
     common
     Kokkos::kokkos
 )
 
-target_include_directories(fft PUBLIC
+target_include_directories(fft INTERFACE
                            $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
                            $<INSTALL_INTERFACE:${INSTALL_INCLUDEDIR}>
 )

diff --git a/fft/src/KokkosFFT.hpp b/fft/src/KokkosFFT.hpp
@@ -10,6 +10,5 @@
 #include "KokkosFFT_Helpers.hpp"
 #include "KokkosFFT_Plans.hpp"
 #include "KokkosFFT_Transform.hpp"
-#include "KokkosFFT_Core.hpp"
 
 #endif
diff --git a/fft/src/KokkosFFT_Cuda_plans.hpp b/fft/src/KokkosFFT_Cuda_plans.hpp
@@ -13,6 +13,26 @@
 
 namespace KokkosFFT {
 namespace Impl {
+
+template <typename ExecutionSpace, typename T,
+          std::enable_if_t<std::is_same_v<ExecutionSpace, Kokkos::Cuda>,
+                           std::nullptr_t> = nullptr>
+void setup() {
+  static bool once = [] {
+    if (!(Kokkos::is_initialized() || Kokkos::is_finalized())) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called before initializing "
+          "Kokkos.\n");
+    }
+    if (Kokkos::is_finalized()) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called after finalizing "
+          "Kokkos.\n");
+    }
+    return true;
+  }();
+}
+
 // 1D transform
 template <typename ExecutionSpace, typename PlanType, typename InViewType,
           typename OutViewType,

diff --git a/fft/src/KokkosFFT_Cuda_types.hpp b/fft/src/KokkosFFT_Cuda_types.hpp
@@ -239,12 +239,8 @@ auto direction_type(Direction direction) {
   return direction == Direction::forward ? CUFFT_FORWARD : CUFFT_INVERSE;
 }
 
-inline void initialize_host() {}
-inline void finalize_host() {}
 #endif
 
-inline void initialize_device() {}
-inline void finalize_device() {}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_FFTW_Types.hpp b/fft/src/KokkosFFT_FFTW_Types.hpp
@@ -138,28 +138,6 @@ struct ScopedFFTWPlan {
   }
 };
 
-#if defined(KOKKOS_ENABLE_OPENMP) || defined(KOKKOS_ENABLE_THREADS)
-inline void initialize_host() {
-  fftwf_init_threads();
-  fftw_init_threads();
-}
-inline void finalize_host() {
-  fftwf_cleanup_threads();
-  fftw_cleanup_threads();
-}
-#else
-inline void initialize_host() {}
-inline void finalize_host() {}
-#endif
-
-// If non of device backend is enabled, then FFTW is responsible
-// for the device cleanup. Otherwise, device backend will cleanup
-#if !(defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || \
-      defined(KOKKOS_ENABLE_SYCL))
-inline void initialize_device() {}
-inline void finalize_device() {}
-#endif
-
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_HIP_plans.hpp b/fft/src/KokkosFFT_HIP_plans.hpp
@@ -13,6 +13,26 @@
 
 namespace KokkosFFT {
 namespace Impl {
+
+template <typename ExecutionSpace, typename T,
+          std::enable_if_t<std::is_same_v<ExecutionSpace, Kokkos::HIP>,
+                           std::nullptr_t> = nullptr>
+void setup() {
+  static bool once = [] {
+    if (!(Kokkos::is_initialized() || Kokkos::is_finalized())) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called before initializing "
+          "Kokkos.\n");
+    }
+    if (Kokkos::is_finalized()) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called after finalizing "
+          "Kokkos.\n");
+    }
+    return true;
+  }();
+}
+
 // 1D transform
 template <typename ExecutionSpace, typename PlanType, typename InViewType,
           typename OutViewType,

diff --git a/fft/src/KokkosFFT_HIP_types.hpp b/fft/src/KokkosFFT_HIP_types.hpp
@@ -239,12 +239,8 @@ auto direction_type(Direction direction) {
   return direction == Direction::forward ? HIPFFT_FORWARD : HIPFFT_BACKWARD;
 }
 
-inline void initialize_host() {}
-inline void finalize_host() {}
 #endif
 
-inline void initialize_device() {}
-inline void finalize_device() {}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_Host_plans.hpp b/fft/src/KokkosFFT_Host_plans.hpp
@@ -12,6 +12,45 @@
 
 namespace KokkosFFT {
 namespace Impl {
+
+template <typename ExecutionSpace, typename T,
+          std::enable_if_t<is_AnyHostSpace_v<ExecutionSpace>, std::nullptr_t> =
+              nullptr>
+void setup() {
+  static bool once = [] {
+    if (!(Kokkos::is_initialized() || Kokkos::is_finalized())) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called before initializing "
+          "Kokkos.\n");
+    }
+    if (Kokkos::is_finalized()) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called after finalizing "
+          "Kokkos.\n");
+    }
+#if defined(KOKKOS_ENABLE_OPENMP) || defined(KOKKOS_ENABLE_THREADS)
+    if constexpr (std::is_same_v<ExecutionSpace,
+                                 Kokkos::DefaultHostExecutionSpace>) {
+      if constexpr (std::is_same_v<T, float>) {
+        fftwf_init_threads();
+      } else {
+        fftw_init_threads();
+      }
+
+      // Register cleanup function as a hook in Kokkos::finalize
+      Kokkos::push_finalize_hook([]() {
+        if constexpr (std::is_same_v<T, float>) {
+          fftwf_cleanup_threads();
+        } else {
+          fftw_cleanup_threads();
+        }
+      });
+    }
+#endif
+    return true;
+  }();
+}
+
 // batched transform, over ND Views
 template <typename ExecutionSpace, typename PlanType, typename InViewType,
           typename OutViewType, std::size_t fft_rank = 1,

diff --git a/fft/src/KokkosFFT_Plans.hpp b/fft/src/KokkosFFT_Plans.hpp
@@ -197,6 +197,7 @@ class Plan {
                        "In-place transform is not supported with reshape. "
                        "Please use out-of-place transform.");
 
+    KokkosFFT::Impl::setup<ExecutionSpace, float_type>();
     m_fft_size = KokkosFFT::Impl::create_plan(
         exec_space, m_plan, in, out, direction, m_axes, s, m_is_inplace);
   }
@@ -260,6 +261,7 @@ class Plan {
                        "In-place transform is not supported with reshape. "
                        "Please use out-of-place transform.");
 
+    KokkosFFT::Impl::setup<ExecutionSpace, float_type>();
     m_fft_size = KokkosFFT::Impl::create_plan(exec_space, m_plan, in, out,
                                               direction, axes, s, m_is_inplace);
   }

diff --git a/fft/src/KokkosFFT_ROCM_plans.hpp b/fft/src/KokkosFFT_ROCM_plans.hpp
@@ -5,6 +5,7 @@
 #ifndef KOKKOSFFT_ROCM_PLANS_HPP
 #define KOKKOSFFT_ROCM_PLANS_HPP
 
+#include <Kokkos_Core.hpp>
 #include "KokkosFFT_ROCM_types.hpp"
 #include "KokkosFFT_Extents.hpp"
 #include "KokkosFFT_traits.hpp"
@@ -14,6 +15,35 @@
 namespace KokkosFFT {
 namespace Impl {
 
+template <typename ExecutionSpace, typename T,
+          std::enable_if_t<std::is_same_v<ExecutionSpace, Kokkos::HIP>,
+                           std::nullptr_t> = nullptr>
+void setup() {
+  static bool once = [] {
+    if (!(Kokkos::is_initialized() || Kokkos::is_finalized())) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called before initializing "
+          "Kokkos.\n");
+    }
+    if (Kokkos::is_finalized()) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called after finalizing "
+          "Kokkos.\n");
+    }
+
+    rocfft_status status = rocfft_setup();
+    if (status != rocfft_status_success) Kokkos::abort("rocfft_setup failed");
+
+    // Register cleanup function as a hook in Kokkos::finalize
+    Kokkos::push_finalize_hook([]() {
+      rocfft_status status = rocfft_cleanup();
+      if (status != rocfft_status_success)
+        Kokkos::abort("rocfft_cleanup failed");
+    });
+    return true;
+  }();
+}
+
 // batched transform, over ND Views
 template <typename ExecutionSpace, typename PlanType, typename InViewType,
           typename OutViewType, std::size_t fft_rank = 1,

diff --git a/fft/src/KokkosFFT_ROCM_types.hpp b/fft/src/KokkosFFT_ROCM_types.hpp
@@ -364,18 +364,8 @@ auto direction_type(Direction direction) {
   return direction == Direction::forward ? ROCFFT_FORWARD : ROCFFT_BACKWARD;
 }
 
-inline void initialize_host() {}
-inline void finalize_host() {}
 #endif
 
-inline void initialize_device() {
-  rocfft_status status = rocfft_setup();
-  if (status != rocfft_status_success) Kokkos::abort("rocfft_setup failed");
-}
-inline void finalize_device() {
-  rocfft_status status = rocfft_cleanup();
-  if (status != rocfft_status_success) Kokkos::abort("rocfft_cleanup failed");
-}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_SYCL_plans.hpp b/fft/src/KokkosFFT_SYCL_plans.hpp
@@ -14,6 +14,27 @@
 
 namespace KokkosFFT {
 namespace Impl {
+
+template <
+    typename ExecutionSpace, typename T,
+    std::enable_if_t<std::is_same_v<ExecutionSpace, Kokkos::Experimental::SYCL>,
+                     std::nullptr_t> = nullptr>
+void setup() {
+  static bool once = [] {
+    if (!(Kokkos::is_initialized() || Kokkos::is_finalized())) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called before initializing "
+          "Kokkos.\n");
+    }
+    if (Kokkos::is_finalized()) {
+      Kokkos::abort(
+          "Error: KokkosFFT APIs must not be called after finalizing "
+          "Kokkos.\n");
+    }
+    return true;
+  }();
+}
+
 // Helper to convert the integer type of vectors
 template <typename InType, typename OutType>
 auto convert_int_type(std::vector<InType>& in) -> std::vector<OutType> {

diff --git a/fft/src/KokkosFFT_SYCL_types.hpp b/fft/src/KokkosFFT_SYCL_types.hpp
@@ -226,12 +226,8 @@ auto direction_type(Direction direction) {
   return direction == Direction::forward ? MKL_FFT_FORWARD : MKL_FFT_BACKWARD;
 }
 
-inline void initialize_host() {}
-inline void finalize_host() {}
 #endif
 
-inline void initialize_device() {}
-inline void finalize_device() {}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/unit_test/CMakeLists.txt b/fft/unit_test/CMakeLists.txt
@@ -8,18 +8,9 @@ add_executable(unit-tests-kokkos-fft-core
     Test_Transform.cpp
 )
 
-add_executable(unit-tests-kokkos-fft-core-no-init
-    Test_MainNoInit.cpp
-    Test_InitializeFinalize.cpp
-)
-
 target_compile_features(unit-tests-kokkos-fft-core PUBLIC cxx_std_17)
 target_link_libraries(unit-tests-kokkos-fft-core PUBLIC KokkosFFT::fft GTest::gtest)
 
-target_compile_features(unit-tests-kokkos-fft-core-no-init PUBLIC cxx_std_17)
-target_link_libraries(unit-tests-kokkos-fft-core-no-init PUBLIC KokkosFFT::fft GTest::gtest)
-
 # Enable GoogleTest
 include(GoogleTest)
 gtest_discover_tests(unit-tests-kokkos-fft-core PROPERTIES DISCOVERY_TIMEOUT 600 DISCOVERY_MODE PRE_TEST)
-gtest_discover_tests(unit-tests-kokkos-fft-core-no-init PROPERTIES DISCOVERY_TIMEOUT 600 DISCOVERY_MODE PRE_TEST)