Merge pull request #16 from CExA-project/add-hfft

Add hfft

common/src/KokkosFFT_default_types.hpp

@@ -34,6 +34,10 @@ namespace KokkosFFT {
   template <std::size_t DIM>
   using axis_type = std::array<int, DIM>;
 
+  // Define type to specify new shape
+  template <std::size_t DIM>
+  using shape_type = std::array<std::size_t, DIM>;
+
   enum class Normalization {
     FORWARD,
     BACKWARD,

common/src/KokkosFFT_normalization.hpp

@@ -53,7 +53,23 @@ namespace Impl {
     auto [coef, to_normalize] = _coefficients(inout, direction, normalization, fft_size);
     if(to_normalize) _normalize(exec_space, inout, coef);
   }
+
+  auto swap_direction(Normalization normalization) {
+    Normalization new_direction = Normalization::FORWARD;
+    switch (normalization) {
+      case Normalization::FORWARD:
+        new_direction = Normalization::BACKWARD;
+        break;
+      case Normalization::BACKWARD:
+        new_direction = Normalization::FORWARD;
+        break;
+      case Normalization::ORTHO:
+        new_direction = Normalization::ORTHO;
+        break;
+    };
+    return new_direction;
+  }
 } // namespace Impl
-}; // namespace KokkosFFT
+} // namespace KokkosFFT
 
 #endif

common/src/KokkosFFT_padding.hpp

