From fede968abda376e5620c0c42ab3c10a64af55263 Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 25 Mar 2024 16:55:13 +0900 Subject: [PATCH 1/5] fix get_modified shape to rely on axis argument --- common/src/KokkosFFT_padding.hpp | 27 +- common/unit_test/Test_Padding.cpp | 862 ++++++++++++++++++++++++------ 2 files changed, 729 insertions(+), 160 deletions(-) diff --git a/common/src/KokkosFFT_padding.hpp b/common/src/KokkosFFT_padding.hpp index cc30e4e5..6278fc2b 100644 --- a/common/src/KokkosFFT_padding.hpp +++ b/common/src/KokkosFFT_padding.hpp @@ -12,17 +12,27 @@ namespace KokkosFFT { namespace Impl { template -auto get_modified_shape(const ViewType& view, shape_type shape) { +auto get_modified_shape(const ViewType& view, shape_type shape, + axis_type axes, bool is_C2R = false) { static_assert(ViewType::rank() >= DIM, "get_modified_shape: Rank of View must be larger " "than or equal to the Rank of new shape"); static_assert(DIM > 0, "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()); + // Convert the input axes to be in the range of [0, rank-1] + std::vector positive_axes; + for (std::size_t i = 0; i < DIM; i++) { + int axis = KokkosFFT::Impl::convert_negative_axis(view, axes.at(i)); + positive_axes.push_back(axis); + } + + // Assert if the elements are overlapped + assert(!KokkosFFT::Impl::has_duplicate_values(positive_axes)); + assert(!KokkosFFT::Impl::is_out_of_range_value_included(positive_axes, rank)); + using full_shape_type = shape_type; full_shape_type modified_shape; for (int i = 0; i < rank; i++) { @@ -31,13 +41,16 @@ auto get_modified_shape(const ViewType& view, shape_type shape) { // Update shapes based on newly given shape for (int i = 0; i < DIM; i++) { + int positive_axis = positive_axes.at(i); assert(shape.at(i) > 0); - modified_shape.at(i) = shape.at(i); + modified_shape.at(positive_axis) = shape.at(i); + } + + if (is_C2R) { + int reshaped_axis = positive_axes.back(); + modified_shape.at(reshaped_axis) = modified_shape.at(reshaped_axis) / 2 + 1; } - // [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; } diff --git a/common/unit_test/Test_Padding.cpp b/common/unit_test/Test_Padding.cpp index 9cdea642..8cf9beeb 100644 --- a/common/unit_test/Test_Padding.cpp +++ b/common/unit_test/Test_Padding.cpp @@ -10,184 +10,740 @@ #include "Test_Utils.hpp" template -using shape_type = std::array; +using shape_type = KokkosFFT::shape_type; -TEST(ModifyShape1D, View1D) { - const int len = 30, len_pad = 32, len_crop = 28; +template +using axes_type = KokkosFFT::axis_type; +class GetModifiedShape1D : public ::testing::TestWithParam {}; +class GetModifiedShape2D : public ::testing::TestWithParam {}; +class GetModifiedShape3D : public ::testing::TestWithParam {}; + +auto get_c2r_shape(std::size_t len, bool is_C2R) { + return is_C2R ? (len / 2 + 1) : len; +} + +void test_reshape1D_1DView(bool is_C2R) { + 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}); + auto shape = KokkosFFT::Impl::get_modified_shape(x, shape_type<1>{len}, + axes_type<1>{-1}, is_C2R); + auto shape_pad = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{len_pad}, axes_type<1>{-1}, is_C2R); + auto shape_crop = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{len_crop}, axes_type<1>{-1}, is_C2R); - shape_type<1> ref_shape = {len}; - shape_type<1> ref_shape_pad = {len_pad}; - shape_type<1> ref_shape_crop = {len_crop}; + shape_type<1> ref_shape = {get_c2r_shape(len, is_C2R)}; + shape_type<1> ref_shape_pad = {get_c2r_shape(len_pad, is_C2R)}; + shape_type<1> ref_shape_crop = {get_c2r_shape(len_crop, is_C2R)}; 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; - +void test_reshape1D_2DView(bool is_C2R) { + const int n0 = 30, n1 = 15; View2D x("x", n0, n1); + constexpr int DIM = 2; - 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> default_shape({n0, n1}); - 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}; + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<2> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); - EXPECT_TRUE(shape == ref_shape); - EXPECT_TRUE(shape_pad == ref_shape_pad); - EXPECT_TRUE(shape_crop == ref_shape_crop); + EXPECT_TRUE(modified_shape == ref_shape); + } + } } -TEST(ModifyShape1D, View3D) { - const int n0 = 30, n0_pad = 32, n0_crop = 28; - const int n1 = 15, n2 = 8; - +void test_reshape1D_3DView(bool is_C2R) { + const int n0 = 30, n1 = 15, n2 = 8; View3D x("x", n0, n1, n2); + constexpr int DIM = 3; - 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> default_shape({n0, n1, n2}); - 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}; + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<3> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); - EXPECT_TRUE(shape == ref_shape); - EXPECT_TRUE(shape_pad == ref_shape_pad); - EXPECT_TRUE(shape_crop == ref_shape_crop); + EXPECT_TRUE(modified_shape == ref_shape); + } + } } -TEST(ModifyShape2D, View2D) { - const int n0 = 30, n0_pad = 32, n0_crop = 28; - const int n1 = 15, n1_pad = 16, n1_crop = 14; +void test_reshape1D_4DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8; + View4D x("x", n0, n1, n2, n3); - View2D x("x", n0, n1); + constexpr int DIM = 4; + shape_type<4> default_shape({n0, n1, n2, n3}); - 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; + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<4> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); + + EXPECT_TRUE(modified_shape == ref_shape); + } + } +} + +void test_reshape1D_5DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3; + View5D x("x", n0, n1, n2, n3, n4); + + constexpr int DIM = 5; + shape_type<5> default_shape({n0, n1, n2, n3, n4}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<5> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); + + EXPECT_TRUE(modified_shape == ref_shape); + } + } +} + +void test_reshape1D_6DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4; + View6D x("x", n0, n1, n2, n3, n4, n5); + + constexpr int DIM = 6; + shape_type<6> default_shape({n0, n1, n2, n3, n4, n5}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<6> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } +} + +void test_reshape1D_7DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4, n6 = 7; + View7D x("x", n0, n1, n2, n3, n4, n5, n6); + + constexpr int DIM = 7; + shape_type<7> default_shape({n0, n1, n2, n3, n4, n5, n6}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<7> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); + + EXPECT_TRUE(modified_shape == ref_shape); + } + } +} + +void test_reshape1D_8DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4, n6 = 7, n7 = 9; + View8D x("x", n0, n1, n2, n3, n4, n5, n6, n7); + + constexpr int DIM = 8; + shape_type<8> default_shape({n0, n1, n2, n3, n4, n5, n6, n7}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int i0 = -1; i0 <= 1; i0++) { + shape_type<8> ref_shape = default_shape; + std::size_t n_new = static_cast(x.extent(axis0) + i0); + ref_shape.at(axis0) = get_c2r_shape(n_new, is_C2R); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, shape_type<1>{n_new}, axes_type<1>{axis0}, is_C2R); + + EXPECT_TRUE(modified_shape == ref_shape); + } + } +} + +TEST_P(GetModifiedShape1D, 1DView) { + bool is_C2R = GetParam(); + test_reshape1D_1DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 2DView) { + bool is_C2R = GetParam(); + test_reshape1D_2DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 3DView) { + bool is_C2R = GetParam(); + test_reshape1D_3DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 4DView) { + bool is_C2R = GetParam(); + test_reshape1D_4DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 5DView) { + bool is_C2R = GetParam(); + test_reshape1D_5DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 6DView) { + bool is_C2R = GetParam(); + test_reshape1D_6DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 7DView) { + bool is_C2R = GetParam(); + test_reshape1D_7DView(is_C2R); +} + +TEST_P(GetModifiedShape1D, 8DView) { + bool is_C2R = GetParam(); + test_reshape1D_8DView(is_C2R); +} + +INSTANTIATE_TEST_SUITE_P(CropOrPad, GetModifiedShape1D, + ::testing::Values(true, false)); + +void test_reshape2D_2DView(bool is_C2R) { + const int n0 = 30, n1 = 15; + View2D x("x", n0, n1); + constexpr int DIM = 2; + shape_type<2> default_shape({n0, n1}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<2> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +void test_reshape2D_3DView(bool is_C2R) { + const int n0 = 30, n1 = 15, n2 = 8; View3D x("x", n0, n1, n2); + constexpr int DIM = 3; + shape_type<3> default_shape({n0, n1, n2}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<3> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +void test_reshape2D_4DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8; + View4D x("x", n0, n1, n2, n3); + constexpr int DIM = 4; + shape_type<4> default_shape({n0, n1, n2, n3}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<4> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +void test_reshape2D_5DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3; + View5D x("x", n0, n1, n2, n3, n4); + constexpr int DIM = 5; + shape_type<5> default_shape({n0, n1, n2, n3, n4}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<5> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +void test_reshape2D_6DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4; + View6D x("x", n0, n1, n2, n3, n4, n5); + constexpr int DIM = 6; + shape_type<6> default_shape({n0, n1, n2, n3, n4, n5}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<6> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +void test_reshape2D_7DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4, n6 = 7; + View7D x("x", n0, n1, n2, n3, n4, n5, n6); + constexpr int DIM = 7; + shape_type<7> default_shape({n0, n1, n2, n3, n4, n5, n6}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<7> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +void test_reshape2D_8DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4, n6 = 7, n7 = 9; + View8D x("x", n0, n1, n2, n3, n4, n5, n6, n7); + constexpr int DIM = 8; + shape_type<8> default_shape({n0, n1, n2, n3, n4, n5, n6, n7}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + if (axis0 == axis1) continue; + axes_type<2> axes({axis0, axis1}); + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type<8> ref_shape = default_shape; + std::size_t n0_new = static_cast(x.extent(axis0) + i0); + std::size_t n1_new = static_cast(x.extent(axis1) + i1); + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = get_c2r_shape(n1_new, is_C2R); + + shape_type<2> new_shape = {n0_new, n1_new}; + + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } +} + +TEST_P(GetModifiedShape2D, 2DView) { + bool is_C2R = GetParam(); + test_reshape2D_2DView(is_C2R); +} + +TEST_P(GetModifiedShape2D, 3DView) { + bool is_C2R = GetParam(); + test_reshape2D_3DView(is_C2R); +} - 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) { +TEST_P(GetModifiedShape2D, 4DView) { + bool is_C2R = GetParam(); + test_reshape2D_4DView(is_C2R); +} + +TEST_P(GetModifiedShape2D, 5DView) { + bool is_C2R = GetParam(); + test_reshape2D_5DView(is_C2R); +} + +TEST_P(GetModifiedShape2D, 6DView) { + bool is_C2R = GetParam(); + test_reshape2D_6DView(is_C2R); +} + +TEST_P(GetModifiedShape2D, 7DView) { + bool is_C2R = GetParam(); + test_reshape2D_7DView(is_C2R); +} + +TEST_P(GetModifiedShape2D, 8DView) { + bool is_C2R = GetParam(); + test_reshape2D_8DView(is_C2R); +} + +INSTANTIATE_TEST_SUITE_P(CropOrPad, GetModifiedShape2D, + ::testing::Values(true, false)); + +void test_reshape3D_3DView(bool is_C2R) { 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); + constexpr int DIM = 3; + shape_type<3> default_shape({n0, n1, n2}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + for (int axis2 = 0; axis2 < DIM; axis2++) { + if (axis0 == axis1 || axis0 == axis2 || axis1 == axis2) continue; + axes_type<3> axes({axis0, axis1, axis2}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + for (int i2 = -1; i2 <= 1; i2++) { + shape_type<3> ref_shape = default_shape; + std::size_t n0_new = + static_cast(x.extent(axis0) + i0); + std::size_t n1_new = + static_cast(x.extent(axis1) + i1); + std::size_t n2_new = + static_cast(x.extent(axis2) + i2); + + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = n1_new; + ref_shape.at(axis2) = get_c2r_shape(n2_new, is_C2R); + + shape_type<3> new_shape = {n0_new, n1_new, n2_new}; + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, new_shape, axes, is_C2R); + + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } + } + } +} - shape_type<3> ref_shape = {n0_new, n1_new, n2_new}; - auto shape = KokkosFFT::Impl::get_modified_shape(x, ref_shape); +void test_reshape3D_4DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8; + View4D x("x", n0, n1, n2, n3); - EXPECT_TRUE(shape == ref_shape); + constexpr int DIM = 4; + shape_type<4> default_shape({n0, n1, n2, n3}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + for (int axis2 = 0; axis2 < DIM; axis2++) { + if (axis0 == axis1 || axis0 == axis2 || axis1 == axis2) continue; + axes_type<3> axes({axis0, axis1, axis2}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + for (int i2 = -1; i2 <= 1; i2++) { + shape_type<4> ref_shape = default_shape; + std::size_t n0_new = + static_cast(x.extent(axis0) + i0); + std::size_t n1_new = + static_cast(x.extent(axis1) + i1); + std::size_t n2_new = + static_cast(x.extent(axis2) + i2); + + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = n1_new; + ref_shape.at(axis2) = get_c2r_shape(n2_new, is_C2R); + + shape_type<3> new_shape = {n0_new, n1_new, n2_new}; + + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } + } + } +} + +void test_reshape3D_5DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3; + View5D x("x", n0, n1, n2, n3, n4); + constexpr int DIM = 5; + shape_type<5> default_shape({n0, n1, n2, n3, n4}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + for (int axis2 = 0; axis2 < DIM; axis2++) { + if (axis0 == axis1 || axis0 == axis2 || axis1 == axis2) continue; + axes_type<3> axes({axis0, axis1, axis2}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + for (int i2 = -1; i2 <= 1; i2++) { + shape_type<5> ref_shape = default_shape; + std::size_t n0_new = + static_cast(x.extent(axis0) + i0); + std::size_t n1_new = + static_cast(x.extent(axis1) + i1); + std::size_t n2_new = + static_cast(x.extent(axis2) + i2); + + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = n1_new; + ref_shape.at(axis2) = get_c2r_shape(n2_new, is_C2R); + + shape_type<3> new_shape = {n0_new, n1_new, n2_new}; + + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } + } + } +} + +void test_reshape3D_6DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4; + View6D x("x", n0, n1, n2, n3, n4, n5); + constexpr int DIM = 6; + shape_type<6> default_shape({n0, n1, n2, n3, n4, n5}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + for (int axis2 = 0; axis2 < DIM; axis2++) { + if (axis0 == axis1 || axis0 == axis2 || axis1 == axis2) continue; + axes_type<3> axes({axis0, axis1, axis2}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + for (int i2 = -1; i2 <= 1; i2++) { + shape_type<6> ref_shape = default_shape; + std::size_t n0_new = + static_cast(x.extent(axis0) + i0); + std::size_t n1_new = + static_cast(x.extent(axis1) + i1); + std::size_t n2_new = + static_cast(x.extent(axis2) + i2); + + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = n1_new; + ref_shape.at(axis2) = get_c2r_shape(n2_new, is_C2R); + + shape_type<3> new_shape = {n0_new, n1_new, n2_new}; + + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } } } } } -TEST(IsCropOrPadNeeded, View1D) { +void test_reshape3D_7DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4, n6 = 7; + View7D x("x", n0, n1, n2, n3, n4, n5, n6); + constexpr int DIM = 7; + shape_type<7> default_shape({n0, n1, n2, n3, n4, n5, n6}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + for (int axis2 = 0; axis2 < DIM; axis2++) { + if (axis0 == axis1 || axis0 == axis2 || axis1 == axis2) continue; + axes_type<3> axes({axis0, axis1, axis2}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + for (int i2 = -1; i2 <= 1; i2++) { + shape_type<7> ref_shape = default_shape; + std::size_t n0_new = + static_cast(x.extent(axis0) + i0); + std::size_t n1_new = + static_cast(x.extent(axis1) + i1); + std::size_t n2_new = + static_cast(x.extent(axis2) + i2); + + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = n1_new; + ref_shape.at(axis2) = get_c2r_shape(n2_new, is_C2R); + + shape_type<3> new_shape = {n0_new, n1_new, n2_new}; + + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } + } + } +} + +void test_reshape3D_8DView(bool is_C2R) { + const int n0 = 5, n1 = 11, n2 = 10, n3 = 8, n4 = 3, n5 = 4, n6 = 7, n7 = 9; + View8D x("x", n0, n1, n2, n3, n4, n5, n6, n7); + constexpr int DIM = 8; + shape_type<8> default_shape({n0, n1, n2, n3, n4, n5, n6, n7}); + + for (int axis0 = 0; axis0 < DIM; axis0++) { + for (int axis1 = 0; axis1 < DIM; axis1++) { + for (int axis2 = 0; axis2 < DIM; axis2++) { + if (axis0 == axis1 || axis0 == axis2 || axis1 == axis2) continue; + axes_type<3> axes({axis0, axis1, axis2}); + + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + for (int i2 = -1; i2 <= 1; i2++) { + shape_type<8> ref_shape = default_shape; + std::size_t n0_new = + static_cast(x.extent(axis0) + i0); + std::size_t n1_new = + static_cast(x.extent(axis1) + i1); + std::size_t n2_new = + static_cast(x.extent(axis2) + i2); + + ref_shape.at(axis0) = n0_new; + ref_shape.at(axis1) = n1_new; + ref_shape.at(axis2) = get_c2r_shape(n2_new, is_C2R); + + shape_type<3> new_shape = {n0_new, n1_new, n2_new}; + + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + x, new_shape, axes, is_C2R); + EXPECT_TRUE(modified_shape == ref_shape); + } + } + } + } + } + } +} + +TEST_P(GetModifiedShape3D, 3DView) { + bool is_C2R = GetParam(); + test_reshape3D_3DView(is_C2R); +} + +TEST_P(GetModifiedShape3D, 4DView) { + bool is_C2R = GetParam(); + test_reshape3D_4DView(is_C2R); +} + +TEST_P(GetModifiedShape3D, 5DView) { + bool is_C2R = GetParam(); + test_reshape3D_5DView(is_C2R); +} + +TEST_P(GetModifiedShape3D, 6DView) { + bool is_C2R = GetParam(); + test_reshape3D_6DView(is_C2R); +} + +TEST_P(GetModifiedShape3D, 7DView) { + bool is_C2R = GetParam(); + test_reshape3D_7DView(is_C2R); +} + +TEST_P(GetModifiedShape3D, 8DView) { + bool is_C2R = GetParam(); + test_reshape3D_8DView(is_C2R); +} + +INSTANTIATE_TEST_SUITE_P(CropOrPad, GetModifiedShape3D, + ::testing::Values(true, false)); + +TEST(IsCropOrPadNeeded, 1DView) { const int len = 30, len_pad = 32, len_crop = 28; View1D x("x", len); @@ -198,7 +754,7 @@ TEST(IsCropOrPadNeeded, View1D) { KokkosFFT::Impl::is_crop_or_pad_needed(x, shape_type<1>{len_crop})); } -TEST(IsCropOrPadNeeded, View2D) { +TEST(IsCropOrPadNeeded, 2DView) { const int n0 = 30, n1 = 15; View2D x("x", n0, n1); @@ -217,7 +773,7 @@ TEST(IsCropOrPadNeeded, View2D) { } } -TEST(IsCropOrPadNeeded, View3D) { +TEST(IsCropOrPadNeeded, 3DView) { const int n0 = 30, n1 = 15, n2 = 8; View3D x("x", n0, n1, n2); @@ -239,7 +795,7 @@ TEST(IsCropOrPadNeeded, View3D) { } } -TEST(CropOrPad1D, View1D) { +TEST(CropOrPad1D, 1DView) { const int len = 30, len_pad = 32, len_crop = 28; View1D x("x", len); @@ -273,7 +829,7 @@ TEST(CropOrPad1D, View1D) { EXPECT_TRUE(allclose(_x_crop, ref_x_crop, 1.e-5, 1.e-12)); } -TEST(CropOrPad1D, View2D) { +TEST(CropOrPad1D, 2DView) { const int n0 = 12, n0_pad = 14, n0_crop = 10; const int n1 = 5, n1_pad = 7, n1_crop = 3; @@ -349,7 +905,7 @@ TEST(CropOrPad1D, View2D) { EXPECT_TRUE(allclose(_x_crop_axis1, ref_x_crop_axis1, 1.e-5, 1.e-12)); } -TEST(CropOrPad1D, View3D) { +TEST(CropOrPad1D, 3DView) { 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; @@ -463,7 +1019,7 @@ TEST(CropOrPad1D, View3D) { EXPECT_TRUE(allclose(_x_crop_axis2, ref_x_crop_axis2, 1.e-5, 1.e-12)); } -TEST(CropOrPad2D, View2D) { +TEST(CropOrPad2D, 2DView) { const int n0 = 12, n0_pad = 14, n0_crop = 10; const int n1 = 5, n1_pad = 7, n1_crop = 3; @@ -572,7 +1128,7 @@ TEST(CropOrPad2D, View2D) { EXPECT_TRUE(allclose(_x_0c_1c, ref_x_0c_1c, 1.e-5, 1.e-12)); } -TEST(CropOrPad2D, View3D) { +TEST(CropOrPad2D, 3DView) { const int n0 = 12, n0_pad = 14, n0_crop = 10; const int n1 = 5, n1_pad = 7, n1_crop = 3; const int n2 = 8; @@ -686,7 +1242,7 @@ TEST(CropOrPad2D, View3D) { EXPECT_TRUE(allclose(_x_0c_1c, ref_x_0c_1c, 1.e-5, 1.e-12)); } -TEST(CropOrPad3D, View3D) { +TEST(CropOrPad3D, 3DView) { const int n0 = 30, n1 = 15, n2 = 8; View3D x("x", n0, n1, n2); @@ -726,7 +1282,7 @@ TEST(CropOrPad3D, View3D) { } } -TEST(CropOrPad4D, View4D) { +TEST(CropOrPad4D, 4DView) { const int n0 = 30, n1 = 15, n2 = 8, n3 = 7; View4D x("x", n0, n1, n2, n3); @@ -773,7 +1329,7 @@ TEST(CropOrPad4D, View4D) { } } -TEST(CropOrPad5D, View5D) { +TEST(CropOrPad5D, 5DView) { const int n0 = 30, n1 = 15, n2 = 8, n3 = 7, n4 = 3; View5D x("x", n0, n1, n2, n3, n4); @@ -825,7 +1381,7 @@ TEST(CropOrPad5D, View5D) { } } -TEST(CropOrPad6D, View6D) { +TEST(CropOrPad6D, 6DView) { const int n0 = 10, n1 = 15, n2 = 8, n3 = 7, n4 = 3, n5 = 4; View6D x("x", n0, n1, n2, n3, n4, n5); @@ -884,7 +1440,7 @@ TEST(CropOrPad6D, View6D) { } } -TEST(CropOrPad7D, View7D) { +TEST(CropOrPad7D, 7DView) { const int n0 = 10, n1 = 15, n2 = 8, n3 = 7, n4 = 3, n5 = 4, n6 = 5; View7D x("x", n0, n1, n2, n3, n4, n5, n6); @@ -948,7 +1504,7 @@ TEST(CropOrPad7D, View7D) { } } -TEST(CropOrPad8D, View8D) { +TEST(CropOrPad8D, 8DView) { const int n0 = 10, n1 = 15, n2 = 8, n3 = 7, n4 = 3, n5 = 4, n6 = 5, n7 = 3; View8D x("x", n0, n1, n2, n3, n4, n5, n6, n7); From 1aff291e51933fe59ea23e71aad6f4d1b68868db Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 25 Mar 2024 16:56:21 +0900 Subject: [PATCH 2/5] Update reshaping inside FFT APIs --- fft/src/KokkosFFT_Transform.hpp | 229 ++++++++++++++------------------ 1 file changed, 99 insertions(+), 130 deletions(-) diff --git a/fft/src/KokkosFFT_Transform.hpp b/fft/src/KokkosFFT_Transform.hpp index fcbf237c..2f3922dc 100644 --- a/fft/src/KokkosFFT_Transform.hpp +++ b/fft/src/KokkosFFT_Transform.hpp @@ -177,9 +177,9 @@ void fft(const ExecutionSpace& exec_space, const InViewType& in, InViewType _in; if (n) { - std::size_t _n = n.value(); - auto modified_shape = - KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + std::size_t _n = n.value(); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + in, shape_type<1>({_n}), axis_type<1>({axis})); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -249,9 +249,9 @@ void fft(const ExecutionSpace& exec_space, const InViewType& in, InViewType _in; if (n) { - std::size_t _n = n.value(); - auto modified_shape = - KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + std::size_t _n = n.value(); + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + in, shape_type<1>({_n}), axis_type<1>({axis})); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -318,11 +318,14 @@ void ifft(const ExecutionSpace& exec_space, const InViewType& in, ExecutionSpace, typename OutViewType::memory_space>::accessible, "ifft: execution_space cannot access data in OutViewType"); + using out_value_type = typename OutViewType::non_const_value_type; + InViewType _in; if (n) { - std::size_t _n = n.value(); - auto modified_shape = - KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + std::size_t _n = n.value(); + bool is_C2R = std::is_floating_point::value; + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + in, shape_type<1>({_n}), axis_type<1>({axis}), is_C2R); /* [FIX THIS] Shallow copy should be sufficient if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { @@ -394,11 +397,14 @@ void ifft(const ExecutionSpace& exec_space, const InViewType& in, ExecutionSpace, typename OutViewType::memory_space>::accessible, "ifft: execution_space cannot access data in OutViewType"); + using out_value_type = typename OutViewType::non_const_value_type; + InViewType _in; if (n) { - std::size_t _n = n.value(); - auto modified_shape = - KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + std::size_t _n = n.value(); + bool is_C2R = std::is_floating_point::value; + auto modified_shape = KokkosFFT::Impl::get_modified_shape( + in, shape_type<1>({_n}), axis_type<1>({axis}), is_C2R); /* [FIX THIS] Shallow copy should be sufficient if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); @@ -575,12 +581,7 @@ void irfft(const ExecutionSpace& exec_space, const InViewType& in, "irfft: InViewType must be complex"); static_assert(std::is_floating_point::value, "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, n); } /// \brief Inverse of rfft with a given @@ -626,12 +627,8 @@ void irfft(const ExecutionSpace& exec_space, const InViewType& in, "irfft: InViewType must be complex"); static_assert(std::is_floating_point::value, "irfft: OutViewType must be real"); - if (n) { - std::size_t _n = n.value() / 2 + 1; - ifft(exec_space, in, out, plan, norm, axis, _n); - } else { - ifft(exec_space, in, out, plan, norm, axis); - } + + ifft(exec_space, in, out, plan, norm, axis, n); } /// \brief One dimensional FFT of a signal that has Hermitian symmetry @@ -865,12 +862,11 @@ void ihfft(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) -template +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}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "fft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -897,9 +893,9 @@ void fft2(const ExecutionSpace& exec_space, const InViewType& in, "fft2: execution_space cannot access data in OutViewType"); InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type<2> zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s, axes); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -936,11 +932,11 @@ void fft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType> void fft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, const PlanType& plan, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - axis_type<2> axes = {-2, -1}, shape_type s = {0}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "fft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -967,9 +963,9 @@ void fft2(const ExecutionSpace& exec_space, const InViewType& in, "fft2: execution_space cannot access data in OutViewType"); InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type<2> zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s, axes); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1005,12 +1001,11 @@ void fft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) -template +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}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "ifft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1036,10 +1031,14 @@ void ifft2(const ExecutionSpace& exec_space, const InViewType& in, ExecutionSpace, typename OutViewType::memory_space>::accessible, "ifft2: execution_space cannot access data in OutViewType"); + using out_value_type = typename OutViewType::non_const_value_type; + InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type<2> zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + bool is_C2R = std::is_floating_point::value; + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(in, s, axes, is_C2R); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1076,11 +1075,11 @@ void ifft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType> void ifft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, const PlanType& plan, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - axis_type<2> axes = {-2, -1}, shape_type s = {0}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "ifft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1106,10 +1105,14 @@ void ifft2(const ExecutionSpace& exec_space, const InViewType& in, ExecutionSpace, typename OutViewType::memory_space>::accessible, "ifft2: execution_space cannot access data in OutViewType"); + using out_value_type = typename OutViewType::non_const_value_type; + InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type<2> zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + bool is_C2R = std::is_floating_point::value; + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(in, s, axes, is_C2R); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1145,12 +1148,11 @@ void ifft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) -template +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}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "rfft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1197,11 +1199,11 @@ void rfft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType> void rfft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, const PlanType& plan, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - axis_type<2> axes = {-2, -1}, shape_type s = {0}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "rfft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1246,12 +1248,11 @@ void rfft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) -template +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}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "irfft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1286,15 +1287,7 @@ void irfft2(const ExecutionSpace& exec_space, const InViewType& in, static_assert(std::is_floating_point::value, "irfft2: OutViewType must be real"); - shape_type zeros = {0}; // default shape means no crop or pad - shape_type _s = {0}; - if (s != zeros) { - for (int i = 0; i < DIM; i++) { - _s.at(i) = s.at(i) / 2 + 1; - } - } - - ifft2(exec_space, in, out, norm, axes, _s); + ifft2(exec_space, in, out, norm, axes, s); } /// \brief Inverse of rfft2 with a given plan @@ -1307,11 +1300,11 @@ void irfft2(const ExecutionSpace& exec_space, const InViewType& in, /// \param axes [in] Axes over which FFT is performed (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType> void irfft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, const PlanType& plan, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - axis_type<2> axes = {-2, -1}, shape_type s = {0}) { + axis_type<2> axes = {-2, -1}, shape_type<2> s = {0}) { static_assert(Kokkos::is_view::value, "irfft2: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1346,15 +1339,7 @@ void irfft2(const ExecutionSpace& exec_space, const InViewType& in, static_assert(std::is_floating_point::value, "irfft2: OutViewType must be real"); - shape_type zeros = {0}; // default shape means no crop or pad - shape_type _s = {0}; - if (s != zeros) { - for (int i = 0; i < DIM; i++) { - _s.at(i) = s.at(i) / 2 + 1; - } - } - - ifft2(exec_space, in, out, plan, norm, axes, _s); + ifft2(exec_space, in, out, plan, norm, axes, s); } // ND FFT @@ -1405,7 +1390,7 @@ void fftn(const ExecutionSpace& exec_space, const InViewType& in, 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); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s, axes); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1441,11 +1426,11 @@ void fftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + std::size_t DIM = 1> void fftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, axis_type axes, + OutViewType& out, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "fftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1472,9 +1457,9 @@ void fftn(const ExecutionSpace& exec_space, const InViewType& in, "fftn: execution_space cannot access data in OutViewType"); InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s, axes); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1511,11 +1496,11 @@ void fftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType, std::size_t DIM = 1> void fftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, const PlanType& plan, axis_type axes, + OutViewType& out, const PlanType& plan, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "fftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1542,9 +1527,9 @@ void fftn(const ExecutionSpace& exec_space, const InViewType& in, "fftn: execution_space cannot access data in OutViewType"); InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s, axes); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1618,7 +1603,7 @@ void ifftn(const ExecutionSpace& exec_space, const InViewType& in, 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); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s, axes); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1654,11 +1639,11 @@ void ifftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + std::size_t DIM = 1> void ifftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, axis_type axes, + OutViewType& out, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "ifftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1684,10 +1669,14 @@ void ifftn(const ExecutionSpace& exec_space, const InViewType& in, ExecutionSpace, typename OutViewType::memory_space>::accessible, "ifftn: execution_space cannot access data in OutViewType"); + using out_value_type = typename OutViewType::non_const_value_type; + InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + bool is_C2R = std::is_floating_point::value; + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(in, s, axes, is_C2R); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1724,11 +1713,11 @@ void ifftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType, std::size_t DIM = 1> void ifftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, const PlanType& plan, axis_type axes, + OutViewType& out, const PlanType& plan, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "ifftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1754,10 +1743,14 @@ void ifftn(const ExecutionSpace& exec_space, const InViewType& in, ExecutionSpace, typename OutViewType::memory_space>::accessible, "ifftn: execution_space cannot access data in OutViewType"); + using out_value_type = typename OutViewType::non_const_value_type; + InViewType _in; - shape_type zeros = {0}; // default shape means no crop or pad + shape_type zeros = {0}; // default shape means no crop or pad if (s != zeros) { - auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + bool is_C2R = std::is_floating_point::value; + auto modified_shape = + KokkosFFT::Impl::get_modified_shape(in, s, axes, is_C2R); if (KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape)) { KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); } else { @@ -1844,11 +1837,11 @@ void rfftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType, std::size_t DIM = 1> void rfftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, const PlanType& plan, axis_type axes, + OutViewType& out, const PlanType& plan, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "rfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1894,11 +1887,11 @@ void rfftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + std::size_t DIM = 1> void rfftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, axis_type axes, + OutViewType& out, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "rfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -1982,14 +1975,6 @@ void irfftn(const ExecutionSpace& exec_space, const InViewType& in, static_assert(std::is_floating_point::value, "irfftn: OutViewType must be real"); - shape_type zeros = {0}; // default shape means no crop or pad - shape_type _s = {0}; - if (s != zeros) { - for (int i = 0; i < DIM; i++) { - _s.at(i) = s.at(i) / 2 + 1; - } - } - ifftn(exec_space, in, out, norm, _s); } @@ -2002,11 +1987,11 @@ void irfftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + std::size_t DIM = 1> void irfftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, axis_type axes, + OutViewType& out, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "irfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -2041,15 +2026,7 @@ void irfftn(const ExecutionSpace& exec_space, const InViewType& in, static_assert(std::is_floating_point::value, "irfftn: OutViewType must be real"); - shape_type zeros = {0}; // default shape means no crop or pad - shape_type _s = {0}; - if (s != zeros) { - for (int i = 0; i < DIM2; i++) { - _s.at(i) = s.at(i) / 2 + 1; - } - } - - ifftn(exec_space, in, out, axes, norm, _s); + ifftn(exec_space, in, out, axes, norm, s); } /// \brief Inverse of rfftn with a given plan @@ -2062,11 +2039,11 @@ void irfftn(const ExecutionSpace& exec_space, const InViewType& in, /// \param norm [in] How the normalization is applied (optional) /// \param s [in] Shape of the transformed axis of the output (optional) template + typename PlanType, std::size_t DIM = 1> void irfftn(const ExecutionSpace& exec_space, const InViewType& in, - OutViewType& out, const PlanType& plan, axis_type axes, + OutViewType& out, const PlanType& plan, axis_type axes, KokkosFFT::Normalization norm = KokkosFFT::Normalization::backward, - shape_type s = {0}) { + shape_type s = {0}) { static_assert(Kokkos::is_view::value, "irfftn: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -2101,15 +2078,7 @@ void irfftn(const ExecutionSpace& exec_space, const InViewType& in, static_assert(std::is_floating_point::value, "irfftn: OutViewType must be real"); - shape_type zeros = {0}; // default shape means no crop or pad - shape_type _s = {0}; - if (s != zeros) { - for (int i = 0; i < DIM2; i++) { - _s.at(i) = s.at(i) / 2 + 1; - } - } - - ifftn(exec_space, in, out, plan, axes, norm, _s); + ifftn(exec_space, in, out, plan, axes, norm, s); } } // namespace KokkosFFT From 1190f42d1b034ceeb4925ba86e222120ccf7f6dc Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 25 Mar 2024 16:56:59 +0900 Subject: [PATCH 3/5] Fix and add tests based on the fix of shape arg --- fft/unit_test/Test_Transform.cpp | 1664 +++++++++++++++++------------- 1 file changed, 935 insertions(+), 729 deletions(-) diff --git a/fft/unit_test/Test_Transform.cpp b/fft/unit_test/Test_Transform.cpp index 32208f9e..d1f12234 100644 --- a/fft/unit_test/Test_Transform.cpp +++ b/fft/unit_test/Test_Transform.cpp @@ -10,6 +10,9 @@ #include "Test_Types.hpp" #include "Test_Utils.hpp" +template +using shape_type = KokkosFFT::shape_type; + /// Kokkos equivalent of fft1 with numpy /// def fft1(x): /// L = len(x) @@ -1041,48 +1044,51 @@ void test_fft1_1dfft_5dview(T atol = 1.e-12) { using ComplexView5DType = Kokkos::View*****, LayoutType, execution_space>; - constexpr int DIM = 5; - std::array shape = {n0, n1, n2, n3, n4}; - ComplexView5DType x("x", n0, n1, n2, n3, n4), - ref_x("ref_x", n0, n1, n2, n3, n4); + constexpr int DIM = 5; + shape_type default_shape({n0, n1, n2, n3, n4}); + ComplexView5DType x("x", n0, n1, n2, n3, n4); for (int axis = 0; axis < DIM; axis++) { - std::array shape_c2r = shape; - shape_c2r.at(axis) = shape_c2r.at(axis) / 2 + 1; - auto [_n0, _n1, _n2, _n3, _n4] = shape_c2r; - ComplexView5DType _x("_x", n0, n1, n2, n3, n4), - out("out", n0, n1, n2, n3, n4), ref_out("ref_out", n0, n1, n2, n3, n4); - RealView5DType xr("xr", n0, n1, n2, n3, n4), - ref_xr("ref_xr", n0, n1, n2, n3, n4), _xr("_xr", n0, n1, n2, n3, n4); - ComplexView5DType outr("outr", _n0, _n1, _n2, _n3, _n4); + for (int i0 = -1; i0 <= 1; i0++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; + const std::size_t n_new = shape.at(axis) + i0; + shape.at(axis) = n_new; + shape_c2r.at(axis) = n_new / 2 + 1; + auto [_n0, _n1, _n2, _n3, _n4] = shape; + auto [_m0, _m1, _m2, _m3, _m4] = shape_c2r; + ComplexView5DType _x("_x", _n0, _n1, _n2, _n3, _n4), + out("out", _n0, _n1, _n2, _n3, _n4), ref_x; + RealView5DType xr("xr", _n0, _n1, _n2, _n3, _n4), + _xr("_xr", _n0, _n1, _n2, _n3, _n4), ref_xr; + ComplexView5DType outr("outr", _m0, _m1, _m2, _m3, _m4); - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Since HIP FFT destructs the input data, we need to keep the input data in - // different place - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - Kokkos::fence(); + Kokkos::fence(); - // Along one axis - // Simple identity tests - KokkosFFT::fft(execution_space(), x, out, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::ifft(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axis); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft(execution_space(), x, out, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::ifft(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axis, n_new); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::irfft(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axis); + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::irfft(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axis, n_new); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } } } @@ -1093,50 +1099,51 @@ void test_fft1_1dfft_6dview(T atol = 1.e-12) { using ComplexView6DType = Kokkos::View******, LayoutType, execution_space>; - constexpr int DIM = 6; - std::array shape = {n0, n1, n2, n3, n4, n5}; - ComplexView6DType x("x", n0, n1, n2, n3, n4, n5), - ref_x("ref_x", n0, n1, n2, n3, n4, n5); + constexpr int DIM = 6; + shape_type default_shape({n0, n1, n2, n3, n4, n5}); + ComplexView6DType x("x", n0, n1, n2, n3, n4, n5); for (int axis = 0; axis < DIM; axis++) { - std::array shape_c2r = shape; - shape_c2r.at(axis) = shape_c2r.at(axis) / 2 + 1; - auto [_n0, _n1, _n2, _n3, _n4, _n5] = shape_c2r; - ComplexView6DType _x("_x", n0, n1, n2, n3, n4, n5), - out("out", n0, n1, n2, n3, n4, n5), - ref_out("ref_out", n0, n1, n2, n3, n4, n5); - RealView6DType xr("xr", n0, n1, n2, n3, n4, n5), - ref_xr("ref_xr", n0, n1, n2, n3, n4, n5), - _xr("_xr", n0, n1, n2, n3, n4, n5); - ComplexView6DType outr("outr", _n0, _n1, _n2, _n3, _n4, _n5); + for (int i0 = -1; i0 <= 1; i0++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; + const std::size_t n_new = shape.at(axis) + i0; + shape.at(axis) = n_new; + shape_c2r.at(axis) = n_new / 2 + 1; + auto [_n0, _n1, _n2, _n3, _n4, _n5] = shape; + auto [_m0, _m1, _m2, _m3, _m4, _m5] = shape_c2r; + ComplexView6DType _x("_x", _n0, _n1, _n2, _n3, _n4, _n5), + out("out", _n0, _n1, _n2, _n3, _n4, _n5), ref_x; + RealView6DType xr("xr", _n0, _n1, _n2, _n3, _n4, _n5), + _xr("_xr", _n0, _n1, _n2, _n3, _n4, _n5), ref_xr; + ComplexView6DType outr("outr", _m0, _m1, _m2, _m3, _m4, _m5); - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Since HIP FFT destructs the input data, we need to keep the input data in - // different place - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - Kokkos::fence(); + Kokkos::fence(); - // Along one axis - // Simple identity tests - KokkosFFT::fft(execution_space(), x, out, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::ifft(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axis); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft(execution_space(), x, out, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::ifft(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axis, n_new); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::irfft(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axis); + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::irfft(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axis, n_new); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } } } @@ -1147,104 +1154,106 @@ void test_fft1_1dfft_7dview(T atol = 1.e-12) { using ComplexView7DType = Kokkos::View*******, LayoutType, execution_space>; - constexpr int DIM = 7; - std::array shape = {n0, n1, n2, n3, n4, n5, n6}; - ComplexView7DType x("x", n0, n1, n2, n3, n4, n5, n6), - ref_x("ref_x", n0, n1, n2, n3, n4, n5, n6); + constexpr int DIM = 7; + shape_type default_shape({n0, n1, n2, n3, n4, n5, n6}); + ComplexView7DType x("x", n0, n1, n2, n3, n4, n5, n6); for (int axis = 0; axis < DIM; axis++) { - std::array shape_c2r = shape; - shape_c2r.at(axis) = shape_c2r.at(axis) / 2 + 1; - auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = shape_c2r; - ComplexView7DType _x("_x", n0, n1, n2, n3, n4, n5, n6), - out("out", n0, n1, n2, n3, n4, n5, n6), - ref_out("ref_out", n0, n1, n2, n3, n4, n5, n6); - RealView7DType xr("xr", n0, n1, n2, n3, n4, n5, n6), - ref_xr("ref_xr", n0, n1, n2, n3, n4, n5, n6), - _xr("_xr", n0, n1, n2, n3, n4, n5, n6); - ComplexView7DType outr("outr", _n0, _n1, _n2, _n3, _n4, _n5, _n6); + for (int i0 = -1; i0 <= 1; i0++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; + const std::size_t n_new = shape.at(axis) + i0; + shape.at(axis) = n_new; + shape_c2r.at(axis) = n_new / 2 + 1; + auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = shape; + auto [_m0, _m1, _m2, _m3, _m4, _m5, _m6] = shape_c2r; + ComplexView7DType _x("_x", _n0, _n1, _n2, _n3, _n4, _n5, _n6), + out("out", _n0, _n1, _n2, _n3, _n4, _n5, _n6), ref_x; + RealView7DType xr("xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6), + _xr("_xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6), ref_xr; + ComplexView7DType outr("outr", _m0, _m1, _m2, _m3, _m4, _m5, _m6); - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Since HIP FFT destructs the input data, we need to keep the input data in - // different place - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - Kokkos::fence(); + Kokkos::fence(); - // Along one axis - // Simple identity tests - KokkosFFT::fft(execution_space(), x, out, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::ifft(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axis); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft(execution_space(), x, out, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::ifft(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axis, n_new); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::irfft(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axis); + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::irfft(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axis, n_new); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } } } template void test_fft1_1dfft_8dview(T atol = 1.e-12) { - const int n0 = 2, n1 = 6, n2 = 8, n3 = 5, n4 = 4, n5 = 3, n6 = 4, n7 = 3; + const int n0 = 4, n1 = 6, n2 = 8, n3 = 5, n4 = 4, n5 = 3, n6 = 4, n7 = 3; using RealView8DType = Kokkos::View; using ComplexView8DType = Kokkos::View********, LayoutType, execution_space>; - constexpr int DIM = 8; - std::array shape = {n0, n1, n2, n3, n4, n5, n6, n7}; - ComplexView8DType x("x", n0, n1, n2, n3, n4, n5, n6, n7), - ref_x("ref_x", n0, n1, n2, n3, n4, n5, n6, n7); + constexpr int DIM = 8; + shape_type default_shape({n0, n1, n2, n3, n4, n5, n6, n7}); + ComplexView8DType x("x", n0, n1, n2, n3, n4, n5, n6, n7); for (int axis = 0; axis < DIM; axis++) { - std::array shape_c2r = shape; - shape_c2r.at(axis) = shape_c2r.at(axis) / 2 + 1; - auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = shape_c2r; - ComplexView8DType _x("_x", n0, n1, n2, n3, n4, n5, n6, n7), - out("out", n0, n1, n2, n3, n4, n5, n6, n7), - ref_out("ref_out", n0, n1, n2, n3, n4, n5, n6, n7); - RealView8DType xr("xr", n0, n1, n2, n3, n4, n5, n6, n7), - ref_xr("ref_xr", n0, n1, n2, n3, n4, n5, n6, n7), - _xr("_xr", n0, n1, n2, n3, n4, n5, n6, n7); - ComplexView8DType outr("outr", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7); + for (int i0 = -1; i0 <= 1; i0++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; + const std::size_t n_new = shape.at(axis) + i0; + shape.at(axis) = n_new; + shape_c2r.at(axis) = n_new / 2 + 1; + auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = shape; + auto [_m0, _m1, _m2, _m3, _m4, _m5, _m6, _m7] = shape_c2r; + ComplexView8DType _x("_x", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), + out("out", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), ref_x; + RealView8DType xr("xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), + _xr("_xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), ref_xr; + ComplexView8DType outr("outr", _m0, _m1, _m2, _m3, _m4, _m5, _m6, _m7); - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Since HIP FFT destructs the input data, we need to keep the input data in - // different place - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - Kokkos::fence(); + Kokkos::fence(); - // Along one axis - // Simple identity tests - KokkosFFT::fft(execution_space(), x, out, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::ifft(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axis); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft(execution_space(), x, out, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::ifft(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axis, n_new); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axis); - KokkosFFT::irfft(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axis); + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axis, n_new); + KokkosFFT::irfft(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axis, n_new); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } } } @@ -1735,6 +1744,92 @@ void test_fft2_2dirfft_2dview() { EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); } +template +void test_fft2_2dfft_2dview_shape(T atol = 1.0e-12) { + const int n0 = 4, n1 = 6; + using RealView2DType = Kokkos::View; + using ComplexView2DType = + Kokkos::View**, LayoutType, execution_space>; + + RealView2DType xr("xr", n0, n1), xr_ref("xr_ref", n0, n1); + ComplexView2DType x("x", n0, n1 / 2 + 1), x_ref("x_ref", n0, n1 / 2 + 1); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(xr, random_pool, 1.0); + Kokkos::fill_random(x, random_pool, I); + Kokkos::deep_copy(xr_ref, xr); + Kokkos::deep_copy(x_ref, x); + + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; + + std::vector shapes0 = {n0 / 2, n0, n0 * 2}; + std::vector shapes1 = {n1 / 2, n1, n1 * 2}; + + for (auto&& shape0 : shapes0) { + for (auto&& shape1 : shapes1) { + // Real to complex + ComplexView2DType outr("outr", shape0, shape1 / 2 + 1), + outr_b("outr_b", shape0, shape1 / 2 + 1), + outr_o("outr_o", shape0, shape1 / 2 + 1), + outr_f("outr_f", shape0, shape1 / 2 + 1); + + shape_type<2> new_shape = {shape0, shape1}; + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::none, axes, new_shape); + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfft2(execution_space(), xr, outr_b, + KokkosFFT::Normalization::backward, axes, new_shape); + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfft2(execution_space(), xr, outr_o, + KokkosFFT::Normalization::ortho, axes, new_shape); + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfft2(execution_space(), xr, outr_f, + KokkosFFT::Normalization::forward, axes, new_shape); + + multiply(outr_o, sqrt(static_cast(shape0 * shape1))); + multiply(outr_f, static_cast(shape0 * shape1)); + + EXPECT_TRUE(allclose(outr_b, outr, 1.e-5, atol)); + EXPECT_TRUE(allclose(outr_o, outr, 1.e-5, atol)); + EXPECT_TRUE(allclose(outr_f, outr, 1.e-5, atol)); + + // Complex to real + RealView2DType out("out", shape0, shape1), out_b("out_b", shape0, shape1), + out_o("out_o", shape0, shape1), out_f("out_f", shape0, shape1); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out, + KokkosFFT::Normalization::none, axes, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out_b, + KokkosFFT::Normalization::backward, axes, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out_o, + KokkosFFT::Normalization::ortho, axes, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out_f, + KokkosFFT::Normalization::forward, axes, new_shape); + + multiply(out_o, sqrt(static_cast(shape0 * shape1))); + multiply(out_b, static_cast(shape0 * shape1)); + + EXPECT_TRUE(allclose(out_b, out, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_o, out, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_f, out, 1.e-5, atol)); + } + } +} + template void test_fft2_2dfft_3dview(T atol = 1.e-12) { const int n0 = 10, n1 = 6, n2 = 8; @@ -1742,8 +1837,8 @@ void test_fft2_2dfft_3dview(T atol = 1.e-12) { using ComplexView3DType = Kokkos::View***, LayoutType, execution_space>; - constexpr int DIM = 3; - std::array shape = {n0, n1, n2}; + constexpr int DIM = 3; + shape_type default_shape({n0, n1, n2}); ComplexView3DType x("x", n0, n1, n2), ref_x("ref_x", n0, n1, n2); using axes_type = KokkosFFT::axis_type<2>; @@ -1756,45 +1851,60 @@ void test_fft2_2dfft_3dview(T atol = 1.e-12) { axes_type axes = {axis0, axis1}; - std::array shape_c2r = shape; - shape_c2r.at(axis1) = shape_c2r.at(axis1) / 2 + 1; + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; - auto [_n0, _n1, _n2] = shape_c2r; + std::size_t n0_new = static_cast(shape.at(axis0) + i0); + std::size_t n1_new = static_cast(shape.at(axis1) + i1); + shape_type<2> new_shape = {n0_new, n1_new}; - ComplexView3DType _x("_x", n0, n1, n2), out("out", n0, n1, n2), - ref_out("ref_out", n0, n1, n2); - RealView3DType xr("xr", n0, n1, n2), ref_xr("ref_xr", n0, n1, n2), - _xr("_xr", n0, n1, n2); - ComplexView3DType outr("outr", _n0, _n1, _n2); + shape.at(axis0) = n0_new; + shape.at(axis1) = n1_new; + shape_c2r.at(axis0) = n0_new; + shape_c2r.at(axis1) = n1_new / 2 + 1; - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + auto [_n0, _n1, _n2] = shape; + auto [_m0, _m1, _m2] = shape_c2r; - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + ComplexView3DType _x("_x", _n0, _n1, _n2), out("out", _n0, _n1, _n2), + ref_x; + RealView3DType xr("xr", _n0, _n1, _n2), _xr("_xr", _n0, _n1, _n2), + ref_xr; + ComplexView3DType outr("outr", _m0, _m1, _m2); - Kokkos::fence(); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Along one axis - // Simple identity tests - KokkosFFT::fft2(execution_space(), x, out, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - KokkosFFT::ifft2(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axes); + Kokkos::fence(); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft2(execution_space(), x, out, + KokkosFFT::Normalization::backward, axes, new_shape); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft2(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::ifft2(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axes, new_shape); - KokkosFFT::irfft2(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axes); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::irfft2(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axes, + new_shape); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } } } } @@ -1806,9 +1916,9 @@ void test_fft2_2dfft_4dview(T atol = 1.e-12) { using ComplexView4DType = Kokkos::View****, LayoutType, execution_space>; - constexpr int DIM = 4; - std::array shape = {n0, n1, n2, n3}; - ComplexView4DType x("x", n0, n1, n2, n3), ref_x("ref_x", n0, n1, n2, n3); + constexpr int DIM = 4; + shape_type default_shape({n0, n1, n2, n3}); + ComplexView4DType x("x", n0, n1, n2, n3); using axes_type = KokkosFFT::axis_type<2>; @@ -1820,44 +1930,60 @@ void test_fft2_2dfft_4dview(T atol = 1.e-12) { axes_type axes = {axis0, axis1}; - std::array shape_c2r = shape; - shape_c2r.at(axis1) = shape_c2r.at(axis1) / 2 + 1; + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; - auto [_n0, _n1, _n2, _n3] = shape_c2r; - ComplexView4DType _x("_x", n0, n1, n2, n3), out("out", n0, n1, n2, n3), - ref_out("ref_out", n0, n1, n2, n3); - RealView4DType xr("xr", n0, n1, n2, n3), ref_xr("ref_xr", n0, n1, n2, n3), - _xr("_xr", n0, n1, n2, n3); - ComplexView4DType outr("outr", _n0, _n1, _n2, _n3); + std::size_t n0_new = static_cast(shape.at(axis0) + i0); + std::size_t n1_new = static_cast(shape.at(axis1) + i1); + shape_type<2> new_shape = {n0_new, n1_new}; - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + shape.at(axis0) = n0_new; + shape.at(axis1) = n1_new; + shape_c2r.at(axis0) = n0_new; + shape_c2r.at(axis1) = n1_new / 2 + 1; - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + auto [_n0, _n1, _n2, _n3] = shape; + auto [_m0, _m1, _m2, _m3] = shape_c2r; - Kokkos::fence(); + ComplexView4DType _x("_x", _n0, _n1, _n2, _n3), + out("out", _n0, _n1, _n2, _n3), ref_x; + RealView4DType xr("xr", _n0, _n1, _n2, _n3), + _xr("_xr", _n0, _n1, _n2, _n3), ref_xr; + ComplexView4DType outr("outr", _m0, _m1, _m2, _m3); - // Along one axis - // Simple identity tests - KokkosFFT::fft2(execution_space(), x, out, - KokkosFFT::Normalization::backward, axes); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - KokkosFFT::ifft2(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + Kokkos::fence(); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft2(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axes); + // Along one axis + // Simple identity tests + KokkosFFT::fft2(execution_space(), x, out, + KokkosFFT::Normalization::backward, axes, new_shape); - KokkosFFT::irfft2(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::ifft2(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axes, new_shape); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::irfft2(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axes, + new_shape); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } } } } @@ -1869,10 +1995,9 @@ void test_fft2_2dfft_5dview(T atol = 1.e-12) { using ComplexView5DType = Kokkos::View*****, LayoutType, execution_space>; - constexpr int DIM = 5; - std::array shape = {n0, n1, n2, n3, n4}; - ComplexView5DType x("x", n0, n1, n2, n3, n4), - ref_x("ref_x", n0, n1, n2, n3, n4); + constexpr int DIM = 5; + shape_type default_shape({n0, n1, n2, n3, n4}); + ComplexView5DType x("x", n0, n1, n2, n3, n4); using axes_type = KokkosFFT::axis_type<2>; @@ -1884,45 +2009,60 @@ void test_fft2_2dfft_5dview(T atol = 1.e-12) { axes_type axes = {axis0, axis1}; - std::array shape_c2r = shape; - shape_c2r.at(axis1) = shape_c2r.at(axis1) / 2 + 1; + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; - auto [_n0, _n1, _n2, _n3, _n4] = shape_c2r; - ComplexView5DType _x("_x", n0, n1, n2, n3, n4), - out("out", n0, n1, n2, n3, n4), - ref_out("ref_out", n0, n1, n2, n3, n4); - RealView5DType xr("xr", n0, n1, n2, n3, n4), - ref_xr("ref_xr", n0, n1, n2, n3, n4), _xr("_xr", n0, n1, n2, n3, n4); - ComplexView5DType outr("outr", _n0, _n1, _n2, _n3, _n4); + std::size_t n0_new = static_cast(shape.at(axis0) + i0); + std::size_t n1_new = static_cast(shape.at(axis1) + i1); + shape_type<2> new_shape = {n0_new, n1_new}; - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + shape.at(axis0) = n0_new; + shape.at(axis1) = n1_new; + shape_c2r.at(axis0) = n0_new; + shape_c2r.at(axis1) = n1_new / 2 + 1; - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + auto [_n0, _n1, _n2, _n3, _n4] = shape; + auto [_m0, _m1, _m2, _m3, _m4] = shape_c2r; - Kokkos::fence(); + ComplexView5DType _x("_x", _n0, _n1, _n2, _n3, _n4), + out("out", _n0, _n1, _n2, _n3, _n4), ref_x; + RealView5DType xr("xr", _n0, _n1, _n2, _n3, _n4), + _xr("_xr", _n0, _n1, _n2, _n3, _n4), ref_xr; + ComplexView5DType outr("outr", _m0, _m1, _m2, _m3, _m4); - // Along one axis - // Simple identity tests - KokkosFFT::fft2(execution_space(), x, out, - KokkosFFT::Normalization::backward, axes); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - KokkosFFT::ifft2(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + Kokkos::fence(); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft2(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axes); + // Along one axis + // Simple identity tests + KokkosFFT::fft2(execution_space(), x, out, + KokkosFFT::Normalization::backward, axes, new_shape); - KokkosFFT::irfft2(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::ifft2(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axes, new_shape); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::irfft2(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axes, + new_shape); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } } } } @@ -1934,10 +2074,9 @@ void test_fft2_2dfft_6dview(T atol = 1.e-12) { using ComplexView6DType = Kokkos::View******, LayoutType, execution_space>; - constexpr int DIM = 6; - std::array shape = {n0, n1, n2, n3, n4, n5}; - ComplexView6DType x("x", n0, n1, n2, n3, n4, n5), - ref_x("ref_x", n0, n1, n2, n3, n4, n5); + constexpr int DIM = 6; + shape_type default_shape({n0, n1, n2, n3, n4, n5}); + ComplexView6DType x("x", n0, n1, n2, n3, n4, n5); using axes_type = KokkosFFT::axis_type<2>; @@ -1949,46 +2088,60 @@ void test_fft2_2dfft_6dview(T atol = 1.e-12) { axes_type axes = {axis0, axis1}; - std::array shape_c2r = shape; - shape_c2r.at(axis1) = shape_c2r.at(axis1) / 2 + 1; + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; - auto [_n0, _n1, _n2, _n3, _n4, _n5] = shape_c2r; - ComplexView6DType _x("_x", n0, n1, n2, n3, n4, n5), - out("out", n0, n1, n2, n3, n4, n5), - ref_out("ref_out", n0, n1, n2, n3, n4, n5); - RealView6DType xr("xr", n0, n1, n2, n3, n4, n5), - ref_xr("ref_xr", n0, n1, n2, n3, n4, n5), - _xr("_xr", n0, n1, n2, n3, n4, n5); - ComplexView6DType outr("outr", _n0, _n1, _n2, _n3, _n4, _n5); + std::size_t n0_new = static_cast(shape.at(axis0) + i0); + std::size_t n1_new = static_cast(shape.at(axis1) + i1); + shape_type<2> new_shape = {n0_new, n1_new}; - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + shape.at(axis0) = n0_new; + shape.at(axis1) = n1_new; + shape_c2r.at(axis0) = n0_new; + shape_c2r.at(axis1) = n1_new / 2 + 1; - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + auto [_n0, _n1, _n2, _n3, _n4, _n5] = shape; + auto [_m0, _m1, _m2, _m3, _m4, _m5] = shape_c2r; - Kokkos::fence(); + ComplexView6DType _x("_x", _n0, _n1, _n2, _n3, _n4, _n5), + out("out", _n0, _n1, _n2, _n3, _n4, _n5), ref_x; + RealView6DType xr("xr", _n0, _n1, _n2, _n3, _n4, _n5), + _xr("_xr", _n0, _n1, _n2, _n3, _n4, _n5), ref_xr; + ComplexView6DType outr("outr", _m0, _m1, _m2, _m3, _m4, _m5); - // Along one axis - // Simple identity tests - KokkosFFT::fft2(execution_space(), x, out, - KokkosFFT::Normalization::backward, axes); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - KokkosFFT::ifft2(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + Kokkos::fence(); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft2(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axes); + // Along one axis + // Simple identity tests + KokkosFFT::fft2(execution_space(), x, out, + KokkosFFT::Normalization::backward, axes, new_shape); - KokkosFFT::irfft2(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::ifft2(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axes, new_shape); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::irfft2(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axes, + new_shape); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } } } } @@ -2000,10 +2153,9 @@ void test_fft2_2dfft_7dview(T atol = 1.e-12) { using ComplexView7DType = Kokkos::View*******, LayoutType, execution_space>; - constexpr int DIM = 7; - std::array shape = {n0, n1, n2, n3, n4, n5, n6}; - ComplexView7DType x("x", n0, n1, n2, n3, n4, n5, n6), - ref_x("ref_x", n0, n1, n2, n3, n4, n5, n6); + constexpr int DIM = 7; + shape_type default_shape({n0, n1, n2, n3, n4, n5, n6}); + ComplexView7DType x("x", n0, n1, n2, n3, n4, n5, n6); using axes_type = KokkosFFT::axis_type<2>; for (int axis0 = 0; axis0 < DIM; axis0++) { @@ -2014,46 +2166,62 @@ void test_fft2_2dfft_7dview(T atol = 1.e-12) { axes_type axes = {axis0, axis1}; - std::array shape_c2r = shape; - shape_c2r.at(axis1) = shape_c2r.at(axis1) / 2 + 1; + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; - auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = shape_c2r; - ComplexView7DType _x("_x", n0, n1, n2, n3, n4, n5, n6), - out("out", n0, n1, n2, n3, n4, n5, n6), - ref_out("ref_out", n0, n1, n2, n3, n4, n5, n6); - RealView7DType xr("xr", n0, n1, n2, n3, n4, n5, n6), - ref_xr("ref_xr", n0, n1, n2, n3, n4, n5, n6), - _xr("_xr", n0, n1, n2, n3, n4, n5, n6); - ComplexView7DType outr("outr", _n0, _n1, _n2, _n3, _n4, _n5, _n6); + std::size_t n0_new = static_cast(shape.at(axis0) + i0); + std::size_t n1_new = static_cast(shape.at(axis1) + i1); + shape_type<2> new_shape = {n0_new, n1_new}; - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + shape.at(axis0) = n0_new; + shape.at(axis1) = n1_new; + shape_c2r.at(axis0) = n0_new; + shape_c2r.at(axis1) = n1_new / 2 + 1; - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = shape; + auto [_m0, _m1, _m2, _m3, _m4, _m5, _m6] = shape_c2r; - Kokkos::fence(); + ComplexView7DType _x("_x", _n0, _n1, _n2, _n3, _n4, _n5, _n6), + out("out", _n0, _n1, _n2, _n3, _n4, _n5, _n6), ref_x; - // Along one axis - // Simple identity tests - KokkosFFT::fft2(execution_space(), x, out, - KokkosFFT::Normalization::backward, axes); + RealView7DType xr("xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6), + _xr("_xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6), ref_xr; - KokkosFFT::ifft2(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axes); + ComplexView7DType outr("outr", _m0, _m1, _m2, _m3, _m4, _m5, _m6); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft2(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - KokkosFFT::irfft2(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axes); + Kokkos::fence(); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft2(execution_space(), x, out, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::ifft2(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axes, new_shape); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::irfft2(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axes, + new_shape); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } } } } @@ -2065,10 +2233,9 @@ void test_fft2_2dfft_8dview(T atol = 1.e-12) { using ComplexView8DType = Kokkos::View********, LayoutType, execution_space>; - constexpr int DIM = 8; - std::array shape = {n0, n1, n2, n3, n4, n5, n6, n7}; - ComplexView8DType x("x", n0, n1, n2, n3, n4, n5, n6, n7), - ref_x("ref_x", n0, n1, n2, n3, n4, n5, n6, n7); + constexpr int DIM = 8; + shape_type default_shape({n0, n1, n2, n3, n4, n5, n6, n7}); + ComplexView8DType x("x", n0, n1, n2, n3, n4, n5, n6, n7); using axes_type = KokkosFFT::axis_type<2>; for (int axis0 = 0; axis0 < DIM; axis0++) { @@ -2079,46 +2246,63 @@ void test_fft2_2dfft_8dview(T atol = 1.e-12) { axes_type axes = {axis0, axis1}; - std::array shape_c2r = shape; - shape_c2r.at(axis1) = shape_c2r.at(axis1) / 2 + 1; + for (int i0 = -1; i0 <= 1; i0++) { + for (int i1 = -1; i1 <= 1; i1++) { + shape_type shape = default_shape; + shape_type shape_c2r = default_shape; - auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = shape_c2r; - ComplexView8DType _x("_x", n0, n1, n2, n3, n4, n5, n6, n7), - out("out", n0, n1, n2, n3, n4, n5, n6, n7), - ref_out("ref_out", n0, n1, n2, n3, n4, n5, n6, n7); - RealView8DType xr("xr", n0, n1, n2, n3, n4, n5, n6, n7), - ref_xr("ref_xr", n0, n1, n2, n3, n4, n5, n6, n7), - _xr("_xr", n0, n1, n2, n3, n4, n5, n6, n7); - ComplexView8DType outr("outr", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7); + std::size_t n0_new = static_cast(shape.at(axis0) + i0); + std::size_t n1_new = static_cast(shape.at(axis1) + i1); + shape_type<2> new_shape = {n0_new, n1_new}; - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - Kokkos::fill_random(xr, random_pool, 1); + shape.at(axis0) = n0_new; + shape.at(axis1) = n1_new; + shape_c2r.at(axis0) = n0_new; + shape_c2r.at(axis1) = n1_new / 2 + 1; - Kokkos::deep_copy(ref_x, x); - Kokkos::deep_copy(ref_xr, xr); + auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = shape; + auto [_m0, _m1, _m2, _m3, _m4, _m5, _m6, _m7] = shape_c2r; - Kokkos::fence(); + ComplexView8DType _x("_x", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), + out("out", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), ref_x; - // Along one axis - // Simple identity tests - KokkosFFT::fft2(execution_space(), x, out, - KokkosFFT::Normalization::backward, axes); + RealView8DType xr("xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), + _xr("_xr", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7), ref_xr; - KokkosFFT::ifft2(execution_space(), out, _x, - KokkosFFT::Normalization::backward, axes); + ComplexView8DType outr("outr", _m0, _m1, _m2, _m3, _m4, _m5, _m6, + _m7); - EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::fill_random(xr, random_pool, 1); - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft2(execution_space(), xr, outr, - KokkosFFT::Normalization::backward, axes); + KokkosFFT::Impl::crop_or_pad(execution_space(), x, ref_x, shape); + KokkosFFT::Impl::crop_or_pad(execution_space(), xr, ref_xr, shape); - KokkosFFT::irfft2(execution_space(), outr, _xr, - KokkosFFT::Normalization::backward, axes); + Kokkos::fence(); - EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + // Along one axis + // Simple identity tests + KokkosFFT::fft2(execution_space(), x, out, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::ifft2(execution_space(), out, _x, + KokkosFFT::Normalization::backward, axes, new_shape); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfft2(execution_space(), xr, outr, + KokkosFFT::Normalization::backward, axes, new_shape); + + KokkosFFT::irfft2(execution_space(), outr, _xr, + KokkosFFT::Normalization::backward, axes, + new_shape); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } } } } @@ -2155,6 +2339,14 @@ TYPED_TEST(FFT2D, IRFFT2_2DView) { test_fft2_2dirfft_2dview(); } +// fft2 on 2D Views with shape argument +TYPED_TEST(FFT2D, 2DFFT_2DView_shape) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + test_fft2_2dfft_2dview_shape(); +} + // batced fft2 on 3D Views TYPED_TEST(FFT2D, FFT_batched_3DView) { using float_type = typename TestFixture::float_type; @@ -2234,27 +2426,8 @@ void test_fftn_2dfft_2dview() { KokkosFFT::fft(execution_space(), out1, out2, KokkosFFT::Normalization::backward, /*axis=*/0); - KokkosFFT::fftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward - KokkosFFT::fftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); - KokkosFFT::fftn(execution_space(), x, out_o, KokkosFFT::Normalization::ortho); - KokkosFFT::fftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); - - multiply(out_o, sqrt(static_cast(n0 * n1))); - multiply(out_f, static_cast(n0 * n1)); - - EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); - - // Same tests with specifying axes - // np.fftn for 2D array is identical to np.fft(np.fft(x, axis=1), axis=0) using axes_type = KokkosFFT::axis_type<2>; axes_type axes = {-2, -1}; - KokkosFFT::fftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::backward KokkosFFT::fftn(execution_space(), x, out_b, axes, @@ -2294,73 +2467,6 @@ void test_fftn_2dfft_2dview() { EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); } -template -void test_fftn_3dfft_3dview(T atol = 1.0e-6) { - const int n0 = 4, n1 = 6, n2 = 8; - using ComplexView3DType = - Kokkos::View***, LayoutType, execution_space>; - - ComplexView3DType x("x", n0, n1, n2); - ComplexView3DType out("out", n0, n1, n2), out1("out1", n0, n1, n2), - out2("out2", n0, n1, n2), out3("out3", n0, n1, n2); - ComplexView3DType out_b("out_b", n0, n1, n2), out_o("out_o", n0, n1, n2), - out_f("out_f", n0, n1, n2); - - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - - Kokkos::fence(); - - // np.fftn for 3D array is identical to np.fft(np.fft(np.fft(x, axis=2), - // axis=1), axis=0) - KokkosFFT::fft(execution_space(), x, out1, KokkosFFT::Normalization::backward, - /*axis=*/2); - KokkosFFT::fft(execution_space(), out1, out2, - KokkosFFT::Normalization::backward, /*axis=*/1); - KokkosFFT::fft(execution_space(), out2, out3, - KokkosFFT::Normalization::backward, /*axis=*/0); - - KokkosFFT::fftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward - KokkosFFT::fftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); - KokkosFFT::fftn(execution_space(), x, out_o, KokkosFFT::Normalization::ortho); - KokkosFFT::fftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); - - multiply(out_o, sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, static_cast(n0 * n1 * n2)); - - EXPECT_TRUE(allclose(out, out3, 1.e-5, atol)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, atol)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, atol)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, atol)); - - // Same tests with specifying axes - // np.fftn for 3D array is identical to np.fft(np.fft(np.fft(x, axis=2), - // axis=1), axis=0) - using axes_type = KokkosFFT::axis_type<3>; - - KokkosFFT::fftn( - execution_space(), x, out, - axes_type{-3, -2, -1}); // default: KokkosFFT::Normalization::backward - KokkosFFT::fftn(execution_space(), x, out_b, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::backward); - KokkosFFT::fftn(execution_space(), x, out_o, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::ortho); - KokkosFFT::fftn(execution_space(), x, out_f, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::forward); - - multiply(out_o, sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, static_cast(n0 * n1 * n2)); - - EXPECT_TRUE(allclose(out, out3, 1.e-5, atol)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, atol)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, atol)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, atol)); -} - template void test_ifftn_2dfft_2dview() { const int n0 = 4, n1 = 6; @@ -2380,33 +2486,14 @@ void test_ifftn_2dfft_2dview() { Kokkos::fence(); // np.ifftn for 2D array is identical to np.ifft(np.ifft(x, axis=1), axis=0) + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; + KokkosFFT::ifft(execution_space(), x, out1, KokkosFFT::Normalization::backward, /*axis=*/1); KokkosFFT::ifft(execution_space(), out1, out2, KokkosFFT::Normalization::backward, /*axis=*/0); - KokkosFFT::ifftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward - KokkosFFT::ifftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); - KokkosFFT::ifftn(execution_space(), x, out_o, - KokkosFFT::Normalization::ortho); - KokkosFFT::ifftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); - - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0 / static_cast(n0 * n1)); - - EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); - - // Same tests with specifying axes - // np.fftn for 2D array is identical to np.fft(np.fft(x, axis=1), axis=0) - using axes_type = KokkosFFT::axis_type<2>; - axes_type axes = {-2, -1}; - KokkosFFT::ifftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::backward KokkosFFT::ifftn(execution_space(), x, out_b, axes, @@ -2445,74 +2532,6 @@ void test_ifftn_2dfft_2dview() { EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); } -template -void test_ifftn_3dfft_3dview() { - const int n0 = 4, n1 = 6, n2 = 8; - using ComplexView3DType = - Kokkos::View***, LayoutType, execution_space>; - - ComplexView3DType x("x", n0, n1, n2); - ComplexView3DType out("out", n0, n1, n2), out1("out1", n0, n1, n2), - out2("out2", n0, n1, n2), out3("out3", n0, n1, n2); - ComplexView3DType out_b("out_b", n0, n1, n2), out_o("out_o", n0, n1, n2), - out_f("out_f", n0, n1, n2); - - const Kokkos::complex I(1.0, 1.0); - Kokkos::Random_XorShift64_Pool<> random_pool(12345); - Kokkos::fill_random(x, random_pool, I); - - Kokkos::fence(); - - // np.ifftn for 3D array is identical to np.ifft(np.ifft(np.ifft(x, axis=2), - // axis=1), axis=0) - KokkosFFT::ifft(execution_space(), x, out1, - KokkosFFT::Normalization::backward, /*axis=*/2); - KokkosFFT::ifft(execution_space(), out1, out2, - KokkosFFT::Normalization::backward, /*axis=*/1); - KokkosFFT::ifft(execution_space(), out2, out3, - KokkosFFT::Normalization::backward, /*axis=*/0); - - KokkosFFT::ifftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward - KokkosFFT::ifftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); - KokkosFFT::ifftn(execution_space(), x, out_o, - KokkosFFT::Normalization::ortho); - KokkosFFT::ifftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); - - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, 1.0 / static_cast(n0 * n1 * n2)); - - EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); - - // Same tests with specifying axes - // np.ifftn for 3D array is identical to np.ifft(np.ifft(np.ifft(x, axis=2), - // axis=1), axis=0) - using axes_type = KokkosFFT::axis_type<3>; - - KokkosFFT::ifftn( - execution_space(), x, out, - axes_type{-3, -2, -1}); // default: KokkosFFT::Normalization::backward - KokkosFFT::ifftn(execution_space(), x, out_b, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::backward); - KokkosFFT::ifftn(execution_space(), x, out_o, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::ortho); - KokkosFFT::ifftn(execution_space(), x, out_f, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::forward); - - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, 1.0 / static_cast(n0 * n1 * n2)); - - EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); -} - template void test_rfftn_2dfft_2dview() { const int n0 = 4, n1 = 6; @@ -2532,25 +2551,27 @@ void test_rfftn_2dfft_2dview() { Kokkos::fence(); // np.rfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; KokkosFFT::rfft(execution_space(), x, out1, KokkosFFT::Normalization::backward, /*axis=*/1); KokkosFFT::fft(execution_space(), out1, out2, KokkosFFT::Normalization::backward, /*axis=*/0); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward + KokkosFFT::rfftn(execution_space(), x, out, + axes); // default: KokkosFFT::Normalization::backward Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_b, + KokkosFFT::rfftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::backward); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_o, + KokkosFFT::rfftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ortho); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_f, + KokkosFFT::rfftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::forward); multiply(out_o, sqrt(static_cast(n0 * n1))); @@ -2561,25 +2582,24 @@ void test_rfftn_2dfft_2dview() { EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); - // Same tests with specifying axes - // np.rfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) - using axes_type = KokkosFFT::axis_type<2>; - axes_type axes = {-2, -1}; + // Reuse plans + KokkosFFT::Impl::Plan rfftn_plan(execution_space(), x, out, + KokkosFFT::Direction::forward, axes); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out, + KokkosFFT::rfftn(execution_space(), x, out, rfftn_plan, axes); // default: KokkosFFT::Normalization::backward Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_b, axes, + KokkosFFT::rfftn(execution_space(), x, out_b, rfftn_plan, axes, KokkosFFT::Normalization::backward); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_o, axes, + KokkosFFT::rfftn(execution_space(), x, out_o, rfftn_plan, axes, KokkosFFT::Normalization::ortho); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_f, axes, + KokkosFFT::rfftn(execution_space(), x, out_f, rfftn_plan, axes, KokkosFFT::Normalization::forward); multiply(out_o, sqrt(static_cast(n0 * n1))); @@ -2589,34 +2609,268 @@ void test_rfftn_2dfft_2dview() { EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); +} + +template +void test_irfftn_2dfft_2dview() { + const int n0 = 4, n1 = 6; + using RealView2DType = Kokkos::View; + using ComplexView2DType = + Kokkos::View**, LayoutType, execution_space>; + + ComplexView2DType x("x", n0, n1 / 2 + 1), x_ref("x_ref", n0, n1 / 2 + 1); + ComplexView2DType out1("out1", n0, n1 / 2 + 1); + RealView2DType out2("out2", n0, n1), out("out", n0, n1); + RealView2DType out_b("out_b", n0, n1), out_o("out_o", n0, n1), + out_f("out_f", n0, n1); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + Kokkos::deep_copy(x_ref, x); + + // np.irfftn for 2D array is identical to np.irfft(np.ifft(x, axis=0), axis=1) + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; + + KokkosFFT::ifft(execution_space(), x, out1, + KokkosFFT::Normalization::backward, /*axis=*/0); + KokkosFFT::irfft(execution_space(), out1, out2, + KokkosFFT::Normalization::backward, /*axis=*/1); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out, + axes); // default: KokkosFFT::Normalization::backward + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_b, axes, + KokkosFFT::Normalization::backward); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_o, axes, + KokkosFFT::Normalization::ortho); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_f, axes, + KokkosFFT::Normalization::forward); + + multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1))); + multiply(out_f, 1.0 / static_cast(n0 * n1)); + + EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); // Reuse plans - KokkosFFT::Impl::Plan rfftn_plan(execution_space(), x, out, - KokkosFFT::Direction::forward, axes); + KokkosFFT::Impl::Plan irfftn_plan(execution_space(), x, out, + KokkosFFT::Direction::backward, axes); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out, irfftn_plan, + axes); // default: KokkosFFT::Normalization::backward + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_b, irfftn_plan, axes, + KokkosFFT::Normalization::backward); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_o, irfftn_plan, axes, + KokkosFFT::Normalization::ortho); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_f, irfftn_plan, axes, + KokkosFFT::Normalization::forward); + + multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1))); + multiply(out_f, 1.0 / static_cast(n0 * n1)); + + EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); +} + +template +void test_fftn_2dfft_2dview_shape(T atol = 1.0e-12) { + const int n0 = 4, n1 = 6; + using RealView2DType = Kokkos::View; + using ComplexView2DType = + Kokkos::View**, LayoutType, execution_space>; + + RealView2DType xr("xr", n0, n1), xr_ref("xr_ref", n0, n1); + ComplexView2DType x("x", n0, n1 / 2 + 1), x_ref("x_ref", n0, n1 / 2 + 1); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(xr, random_pool, 1.0); + Kokkos::fill_random(x, random_pool, I); + Kokkos::deep_copy(xr_ref, xr); + Kokkos::deep_copy(x_ref, x); + + // np.irfftn for 2D array is identical to np.irfft(np.ifft(x, axis=0), axis=1) + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; + + std::vector shapes0 = {n0 / 2, n0, n0 * 2}; + std::vector shapes1 = {n1 / 2, n1, n1 * 2}; + + for (auto&& shape0 : shapes0) { + for (auto&& shape1 : shapes1) { + // Real to complex + ComplexView2DType outr("outr", shape0, shape1 / 2 + 1), + outr_b("outr_b", shape0, shape1 / 2 + 1), + outr_o("outr_o", shape0, shape1 / 2 + 1), + outr_f("outr_f", shape0, shape1 / 2 + 1); + + shape_type<2> new_shape = {shape0, shape1}; + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr, axes, + KokkosFFT::Normalization::none, new_shape); + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr_b, axes, + KokkosFFT::Normalization::backward, new_shape); + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr_o, axes, + KokkosFFT::Normalization::ortho, new_shape); + + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr_f, axes, + KokkosFFT::Normalization::forward, new_shape); + + multiply(outr_o, sqrt(static_cast(shape0 * shape1))); + multiply(outr_f, static_cast(shape0 * shape1)); + + EXPECT_TRUE(allclose(outr_b, outr, 1.e-5, atol)); + EXPECT_TRUE(allclose(outr_o, outr, 1.e-5, atol)); + EXPECT_TRUE(allclose(outr_f, outr, 1.e-5, atol)); + + // Complex to real + RealView2DType out("out", shape0, shape1), out_b("out_b", shape0, shape1), + out_o("out_o", shape0, shape1), out_f("out_f", shape0, shape1); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out, axes, + KokkosFFT::Normalization::none, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_b, axes, + KokkosFFT::Normalization::backward, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_o, axes, + KokkosFFT::Normalization::ortho, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_f, axes, + KokkosFFT::Normalization::forward, new_shape); + + multiply(out_o, sqrt(static_cast(shape0 * shape1))); + multiply(out_b, static_cast(shape0 * shape1)); + + EXPECT_TRUE(allclose(out_b, out, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_o, out, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_f, out, 1.e-5, atol)); + } + } +} + +template +void test_fftn_3dfft_3dview(T atol = 1.0e-6) { + const int n0 = 4, n1 = 6, n2 = 8; + using ComplexView3DType = + Kokkos::View***, LayoutType, execution_space>; + + ComplexView3DType x("x", n0, n1, n2); + ComplexView3DType out("out", n0, n1, n2), out1("out1", n0, n1, n2), + out2("out2", n0, n1, n2), out3("out3", n0, n1, n2); + ComplexView3DType out_b("out_b", n0, n1, n2), out_o("out_o", n0, n1, n2), + out_f("out_f", n0, n1, n2); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + + Kokkos::fence(); + + // np.fftn for 3D array is identical to np.fft(np.fft(np.fft(x, axis=2), + // axis=1), axis=0) + using axes_type = KokkosFFT::axis_type<3>; + axes_type axes = {-3, -2, -1}; + + KokkosFFT::fft(execution_space(), x, out1, KokkosFFT::Normalization::backward, + /*axis=*/2); + KokkosFFT::fft(execution_space(), out1, out2, + KokkosFFT::Normalization::backward, /*axis=*/1); + KokkosFFT::fft(execution_space(), out2, out3, + KokkosFFT::Normalization::backward, /*axis=*/0); + + KokkosFFT::fftn(execution_space(), x, out, + axes); // default: KokkosFFT::Normalization::backward + KokkosFFT::fftn(execution_space(), x, out_b, axes, + KokkosFFT::Normalization::backward); + KokkosFFT::fftn(execution_space(), x, out_o, axes, + KokkosFFT::Normalization::ortho); + KokkosFFT::fftn(execution_space(), x, out_f, axes, + KokkosFFT::Normalization::forward); + + multiply(out_o, sqrt(static_cast(n0 * n1 * n2))); + multiply(out_f, static_cast(n0 * n1 * n2)); + + EXPECT_TRUE(allclose(out, out3, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_b, out3, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_o, out3, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_f, out3, 1.e-5, atol)); +} + +template +void test_ifftn_3dfft_3dview() { + const int n0 = 4, n1 = 6, n2 = 8; + using ComplexView3DType = + Kokkos::View***, LayoutType, execution_space>; + + ComplexView3DType x("x", n0, n1, n2); + ComplexView3DType out("out", n0, n1, n2), out1("out1", n0, n1, n2), + out2("out2", n0, n1, n2), out3("out3", n0, n1, n2); + ComplexView3DType out_b("out_b", n0, n1, n2), out_o("out_o", n0, n1, n2), + out_f("out_f", n0, n1, n2); + + const Kokkos::complex I(1.0, 1.0); + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, I); + + Kokkos::fence(); + + // np.ifftn for 3D array is identical to np.ifft(np.ifft(np.ifft(x, axis=2), + // axis=1), axis=0) + using axes_type = KokkosFFT::axis_type<3>; + axes_type axes = {-3, -2, -1}; + + KokkosFFT::ifft(execution_space(), x, out1, + KokkosFFT::Normalization::backward, /*axis=*/2); + KokkosFFT::ifft(execution_space(), out1, out2, + KokkosFFT::Normalization::backward, /*axis=*/1); + KokkosFFT::ifft(execution_space(), out2, out3, + KokkosFFT::Normalization::backward, /*axis=*/0); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out, rfftn_plan, + KokkosFFT::ifftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::backward - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_b, rfftn_plan, axes, + KokkosFFT::ifftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::backward); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_o, rfftn_plan, axes, + KokkosFFT::ifftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ortho); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_f, rfftn_plan, axes, + KokkosFFT::ifftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::forward); - multiply(out_o, sqrt(static_cast(n0 * n1))); - multiply(out_f, static_cast(n0 * n1)); + multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1 * n2))); + multiply(out_f, 1.0 / static_cast(n0 * n1 * n2)); - EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); } template @@ -2640,6 +2894,9 @@ void test_rfftn_3dfft_3dview() { // np.rfftn for 3D array is identical to np.fft(np.fft(np.rfft(x, axis=2), // axis=1), axis=0) + using axes_type = KokkosFFT::axis_type<3>; + axes_type axes = {-3, -2, -1}; + KokkosFFT::rfft(execution_space(), x, out1, KokkosFFT::Normalization::backward, /*axis=*/2); KokkosFFT::fft(execution_space(), out1, out2, @@ -2648,49 +2905,19 @@ void test_rfftn_3dfft_3dview() { KokkosFFT::Normalization::backward, /*axis=*/0); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_o, - KokkosFFT::Normalization::ortho); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); - - multiply(out_o, sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, static_cast(n0 * n1 * n2)); - - EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); - - // Same tests with specifying axes - // np.rfftn for 3D array is identical to np.fft(np.fft(np.rfft(x, axis=2), - // axis=1), axis=0) - using axes_type = KokkosFFT::axis_type<3>; - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn( - execution_space(), x, out, - axes_type{-3, -2, -1}); // default: KokkosFFT::Normalization::backward + KokkosFFT::rfftn(execution_space(), x, out, + axes); // default: KokkosFFT::Normalization::backward Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_b, axes_type{-3, -2, -1}, + KokkosFFT::rfftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::backward); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_o, axes_type{-3, -2, -1}, + KokkosFFT::rfftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ortho); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_f, axes_type{-3, -2, -1}, + KokkosFFT::rfftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::forward); multiply(out_o, sqrt(static_cast(n0 * n1 * n2))); @@ -2703,57 +2930,36 @@ void test_rfftn_3dfft_3dview() { } template -void test_irfftn_2dfft_2dview() { - const int n0 = 4, n1 = 6; - using RealView2DType = Kokkos::View; - using ComplexView2DType = - Kokkos::View**, LayoutType, execution_space>; +void test_irfftn_3dfft_3dview() { + const int n0 = 4, n1 = 6, n2 = 8; + using RealView3DType = Kokkos::View; + using ComplexView3DType = + Kokkos::View***, LayoutType, execution_space>; - ComplexView2DType x("x", n0, n1 / 2 + 1), x_ref("x_ref", n0, n1 / 2 + 1); - ComplexView2DType out1("out1", n0, n1 / 2 + 1); - RealView2DType out2("out2", n0, n1), out("out", n0, n1); - RealView2DType out_b("out_b", n0, n1), out_o("out_o", n0, n1), - out_f("out_f", n0, n1); + ComplexView3DType x("x", n0, n1, n2 / 2 + 1), + x_ref("x_ref", n0, n1, n2 / 2 + 1); + ComplexView3DType out1("out1", n0, n1, n2 / 2 + 1), + out2("out2", n0, n1, n2 / 2 + 1); + RealView3DType out("out", n0, n1, n2), out3("out3", n0, n1, n2); + RealView3DType out_b("out_b", n0, n1, n2), out_o("out_o", n0, n1, n2), + out_f("out_f", n0, n1, n2); const Kokkos::complex I(1.0, 1.0); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, I); Kokkos::deep_copy(x_ref, x); - // np.irfftn for 2D array is identical to np.irfft(np.ifft(x, axis=0), axis=1) + // np.irfftn for 3D array is identical to np.irfft(np.ifft(np.ifft(x, axis=0), + // axis=1), axis=2) + using axes_type = KokkosFFT::axis_type<3>; + axes_type axes = {-3, -2, -1}; + KokkosFFT::ifft(execution_space(), x, out1, KokkosFFT::Normalization::backward, /*axis=*/0); - KokkosFFT::irfft(execution_space(), out1, out2, - KokkosFFT::Normalization::backward, /*axis=*/1); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_o, - KokkosFFT::Normalization::ortho); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); - - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0 / static_cast(n0 * n1)); - - EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); - - // Same tests with specifying axes - // np.irfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) - using axes_type = KokkosFFT::axis_type<2>; - axes_type axes = {-2, -1}; + KokkosFFT::ifft(execution_space(), out1, out2, + KokkosFFT::Normalization::backward, /*axis=*/1); + KokkosFFT::irfft(execution_space(), out2, out3, + KokkosFFT::Normalization::backward, /*axis=*/2); Kokkos::deep_copy(x, x_ref); KokkosFFT::irfftn(execution_space(), x, out, @@ -2771,124 +2977,107 @@ void test_irfftn_2dfft_2dview() { KokkosFFT::irfftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::forward); - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0 / static_cast(n0 * n1)); - - EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); - - // Reuse plans - KokkosFFT::Impl::Plan irfftn_plan(execution_space(), x, out, - KokkosFFT::Direction::backward, axes); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out, irfftn_plan, - axes); // default: KokkosFFT::Normalization::backward - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_b, irfftn_plan, axes, - KokkosFFT::Normalization::backward); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_o, irfftn_plan, axes, - KokkosFFT::Normalization::ortho); - - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_f, irfftn_plan, axes, - KokkosFFT::Normalization::forward); - - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0 / static_cast(n0 * n1)); + multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1 * n2))); + multiply(out_f, 1.0 / static_cast(n0 * n1 * n2)); - EXPECT_TRUE(allclose(out, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out2, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out2, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); + EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); } template -void test_irfftn_3dfft_3dview() { +void test_fftn_3dfft_3dview_shape(T atol = 1.0e-12) { const int n0 = 4, n1 = 6, n2 = 8; using RealView3DType = Kokkos::View; using ComplexView3DType = Kokkos::View***, LayoutType, execution_space>; + RealView3DType xr("xr", n0, n1, n2), xr_ref("xr_ref", n0, n1, n2); ComplexView3DType x("x", n0, n1, n2 / 2 + 1), x_ref("x_ref", n0, n1, n2 / 2 + 1); - ComplexView3DType out1("out1", n0, n1, n2 / 2 + 1), - out2("out2", n0, n1, n2 / 2 + 1); - RealView3DType out("out", n0, n1, n2), out3("out3", n0, n1, n2); - RealView3DType out_b("out_b", n0, n1, n2), out_o("out_o", n0, n1, n2), - out_f("out_f", n0, n1, n2); const Kokkos::complex I(1.0, 1.0); Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(xr, random_pool, 1.0); Kokkos::fill_random(x, random_pool, I); + Kokkos::deep_copy(xr_ref, xr); Kokkos::deep_copy(x_ref, x); // np.irfftn for 3D array is identical to np.irfft(np.ifft(np.ifft(x, axis=0), // axis=1), axis=2) - KokkosFFT::ifft(execution_space(), x, out1, - KokkosFFT::Normalization::backward, /*axis=*/0); - KokkosFFT::ifft(execution_space(), out1, out2, - KokkosFFT::Normalization::backward, /*axis=*/1); - KokkosFFT::irfft(execution_space(), out2, out3, - KokkosFFT::Normalization::backward, /*axis=*/2); + using axes_type = KokkosFFT::axis_type<3>; + axes_type axes = {-3, -2, -1}; - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, - out); // default: KokkosFFT::Normalization::backward + std::vector shapes0 = {n0 / 2, n0, n0 * 2}; + std::vector shapes1 = {n1 / 2, n1, n1 * 2}; + std::vector shapes2 = {n2 / 2, n2, n2 * 2}; - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_b, - KokkosFFT::Normalization::backward); + for (auto&& shape0 : shapes0) { + for (auto&& shape1 : shapes1) { + for (auto&& shape2 : shapes2) { + shape_type<3> new_shape = {shape0, shape1, shape2}; - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_o, - KokkosFFT::Normalization::ortho); + // Real to comple + ComplexView3DType outr("outr", shape0, shape1, shape2 / 2 + 1), + outr_b("outr_b", shape0, shape1, shape2 / 2 + 1), + outr_o("outr_o", shape0, shape1, shape2 / 2 + 1), + outr_f("outr_f", shape0, shape1, shape2 / 2 + 1); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_f, - KokkosFFT::Normalization::forward); + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr, axes, + KokkosFFT::Normalization::none, new_shape); - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, 1.0 / static_cast(n0 * n1 * n2)); + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr_b, axes, + KokkosFFT::Normalization::backward, new_shape); - EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr_o, axes, + KokkosFFT::Normalization::ortho, new_shape); - // Same tests with specifying axes - // np.irfftn for 3D array is identical to np.irfft(np.ifft(np.ifft(x, axis=0), - // axis=1), axis=2) - using axes_type = KokkosFFT::axis_type<3>; + Kokkos::deep_copy(xr, xr_ref); + KokkosFFT::rfftn(execution_space(), xr, outr_f, axes, + KokkosFFT::Normalization::forward, new_shape); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn( - execution_space(), x, out, - axes_type{-3, -2, -1}); // default: KokkosFFT::Normalization::backward + multiply(outr_o, sqrt(static_cast(shape0 * shape1 * shape2))); + multiply(outr_f, static_cast(shape0 * shape1 * shape2)); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_b, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::backward); + EXPECT_TRUE(allclose(outr_b, outr, 1.e-5, atol)); + EXPECT_TRUE(allclose(outr_o, outr, 1.e-5, atol)); + EXPECT_TRUE(allclose(outr_f, outr, 1.e-5, atol)); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_o, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::ortho); + // Complex to real + RealView3DType out("out", shape0, shape1, shape2), + out_b("out_b", shape0, shape1, shape2), + out_o("out_o", shape0, shape1, shape2), + out_f("out_f", shape0, shape1, shape2); - Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_f, axes_type{-3, -2, -1}, - KokkosFFT::Normalization::forward); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out, axes, + KokkosFFT::Normalization::none, new_shape); - multiply(out_o, 1.0 / sqrt(static_cast(n0 * n1 * n2))); - multiply(out_f, 1.0 / static_cast(n0 * n1 * n2)); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_b, axes, + KokkosFFT::Normalization::backward, new_shape); - EXPECT_TRUE(allclose(out, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_b, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_o, out3, 1.e-5, 1.e-6)); - EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_o, axes, + KokkosFFT::Normalization::ortho, new_shape); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_f, axes, + KokkosFFT::Normalization::forward, new_shape); + + multiply(out_o, sqrt(static_cast(shape0 * shape1 * shape2))); + multiply(out_b, static_cast(shape0 * shape1 * shape2)); + + EXPECT_TRUE(allclose(out_b, out, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_o, out, 1.e-5, atol)); + EXPECT_TRUE(allclose(out_f, out, 1.e-5, atol)); + } + } + } } template @@ -3269,53 +3458,61 @@ TYPED_TEST(FFTND, 2DFFT_2DView) { test_fftn_2dfft_2dview(); } -// fftn on 3D Views -TYPED_TEST(FFTND, 3DFFT_3DView) { +// ifftn on 2D Views +TYPED_TEST(FFTND, 2DIFFT_2DView) { using float_type = typename TestFixture::float_type; using layout_type = typename TestFixture::layout_type; - float_type atol = std::is_same_v ? 5.0e-5 : 1.0e-6; - test_fftn_3dfft_3dview(atol); + test_ifftn_2dfft_2dview(); } -// ifftn on 2D Views -TYPED_TEST(FFTND, 2DIFFT_2DView) { +// rfftn on 2D Views +TYPED_TEST(FFTND, 2DRFFT_2DView) { using float_type = typename TestFixture::float_type; using layout_type = typename TestFixture::layout_type; - test_ifftn_2dfft_2dview(); + test_rfftn_2dfft_2dview(); } -// ifftn on 3D Views -TYPED_TEST(FFTND, 3DIFFT_3DView) { +// irfftn on 2D Views +TYPED_TEST(FFTND, 2DIRFFT_2DView) { using float_type = typename TestFixture::float_type; using layout_type = typename TestFixture::layout_type; - test_ifftn_3dfft_3dview(); + test_irfftn_2dfft_2dview(); } -// rfftn on 2D Views -TYPED_TEST(FFTND, 2DRFFT_2DView) { +// fftn on 2D Views with shape argument +TYPED_TEST(FFTND, 2DFFT_2DView_shape) { using float_type = typename TestFixture::float_type; using layout_type = typename TestFixture::layout_type; - test_rfftn_2dfft_2dview(); + test_fftn_2dfft_2dview_shape(); } -// rfftn on 3D Views -TYPED_TEST(FFTND, 3DRFFT_3DView) { +// fftn on 3D Views +TYPED_TEST(FFTND, 3DFFT_3DView) { using float_type = typename TestFixture::float_type; using layout_type = typename TestFixture::layout_type; - test_rfftn_3dfft_3dview(); + float_type atol = std::is_same_v ? 5.0e-5 : 1.0e-10; + test_fftn_3dfft_3dview(atol); } -// irfftn on 2D Views -TYPED_TEST(FFTND, 2DIRFFT_2DView) { +// ifftn on 3D Views +TYPED_TEST(FFTND, 3DIFFT_3DView) { using float_type = typename TestFixture::float_type; using layout_type = typename TestFixture::layout_type; - test_irfftn_2dfft_2dview(); + test_ifftn_3dfft_3dview(); +} + +// rfftn on 3D Views +TYPED_TEST(FFTND, 3DRFFT_3DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + test_rfftn_3dfft_3dview(); } // irfftn on 3D Views @@ -3325,6 +3522,15 @@ TYPED_TEST(FFTND, 3DIRFFT_3DView) { test_irfftn_3dfft_3dview(); } + +// fftn on 3D Views with shape argument +TYPED_TEST(FFTND, 3DFFT_3DView_shape) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + test_fftn_3dfft_3dview_shape(); +} + // batched fftn on 4D Views TYPED_TEST(FFTND, 3DFFT_batched_4DView) { using float_type = typename TestFixture::float_type; From ef451fa5c46ade6ddc9d56dbecf7f743458c4d8e Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 25 Mar 2024 16:57:31 +0900 Subject: [PATCH 4/5] specify args in examples of ND FFT --- examples/03_NDFFT/03_NDFFT.cpp | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/03_NDFFT/03_NDFFT.cpp b/examples/03_NDFFT/03_NDFFT.cpp index 3c42dd7e..ea282d52 100644 --- a/examples/03_NDFFT/03_NDFFT.cpp +++ b/examples/03_NDFFT/03_NDFFT.cpp @@ -10,6 +10,8 @@ using execution_space = Kokkos::DefaultExecutionSpace; template using View3D = Kokkos::View; +template +using axis_type = KokkosFFT::axis_type; int main(int argc, char* argv[]) { Kokkos::initialize(argc, argv); @@ -27,8 +29,8 @@ int main(int argc, char* argv[]) { Kokkos::fill_random(exec, xc2c, random_pool, I); exec.fence(); - KokkosFFT::fftn(exec, xc2c, xc2c_hat); - KokkosFFT::ifftn(exec, xc2c_hat, xc2c_inv); + KokkosFFT::fftn(exec, xc2c, xc2c_hat, axis_type<3>{-3, -2, -1}); + KokkosFFT::ifftn(exec, xc2c_hat, xc2c_inv, axis_type<3>{-3, -2, -1}); exec.fence(); // 3D R2C FFT @@ -37,7 +39,7 @@ int main(int argc, char* argv[]) { Kokkos::fill_random(exec, xr2c, random_pool, 1); exec.fence(); - KokkosFFT::rfftn(exec, xr2c, xr2c_hat); + KokkosFFT::rfftn(exec, xr2c, xr2c_hat, axis_type<3>{-3, -2, -1}); exec.fence(); // 3D C2R FFT @@ -46,7 +48,7 @@ int main(int argc, char* argv[]) { Kokkos::fill_random(exec, xc2r, random_pool, I); exec.fence(); - KokkosFFT::irfftn(exec, xc2r, xc2r_hat); + KokkosFFT::irfftn(exec, xc2r, xc2r_hat, axis_type<3>{-3, -2, -1}); exec.fence(); } Kokkos::finalize(); From e929a4c6d06c0b1818207c1e972f4ae50922965d Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 25 Mar 2024 17:37:54 +0900 Subject: [PATCH 5/5] fix argument --- fft/src/KokkosFFT_Transform.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fft/src/KokkosFFT_Transform.hpp b/fft/src/KokkosFFT_Transform.hpp index 2f3922dc..66953be8 100644 --- a/fft/src/KokkosFFT_Transform.hpp +++ b/fft/src/KokkosFFT_Transform.hpp @@ -1975,7 +1975,7 @@ void irfftn(const ExecutionSpace& exec_space, const InViewType& in, static_assert(std::is_floating_point::value, "irfftn: OutViewType must be real"); - ifftn(exec_space, in, out, norm, _s); + ifftn(exec_space, in, out, norm, s); } /// \brief Inverse of rfftn