Support in place transform (#149)

* fix: conflict * compute extents for in-place transform * Add new line to tests * fix: conflicts * fix conflicts * fix: conflicts * fix: conflicts * cleanup in-place tests * fix: normalization for C2R * Add fence as a temporary solution * prefer static cast for casting to void* * Use Kokkos::numbers::pi instead of M_PI * Do not use Unmanaged traits explicitly * format * use z for imaginary unit * fix Kokkos::numbers::pi usage * Use pi in double precision * fix test utils based on review --------- Co-authored-by: Yuuichi Asahi <[email protected]>
kokkos · Oct 24, 2024 · d620636 · d620636
1 parent 51a9895
commit d620636
Show file tree

Hide file tree

Showing 11 changed files with 391 additions and 135 deletions.
diff --git a/common/src/KokkosFFT_layouts.hpp b/common/src/KokkosFFT_layouts.hpp
@@ -21,7 +21,8 @@ namespace Impl {
 */
 template <typename InViewType, typename OutViewType, std::size_t DIM = 1>
 auto get_extents(const InViewType& in, const OutViewType& out,
-                 axis_type<DIM> axes, shape_type<DIM> shape = {0}) {
+                 axis_type<DIM> axes, shape_type<DIM> shape = {},
+                 bool is_inplace = false) {
   using in_value_type     = typename InViewType::non_const_value_type;
   using out_value_type    = typename OutViewType::non_const_value_type;
   using array_layout_type = typename InViewType::array_layout;
@@ -65,20 +66,28 @@ auto get_extents(const InViewType& in, const OutViewType& out,
 
   if constexpr (is_real_v<in_value_type>) {
     // Then R2C
-    KOKKOSFFT_THROW_IF(
-        _out_extents.at(inner_most_axis) !=
-            _in_extents.at(inner_most_axis) / 2 + 1,
-        "For R2C, the 'output extent' of transform must be equal to "
-        "'input extent'/2 + 1");
+    if (is_inplace) {
+      _in_extents.at(inner_most_axis) = _out_extents.at(inner_most_axis) * 2;
+    } else {
+      KOKKOSFFT_THROW_IF(
+          _out_extents.at(inner_most_axis) !=
+              _in_extents.at(inner_most_axis) / 2 + 1,
+          "For R2C, the 'output extent' of transform must be equal to "
+          "'input extent'/2 + 1");
+    }
   }
 
   if constexpr (is_real_v<out_value_type>) {
     // Then C2R
-    KOKKOSFFT_THROW_IF(
-        _in_extents.at(inner_most_axis) !=
-            _out_extents.at(inner_most_axis) / 2 + 1,
-        "For C2R, the 'input extent' of transform must be equal to "
-        "'output extent' / 2 + 1");
+    if (is_inplace) {
+      _out_extents.at(inner_most_axis) = _in_extents.at(inner_most_axis) * 2;
+    } else {
+      KOKKOSFFT_THROW_IF(
+          _in_extents.at(inner_most_axis) !=
+              _out_extents.at(inner_most_axis) / 2 + 1,
+          "For C2R, the 'input extent' of transform must be equal to "
+          "'output extent' / 2 + 1");
+    }
   }
 
   if constexpr (std::is_same_v<array_layout_type, Kokkos::LayoutLeft>) {

diff --git a/common/src/KokkosFFT_utils.hpp b/common/src/KokkosFFT_utils.hpp
@@ -17,6 +17,11 @@
 namespace KokkosFFT {
 namespace Impl {
 
+template <typename ScalarType1, typename ScalarType2>
+bool are_aliasing(const ScalarType1* ptr1, const ScalarType2* ptr2) {
+  return (static_cast<const void*>(ptr1) == static_cast<const void*>(ptr2));
+}
+
 template <typename ViewType>
 auto convert_negative_axis(ViewType, int _axis = -1) {
   static_assert(Kokkos::is_view_v<ViewType>,

diff --git a/common/unit_test/Test_Utils.cpp b/common/unit_test/Test_Utils.cpp
@@ -11,6 +11,13 @@ using test_types = ::testing::Types<Kokkos::LayoutLeft, Kokkos::LayoutRight>;
 // Int like types
 using base_int_types = ::testing::Types<int, std::size_t>;
 
+// value type combinations that are tested
+using paired_scalar_types = ::testing::Types<
+    std::pair<float, float>, std::pair<float, Kokkos::complex<float>>,
+    std::pair<Kokkos::complex<float>, Kokkos::complex<float>>,
+    std::pair<double, double>, std::pair<double, Kokkos::complex<double>>,
+    std::pair<Kokkos::complex<double>, Kokkos::complex<double>>>;
+
 // Basically the same fixtures, used for labeling tests
 template <typename T>
 struct ConvertNegativeAxis : public ::testing::Test {
@@ -30,9 +37,16 @@ struct ContainerTypes : public ::testing::Test {
   using array_type                  = std::array<T, rank>;
 };
 
+template <typename T>
+struct PairedScalarTypes : public ::testing::Test {
+  using value_type1 = typename T::first_type;
+  using value_type2 = typename T::second_type;
+};
+
 TYPED_TEST_SUITE(ConvertNegativeAxis, test_types);
 TYPED_TEST_SUITE(ConvertNegativeShift, test_types);
 TYPED_TEST_SUITE(ContainerTypes, base_int_types);
+TYPED_TEST_SUITE(PairedScalarTypes, paired_scalar_types);
 
 // Tests for convert_negative_axes over ND views
 template <typename LayoutType>
@@ -550,3 +564,25 @@ TEST(ToArray, lvalue) {
 TEST(ToArray, rvalue) {
   ASSERT_EQ(KokkosFFT::Impl::to_array(std::array{1, 2}), (Kokkos::Array{1, 2}));
 }
+
+template <typename ValueType1, typename ValueType2>
+void test_are_pointers_aliasing() {
+  using View1 = Kokkos::View<ValueType1*, execution_space>;
+  using View2 = Kokkos::View<ValueType2*, execution_space>;
+
+  const int n1 = 10;
+  // sizeof ValeuType2 is larger or equal to ValueType1
+  const int n2 = sizeof(ValueType1) == sizeof(ValueType2) ? n1 : n1 / 2 + 1;
+  View1 view1("view1", n1);
+  View2 view2("view2", n1);
+  View2 uview2(reinterpret_cast<ValueType2*>(view1.data()), n2);
+
+  EXPECT_TRUE(KokkosFFT::Impl::are_aliasing(view1.data(), uview2.data()));
+  EXPECT_FALSE(KokkosFFT::Impl::are_aliasing(view1.data(), view2.data()));
+}
+
+TYPED_TEST(PairedScalarTypes, are_pointers_aliasing) {
+  using value_type1 = typename TestFixture::value_type1;
+  using value_type2 = typename TestFixture::value_type2;
+  test_are_pointers_aliasing<value_type1, value_type2>();
+}
diff --git a/fft/src/KokkosFFT_Cuda_plans.hpp b/fft/src/KokkosFFT_Cuda_plans.hpp
@@ -22,7 +22,8 @@ template <typename ExecutionSpace, typename PlanType, typename InViewType,
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
-                 Direction /*direction*/, axis_type<1> axes, shape_type<1> s) {
+                 Direction /*direction*/, axis_type<1> axes, shape_type<1> s,
+                 bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -44,7 +45,7 @@ auto create_plan(const ExecutionSpace& exec_space,
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   const int nx = fft_extents.at(0);
   int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1,
                                  std::multiplies<>());
@@ -64,7 +65,8 @@ template <typename ExecutionSpace, typename PlanType, typename InViewType,
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
-                 Direction /*direction*/, axis_type<2> axes, shape_type<2> s) {
+                 Direction /*direction*/, axis_type<2> axes, shape_type<2> s,
+                 bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -86,7 +88,7 @@ auto create_plan(const ExecutionSpace& exec_space,
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   const int nx = fft_extents.at(0), ny = fft_extents.at(1);
   int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1,
                                  std::multiplies<>());
@@ -106,7 +108,8 @@ template <typename ExecutionSpace, typename PlanType, typename InViewType,
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
-                 Direction /*direction*/, axis_type<3> axes, shape_type<3> s) {
+                 Direction /*direction*/, axis_type<3> axes, shape_type<3> s,
+                 bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -128,7 +131,7 @@ auto create_plan(const ExecutionSpace& exec_space,
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
 
   const int nx = fft_extents.at(0), ny = fft_extents.at(1),
             nz = fft_extents.at(2);
@@ -151,7 +154,7 @@ auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
                  Direction /*direction*/, axis_type<fft_rank> axes,
-                 shape_type<fft_rank> s) {
+                 shape_type<fft_rank> s, bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -172,7 +175,7 @@ auto create_plan(const ExecutionSpace& exec_space,
       KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                       out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   int idist    = std::accumulate(in_extents.begin(), in_extents.end(), 1,
                                  std::multiplies<>());
   int odist    = std::accumulate(out_extents.begin(), out_extents.end(), 1,

diff --git a/fft/src/KokkosFFT_HIP_plans.hpp b/fft/src/KokkosFFT_HIP_plans.hpp
@@ -22,7 +22,8 @@ template <typename ExecutionSpace, typename PlanType, typename InViewType,
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
-                 Direction /*direction*/, axis_type<1> axes, shape_type<1> s) {
+                 Direction /*direction*/, axis_type<1> axes, shape_type<1> s,
+                 bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -44,7 +45,7 @@ auto create_plan(const ExecutionSpace& exec_space,
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   const int nx = fft_extents.at(0);
   int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1,
                                  std::multiplies<>());
@@ -64,7 +65,8 @@ template <typename ExecutionSpace, typename PlanType, typename InViewType,
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
-                 Direction /*direction*/, axis_type<2> axes, shape_type<2> s) {
+                 Direction /*direction*/, axis_type<2> axes, shape_type<2> s,
+                 bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -86,7 +88,7 @@ auto create_plan(const ExecutionSpace& exec_space,
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   const int nx = fft_extents.at(0), ny = fft_extents.at(1);
   int fft_size = std::accumulate(fft_extents.begin(), fft_extents.end(), 1,
                                  std::multiplies<>());
@@ -106,7 +108,8 @@ template <typename ExecutionSpace, typename PlanType, typename InViewType,
 auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
-                 Direction /*direction*/, axis_type<3> axes, shape_type<3> s) {
+                 Direction /*direction*/, axis_type<3> axes, shape_type<3> s,
+                 bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -128,7 +131,7 @@ auto create_plan(const ExecutionSpace& exec_space,
   auto type = KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                               out_value_type>::type();
   [[maybe_unused]] auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
 
   const int nx = fft_extents.at(0), ny = fft_extents.at(1),
             nz = fft_extents.at(2);
@@ -151,7 +154,7 @@ auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
                  Direction /*direction*/, axis_type<fft_rank> axes,
-                 shape_type<fft_rank> s) {
+                 shape_type<fft_rank> s, bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -172,7 +175,7 @@ auto create_plan(const ExecutionSpace& exec_space,
       KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                       out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   int idist    = std::accumulate(in_extents.begin(), in_extents.end(), 1,
                                  std::multiplies<>());
   int odist    = std::accumulate(out_extents.begin(), out_extents.end(), 1,

diff --git a/fft/src/KokkosFFT_Host_plans.hpp b/fft/src/KokkosFFT_Host_plans.hpp
@@ -38,7 +38,7 @@ auto create_plan(const ExecutionSpace& exec_space,
                  std::unique_ptr<PlanType>& plan, const InViewType& in,
                  const OutViewType& out, BufferViewType&, InfoType&,
                  Direction direction, axis_type<fft_rank> axes,
-                 shape_type<fft_rank> s) {
+                 shape_type<fft_rank> s, bool is_inplace) {
   static_assert(
       KokkosFFT::Impl::are_operatable_views_v<ExecutionSpace, InViewType,
                                               OutViewType>,
@@ -64,7 +64,7 @@ auto create_plan(const ExecutionSpace& exec_space,
       KokkosFFT::Impl::transform_type<ExecutionSpace, in_value_type,
                                       out_value_type>::type();
   auto [in_extents, out_extents, fft_extents, howmany] =
-      KokkosFFT::Impl::get_extents(in, out, axes, s);
+      KokkosFFT::Impl::get_extents(in, out, axes, s, is_inplace);
   int idist    = std::accumulate(in_extents.begin(), in_extents.end(), 1,
                                  std::multiplies<>());
   int odist    = std::accumulate(out_extents.begin(), out_extents.end(), 1,