Introduce global setup/cleanup in APIs

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,

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() {
+  [[maybe_unused]] 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,

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() {
+  [[maybe_unused]] 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,

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);
   }

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() {
+  [[maybe_unused]] 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,

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> {

fft/unit_test/CMakeLists.txt

@@ -9,7 +9,6 @@ add_executable(unit-tests-kokkos-fft-core
 )
 
 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)
 
 # Enable GoogleTest

fft/unit_test/Test_Main.cpp

@@ -3,25 +3,16 @@
 //
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include <Kokkos_Core.hpp>
 #include <gtest/gtest.h>
-
-namespace testing::internal {
-// accessing gtest internals is not very clean, but gtest provides no public
-// access...
-extern bool g_help_flag;
-}  // namespace testing::internal
+#include <Kokkos_Core.hpp>
 
 int main(int argc, char* argv[]) {
   ::testing::InitGoogleTest(&argc, argv);
+
   int result = 0;
-  if (::testing::GTEST_FLAG(list_tests) || ::testing::internal::g_help_flag) {
-    result = RUN_ALL_TESTS();
-  } else {
-    Kokkos::initialize(argc, argv);
-    result = RUN_ALL_TESTS();
-    Kokkos::finalize();
-  }
+  Kokkos::initialize(argc, argv);
+  result = RUN_ALL_TESTS();
+  Kokkos::finalize();
 
   return result;
 }