add device specific setup functions

yasahi-hpc · Jan 13, 2025 · e0abd00 · e0abd00
1 parent 673b484
commit e0abd00
Show file tree

Hide file tree

Showing 9 changed files with 73 additions and 18 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 INTERFACE)
+add_library(fft STATIC KokkosFFT_Core.cpp)
 
 target_link_libraries(fft
-    INTERFACE
+    PUBLIC
     common
     Kokkos::kokkos
 )
 
-target_include_directories(fft INTERFACE
+target_include_directories(fft PUBLIC
                            $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
                            $<INSTALL_INTERFACE:${INSTALL_INCLUDEDIR}>
 )

diff --git a/fft/src/KokkosFFT.hpp b/fft/src/KokkosFFT.hpp
@@ -10,5 +10,6 @@
 #include "KokkosFFT_Helpers.hpp"
 #include "KokkosFFT_Plans.hpp"
 #include "KokkosFFT_Transform.hpp"
+#include "KokkosFFT_Core.hpp"
 
 #endif
diff --git a/fft/src/KokkosFFT_Cuda_types.hpp b/fft/src/KokkosFFT_Cuda_types.hpp
@@ -238,7 +238,13 @@ template <typename ExecutionSpace>
 auto direction_type(Direction direction) {
   return direction == Direction::forward ? CUFFT_FORWARD : CUFFT_INVERSE;
 }
+
+static void initialize_host() {}
+static void finalize_host() {}
 #endif
+
+static void initialize_device() {}
+static void finalize_device() {}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_FFTW_Types.hpp b/fft/src/KokkosFFT_FFTW_Types.hpp
@@ -138,6 +138,28 @@ struct ScopedFFTWPlan {
   }
 };
 
+#if defined(KOKKOS_ENABLE_OPENMP) || defined(KOKKOS_ENABLE_THREADS)
+static void initialize_host() {
+  fftwf_init_threads();
+  fftw_init_threads();
+}
+static void finalize_host() {
+  fftwf_cleanup_threads();
+  fftw_cleanup_threads();
+}
+#else
+static void initialize_host() {}
+static 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))
+static void initialize_device() {}
+static void finalize_device() {}
+#endif
+
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_HIP_types.hpp b/fft/src/KokkosFFT_HIP_types.hpp
@@ -238,7 +238,13 @@ template <typename ExecutionSpace>
 auto direction_type(Direction direction) {
   return direction == Direction::forward ? HIPFFT_FORWARD : HIPFFT_BACKWARD;
 }
+
+static void initialize_host() {}
+static void finalize_host() {}
 #endif
+
+static void initialize_device() {}
+static void finalize_device() {}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/src/KokkosFFT_ROCM_types.hpp b/fft/src/KokkosFFT_ROCM_types.hpp
@@ -363,7 +363,19 @@ template <typename ExecutionSpace>
 auto direction_type(Direction direction) {
   return direction == Direction::forward ? ROCFFT_FORWARD : ROCFFT_BACKWARD;
 }
+
+static void initialize_host() {}
+static void finalize_host() {}
 #endif
+
+static void initialize_device() {
+  rocfft_status status = rocfft_setup();
+  if (status != rocfft_status_success) Kokkos::abort("rocfft_setup failed");
+}
+static 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_types.hpp b/fft/src/KokkosFFT_SYCL_types.hpp
@@ -225,7 +225,13 @@ template <typename ExecutionSpace>
 auto direction_type(Direction direction) {
   return direction == Direction::forward ? MKL_FFT_FORWARD : MKL_FFT_BACKWARD;
 }
+
+static void initialize_host() {}
+static void finalize_host() {}
 #endif
+
+static void initialize_device() {}
+static void finalize_device() {}
 }  // namespace Impl
 }  // namespace KokkosFFT
 

diff --git a/fft/unit_test/CMakeLists.txt b/fft/unit_test/CMakeLists.txt
@@ -8,10 +8,18 @@ add_executable(unit-tests-kokkos-fft-core
     Test_Transform.cpp
 )
 
-target_compile_features(unit-tests-kokkos-fft-core PUBLIC cxx_std_17)
+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)
diff --git a/fft/unit_test/Test_Main.cpp b/fft/unit_test/Test_Main.cpp
@@ -3,25 +3,19 @@
 //
 // 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>
+#include <KokkosFFT_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);
+  KokkosFFT::initialize();
+  result = RUN_ALL_TESTS();
+  KokkosFFT::finalize();
+  Kokkos::finalize();
 
   return result;
 }