@@ -0,0 +1,153 @@
+#ifndef KOKKOSFFT_PADDING_HPP
+#define KOKKOSFFT_PADDING_HPP
+
+#include <tuple>
+#include "KokkosFFT_default_types.hpp"
+#include "KokkosFFT_utils.hpp"
+
+namespace KokkosFFT {
+namespace Impl {
+  template <typename ViewType, std::size_t DIM>
+  auto get_modified_shape(const ViewType& view, shape_type<DIM> shape) {
+    static_assert(ViewType::rank() >= DIM,
+                  "KokkosFFT::get_modified_shape: Rank of View must be larger than or equal to the Rank of new shape");
+    static_assert(DIM > 0,
+                  "KokkosFFT::get_modified_shape: Rank of FFT axes must be larger than or equal to 1");
+
+    // [TO DO] Add a is_C2R arg. If is_C2R is true, then shape should be shape/2+1
+    constexpr int rank = static_cast<int>( ViewType::rank() );
+
+    using full_shape_type = shape_type<rank>;
+    full_shape_type modified_shape;
+    for(int i=0; i<rank; i++) {
+      modified_shape.at(i) = view.extent(i);
+    }
+
+    // Update shapes based on newly given shape
+    for(int i=0; i<DIM; i++) {
+      assert(shape.at(i) > 0);
+      modified_shape.at(i) = shape.at(i);
+    }
+
+    // [TO DO] may return, is_modification_needed if modified_shape is not equal view.extents()
+    // May be implement other function. is_crop_or_pad_needed(view, shape);
+    return modified_shape;
+  }
+
+  template <typename ViewType, std::size_t DIM>
+  auto is_crop_or_pad_needed(const ViewType& view, const shape_type<DIM>& modified_shape) {
+    static_assert(ViewType::rank() == DIM,
+                  "KokkosFFT::_crop_or_pad: Rank of View must be equal to Rank of extended shape.");
+
+    // [TO DO] Add a is_C2R arg. If is_C2R is true, then shape should be shape/2+1
+    constexpr int rank = static_cast<int>( ViewType::rank() );
+
+    bool not_same = false;
+    for(int i=0; i<rank; i++) {
+      if(modified_shape.at(i) != view.extent(i)) {
+        not_same = true;
+        break;
+      }
+    }
+
+    return not_same;
+  }
+
+  template <typename ExecutionSpace, typename ViewType>
+  void _crop_or_pad(const ExecutionSpace& exec_space,
+                    const ViewType& in,
+                    ViewType& out,
+                    shape_type<1> s) {
+    constexpr std::size_t DIM = 1;
+    static_assert(ViewType::rank() == DIM,
+                  "KokkosFFT::_crop_or_pad: Rank of View must be equal to Rank of extended shape.");
+
+    auto _n0 = s.at(0);
+    out = ViewType("out", _n0);
+
+    auto n0 = std::min(_n0, in.extent(0));
+
+    Kokkos::parallel_for(
+      Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<std::size_t>>(exec_space, 0, n0),
+      KOKKOS_LAMBDA (int i0) {
+        out(i0) = in(i0);
+      }
+    );
+  }
+
+  template <typename ExecutionSpace, typename ViewType>
+  void _crop_or_pad(const ExecutionSpace& exec_space,
+                    const ViewType& in,
+                    ViewType& out,
+                    shape_type<2> s) {
+    constexpr std::size_t DIM = 2;
+    static_assert(ViewType::rank() == DIM,
+                  "KokkosFFT::_crop_or_pad: Rank of View must be equal to Rank of extended shape.");
+
+    auto [_n0, _n1] = s;
+    out = ViewType("out", _n0, _n1);
+
+    int n0 = std::min(_n0, in.extent(0));
+    int n1 = std::min(_n1, in.extent(1));
+
+    using range_type = Kokkos::MDRangePolicy<ExecutionSpace,  Kokkos::Rank<2, Kokkos::Iterate::Default, Kokkos::Iterate::Default> >;
+    using tile_type = typename range_type::tile_type;
+    using point_type = typename range_type::point_type;
+
+    range_type range(
+      point_type{{0, 0}},
+      point_type{{n0, n1}},
+      tile_type{{4, 4}} // [TO DO] Choose optimal tile sizes for each device
+    );
+
+    Kokkos::parallel_for(range,
+      KOKKOS_LAMBDA (int i0, int i1) {
+        out(i0, i1) = in(i0, i1);
+      }
+    );
+  }
+
+  template <typename ExecutionSpace, typename ViewType>
+  void _crop_or_pad(const ExecutionSpace& exec_space,
+                    const ViewType& in,
+                    ViewType& out,
+                    shape_type<3> s) {
+    constexpr std::size_t DIM = 3;
+    static_assert(ViewType::rank() == DIM,
+                  "KokkosFFT::_crop_or_pad: Rank of View must be equal to Rank of extended shape.");
+
+    auto [_n0, _n1, _n2] = s;
+    out = ViewType("out", _n0, _n1, _n2);
+
+    int n0 = std::min(_n0, in.extent(0));
+    int n1 = std::min(_n1, in.extent(1));
+    int n2 = std::min(_n2, in.extent(2));
+
+    using range_type = Kokkos::MDRangePolicy<ExecutionSpace,  Kokkos::Rank<3, Kokkos::Iterate::Default, Kokkos::Iterate::Default> >;
+    using tile_type = typename range_type::tile_type;
+    using point_type = typename range_type::point_type;
+
+    range_type range(
+      point_type{{0, 0, 0}},
+      point_type{{n0, n1, n2}},
+      tile_type{{4, 4, 4}} // [TO DO] Choose optimal tile sizes for each device
+    );
+
+    Kokkos::parallel_for(range,
+      KOKKOS_LAMBDA (int i0, int i1, int i2) {
+        out(i0, i1, i2) = in(i0, i1, i2);
+      }
+    );
+  }
+
+  template <typename ExecutionSpace, typename ViewType, std::size_t DIM = 1>
+  void crop_or_pad(const ExecutionSpace& exec_space,
+                   const ViewType& in,
+                   ViewType& out,
+                   shape_type<DIM> s) {
+    _crop_or_pad(exec_space, in, out, s);
+  }
+} // namespace Impl
+} // namespace KokkosFFT
+
+#endif

common/src/KokkosFFT_utils.hpp

