diff --git a/common/src/KokkosFFT_default_types.hpp b/common/src/KokkosFFT_default_types.hpp index fb05d45a..3f2b6ec6 100644 --- a/common/src/KokkosFFT_default_types.hpp +++ b/common/src/KokkosFFT_default_types.hpp @@ -34,6 +34,10 @@ namespace KokkosFFT { template using axis_type = std::array; + // Define type to specify new shape + template + using shape_type = std::array; + enum class Normalization { FORWARD, BACKWARD, diff --git a/common/src/KokkosFFT_normalization.hpp b/common/src/KokkosFFT_normalization.hpp index 6c5ac9e3..d7f0e189 100644 --- a/common/src/KokkosFFT_normalization.hpp +++ b/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 \ No newline at end of file diff --git a/common/src/KokkosFFT_padding.hpp b/common/src/KokkosFFT_padding.hpp new file mode 100644 index 00000000..044b9d45 --- /dev/null +++ b/common/src/KokkosFFT_padding.hpp @@ -0,0 +1,153 @@ +#ifndef KOKKOSFFT_PADDING_HPP +#define KOKKOSFFT_PADDING_HPP + +#include +#include "KokkosFFT_default_types.hpp" +#include "KokkosFFT_utils.hpp" + +namespace KokkosFFT { +namespace Impl { + template + auto get_modified_shape(const ViewType& view, shape_type 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( ViewType::rank() ); + + using full_shape_type = shape_type; + full_shape_type modified_shape; + for(int i=0; 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 + auto is_crop_or_pad_needed(const ViewType& view, const shape_type& 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( ViewType::rank() ); + + bool not_same = false; + for(int i=0; i + 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>(exec_space, 0, n0), + KOKKOS_LAMBDA (int i0) { + out(i0) = in(i0); + } + ); + } + + template + 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 >; + 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 + 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 >; + 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 + void crop_or_pad(const ExecutionSpace& exec_space, + const ViewType& in, + ViewType& out, + shape_type s) { + _crop_or_pad(exec_space, in, out, s); + } +} // namespace Impl +} // namespace KokkosFFT + +#endif \ No newline at end of file diff --git a/common/src/KokkosFFT_utils.hpp b/common/src/KokkosFFT_utils.hpp index 09d49b0e..f138d0fa 100644 --- a/common/src/KokkosFFT_utils.hpp +++ b/common/src/KokkosFFT_utils.hpp @@ -28,6 +28,16 @@ namespace Impl { template struct is_complex> : std::true_type {}; + template = nullptr> + struct complex_view_type { + using value_type = typename ViewType::non_const_value_type; + using float_type = KokkosFFT::Impl::real_type_t; + using complex_type = Kokkos::complex; + using array_layout_type = typename ViewType::array_layout; + using type = Kokkos::View; + }; + template auto convert_negative_axis(const ViewType& view, int _axis=-1) { static_assert(Kokkos::is_view::value, @@ -38,6 +48,27 @@ namespace Impl { return axis; } + template + auto convert_negative_shift(const ViewType& view, int _shift, int _axis) { + static_assert(Kokkos::is_view::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({shift0, shift1, shift2}); + } + template void permute(std::vector& values, const std::vector& indices) { std::vector out; @@ -59,6 +90,17 @@ namespace Impl { return set_values.size() < values.size(); } + template , std::nullptr_t> = nullptr> + bool is_out_of_range_value_included(const std::array& values, IntType max) { + bool is_included = false; + for(auto value: values) { + is_included = value >= max; + } + return is_included; + } + template bool is_transpose_needed(std::array map) { std::array contiguous_map; @@ -91,7 +133,56 @@ namespace Impl { [=](const T sequence) -> T {return start + sequence;}); return sequence; } + + template + inline std::vector arange(const ElementType start, + const ElementType stop, + const ElementType step=1 + ) { + const size_t length = ceil((stop - start) / step); + std::vector result; + ElementType delta = (stop - start) / length; + + // thrust::sequence + for(auto i=0; i + 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::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>(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 \ No newline at end of file diff --git a/common/unit_test/CMakeLists.txt b/common/unit_test/CMakeLists.txt index edb641fd..565468f6 100644 --- a/common/unit_test/CMakeLists.txt +++ b/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) diff --git a/common/unit_test/Test_Padding.cpp b/common/unit_test/Test_Padding.cpp new file mode 100644 index 00000000..3eb66e10 --- /dev/null +++ b/common/unit_test/Test_Padding.cpp @@ -0,0 +1,643 @@ +#include +#include +#include +#include "KokkosFFT_padding.hpp" +#include "Test_Types.hpp" +#include "Test_Utils.hpp" + +template +using shape_type = std::array; + +TEST(ModifyShape1D, View1D) { + const int len = 30, len_pad = 32, len_crop = 28; + + View1D x("x", len); + + auto shape = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{len}); + auto shape_pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{len_pad}); + auto shape_crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{len_crop}); + + shape_type<1> ref_shape = {len}; + shape_type<1> ref_shape_pad = {len_pad}; + shape_type<1> ref_shape_crop = {len_crop}; + + EXPECT_TRUE( shape == ref_shape ); + EXPECT_TRUE( shape_pad == ref_shape_pad ); + EXPECT_TRUE( shape_crop == ref_shape_crop ); +} + +TEST(ModifyShape1D, View2D) { + const int n0 = 30, n0_pad = 32, n0_crop = 28; + const int n1 = 15; + + View2D x("x", n0, n1); + + auto shape = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{n0}); + auto shape_pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{n0_pad}); + auto shape_crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{n0_crop}); + + shape_type<2> ref_shape = {n0, n1}; + shape_type<2> ref_shape_pad = {n0_pad, n1}; + shape_type<2> ref_shape_crop = {n0_crop, n1}; + + EXPECT_TRUE( shape == ref_shape ); + EXPECT_TRUE( shape_pad == ref_shape_pad ); + EXPECT_TRUE( shape_crop == ref_shape_crop ); +} + +TEST(ModifyShape1D, View3D) { + const int n0 = 30, n0_pad = 32, n0_crop = 28; + const int n1 = 15, n2 = 8; + + View3D x("x", n0, n1, n2); + + auto shape = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{n0}); + auto shape_pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{n0_pad}); + auto shape_crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{n0_crop}); + + shape_type<3> ref_shape = {n0, n1, n2}; + shape_type<3> ref_shape_pad = {n0_pad, n1, n2}; + shape_type<3> ref_shape_crop = {n0_crop, n1, n2}; + + EXPECT_TRUE( shape == ref_shape ); + EXPECT_TRUE( shape_pad == ref_shape_pad ); + EXPECT_TRUE( shape_crop == ref_shape_crop ); +} + +TEST(ModifyShape2D, View2D) { + const int n0 = 30, n0_pad = 32, n0_crop = 28; + const int n1 = 15, n1_pad = 16, n1_crop = 14; + + View2D x("x", n0, n1); + + auto shape_n0_n1 = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0, n1}); + auto shape_n0_n1pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0, n1_pad}); + auto shape_n0_n1crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0, n1_crop}); + auto shape_n0pad_n1 = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_pad, n1}); + auto shape_n0pad_n1pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_pad, n1_pad}); + auto shape_n0pad_n1crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_pad, n1_crop}); + auto shape_n0crop_n1 = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_crop, n1}); + auto shape_n0crop_n1pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_crop, n1_pad}); + auto shape_n0crop_n1crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_crop, n1_crop}); + + shape_type<2> ref_shape_n0_n1 = {n0, n1}; + shape_type<2> ref_shape_n0_n1pad = {n0, n1_pad}; + shape_type<2> ref_shape_n0_n1crop = {n0, n1_crop}; + shape_type<2> ref_shape_n0pad_n1 = {n0_pad, n1}; + shape_type<2> ref_shape_n0pad_n1pad = {n0_pad, n1_pad}; + shape_type<2> ref_shape_n0pad_n1crop = {n0_pad, n1_crop}; + shape_type<2> ref_shape_n0crop_n1 = {n0_crop, n1}; + shape_type<2> ref_shape_n0crop_n1pad = {n0_crop, n1_pad}; + shape_type<2> ref_shape_n0crop_n1crop = {n0_crop, n1_crop}; + + EXPECT_TRUE( shape_n0_n1 == ref_shape_n0_n1 ); + EXPECT_TRUE( shape_n0_n1pad == ref_shape_n0_n1pad ); + EXPECT_TRUE( shape_n0_n1crop == ref_shape_n0_n1crop ); + EXPECT_TRUE( shape_n0pad_n1 == ref_shape_n0pad_n1 ); + EXPECT_TRUE( shape_n0pad_n1pad == ref_shape_n0pad_n1pad ); + EXPECT_TRUE( shape_n0pad_n1crop == ref_shape_n0pad_n1crop ); + EXPECT_TRUE( shape_n0crop_n1 == ref_shape_n0crop_n1 ); + EXPECT_TRUE( shape_n0crop_n1pad == ref_shape_n0crop_n1pad ); + EXPECT_TRUE( shape_n0crop_n1crop == ref_shape_n0crop_n1crop ); +} + +TEST(ModifyShape2D, View3D) { + const int n0 = 30, n0_pad = 32, n0_crop = 28; + const int n1 = 15, n1_pad = 16, n1_crop = 14; + const int n2 = 8; + + View3D x("x", n0, n1, n2); + + auto shape_n0_n1 = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0, n1}); + auto shape_n0_n1pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0, n1_pad}); + auto shape_n0_n1crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0, n1_crop}); + auto shape_n0pad_n1 = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_pad, n1}); + auto shape_n0pad_n1pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_pad, n1_pad}); + auto shape_n0pad_n1crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_pad, n1_crop}); + auto shape_n0crop_n1 = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_crop, n1}); + auto shape_n0crop_n1pad = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_crop, n1_pad}); + auto shape_n0crop_n1crop = KokkosFFT::Impl::get_modified_shape(x, shape_type<2>{n0_crop, n1_crop}); + + shape_type<3> ref_shape_n0_n1 = {n0, n1, n2}; + shape_type<3> ref_shape_n0_n1pad = {n0, n1_pad, n2}; + shape_type<3> ref_shape_n0_n1crop = {n0, n1_crop, n2}; + shape_type<3> ref_shape_n0pad_n1 = {n0_pad, n1, n2}; + shape_type<3> ref_shape_n0pad_n1pad = {n0_pad, n1_pad, n2}; + shape_type<3> ref_shape_n0pad_n1crop = {n0_pad, n1_crop, n2}; + shape_type<3> ref_shape_n0crop_n1 = {n0_crop, n1, n2}; + shape_type<3> ref_shape_n0crop_n1pad = {n0_crop, n1_pad, n2}; + shape_type<3> ref_shape_n0crop_n1crop = {n0_crop, n1_crop, n2}; + + EXPECT_TRUE( shape_n0_n1 == ref_shape_n0_n1 ); + EXPECT_TRUE( shape_n0_n1pad == ref_shape_n0_n1pad ); + EXPECT_TRUE( shape_n0_n1crop == ref_shape_n0_n1crop ); + EXPECT_TRUE( shape_n0pad_n1 == ref_shape_n0pad_n1 ); + EXPECT_TRUE( shape_n0pad_n1pad == ref_shape_n0pad_n1pad ); + EXPECT_TRUE( shape_n0pad_n1crop == ref_shape_n0pad_n1crop ); + EXPECT_TRUE( shape_n0crop_n1 == ref_shape_n0crop_n1 ); + EXPECT_TRUE( shape_n0crop_n1pad == ref_shape_n0crop_n1pad ); + EXPECT_TRUE( shape_n0crop_n1crop == ref_shape_n0crop_n1crop ); +} + +TEST(ModifyShape3D, View3D) { + const int n0 = 30, n1 = 15, n2 = 8; + View3D x("x", n0, n1, n2); + + for(int i0=-1; i0<=1; i0++) { + for(int i1=-1; i1<=1; i1++) { + for(int i2=-1; i2<=1; i2++) { + std::size_t n0_new = static_cast(n0+i0); + std::size_t n1_new = static_cast(n1+i1); + std::size_t n2_new = static_cast(n2+i2); + + shape_type<3> ref_shape = {n0_new, n1_new, n2_new}; + auto shape = KokkosFFT::Impl::get_modified_shape(x, ref_shape); + + EXPECT_TRUE( shape == ref_shape ); + } + } + } +} + +TEST(IsCropOrPadNeeded, View1D) { + const int len = 30, len_pad = 32, len_crop = 28; + View1D x("x", len); + + EXPECT_FALSE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_type<1>{len}) ); + EXPECT_TRUE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_type<1>{len_pad}) ); + EXPECT_TRUE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_type<1>{len_crop}) ); +} + +TEST(IsCropOrPadNeeded, View2D) { + const int n0 = 30, n1 = 15; + View2D x("x", n0, n1); + + for(int i0=-1; i0<=1; i0++) { + for(int i1=-1; i1<=1; i1++) { + std::size_t n0_new = static_cast(n0+i0); + std::size_t n1_new = static_cast(n1+i1); + + shape_type<2> shape_new = {n0_new, n1_new}; + if(i0 == 0 && i1 == 0) { + EXPECT_FALSE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_new) ); + } else { + EXPECT_TRUE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_new) ); + } + } + } +} + +TEST(IsCropOrPadNeeded, View3D) { + const int n0 = 30, n1 = 15, n2 = 8; + View3D x("x", n0, n1, n2); + + for(int i0=-1; i0<=1; i0++) { + for(int i1=-1; i1<=1; i1++) { + for(int i2=-1; i2<=1; i2++) { + std::size_t n0_new = static_cast(n0+i0); + std::size_t n1_new = static_cast(n1+i1); + std::size_t n2_new = static_cast(n2+i2); + + shape_type<3> shape_new = {n0_new, n1_new, n2_new}; + if(i0 == 0 && i1 == 0 && i2 == 0) { + EXPECT_FALSE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_new) ); + } else { + EXPECT_TRUE( KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_new) ); + } + } + } + } +} + +TEST(CropOrPad1D, View1D) { + const int len = 30, len_pad = 32, len_crop = 28; + + View1D x("x", len); + View1D _x, _x_pad, _x_crop; + View1D ref_x("ref_x", len), ref_x_pad("ref_x_pad", len_pad), ref_x_crop("ref_x_crop", len_crop); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1.0); + + Kokkos::deep_copy(ref_x, x); + + // Copying the first len elements, others are initialized with zeros + auto range0 = std::pair(0, len); + auto sub_ref_x_pad = Kokkos::subview(ref_x_pad, range0); + Kokkos::deep_copy(sub_ref_x_pad, x); + + // Copying the cropped part + auto range1 = std::pair(0, len_crop); + auto sub_x = Kokkos::subview(x, range1); + Kokkos::deep_copy(ref_x_crop, sub_x); + + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x, shape_type<1>{len}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_pad, shape_type<1>{len_pad}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_crop, shape_type<1>{len_crop}); + + EXPECT_TRUE( allclose(_x, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_pad, ref_x_pad, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_crop, ref_x_crop, 1.e-5, 1.e-12) ); +} + +TEST(CropOrPad1D, View2D) { + const int n0 = 12, n0_pad = 14, n0_crop = 10; + const int n1 = 5, n1_pad = 7, n1_crop = 3; + + View2D x("x", n0, n1); + View2D _x, _x_pad_axis0, _x_pad_axis1, _x_crop_axis0, _x_crop_axis1; + View2D ref_x("ref_x", n0, n1), ref_x_pad_axis0("ref_x_pad_axis0", n0_pad, n1), ref_x_crop_axis0("ref_x_crop_axis0", n0_crop, n1); + View2D ref_x_pad_axis1("ref_x_pad_axis1", n0, n1_pad), ref_x_crop_axis1("ref_x_cro_axis1", n0, n1_crop); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1.0); + + Kokkos::deep_copy(ref_x, x); + + auto h_x = Kokkos::create_mirror_view(x); + auto h_ref_x_pad_axis0 = Kokkos::create_mirror_view(ref_x_pad_axis0); + auto h_ref_x_crop_axis0 = Kokkos::create_mirror_view(ref_x_crop_axis0); + auto h_ref_x_pad_axis1 = Kokkos::create_mirror_view(ref_x_pad_axis1); + auto h_ref_x_crop_axis1 = Kokkos::create_mirror_view(ref_x_crop_axis1); + Kokkos::deep_copy(h_x, x); + Kokkos::deep_copy(h_ref_x_pad_axis0, ref_x_pad_axis0); + Kokkos::deep_copy(h_ref_x_crop_axis0, ref_x_crop_axis0); + Kokkos::deep_copy(h_ref_x_pad_axis1, ref_x_pad_axis1); + Kokkos::deep_copy(h_ref_x_crop_axis1, ref_x_crop_axis1); + + // Pad or crop along axis 0 + for(int i1=0; i1{n0, n1}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_pad_axis0, shape_type<2>{n0_pad, n1}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_crop_axis0, shape_type<2>{n0_crop, n1}); + + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_pad_axis1, shape_type<2>{n0, n1_pad}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_crop_axis1, shape_type<2>{n0, n1_crop}); + + EXPECT_TRUE( allclose(_x, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_pad_axis0, ref_x_pad_axis0, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_crop_axis0, ref_x_crop_axis0, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_pad_axis1, ref_x_pad_axis1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_crop_axis1, ref_x_crop_axis1, 1.e-5, 1.e-12) ); +} + +TEST(CropOrPad1D, View3D) { + const int n0 = 12, n0_pad = 14, n0_crop = 10; + const int n1 = 5, n1_pad = 7, n1_crop = 3; + const int n2 = 8, n2_pad = 10, n2_crop = 6; + + View3D x("x", n0, n1, n2); + View3D _x, _x_pad_axis0, _x_crop_axis0, _x_pad_axis1, _x_crop_axis1, _x_pad_axis2, _x_crop_axis2; + View3D ref_x("ref_x", n0, n1, n2), ref_x_pad_axis0("ref_x_pad_axis0", n0_pad, n1, n2), ref_x_crop_axis0("ref_x_crop_axis0", n0_crop, n1, n2); + View3D ref_x_pad_axis1("ref_x_pad_axis1", n0, n1_pad, n2), ref_x_crop_axis1("ref_x_cro_axis1", n0, n1_crop, n2); + View3D ref_x_pad_axis2("ref_x_pad_axis2", n0, n1, n2_pad), ref_x_crop_axis2("ref_x_cro_axis2", n0, n1, n2_crop); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1.0); + + Kokkos::deep_copy(ref_x, x); + + auto h_x = Kokkos::create_mirror_view(x); + auto h_ref_x_pad_axis0 = Kokkos::create_mirror_view(ref_x_pad_axis0); + auto h_ref_x_crop_axis0 = Kokkos::create_mirror_view(ref_x_crop_axis0); + auto h_ref_x_pad_axis1 = Kokkos::create_mirror_view(ref_x_pad_axis1); + auto h_ref_x_crop_axis1 = Kokkos::create_mirror_view(ref_x_crop_axis1); + auto h_ref_x_pad_axis2 = Kokkos::create_mirror_view(ref_x_pad_axis2); + auto h_ref_x_crop_axis2 = Kokkos::create_mirror_view(ref_x_crop_axis2); + Kokkos::deep_copy(h_x, x); + Kokkos::deep_copy(h_ref_x_pad_axis0, ref_x_pad_axis0); + Kokkos::deep_copy(h_ref_x_crop_axis0, ref_x_crop_axis0); + Kokkos::deep_copy(h_ref_x_pad_axis1, ref_x_pad_axis1); + Kokkos::deep_copy(h_ref_x_crop_axis1, ref_x_crop_axis1); + Kokkos::deep_copy(h_ref_x_pad_axis2, ref_x_pad_axis2); + Kokkos::deep_copy(h_ref_x_crop_axis2, ref_x_crop_axis2); + + // Pad or crop along axis 0 + for(int i2=0; i2{n0, n1, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_pad_axis0, shape_type<3>{n0_pad, n1, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_crop_axis0, shape_type<3>{n0_crop, n1, n2}); + + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_pad_axis1, shape_type<3>{n0, n1_pad, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_crop_axis1, shape_type<3>{n0, n1_crop, n2}); + + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_pad_axis2, shape_type<3>{n0, n1, n2_pad}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_crop_axis2, shape_type<3>{n0, n1, n2_crop}); + + EXPECT_TRUE( allclose(_x, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_pad_axis0, ref_x_pad_axis0, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_crop_axis0, ref_x_crop_axis0, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_pad_axis1, ref_x_pad_axis1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_crop_axis1, ref_x_crop_axis1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_pad_axis2, ref_x_pad_axis2, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_crop_axis2, ref_x_crop_axis2, 1.e-5, 1.e-12) ); +} + +TEST(CropOrPad2D, View2D) { + const int n0 = 12, n0_pad = 14, n0_crop = 10; + const int n1 = 5, n1_pad = 7, n1_crop = 3; + + View2D x("x", n0, n1); + View2D _x, _x_0_1p, _x_0_1c, _x_0p_1, _x_0p_1p, _x_0p_1c, _x_0c_1, _x_0c_1p, _x_0c_1c; + View2D ref_x("ref_x", n0, n1), ref_x_0_1p("ref_x_0_1p", n0, n1_pad), ref_x_0_1c("ref_x_0_1c", n0, n1_crop); + View2D ref_x_0p_1("ref_x_0p_1", n0_pad, n1), ref_x_0p_1p("ref_x_0p_1p", n0_pad, n1_pad), ref_x_0p_1c("ref_x_0p_1c", n0_pad, n1_crop); + View2D ref_x_0c_1("ref_x_0c_1", n0_crop, n1), ref_x_0c_1p("ref_x_0c_1p", n0_crop, n1_pad), ref_x_0c_1c("ref_x_0c_1c", n0_crop, n1_crop); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1.0); + + Kokkos::deep_copy(ref_x, x); + + auto h_x = Kokkos::create_mirror_view(x); + auto h_ref_x_0_1p = Kokkos::create_mirror_view(ref_x_0_1p); + auto h_ref_x_0_1c = Kokkos::create_mirror_view(ref_x_0_1c); + + auto h_ref_x_0p_1 = Kokkos::create_mirror_view(ref_x_0p_1); + auto h_ref_x_0p_1p = Kokkos::create_mirror_view(ref_x_0p_1p); + auto h_ref_x_0p_1c = Kokkos::create_mirror_view(ref_x_0p_1c); + auto h_ref_x_0c_1 = Kokkos::create_mirror_view(ref_x_0c_1); + auto h_ref_x_0c_1p = Kokkos::create_mirror_view(ref_x_0c_1p); + auto h_ref_x_0c_1c = Kokkos::create_mirror_view(ref_x_0c_1c); + + Kokkos::deep_copy(h_x, x); + Kokkos::deep_copy(h_ref_x_0_1p, ref_x_0_1p); + Kokkos::deep_copy(h_ref_x_0_1c, ref_x_0_1c); + Kokkos::deep_copy(h_ref_x_0p_1, ref_x_0p_1); + Kokkos::deep_copy(h_ref_x_0p_1p, ref_x_0p_1p); + Kokkos::deep_copy(h_ref_x_0p_1c, ref_x_0p_1c); + Kokkos::deep_copy(h_ref_x_0c_1, ref_x_0c_1); + Kokkos::deep_copy(h_ref_x_0c_1p, ref_x_0c_1p); + Kokkos::deep_copy(h_ref_x_0c_1c, ref_x_0c_1c); + + // Along axis 0 + for(int i1=0; i1{n0, n1}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0_1p, shape_type<2>{n0, n1_pad}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0_1c, shape_type<2>{n0, n1_crop}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0p_1, shape_type<2>{n0_pad, n1}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0p_1p, shape_type<2>{n0_pad, n1_pad}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0p_1c, shape_type<2>{n0_pad, n1_crop}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0c_1, shape_type<2>{n0_crop, n1}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0c_1p, shape_type<2>{n0_crop, n1_pad}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0c_1c, shape_type<2>{n0_crop, n1_crop}); + + EXPECT_TRUE( allclose(_x, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0_1p, ref_x_0_1p, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0_1c, ref_x_0_1c, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0p_1, ref_x_0p_1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0p_1p, ref_x_0p_1p, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0p_1c, ref_x_0p_1c, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0c_1, ref_x_0c_1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0c_1p, ref_x_0c_1p, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0c_1c, ref_x_0c_1c, 1.e-5, 1.e-12) ); +} + +TEST(CropOrPad2D, View3D) { + const int n0 = 12, n0_pad = 14, n0_crop = 10; + const int n1 = 5, n1_pad = 7, n1_crop = 3; + const int n2 = 8; + + View3D x("x", n0, n1, n2); + View3D _x, _x_0_1p, _x_0_1c, _x_0p_1, _x_0p_1p, _x_0p_1c, _x_0c_1, _x_0c_1p, _x_0c_1c; + View3D ref_x("ref_x", n0, n1, n2), ref_x_0_1p("ref_x_0_1p", n0, n1_pad, n2), ref_x_0_1c("ref_x_0_1c", n0, n1_crop, n2); + View3D ref_x_0p_1("ref_x_0p_1", n0_pad, n1, n2), ref_x_0p_1p("ref_x_0p_1p", n0_pad, n1_pad, n2), ref_x_0p_1c("ref_x_0p_1c", n0_pad, n1_crop, n2); + View3D ref_x_0c_1("ref_x_0c_1", n0_crop, n1, n2), ref_x_0c_1p("ref_x_0c_1p", n0_crop, n1_pad, n2), ref_x_0c_1c("ref_x_0c_1c", n0_crop, n1_crop, n2); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1.0); + + Kokkos::deep_copy(ref_x, x); + + auto h_x = Kokkos::create_mirror_view(x); + auto h_ref_x_0_1p = Kokkos::create_mirror_view(ref_x_0_1p); + auto h_ref_x_0_1c = Kokkos::create_mirror_view(ref_x_0_1c); + + auto h_ref_x_0p_1 = Kokkos::create_mirror_view(ref_x_0p_1); + auto h_ref_x_0p_1p = Kokkos::create_mirror_view(ref_x_0p_1p); + auto h_ref_x_0p_1c = Kokkos::create_mirror_view(ref_x_0p_1c); + auto h_ref_x_0c_1 = Kokkos::create_mirror_view(ref_x_0c_1); + auto h_ref_x_0c_1p = Kokkos::create_mirror_view(ref_x_0c_1p); + auto h_ref_x_0c_1c = Kokkos::create_mirror_view(ref_x_0c_1c); + + Kokkos::deep_copy(h_x, x); + Kokkos::deep_copy(h_ref_x_0_1p, ref_x_0_1p); + Kokkos::deep_copy(h_ref_x_0_1c, ref_x_0_1c); + Kokkos::deep_copy(h_ref_x_0p_1, ref_x_0p_1); + Kokkos::deep_copy(h_ref_x_0p_1p, ref_x_0p_1p); + Kokkos::deep_copy(h_ref_x_0p_1c, ref_x_0p_1c); + Kokkos::deep_copy(h_ref_x_0c_1, ref_x_0c_1); + Kokkos::deep_copy(h_ref_x_0c_1p, ref_x_0c_1p); + Kokkos::deep_copy(h_ref_x_0c_1c, ref_x_0c_1c); + + // Along axis 0 + for(int i2=0; i2{n0, n1, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0_1p, shape_type<3>{n0, n1_pad, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0_1c, shape_type<3>{n0, n1_crop, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0p_1, shape_type<3>{n0_pad, n1, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0p_1p, shape_type<3>{n0_pad, n1_pad, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0p_1c, shape_type<3>{n0_pad, n1_crop, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0c_1, shape_type<3>{n0_crop, n1, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0c_1p, shape_type<3>{n0_crop, n1_pad, n2}); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x_0c_1c, shape_type<3>{n0_crop, n1_crop, n2}); + + EXPECT_TRUE( allclose(_x, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0_1p, ref_x_0_1p, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0_1c, ref_x_0_1c, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0p_1, ref_x_0p_1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0p_1p, ref_x_0p_1p, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0p_1c, ref_x_0p_1c, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0c_1, ref_x_0c_1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0c_1p, ref_x_0c_1p, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_x_0c_1c, ref_x_0c_1c, 1.e-5, 1.e-12) ); +} + +TEST(CropOrPad3D, View3D) { + const int n0 = 30, n1 = 15, n2 = 8; + View3D x("x", n0, n1, n2); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1.0); + + auto h_x = Kokkos::create_mirror_view(x); + Kokkos::deep_copy(h_x, x); + for(int i0=-1; i0<=1; i0++) { + for(int i1=-1; i1<=1; i1++) { + for(int i2=-1; i2<=1; i2++) { + std::size_t n0_new = static_cast(n0+i0); + std::size_t n1_new = static_cast(n1+i1); + std::size_t n2_new = static_cast(n2+i2); + shape_type<3> shape_new = {n0_new, n1_new, n2_new}; + + View3D _x; + View3D ref_x("ref_x", n0_new, n1_new, n2_new); + + auto h_ref_x = Kokkos::create_mirror_view(ref_x); + for(int i2=0; i2= h_ref_x.extent(0) || i1 >= h_ref_x.extent(1) || i2 >= h_ref_x.extent(2)) continue; + h_ref_x(i0, i1, i2) = h_x(i0, i1, i2); + } + } + } + + Kokkos::deep_copy(ref_x, h_ref_x); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, _x, shape_new); + EXPECT_TRUE( allclose(_x, ref_x, 1.e-5, 1.e-12) ); + } + } + } +} \ No newline at end of file diff --git a/common/unit_test/Test_Utils.cpp b/common/unit_test/Test_Utils.cpp index 19c757ac..82cc1e83 100644 --- a/common/unit_test/Test_Utils.cpp +++ b/common/unit_test/Test_Utils.cpp @@ -13,7 +13,13 @@ struct ConvertNegativeAxis : public ::testing::Test { using layout_type = T; }; +template +struct ConvertNegativeShift : public ::testing::Test { + using layout_type = T; +}; + TYPED_TEST_SUITE(ConvertNegativeAxis, test_types); +TYPED_TEST_SUITE(ConvertNegativeShift, test_types); // Tests for convert_negative_axes over ND views template @@ -135,6 +141,57 @@ TYPED_TEST(ConvertNegativeAxis, 4DView) { test_convert_negative_axes_4d(); } +// Tests for convert_negative_shift over ND views +template +void test_convert_negative_shift_1d() { + const int n0_odd = 29, n0_even = 30; + const int shift = 5; + using RealView1Dtype = Kokkos::View; + RealView1Dtype x_odd("x_odd", n0_odd), x_even("x_even", n0_even); + + auto [shift_5_0_odd, shift_5_1_odd, shift_5_2_odd] = KokkosFFT::Impl::convert_negative_shift(x_odd, shift, 0); + auto [shift_5_0_even, shift_5_1_even, shift_5_2_even] = KokkosFFT::Impl::convert_negative_shift(x_even, shift, 0); + auto [shift_0_0_odd, shift_0_1_odd, shift_0_2_odd] = KokkosFFT::Impl::convert_negative_shift(x_odd, 0, 0); + auto [shift_0_0_even, shift_0_1_even, shift_0_2_even] = KokkosFFT::Impl::convert_negative_shift(x_even, 0, 0); + auto [shift_m5_0_odd, shift_m5_1_odd, shift_m5_2_odd] = KokkosFFT::Impl::convert_negative_shift(x_odd, -shift, 0); + auto [shift_m5_0_even, shift_m5_1_even, shift_m5_2_even] = KokkosFFT::Impl::convert_negative_shift(x_even, -shift, 0); + + int ref_shift_5_0_odd = shift, ref_shift_5_1_odd = shift+1, ref_shift_5_2_odd = 0; + int ref_shift_5_0_even = shift, ref_shift_5_1_even = shift, ref_shift_5_2_even = 0; + int ref_shift_0_0_odd = 0, ref_shift_0_1_odd = 0, ref_shift_0_2_odd = n0_odd / 2; + int ref_shift_0_0_even = 0, ref_shift_0_1_even = 0, ref_shift_0_2_even = n0_even / 2; + int ref_shift_m5_0_odd = shift+1, ref_shift_m5_1_odd = shift, ref_shift_m5_2_odd = 0; + int ref_shift_m5_0_even = shift, ref_shift_m5_1_even = shift, ref_shift_m5_2_even = 0; + + EXPECT_EQ( shift_5_0_odd, ref_shift_5_0_odd ); + EXPECT_EQ( shift_5_0_even, ref_shift_5_0_even ); + EXPECT_EQ( shift_0_0_odd, ref_shift_0_0_odd ); + EXPECT_EQ( shift_0_0_even, ref_shift_0_0_even ); + EXPECT_EQ( shift_m5_0_odd, ref_shift_m5_0_odd ); + EXPECT_EQ( shift_m5_0_even, ref_shift_m5_0_even ); + + EXPECT_EQ( shift_5_1_odd, ref_shift_5_1_odd ); + EXPECT_EQ( shift_5_1_even, ref_shift_5_1_even ); + EXPECT_EQ( shift_0_1_odd, ref_shift_0_1_odd ); + EXPECT_EQ( shift_0_1_even, ref_shift_0_1_even ); + EXPECT_EQ( shift_m5_1_odd, ref_shift_m5_1_odd ); + EXPECT_EQ( shift_m5_1_even, ref_shift_m5_1_even ); + + EXPECT_EQ( shift_5_2_odd, ref_shift_5_2_odd ); + EXPECT_EQ( shift_5_2_even, ref_shift_5_2_even ); + EXPECT_EQ( shift_0_2_odd, ref_shift_0_2_odd ); + EXPECT_EQ( shift_0_2_even, ref_shift_0_2_even ); + EXPECT_EQ( shift_m5_2_odd, ref_shift_m5_2_odd ); + EXPECT_EQ( shift_m5_2_even, ref_shift_m5_2_even ); +} + +// Tests for 1D View +TYPED_TEST(ConvertNegativeShift, 1DView) { + using layout_type = typename TestFixture::layout_type; + + test_convert_negative_shift_1d(); +} + TEST(IsTransposeNeeded, 1Dto3D) { std::array map1D ={0}; EXPECT_FALSE( KokkosFFT::Impl::is_transpose_needed(map1D) ); @@ -176,4 +233,13 @@ TEST(GetIndex, Vectors) { KokkosFFT::Impl::get_index(v, 5), std::runtime_error ); +} + +TEST(IsOutOfRangeValueIncluded, Array) { + std::array v = {0, 1, 2, 3}; + + EXPECT_TRUE( KokkosFFT::Impl::is_out_of_range_value_included(v, 2) ); + EXPECT_TRUE( KokkosFFT::Impl::is_out_of_range_value_included(v, 3) ); + EXPECT_FALSE( KokkosFFT::Impl::is_out_of_range_value_included(v, 4) ); + EXPECT_FALSE( KokkosFFT::Impl::is_out_of_range_value_included(v, 5) ); } \ No newline at end of file diff --git a/fft/src/KokkosFFT_Transform.hpp b/fft/src/KokkosFFT_Transform.hpp index bf41c64e..f4e9eaec 100644 --- a/fft/src/KokkosFFT_Transform.hpp +++ b/fft/src/KokkosFFT_Transform.hpp @@ -6,6 +6,7 @@ #include "KokkosFFT_utils.hpp" #include "KokkosFFT_normalization.hpp" #include "KokkosFFT_transpose.hpp" +#include "KokkosFFT_padding.hpp" #include "KokkosFFT_Plans.hpp" #if defined(KOKKOS_ENABLE_CUDA) @@ -66,7 +67,11 @@ namespace Impl { namespace KokkosFFT { template - void fft(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, int axis=-1) { + void fft(const ExecutionSpace& exec_space, + const InViewType& in, OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { static_assert(Kokkos::is_view::value, "KokkosFFT::fft: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -84,12 +89,25 @@ namespace KokkosFFT { "KokkosFFT::fft: execution_space cannot access data in OutViewType" ); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_FORWARD, axis); + InViewType _in; + if(n) { + std::size_t _n = n.value(); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_FORWARD, axis); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); @@ -97,23 +115,43 @@ namespace KokkosFFT { KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_fft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); } } template - void ifft(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, int axis=-1) { + void ifft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { static_assert(Kokkos::is_view::value, "KokkosFFT::ifft: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, "KokkosFFT::ifft: OutViewType is not a Kokkos::View."); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_BACKWARD, axis); + InViewType _in; + // [TO DO] Modify crop_or_pad to perform the following lines + // KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, n); + if(n) { + std::size_t _n = n.value(); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_BACKWARD, axis); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); @@ -121,12 +159,17 @@ namespace KokkosFFT { KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_ifft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); } } template - void rfft(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, int axis=-1) { + void rfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { static_assert(Kokkos::is_view::value, "KokkosFFT::rfft: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -140,11 +183,16 @@ namespace KokkosFFT { static_assert(KokkosFFT::Impl::is_complex::value, "KokkosFFT::rfft: OutViewType must be complex"); - fft(exec_space, in, out, norm, axis); + fft(exec_space, in, out, norm, axis, n); } template - void irfft(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, int axis=-1) { + void irfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { static_assert(Kokkos::is_view::value, "KokkosFFT::irfft: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -157,58 +205,156 @@ namespace KokkosFFT { "KokkosFFT::irfft: InViewType must be complex"); static_assert(std::is_floating_point::value, "KokkosFFT::irfft: OutViewType must be real"); + if(n) { + std::size_t _n = n.value() / 2 + 1; + ifft(exec_space, in, out, norm, axis, _n); + } else { + ifft(exec_space, in, out, norm, axis); + } + } - ifft(exec_space, in, out, norm, axis); + template + void hfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::hfft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::hfft: OutViewType is not a Kokkos::View."); + + // [TO DO] + // allow real type as input, need to obtain complex view type from in view type + 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::value, + "KokkosFFT::hfft: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::hfft: OutViewType must be real"); + auto new_norm = KokkosFFT::Impl::swap_direction(norm); + //using ComplexViewType = typename KokkosFFT::Impl::complex_view_type::type; + //ComplexViewType in_conj; + InViewType in_conj; + KokkosFFT::Impl::conjugate(exec_space, in, in_conj); + irfft(exec_space, in_conj, out, new_norm, axis, n); } template - void fft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, axis_type<2> axes={-2, -1}) { + void ihfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ihfft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ihfft: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + static_assert(std::is_floating_point::value, + "KokkosFFT::rfft: InViewType must be real"); + static_assert(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::ihfft: OutViewType must be complex"); + + auto new_norm = KokkosFFT::Impl::swap_direction(norm); + OutViewType out_conj; + rfft(exec_space, in, out, new_norm, axis, n); + KokkosFFT::Impl::conjugate(exec_space, out, out_conj); + out = out_conj; + } + + template + void fft2(const ExecutionSpace& exec_space, + const InViewType& in, OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::fft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, "KokkosFFT::fft2: OutViewType is not a Kokkos::View."); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_FORWARD, axes); + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_FORWARD, axes); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_fft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); } } - template - void ifft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, axis_type<2> axes={-2, -1}) { + template + void ifft2(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::ifft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, "KokkosFFT::ifft2: OutViewType is not a Kokkos::View."); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_BACKWARD, axes); + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_BACKWARD, axes); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_ifft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); } } - template - void rfft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, axis_type<2> axes={-2, -1}) { + template + void rfft2(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::rfft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -222,11 +368,16 @@ namespace KokkosFFT { static_assert(KokkosFFT::Impl::is_complex::value, "KokkosFFT::rfft2: OutViewType must be complex"); - fft2(exec_space, in, out, norm, axes); + fft2(exec_space, in, out, norm, axes, s); } - template - void irfft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, axis_type<2> axes={-2, -1}) { + template + void irfft2(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::irfft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -240,11 +391,23 @@ namespace KokkosFFT { static_assert(std::is_floating_point::value, "KokkosFFT::irfft2: OutViewType must be real"); - ifft2(exec_space, in, out, norm, axes); + shape_type zeros = {0}; // default shape means no crop or pad + shape_type _s = {0}; + if( s != zeros ) { + for(int i=0; i - void fftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void fftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::fftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -255,47 +418,82 @@ namespace KokkosFFT { constexpr int start = -static_cast(rank); axis_type axes = KokkosFFT::Impl::index_sequence(start); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_FORWARD, axes); + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_FORWARD, axes); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_fft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); } } - template - void fftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, axis_type axes, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void fftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::fftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, "KokkosFFT::fftn: OutViewType is not a Kokkos::View."); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_FORWARD, axes); + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_FORWARD, axes); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_fft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); } } - template - void ifftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void ifftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::ifftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -306,47 +504,82 @@ namespace KokkosFFT { constexpr int start = -static_cast(rank); axis_type axes = KokkosFFT::Impl::index_sequence(start); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_BACKWARD, axes); + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_BACKWARD, axes); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_ifft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); } } - template - void ifftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, axis_type axes, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void ifftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::ifftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, "KokkosFFT::ifftn: OutViewType is not a Kokkos::View."); - KokkosFFT::Impl::Plan plan(exec_space, in, out, KOKKOS_FFT_BACKWARD, axes); + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + KokkosFFT::Impl::Plan plan(exec_space, _in, out, KOKKOS_FFT_BACKWARD, axes); if(plan.is_transpose_needed()) { InViewType in_T; OutViewType out_T; - KokkosFFT::Impl::transpose(exec_space, in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); } else { - KokkosFFT::Impl::_ifft(exec_space, plan, in, out, norm); + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); } } - template - void rfftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void rfftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::rfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -360,11 +593,16 @@ namespace KokkosFFT { static_assert(KokkosFFT::Impl::is_complex::value, "KokkosFFT::rfftn: OutViewType must be complex"); - fftn(exec_space, in, out, norm); + fftn(exec_space, in, out, norm, s); } - template - void rfftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, axis_type axes, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void rfftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::rfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -378,11 +616,15 @@ namespace KokkosFFT { static_assert(KokkosFFT::Impl::is_complex::value, "KokkosFFT::rfftn: OutViewType must be complex"); - fftn(exec_space, in, out, axes, norm); + fftn(exec_space, in, out, axes, norm, s); } - template - void irfftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void irfftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::irfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -396,11 +638,24 @@ namespace KokkosFFT { static_assert(std::is_floating_point::value, "KokkosFFT::irfftn: OutViewType must be real"); - ifftn(exec_space, in, out, norm); + shape_type zeros = {0}; // default shape means no crop or pad + shape_type _s = {0}; + if( s != zeros ) { + for(int i=0; i - void irfftn(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, axis_type axes, KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD) { + template + void irfftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { static_assert(Kokkos::is_view::value, "KokkosFFT::irfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -414,7 +669,15 @@ namespace KokkosFFT { static_assert(std::is_floating_point::value, "KokkosFFT::irfftn: OutViewType must be real"); - ifftn(exec_space, in, out, axes, norm); + shape_type zeros = {0}; // default shape means no crop or pad + shape_type _s = {0}; + if( s != zeros ) { + for(int i=0; i +void test_fft1_1dhfft_1dview() { + const int len_herm = 16, len = len_herm * 2 - 2; + using RealView1DType = Kokkos::View; + using ComplexView1DType = Kokkos::View*, LayoutType, execution_space>; + + ComplexView1DType x_herm("x_herm", len_herm), x_herm_ref("x_herm_ref", len_herm); + ComplexView1DType x("x", len), ref("ref", len); + RealView1DType out("out", len); + RealView1DType out_b("out_b", len), out_o("out_o", len), out_f("out_f", len); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x_herm, random_pool, I); + + auto h_x = Kokkos::create_mirror_view(x); + auto h_x_herm = Kokkos::create_mirror_view(x_herm); + Kokkos::deep_copy(h_x_herm, x_herm); + + auto last = h_x_herm.extent(0) - 1; + h_x_herm(0) = h_x_herm(0).real(); + h_x_herm(last) = h_x_herm(last).real(); + + for(int i=0; i0; i--) { + h_x(len-i) = Kokkos::conj(h_x_herm(i)); + } + + Kokkos::deep_copy(x_herm, h_x_herm); + Kokkos::deep_copy(x_herm_ref, h_x_herm); + Kokkos::deep_copy(x, h_x); + + Kokkos::fence(); + + KokkosFFT::fft(execution_space(), x, ref); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out_b, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out_o, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out_f, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, sqrt(static_cast(len))); + multiply(out_f, static_cast(len)); + + EXPECT_TRUE( allclose(out, ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out, 1.e-5, 1.e-6) ); +} + +template +void test_fft1_1dihfft_1dview() { + const int len_herm = 16, len = len_herm * 2 - 2; + using RealView1DType = Kokkos::View; + using ComplexView1DType = Kokkos::View*, LayoutType, execution_space>; + + ComplexView1DType x_herm("x_herm", len_herm), x_herm_ref("x_herm_ref", len_herm); + RealView1DType out1("out1", len); + RealView1DType out1_b("out1_b", len), out1_o("out1_o", len), out1_f("out1_f", len); + ComplexView1DType out2("out2", len_herm), out2_b("out2_b", len_herm), out2_o("out2_o", len_herm), out2_f("out2_f", len_herm); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x_herm, random_pool, I); + + auto h_x_herm = Kokkos::create_mirror_view(x_herm); + Kokkos::deep_copy(h_x_herm, x_herm); + + auto last = h_x_herm.extent(0) - 1; + h_x_herm(0) = h_x_herm(0).real(); + h_x_herm(last) = h_x_herm(last).real(); + + Kokkos::deep_copy(x_herm, h_x_herm); + Kokkos::deep_copy(x_herm_ref, h_x_herm); + Kokkos::fence(); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::ihfft(execution_space(), out1, out2); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1_b, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::ihfft(execution_space(), out1_b, out2_b, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1_o, KokkosFFT::Normalization::ORTHO); + KokkosFFT::ihfft(execution_space(), out1_o, out2_o, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1_f, KokkosFFT::Normalization::FORWARD); + KokkosFFT::ihfft(execution_space(), out1_f, out2_f, KokkosFFT::Normalization::FORWARD); + + EXPECT_TRUE( allclose(out2, x_herm_ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out2_b, x_herm_ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out2_o, x_herm_ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out2_f, x_herm_ref, 1.e-5, 1.e-6) ); +} + template void test_fft1_1dfft_2dview(T atol=1.e-12) { const int n0 = 10, n1 = 12; @@ -401,6 +510,22 @@ TYPED_TEST(FFT1D, IFFT_1DView) { test_fft1_1difft_1dview(); } +// hfft on 1D Views +TYPED_TEST(FFT1D, HFFT_1DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + test_fft1_1dhfft_1dview(); +} + +// ihfft on 1D Views +TYPED_TEST(FFT1D, IHFFT_1DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + test_fft1_1dihfft_1dview(); +} + // batced fft1 on 2D Views TYPED_TEST(FFT1D, FFT_batched_2DView) { using float_type = typename TestFixture::float_type; @@ -1076,7 +1201,7 @@ TYPED_TEST(FFTND, 3DFFT_3DView) { using layout_type = typename TestFixture::layout_type; float_type atol = std::is_same_v ? 5.0e-5 : 1.0e-6; - test_fftn_3dfft_3dview(); + test_fftn_3dfft_3dview(atol); } // ifftn on 2D Views diff --git a/fft/unit_test/Test_Utils.hpp b/fft/unit_test/Test_Utils.hpp index 712726ad..997a4a8b 100644 --- a/fft/unit_test/Test_Utils.hpp +++ b/fft/unit_test/Test_Utils.hpp @@ -5,13 +5,13 @@ #include #include "Test_Types.hpp" -template -bool allclose(const ViewType& a, - const ViewType& b, +template +bool allclose(const AViewType& a, + const BViewType& b, double rtol=1.e-5, double atol=1.e-8 ) { - constexpr std::size_t rank = ViewType::rank; + constexpr std::size_t rank = AViewType::rank; for(std::size_t i=0; i