Disallow plan creation if execution space is not appropriate

yasahi-hpc · Aug 26, 2024 · 771f97d · 771f97d
1 parent f3bf5ad
commit 771f97d
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 33 deletions.
diff --git a/fft/src/KokkosFFT_Host_plans.hpp b/fft/src/KokkosFFT_Host_plans.hpp
@@ -32,9 +32,8 @@ void init_threads([[maybe_unused]] const ExecutionSpace& exec_space) {
 template <typename ExecutionSpace, typename PlanType, typename InViewType,
           typename OutViewType, typename BufferViewType, typename InfoType,
           std::size_t fft_rank = 1,
-          std::enable_if_t<
-              is_HostSpace_v<ExecutionSpace>,
-              std::nullptr_t> = nullptr>
+          std::enable_if_t<is_AnyHostSpace_v<ExecutionSpace>, std::nullptr_t> =
+              nullptr>
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
@@ -109,9 +108,8 @@ auto create_plan(const ExecutionSpace& exec_space,
 }
 
 template <typename ExecutionSpace, typename PlanType, typename InfoType,
-          std::enable_if_t<
-              is_HostSpace_v<ExecutionSpace>,
-              std::nullptr_t> = nullptr>
+          std::enable_if_t<is_AnyHostSpace_v<ExecutionSpace>, std::nullptr_t> =
+              nullptr>
 void destroy_plan_and_info(std::unique_ptr<PlanType>& plan, InfoType&) {
   if constexpr (std::is_same_v<PlanType, fftwf_plan>) {
     fftwf_destroy_plan(*plan);

diff --git a/fft/src/KokkosFFT_Plans.hpp b/fft/src/KokkosFFT_Plans.hpp
@@ -155,10 +155,11 @@ class Plan {
   /// \param axis [in] Axis over which FFT is performed
   /// \param n [in] Length of the transformed axis of the output (optional)
   //
-  explicit Plan(const ExecutionSpace& exec_space, InViewType& in,
-                OutViewType& out, KokkosFFT::Direction direction, int axis,
-                std::optional<std::size_t> n = std::nullopt,
-                std::enable_if_t<is_AllowedSpace_v<execSpace>, std::nullptr_t> = nullptr)
+  explicit Plan(
+      const ExecutionSpace& exec_space, InViewType& in, OutViewType& out,
+      KokkosFFT::Direction direction, int axis,
+      std::optional<std::size_t> n = std::nullopt,
+      std::enable_if_t<is_AllowedSpace_v<execSpace>, std::nullptr_t> = nullptr)
       : m_exec_space(exec_space), m_axes({axis}), m_direction(direction) {
     static_assert(
         KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
@@ -211,10 +212,11 @@ class Plan {
   /// \param axes [in] Axes over which FFT is performed
   /// \param s [in] Shape of the transformed axis of the output (optional)
   //
-  explicit Plan(const ExecutionSpace& exec_space, InViewType& in,
-                OutViewType& out, KokkosFFT::Direction direction,
-                axis_type<DIM> axes, shape_type<DIM> s = {0},
-                std::enable_if_t<is_AllowedSpace_v<execSpace>, std::nullptr_t> = nullptr)
+  explicit Plan(
+      const ExecutionSpace& exec_space, InViewType& in, OutViewType& out,
+      KokkosFFT::Direction direction, axis_type<DIM> axes,
+      shape_type<DIM> s                                              = {0},
+      std::enable_if_t<is_AllowedSpace_v<execSpace>, std::nullptr_t> = nullptr)
       : m_exec_space(exec_space), m_axes(axes), m_direction(direction) {
     static_assert(
         KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
@@ -261,7 +263,7 @@ class Plan {
                                                                         m_info);
   }
 
-  Plan() = delete;
+  Plan()            = delete;
   Plan(const Plan&) = delete;
   Plan& operator=(const Plan&) = delete;
   Plan& operator=(Plan&&) = delete;

diff --git a/fft/unit_test/Test_Plans.cpp b/fft/unit_test/Test_Plans.cpp
@@ -15,9 +15,9 @@ using test_types = ::testing::Types<std::pair<float, Kokkos::LayoutLeft>,
                                     std::pair<double, Kokkos::LayoutRight> >;
 
 #if defined(KOKKOS_ENABLE_SERIAL)
-using execution_spaces = ::testing::Types<Kokkos::Serial,
-                                          Kokkos::DefaultHostExecutionSpace,
-                                          Kokkos::DefaultExecutionSpace>;
+using execution_spaces =
+    ::testing::Types<Kokkos::Serial, Kokkos::DefaultHostExecutionSpace,
+                     Kokkos::DefaultExecutionSpace>;
 #else
 using execution_spaces = ::testing::Types<Kokkos::DefaultHostExecutionSpace,
                                           Kokkos::DefaultExecutionSpace>;
@@ -59,24 +59,35 @@ TYPED_TEST_SUITE(Plans3D, test_types);
 // Tests for execution space
 template <typename ExecutionSpace>
 void test_plan_constructible() {
-  using ValueType = double;
+  using ValueType      = double;
   using RealView1DType = Kokkos::View<ValueType*, ExecutionSpace>;
   using ComplexView1DType =
       Kokkos::View<Kokkos::complex<ValueType>*, ExecutionSpace>;
-  using PlanType = KokkosFFT::Impl::Plan<ExecutionSpace, RealView1DType, ComplexView1DType>;
-
-  #if !defined(ENABLE_HOST_AND_DEVICE) && (defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL))
-    // For GPUs without ENABLE_HOST_AND_DEVICE, a plan can be constructible from Kokkos::DefaultExecutionSpace only
-    if constexpr (std::is_same_v<ExecutionSpace, Kokkos::DefaultExecutionSpace>) {
-      static_assert(std::is_constructible_v<PlanType, const ExecutionSpace&, RealView1DType&, ComplexView1DType&, KokkosFFT::Direction, int>);
-    } else {
-      static_assert(!std::is_constructible_v<PlanType, const ExecutionSpace&, RealView1DType&, ComplexView1DType&, KokkosFFT::Direction, int>);
-    }
-  #else
-    // For CPUs or GPUs with ENABLE_HOST_AND_DEVICE, 
-    // a plan can be constructible from Kokkos::DefaultExecutionSpace, Kokkos::DefaultHostExecutionSpace or Kokkos::Serial (if enabled)
-    static_assert(std::is_constructible_v<PlanType, const ExecutionSpace&, RealView1DType&, ComplexView1DType&, KokkosFFT::Direction, int>);
-  #endif
+  using PlanType =
+      KokkosFFT::Impl::Plan<ExecutionSpace, RealView1DType, ComplexView1DType>;
+
+#if !defined(ENABLE_HOST_AND_DEVICE) &&                           \
+    (defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || \
+     defined(KOKKOS_ENABLE_SYCL))
+  // For GPUs without ENABLE_HOST_AND_DEVICE, a plan can be constructible from
+  // Kokkos::DefaultExecutionSpace only
+  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::DefaultExecutionSpace>) {
+    static_assert(std::is_constructible_v<PlanType, const ExecutionSpace&,
+                                          RealView1DType&, ComplexView1DType&,
+                                          KokkosFFT::Direction, int>);
+  } else {
+    static_assert(!std::is_constructible_v<PlanType, const ExecutionSpace&,
+                                           RealView1DType&, ComplexView1DType&,
+                                           KokkosFFT::Direction, int>);
+  }
+#else
+  // For CPUs or GPUs with ENABLE_HOST_AND_DEVICE,
+  // a plan can be constructible from Kokkos::DefaultExecutionSpace,
+  // Kokkos::DefaultHostExecutionSpace or Kokkos::Serial (if enabled)
+  static_assert(
+      std::is_constructible_v<PlanType, const ExecutionSpace&, RealView1DType&,
+                              ComplexView1DType&, KokkosFFT::Direction, int>);
+#endif
 }
 
 // Tests for 1D FFT plan on 1D View