@@ -28,6 +28,16 @@ namespace Impl {
   template <typename T>
   struct is_complex<Kokkos::complex<T>> : std::true_type {};
 
+  template <typename ExecutionSpace, typename ViewType,
+            std::enable_if_t<ViewType::rank()==1, std::nullptr_t> = nullptr>
+  struct complex_view_type {
+    using value_type = typename ViewType::non_const_value_type;
+    using float_type = KokkosFFT::Impl::real_type_t<value_type>;
+    using complex_type = Kokkos::complex<float_type>;
+    using array_layout_type = typename ViewType::array_layout;
+    using type = Kokkos::View<complex_type*, array_layout_type, ExecutionSpace>;
+  };
+
   template <typename ViewType>
   auto convert_negative_axis(const ViewType& view, int _axis=-1) {
     static_assert(Kokkos::is_view<ViewType>::value,
@@ -38,6 +48,27 @@ namespace Impl {
     return axis;
   }
 
+  template <typename ViewType>
+  auto convert_negative_shift(const ViewType& view, int _shift, int _axis) {
+    static_assert(Kokkos::is_view<ViewType>::value,
+                  "KokkosFFT::convert_negative_shift: ViewType is not a Kokkos::View.");
+    int axis = convert_negative_axis(view, _axis);
+    int extent = view.extent(axis);
+    int shift0 = 0, shift1 = 0, shift2 = extent/2;
+
+    if(_shift < 0) {
+      shift0 = -_shift + extent%2; // add 1 for odd case
+      shift1 = -_shift;
+      shift2 = 0;
+    } else if(_shift>0){
+      shift0 = _shift;
+      shift1 = _shift + extent%2; // add 1 for odd case
+      shift2 = 0;
+    }
+
+    return std::tuple<int, int, int>({shift0, shift1, shift2});
+  }
+
   template <typename T>
   void permute(std::vector<T>& values, const std::vector<std::size_t>& indices) {
     std::vector<T> out;
@@ -59,6 +90,17 @@ namespace Impl {
     return set_values.size() < values.size();
   }
 
+  template <typename IntType,
+            std::size_t DIM=1,
+            std::enable_if_t<std::is_integral_v<IntType>, std::nullptr_t> = nullptr>
+  bool is_out_of_range_value_included(const std::array<IntType, DIM>& values, IntType max) {
+    bool is_included = false;
+    for(auto value: values) {
+      is_included = value >= max;
+    }
+    return is_included;
+  }
+
   template <std::size_t DIM=1>
   bool is_transpose_needed(std::array<int, DIM> map) {
     std::array<int, DIM> contiguous_map;
@@ -91,7 +133,56 @@ namespace Impl {
                    [=](const T sequence) -> T {return start + sequence;});
     return sequence;
   }
+
+  template <typename ElementType>
+  inline std::vector<ElementType> arange(const ElementType start,
+                                         const ElementType stop,
+                                         const ElementType step=1
+                                        ) {
+    const size_t length = ceil((stop - start) / step);
+    std::vector<ElementType> result;
+    ElementType delta = (stop - start) / length;
+
+    // thrust::sequence
+    for(auto i=0; i<length; i++) {
+      ElementType value = start + delta*i;
+      result.push_back(value);
+    }
+    return result;
+  }
+
+  template <typename ExecutionSpace, typename InViewType, typename OutViewType>
+  void conjugate(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out) {
+    using in_value_type = typename InViewType::non_const_value_type;
+    using out_value_type = typename OutViewType::non_const_value_type;
+
+    static_assert(KokkosFFT::Impl::is_complex<out_value_type>::value,
+                  "KokkosFFT::Impl::conjugate: OutViewType must be complex");
+    std::size_t size = in.size();
+
+    // [TO DO] is there a way to get device mirror?
+    if constexpr(InViewType::rank() == 1) {
+      out = OutViewType("out", in.extent(0));
+    } else if constexpr(InViewType::rank() == 2) {
+      out = OutViewType("out", in.extent(0), in.extent(1));
+    } else if constexpr(InViewType::rank() == 3) {
+      out = OutViewType("out", in.extent(0), in.extent(1), in.extent(2));
+    } else if constexpr(InViewType::rank() == 4) {
+      out = OutViewType("out", in.extent(0), in.extent(1), in.extent(2), in.extent(3));
+    }
+
+    auto* in_data = in.data();
+    auto* out_data = out.data();
+
+    Kokkos::parallel_for(
+      Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<std::size_t>>(exec_space, 0, size),
+      KOKKOS_LAMBDA(const int& i) {
+        out_value_type tmp = in_data[i];
+        out_data[i] = Kokkos::conj(in_data[i]);
+      }
+    );
+  }
 } // namespace Impl
-}; // namespace KokkosFFT
+} // namespace KokkosFFT
 
 #endif

common/unit_test/CMakeLists.txt

@@ -4,6 +4,7 @@ add_executable(unit-tests-kokkos-fft-common
     Test_Normalization.cpp
     Test_Transpose.cpp
     Test_Layouts.cpp
+    Test_Padding.cpp
 )
 
 target_compile_features(unit-tests-kokkos-fft-common PUBLIC cxx_std_17)