diff --git a/common/src/KokkosFFT_utils.hpp b/common/src/KokkosFFT_utils.hpp index 768f9293..2c1d0010 100644 --- a/common/src/KokkosFFT_utils.hpp +++ b/common/src/KokkosFFT_utils.hpp @@ -77,6 +77,21 @@ namespace KokkosFFT { return index; } + + template + std::array make_sequence_array(std::index_sequence) { + return std::array{ {I...} }; + } + + template + std::array index_sequence(T const& start) { + auto sequence = make_sequence_array(std::make_index_sequence()); + std::transform(sequence.begin(), sequence.end(), sequence.begin(), + [=](const T sequence) -> T {return start + sequence;}); + return sequence; + } + + }; #endif \ No newline at end of file diff --git a/fft/src/KokkosFFT_OpenMP_plans.hpp b/fft/src/KokkosFFT_OpenMP_plans.hpp index 6272e78e..b31572cb 100644 --- a/fft/src/KokkosFFT_OpenMP_plans.hpp +++ b/fft/src/KokkosFFT_OpenMP_plans.hpp @@ -175,6 +175,7 @@ namespace KokkosFFT { istride, idist, odata, + out_extents.data(), ostride, odist, FFTW_ESTIMATE diff --git a/fft/src/KokkosFFT_Transform.hpp b/fft/src/KokkosFFT_Transform.hpp index cb56d2da..4b785cbb 100644 --- a/fft/src/KokkosFFT_Transform.hpp +++ b/fft/src/KokkosFFT_Transform.hpp @@ -390,6 +390,314 @@ namespace KokkosFFT { _irfft2(plan, in, out, norm, axes); } } +} + +namespace KokkosFFT { + template + void _fftn(PlanType& plan, const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + auto* idata = reinterpret_cast::type*>(in.data()); + auto* odata = reinterpret_cast::type*>(out.data()); + + _exec(plan.plan(), idata, odata, KOKKOS_FFT_FORWARD); + normalize(out, KOKKOS_FFT_FORWARD, norm, plan.fft_size()); + } + + template + void _ifftn(PlanType& plan, const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + auto* idata = reinterpret_cast::type*>(in.data()); + auto* odata = reinterpret_cast::type*>(out.data()); + + _exec(plan.plan(), idata, odata, KOKKOS_FFT_BACKWARD); + normalize(out, KOKKOS_FFT_BACKWARD, norm, plan.fft_size()); + } + + template + void _rfftn(PlanType& plan, const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(std::is_floating_point::value, + "KokkosFFT::rfftn: InViewType must be real"); + static_assert(is_complex::value, + "KokkosFFT::rfftn: OutViewType must be complex"); + + _fftn(plan, in, out, norm); + } + + template + void _irfftn(PlanType& plan, const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(is_complex::value, + "KokkosFFT::irfftn: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::irfftn: OutViewType must be real"); + + _ifftn(plan, in, out, norm); + } + + template + void fftn(const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: OutViewType is not a Kokkos::View."); + + // Create a default sequence of axes {-rank, -(rank-1), ..., -1} + constexpr std::size_t rank = InViewType::rank(); + constexpr int start = -static_cast(rank); + axis_type axes = index_sequence(start); + + Plan plan(in, out, KOKKOS_FFT_FORWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _fftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _fftn(plan, in, out, norm); + } + } + + template + void fftn(const InViewType& in, OutViewType& out, axis_type axes, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: OutViewType is not a Kokkos::View."); + + Plan plan(in, out, KOKKOS_FFT_FORWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _fftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _fftn(plan, in, out, norm); + } + } + + template + void ifftn(const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: OutViewType is not a Kokkos::View."); + + // Create a default sequence of axes {-rank, -(rank-1), ..., -1} + constexpr std::size_t rank = InViewType::rank(); + constexpr int start = -static_cast(rank); + axis_type axes = index_sequence(start); + + Plan plan(in, out, KOKKOS_FFT_BACKWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _ifftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _ifftn(plan, in, out, norm); + } + } + + template + void ifftn(const InViewType& in, OutViewType& out, axis_type axes, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: OutViewType is not a Kokkos::View."); + + Plan plan(in, out, KOKKOS_FFT_BACKWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _ifftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _ifftn(plan, in, out, norm); + } + } + + template + void rfftn(const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(std::is_floating_point::value, + "KokkosFFT::rfftn: InViewType must be real"); + static_assert(is_complex::value, + "KokkosFFT::rfftn: OutViewType must be complex"); + + // Create a default sequence of axes {-rank, -(rank-1), ..., -1} + constexpr std::size_t rank = InViewType::rank(); + constexpr int start = -static_cast(rank); + axis_type axes = index_sequence(start); + + Plan plan(in, out, KOKKOS_FFT_FORWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _rfftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _rfftn(plan, in, out, norm); + } + } + + template + void rfftn(const InViewType& in, OutViewType& out, axis_type axes, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(std::is_floating_point::value, + "KokkosFFT::rfftn: InViewType must be real"); + static_assert(is_complex::value, + "KokkosFFT::rfftn: OutViewType must be complex"); + + Plan plan(in, out, KOKKOS_FFT_FORWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _rfftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _rfftn(plan, in, out, norm); + } + } + + template + void irfftn(const InViewType& in, OutViewType& out, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(is_complex::value, + "KokkosFFT::irfftn: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::irfftn: OutViewType must be real"); + + // Create a default sequence of axes {-rank, -(rank-1), ..., -1} + constexpr std::size_t rank = InViewType::rank(); + constexpr int start = -static_cast(rank); + axis_type axes = index_sequence(start); + + Plan plan(in, out, KOKKOS_FFT_BACKWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _irfftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _irfftn(plan, in, out, norm); + } + } + + template + void irfftn(const InViewType& in, OutViewType& out, axis_type axes, FFT_Normalization norm=FFT_Normalization::BACKWARD) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(is_complex::value, + "KokkosFFT::irfftn: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::irfftn: OutViewType must be real"); + + Plan plan(in, out, KOKKOS_FFT_BACKWARD, axes); + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::transpose(in, in_T, plan.map()); + KokkosFFT::transpose(out, out_T, plan.map()); + + _irfftn(plan, in_T, out_T, norm); + + KokkosFFT::transpose(out_T, out, plan.map_inv()); + } else { + _irfftn(plan, in, out, norm); + } + } }; #endif \ No newline at end of file diff --git a/fft/unit_test/Test_Transform.cpp b/fft/unit_test/Test_Transform.cpp index 21fb1fbf..f1249d0a 100644 --- a/fft/unit_test/Test_Transform.cpp +++ b/fft/unit_test/Test_Transform.cpp @@ -84,17 +84,22 @@ void ifft1(ViewType& in, ViewType& out) { ); } -TEST(FFT1D, Identity) { - const int maxlen = 8; +template +void test_fft1_identity(T atol=1.0e-12) { + const int maxlen = 30; + using RealView1DType = Kokkos::View; + using ComplexView1DType = Kokkos::View*, LayoutType, execution_space>; - View1D > a("a", maxlen), _a("_a", maxlen); - View1D > out("out", maxlen), outr("outr", maxlen/2+1); - View1D ar("ar", maxlen), _ar("_ar", maxlen); + ComplexView1DType a("a", maxlen), _a("_a", maxlen), a_ref("a_ref", maxlen); + ComplexView1DType out("out", maxlen), outr("outr", maxlen/2+1); + RealView1DType ar("ar", maxlen), _ar("_ar", maxlen), ar_ref("ar_ref", maxlen); - const Kokkos::complex I(1.0, 1.0); + const Kokkos::complex I(1.0, 1.0); Kokkos::Random_XorShift64_Pool<> random_pool(/*seed=*/12345); Kokkos::fill_random(a, random_pool, I); Kokkos::fill_random(ar, random_pool, 1.0); + Kokkos::deep_copy(a_ref, a); + Kokkos::deep_copy(ar_ref, ar); Kokkos::fence(); @@ -104,17 +109,19 @@ TEST(FFT1D, Identity) { KokkosFFT::rfft(ar, outr); KokkosFFT::irfft(outr, _ar); - EXPECT_TRUE( allclose(_a, a, 1.e-5, 1.e-12) ); - EXPECT_TRUE( allclose(_ar, ar, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(_a, a_ref, 1.e-5, atol) ); + EXPECT_TRUE( allclose(_ar, ar_ref, 1.e-5, atol) ); } -TEST(FFT1D, FFT) { +template +void test_fft1_1dfft_1dview() { const int len = 30; + using ComplexView1DType = Kokkos::View*, LayoutType, execution_space>; - View1D > x("x", len), out("out", len), ref("ref", len); - View1D > out_b("out_b", len), out_o("out_o", len), out_f("out_f", len); + ComplexView1DType x("x", len), out("out", len), ref("ref", len); + ComplexView1DType out_b("out_b", len), out_o("out_o", len), out_f("out_f", len); - const Kokkos::complex I(1.0, 1.0); + const Kokkos::complex I(1.0, 1.0); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, I); @@ -126,8 +133,8 @@ TEST(FFT1D, FFT) { KokkosFFT::fft(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); fft1(x, ref); - multiply(out_o, sqrt(static_cast(len))); - multiply(out_f, static_cast(len)); + multiply(out_o, sqrt(static_cast(len))); + multiply(out_f, static_cast(len)); EXPECT_TRUE( allclose(out, ref, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_b, ref, 1.e-5, 1.e-6) ); @@ -135,13 +142,15 @@ TEST(FFT1D, FFT) { EXPECT_TRUE( allclose(out_f, ref, 1.e-5, 1.e-6) ); } -TEST(FFT1D, IFFT) { +template +void test_fft1_1difft_1dview() { const int len = 30; + using ComplexView1DType = Kokkos::View*, LayoutType, execution_space>; - View1D > x("x", len), out("out", len), ref("ref", len); - View1D > out_b("out_b", len), out_o("out_o", len), out_f("out_f", len); + ComplexView1DType x("x", len), out("out", len), ref("ref", len); + ComplexView1DType out_b("out_b", len), out_o("out_o", len), out_f("out_f", len); - const Kokkos::complex I(1.0, 1.0); + const Kokkos::complex I(1.0, 1.0); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, I); @@ -153,9 +162,9 @@ TEST(FFT1D, IFFT) { KokkosFFT::ifft(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); ifft1(x, ref); - multiply(out_o, sqrt(static_cast(len))); - multiply(out_b, static_cast(len)); - multiply(out, static_cast(len)); + multiply(out_o, sqrt(static_cast(len))); + multiply(out_b, static_cast(len)); + multiply(out, static_cast(len)); EXPECT_TRUE( allclose(out, ref, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_b, ref, 1.e-5, 1.e-6) ); @@ -163,87 +172,22 @@ TEST(FFT1D, IFFT) { EXPECT_TRUE( allclose(out_f, ref, 1.e-5, 1.e-6) ); } -TEST(FFT1D, 1DbatchedFFT_2DLeftView) { +template +void test_fft1_1dfft_2dview(T atol=1.e-12) { const int n0 = 10, n1 = 12; + using RealView2DType = Kokkos::View; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - LeftView2D > x("x", n0, n1), ref_x("ref_x", n0, n1); - LeftView2D > x_axis0("x_axis0", n0, n1), x_axis1("x_axis1", n0, n1); - LeftView2D > out_axis0("out_axis0", n0, n1), ref_out_axis0("ref_out_axis0", n0, n1); - LeftView2D > out_axis1("out_axis1", n0, n1), ref_out_axis1("ref_out_axis1", n0, n1); + ComplexView2DType x("x", n0, n1), ref_x("ref_x", n0, n1); + ComplexView2DType x_axis0("x_axis0", n0, n1), x_axis1("x_axis1", n0, n1); + ComplexView2DType out_axis0("out_axis0", n0, n1), ref_out_axis0("ref_out_axis0", n0, n1); + ComplexView2DType out_axis1("out_axis1", n0, n1), ref_out_axis1("ref_out_axis1", n0, n1); - LeftView2D xr("xr", n0, n1), ref_xr("ref_xr", n0, n1); - LeftView2D xr_axis0("xr_axis0", n0, n1), xr_axis1("xr_axis1", n0, n1); - LeftView2D > outr_axis0("outr_axis0", n0/2+1, n1), outr_axis1("outr_axis1", n0, n1/2+1); + RealView2DType xr("xr", n0, n1), ref_xr("ref_xr", n0, n1); + RealView2DType xr_axis0("xr_axis0", n0, n1), xr_axis1("xr_axis1", n0, n1); + ComplexView2DType outr_axis0("outr_axis0", n0/2+1, n1), outr_axis1("outr_axis1", n0, n1/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); - - // 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); - - Kokkos::fence(); - - // Along axis 0 - // Perform batched 1D (along 0th axis) FFT sequentially - for(int i1=0; i1 > x("x", n0, n1), ref_x("ref_x", n0, n1); - RightView2D > x_axis0("x_axis0", n0, n1), x_axis1("x_axis1", n0, n1); - RightView2D > out_axis0("out_axis0", n0, n1), ref_out_axis0("ref_out_axis0", n0, n1); - RightView2D > out_axis1("out_axis1", n0, n1), ref_out_axis1("ref_out_axis1", n0, n1); - - RightView2D xr("xr", n0, n1), ref_xr("ref_xr", n0, n1); - RightView2D xr_axis0("xr_axis0", n0, n1), xr_axis1("xr_axis1", n0, n1); - RightView2D > outr_axis0("outr_axis0", n0/2+1, n1), outr_axis1("outr_axis1", n0, n1/2+1); - - const Kokkos::complex I(1.0, 1.0); + 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); @@ -263,16 +207,16 @@ TEST(FFT1D, 1DbatchedFFT_2DRightView) { } KokkosFFT::fft(x, out_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(out_axis0, ref_out_axis0, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(out_axis0, ref_out_axis0, 1.e-5, atol) ); KokkosFFT::ifft(out_axis0, x_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(x_axis0, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(x_axis0, ref_x, 1.e-5, atol) ); // Simple identity tests for r2c and c2r transforms KokkosFFT::rfft(xr, outr_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); KokkosFFT::irfft(outr_axis0, xr_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(xr_axis0, ref_xr, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(xr_axis0, ref_xr, 1.e-5, atol) ); // Recover input from reference Kokkos::deep_copy(x, ref_x); @@ -287,32 +231,35 @@ TEST(FFT1D, 1DbatchedFFT_2DRightView) { } KokkosFFT::fft(x, out_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(out_axis1, ref_out_axis1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(out_axis1, ref_out_axis1, 1.e-5, atol) ); KokkosFFT::ifft(out_axis1, x_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(x_axis1, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(x_axis1, ref_x, 1.e-5, atol) ); // Simple identity tests for r2c and c2r transforms KokkosFFT::rfft(xr, outr_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); KokkosFFT::irfft(outr_axis1, xr_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(xr_axis1, ref_xr, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(xr_axis1, ref_xr, 1.e-5, atol) ); } -TEST(FFT1D, 1DbatchedFFT_3DLeftView) { +template +void test_fft1_1dfft_3dview(T atol=1.e-12) { const int n0 = 10, n1 = 12, n2 = 8; + using RealView3DType = Kokkos::View; + using ComplexView3DType = Kokkos::View***, LayoutType, execution_space>; - LeftView3D > x("x", n0, n1, n2), ref_x("ref_x", n0, n1, n2); - LeftView3D > x_axis0("x_axis0", n0, n1, n2), x_axis1("x_axis1", n0, n1, n2), x_axis2("x_axis2", n0, n1, n2); - LeftView3D > out_axis0("out_axis0", n0, n1, n2), ref_out_axis0("ref_out_axis0", n0, n1, n2); - LeftView3D > out_axis1("out_axis1", n0, n1, n2), ref_out_axis1("ref_out_axis1", n0, n1, n2); - LeftView3D > out_axis2("out_axis2", n0, n1, n2), ref_out_axis2("ref_out_axis2", n0, n1, n2); + ComplexView3DType x("x", n0, n1, n2), ref_x("ref_x", n0, n1, n2); + ComplexView3DType x_axis0("x_axis0", n0, n1, n2), x_axis1("x_axis1", n0, n1, n2), x_axis2("x_axis2", n0, n1, n2); + ComplexView3DType out_axis0("out_axis0", n0, n1, n2), ref_out_axis0("ref_out_axis0", n0, n1, n2); + ComplexView3DType out_axis1("out_axis1", n0, n1, n2), ref_out_axis1("ref_out_axis1", n0, n1, n2); + ComplexView3DType out_axis2("out_axis2", n0, n1, n2), ref_out_axis2("ref_out_axis2", n0, n1, n2); - LeftView3D xr("xr", n0, n1, n2), ref_xr("ref_xr", n0, n1, n2); - LeftView3D xr_axis0("xr_axis0", n0, n1, n2), xr_axis1("xr_axis1", n0, n1, n2), xr_axis2("xr_axis2", n0, n1, n2); - LeftView3D > outr_axis0("outr_axis0", n0/2+1, n1, n2), outr_axis1("outr_axis1", n0, n1/2+1, n2), outr_axis2("outr_axis2", n0, n1, n2/2+1); + RealView3DType xr("xr", n0, n1, n2), ref_xr("ref_xr", n0, n1, n2); + RealView3DType xr_axis0("xr_axis0", n0, n1, n2), xr_axis1("xr_axis1", n0, n1, n2), xr_axis2("xr_axis2", n0, n1, n2); + ComplexView3DType outr_axis0("outr_axis0", n0/2+1, n1, n2), outr_axis1("outr_axis1", n0, n1/2+1, n2), outr_axis2("outr_axis2", n0, n1, n2/2+1); - const Kokkos::complex I(1.0, 1.0); + 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); @@ -334,16 +281,16 @@ TEST(FFT1D, 1DbatchedFFT_3DLeftView) { } KokkosFFT::fft(x, out_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(out_axis0, ref_out_axis0, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(out_axis0, ref_out_axis0, 1.e-5, atol) ); KokkosFFT::ifft(out_axis0, x_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(x_axis0, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(x_axis0, ref_x, 1.e-5, atol) ); // Simple identity tests for r2c and c2r transforms KokkosFFT::rfft(xr, outr_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); KokkosFFT::irfft(outr_axis0, xr_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(xr_axis0, ref_xr, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(xr_axis0, ref_xr, 1.e-5, atol) ); // Recover input from reference Kokkos::deep_copy(x, ref_x); @@ -360,16 +307,16 @@ TEST(FFT1D, 1DbatchedFFT_3DLeftView) { } KokkosFFT::fft(x, out_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(out_axis1, ref_out_axis1, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(out_axis1, ref_out_axis1, 1.e-5, atol) ); KokkosFFT::ifft(out_axis1, x_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(x_axis1, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(x_axis1, ref_x, 1.e-5, atol) ); // Simple identity tests for r2c and c2r transforms KokkosFFT::rfft(xr, outr_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); KokkosFFT::irfft(outr_axis1, xr_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(xr_axis1, ref_xr, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(xr_axis1, ref_xr, 1.e-5, atol) ); // Recover input from reference Kokkos::deep_copy(x, ref_x); @@ -386,125 +333,114 @@ TEST(FFT1D, 1DbatchedFFT_3DLeftView) { } KokkosFFT::fft(x, out_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - EXPECT_TRUE( allclose(out_axis2, ref_out_axis2, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(out_axis2, ref_out_axis2, 1.e-5, atol) ); KokkosFFT::ifft(out_axis2, x_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - EXPECT_TRUE( allclose(x_axis2, ref_x, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(x_axis2, ref_x, 1.e-5, atol) ); // Simple identity tests for r2c and c2r transforms KokkosFFT::rfft(xr, outr_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); KokkosFFT::irfft(outr_axis2, xr_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - EXPECT_TRUE( allclose(xr_axis2, ref_xr, 1.e-5, 1.e-12) ); + EXPECT_TRUE( allclose(xr_axis2, ref_xr, 1.e-5, atol) ); } -TEST(FFT1D, 1DbatchedFFT_3DRightView) { - const int n0 = 10, n1 = 12, n2 = 8; - - RightView3D > x("x", n0, n1, n2), ref_x("ref_x", n0, n1, n2); - RightView3D > x_axis0("x_axis0", n0, n1, n2), x_axis1("x_axis1", n0, n1, n2), x_axis2("x_axis2", n0, n1, n2); - RightView3D > out_axis0("out_axis0", n0, n1, n2), ref_out_axis0("ref_out_axis0", n0, n1, n2); - RightView3D > out_axis1("out_axis1", n0, n1, n2), ref_out_axis1("ref_out_axis1", n0, n1, n2); - RightView3D > out_axis2("out_axis2", n0, n1, n2), ref_out_axis2("ref_out_axis2", n0, n1, n2); - - RightView3D xr("xr", n0, n1, n2), ref_xr("ref_xr", n0, n1, n2); - RightView3D xr_axis0("xr_axis0", n0, n1, n2), xr_axis1("xr_axis1", n0, n1, n2), xr_axis2("xr_axis2", n0, n1, n2); - RightView3D > outr_axis0("outr_axis0", n0/2+1, n1, n2), outr_axis1("outr_axis1", n0, n1/2+1, n2), outr_axis2("outr_axis2", n0, n1, n2/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); +// Identity tests of fft on 1D Views +TEST(FFT1D, Identity_1DLeftViewFloat) { + test_fft1_identity(1.0e-6); +} - // 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); +TEST(FFT1D, Identity_1DLeftViewDouble) { + test_fft1_identity(1.0e-12); +} - Kokkos::fence(); +TEST(FFT1D, Identity_1DRightViewFloat) { + test_fft1_identity(1.0e-6); +} - // Along axis 0 (transpose neeed) - // Perform batched 1D (along 0th axis) FFT sequentially - for(int i2=0; i2(1.0e-12); +} - KokkosFFT::fft(x, out_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(out_axis0, ref_out_axis0, 1.e-5, 1.e-12) ); +// fft on 1D Views +TEST(FFT1D, 1DFFT_1DLeftViewFloat) { + test_fft1_1dfft_1dview(); +} - KokkosFFT::ifft(out_axis0, x_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - EXPECT_TRUE( allclose(x_axis0, ref_x, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DFFT_1DLeftViewDouble) { + test_fft1_1dfft_1dview(); +} - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(xr, outr_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - KokkosFFT::irfft(outr_axis0, xr_axis0, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); +TEST(FFT1D, 1DFFT_1DRightViewFloat) { + test_fft1_1dfft_1dview(); +} - EXPECT_TRUE( allclose(xr_axis0, ref_xr, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DFFT_1DRightViewDouble) { + test_fft1_1dfft_1dview(); +} - // Recover input from reference - Kokkos::deep_copy(x, ref_x); - Kokkos::deep_copy(xr, ref_xr); +// ifft on 1D Views +TEST(FFT1D, 1DIFFT_1DLeftViewFloat) { + test_fft1_1difft_1dview(); +} - // Along axis 1 (transpose neeed) - // Perform batched 1D (along 1st axis) FFT sequentially - for(int i2=0; i2(); +} - KokkosFFT::fft(x, out_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(out_axis1, ref_out_axis1, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DIFFT_1DRightViewFloat) { + test_fft1_1difft_1dview(); +} - KokkosFFT::ifft(out_axis1, x_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - EXPECT_TRUE( allclose(x_axis1, ref_x, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DIFFT_1DRightViewDouble) { + test_fft1_1difft_1dview(); +} - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(xr, outr_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - KokkosFFT::irfft(outr_axis1, xr_axis1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); +// batced fft1 on 2D Views +TEST(FFT1D, 1DBatchedFFT_2DLeftViewFloat) { + test_fft1_1dfft_2dview(1.0e-6); +} - EXPECT_TRUE( allclose(xr_axis1, ref_xr, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DBatchedFFT_2DLeftViewDouble) { + test_fft1_1dfft_2dview(1.0e-12); +} - // Recover input from reference - Kokkos::deep_copy(x, ref_x); - Kokkos::deep_copy(xr, ref_xr); +TEST(FFT1D, 1DBatchedFFT_2DRightViewFloat) { + test_fft1_1dfft_2dview(1.0e-6); +} - // Along axis 2 - // Perform batched 1D (along 2nd axis) FFT sequentially - for(int i1=0; i1(1.0e-12); +} - KokkosFFT::fft(x, out_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - EXPECT_TRUE( allclose(out_axis2, ref_out_axis2, 1.e-5, 1.e-12) ); +// batced fft1 on 3D Views +TEST(FFT1D, 1DBatchedFFT_3DLeftViewFloat) { + test_fft1_1dfft_3dview(1.0e-6); +} - KokkosFFT::ifft(out_axis2, x_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - EXPECT_TRUE( allclose(x_axis2, ref_x, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DBatchedFFT_3DLeftViewDouble) { + test_fft1_1dfft_3dview(1.0e-12); +} - // Simple identity tests for r2c and c2r transforms - KokkosFFT::rfft(xr, outr_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - KokkosFFT::irfft(outr_axis2, xr_axis2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); +TEST(FFT1D, 1DBatchedFFT_3DRightViewFloat) { + test_fft1_1dfft_3dview(1.0e-6); +} - EXPECT_TRUE( allclose(xr_axis2, ref_xr, 1.e-5, 1.e-12) ); +TEST(FFT1D, 1DBatchedFFT_3DRightViewDouble) { + test_fft1_1dfft_3dview(1.0e-12); } -TEST(FFT2D, 2DFFT_2DLeftView) { +// Tests for FFT2 +template +void test_fft2_2dfft_2dview() { const int n0 = 4, n1 = 6; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - LeftView2D > x("x", n0, n1); - LeftView2D > out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); - LeftView2D > out_b("out_b", n0, n1), out_o("out_o", n0, n1), out_f("out_f", n0, n1); + ComplexView2DType x("x", n0, n1); + ComplexView2DType out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); + ComplexView2DType 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); + const Kokkos::complex I(1.0, 1.0); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, I); @@ -519,8 +455,8 @@ TEST(FFT2D, 2DFFT_2DLeftView) { KokkosFFT::fft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); KokkosFFT::fft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); - multiply(out_o, sqrt(static_cast(n0 * n1))); - multiply(out_f, static_cast(n0 * n1)); + 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) ); @@ -528,30 +464,72 @@ TEST(FFT2D, 2DFFT_2DLeftView) { EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } -TEST(FFT2D, 2DFFT_2DRightView) { +template +void test_fft2_2difft_2dview() { const int n0 = 4, n1 = 6; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - RightView2D > x("x", n0, n1); - RightView2D > out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); - RightView2D > out_b("out_b", n0, n1), out_o("out_o", n0, n1), out_f("out_f", n0, n1); + ComplexView2DType x("x", n0, n1); + ComplexView2DType out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); + ComplexView2DType 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); + 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.fft2 is identical to np.fft(np.fft(x, axis=1), axis=0) - KokkosFFT::fft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + // np.ifft2 is identical to np.ifft(np.ifft(x, axis=1), axis=0) + KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + KokkosFFT::ifft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + + KokkosFFT::ifft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::ifft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::ifft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::ifft2(x, out_f, KokkosFFT::FFT_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_fft2_2drfft_2dview() { + const int n0 = 4, n1 = 6; + using RealView2DType = Kokkos::View; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; + + RealView2DType x("x", n0, n1), x_ref("x_ref", n0, n1); + ComplexView2DType out("out", n0, n1/2+1), out1("out1", n0, n1/2+1), out2("out2", n0, n1/2+1); + ComplexView2DType out_b("out_b", n0, n1/2+1), out_o("out_o", n0, n1/2+1), out_f("out_f", n0, n1/2+1); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(x, random_pool, 1); + Kokkos::deep_copy(x_ref, x); + Kokkos::fence(); + + // np.rfft2 is identical to np.fft(np.rfft(x, axis=1), axis=0) + KokkosFFT::rfft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); KokkosFFT::fft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - KokkosFFT::fft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD - KokkosFFT::fft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); - KokkosFFT::fft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); - KokkosFFT::fft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); - multiply(out_o, sqrt(static_cast(n0 * n1))); - multiply(out_f, static_cast(n0 * n1)); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(x, out_f, KokkosFFT::FFT_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) ); @@ -559,30 +537,159 @@ TEST(FFT2D, 2DFFT_2DRightView) { EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } -TEST(FFT2D, 2DIFFT_2DLeftView) { +template +void test_fft2_2dirfft_2dview() { const int n0 = 4, n1 = 6; + using RealView2DType = Kokkos::View; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - LeftView2D > x("x", n0, n1); - LeftView2D > out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); - LeftView2D > out_b("out_b", n0, n1), out_o("out_o", n0, n1), out_f("out_f", n0, n1); + 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); + 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.irfft2 is identical to np.irfft(np.ifft(x, axis=0), axis=1) + KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, 0); + KokkosFFT::irfft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, 1); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(x, out_f, KokkosFFT::FFT_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) ); +} + +// fft2 on 2D Views +TEST(FFT2D, 2DFFT_2DLeftViewFloat) { + test_fft2_2dfft_2dview(); +} + +TEST(FFT2D, 2DFFT_2DLeftViewDouble) { + test_fft2_2dfft_2dview(); +} + +TEST(FFT2D, 2DFFT_2DRightViewFloat) { + test_fft2_2dfft_2dview(); +} + +TEST(FFT2D, 2DFFT_2DRightViewDouble) { + test_fft2_2dfft_2dview(); +} + +// ifft2 on 2D Views +TEST(FFT2D, 2DIFFT_2DLeftViewFloat) { + test_fft2_2difft_2dview(); +} + +TEST(FFT2D, 2DIFFT_2DLeftViewDouble) { + test_fft2_2difft_2dview(); +} + +TEST(FFT2D, 2DIFFT_2DRightViewFloat) { + test_fft2_2difft_2dview(); +} + +TEST(FFT2D, 2DIFFT_2DRightViewDouble) { + test_fft2_2difft_2dview(); +} + +// rfft2 on 2D Views +TEST(FFT2D, 2DRFFT_2DLeftViewFloat) { + test_fft2_2drfft_2dview(); +} + +TEST(FFT2D, 2DRFFT_2DLeftViewDouble) { + test_fft2_2drfft_2dview(); +} + +TEST(FFT2D, 2DRFFT_2DRightViewFloat) { + test_fft2_2drfft_2dview(); +} + +TEST(FFT2D, 2DRFFT_2DRightViewDouble) { + test_fft2_2drfft_2dview(); +} + +// irfft2 on 2D Views +TEST(FFT2D, 2DIRFFT_2DLeftViewFloat) { + test_fft2_2dirfft_2dview(); +} + +TEST(FFT2D, 2DIRFFT_2DLeftViewDouble) { + test_fft2_2dirfft_2dview(); +} + +TEST(FFT2D, 2DIRFFT_2DRightViewFloat) { + test_fft2_2dirfft_2dview(); +} + +TEST(FFT2D, 2DIRFFT_2DRightViewDouble) { + test_fft2_2dirfft_2dview(); +} + +// Tests for FFTN +template +void test_fftn_2dfft_2dview() { + const int n0 = 4, n1 = 6; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; + + ComplexView2DType x("x", n0, n1); + ComplexView2DType out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); + ComplexView2DType 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::fence(); - // np.ifft2 is identical to np.ifft(np.ifft(x, axis=1), axis=0) - KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - KokkosFFT::ifft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + // np.fftn for 2D array is identical to np.fft(np.fft(x, axis=1), axis=0) + KokkosFFT::fft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + KokkosFFT::fft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - KokkosFFT::ifft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD - KokkosFFT::ifft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); - KokkosFFT::ifft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); - KokkosFFT::ifft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + KokkosFFT::fftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::fftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::fftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::fftn(x, out_f, KokkosFFT::FFT_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>; + + KokkosFFT::fftn(x, out, axes_type{-2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::fftn(x, out_b, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::fftn(x, out_o, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::fftn(x, out_f, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::FORWARD); - multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0/static_cast(n0 * n1)); + 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) ); @@ -590,30 +697,100 @@ TEST(FFT2D, 2DIFFT_2DLeftView) { EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } -TEST(FFT2D, 2DIFFT_2DRightView) { +template +void test_fftn_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.fftn for 3D array is identical to np.fft(np.fft(np.fft(x, axis=2), axis=1), axis=0) + KokkosFFT::fft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); + KokkosFFT::fft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + KokkosFFT::fft(out2, out3, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + + KokkosFFT::fftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::fftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::fftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::fftn(x, out_f, KokkosFFT::FFT_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.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(x, out, axes_type{-3, -2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::fftn(x, out_b, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::fftn(x, out_o, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::fftn(x, out_f, axes_type{-3, -2, -1}, KokkosFFT::FFT_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) ); +} + +template +void test_ifftn_2dfft_2dview() { const int n0 = 4, n1 = 6; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - RightView2D > x("x", n0, n1); - RightView2D > out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); - RightView2D > out_b("out_b", n0, n1), out_o("out_o", n0, n1), out_f("out_f", n0, n1); + ComplexView2DType x("x", n0, n1); + ComplexView2DType out("out", n0, n1), out1("out1", n0, n1), out2("out2", n0, n1); + ComplexView2DType 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); + 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.ifft2 is identical to np.ifft(np.ifft(x, axis=1), axis=0) + // np.ifftn for 2D array is identical to np.ifft(np.ifft(x, axis=1), axis=0) KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); KokkosFFT::ifft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); - KokkosFFT::ifft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD - KokkosFFT::ifft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); - KokkosFFT::ifft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); - KokkosFFT::ifft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + KokkosFFT::ifftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::ifftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::ifftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::ifftn(x, out_f, KokkosFFT::FFT_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>; - multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0/static_cast(n0 * n1)); + KokkosFFT::ifftn(x, out, axes_type{-2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::ifftn(x, out_b, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::ifftn(x, out_o, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::ifftn(x, out_f, axes_type{-2, -1}, KokkosFFT::FFT_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) ); @@ -621,36 +798,114 @@ TEST(FFT2D, 2DIFFT_2DRightView) { EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } -TEST(FFT2D, 2DRFFT_2DLeftView) { +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(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); + KokkosFFT::ifft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + KokkosFFT::ifft(out2, out3, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + + KokkosFFT::ifftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::ifftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::ifftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::ifftn(x, out_f, KokkosFFT::FFT_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(x, out, axes_type{-3, -2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::ifftn(x, out_b, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::ifftn(x, out_o, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::ifftn(x, out_f, axes_type{-3, -2, -1}, KokkosFFT::FFT_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; + using RealView2DType = Kokkos::View; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - LeftView2D x("x", n0, n1), x_ref("x_ref", n0, n1); - LeftView2D > out("out", n0, n1/2+1), out1("out1", n0, n1/2+1), out2("out2", n0, n1/2+1); - LeftView2D > out_b("out_b", n0, n1/2+1), out_o("out_o", n0, n1/2+1), out_f("out_f", n0, n1/2+1); + RealView2DType x("x", n0, n1), x_ref("x_ref", n0, n1); + ComplexView2DType out("out", n0, n1/2+1), out1("out1", n0, n1/2+1), out2("out2", n0, n1/2+1); + ComplexView2DType out_b("out_b", n0, n1/2+1), out_o("out_o", n0, n1/2+1), out_f("out_f", n0, n1/2+1); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, 1); Kokkos::deep_copy(x_ref, x); Kokkos::fence(); - // np.rfft2 is identical to np.fft(np.rfft(x, axis=1), axis=0) + // np.rfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) KokkosFFT::rfft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); KokkosFFT::fft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::rfftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::rfftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::rfftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + KokkosFFT::rfftn(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); - multiply(out_o, sqrt(static_cast(n0 * n1))); - multiply(out_f, static_cast(n0 * n1)); + 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.rfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) + using axes_type = KokkosFFT::axis_type<2>; + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out, axes_type{-2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out_b, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out_o, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out_f, axes_type{-2, -1}, KokkosFFT::FFT_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) ); @@ -658,76 +913,129 @@ TEST(FFT2D, 2DRFFT_2DLeftView) { EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } -TEST(FFT2D, 2DRFFT_2DRightView) { - const int n0 = 4, n1 = 6; +template +void test_rfftn_3dfft_3dview() { + const int n0 = 4, n1 = 6, n2 = 8; + using RealView3DType = Kokkos::View; + using ComplexView3DType = Kokkos::View***, LayoutType, execution_space>; - RightView2D x("x", n0, n1), x_ref("x_ref", n0, n1); - RightView2D > out("out", n0, n1/2+1), out1("out1", n0, n1/2+1), out2("out2", n0, n1/2+1); - RightView2D > out_b("out_b", n0, n1/2+1), out_o("out_o", n0, n1/2+1), out_f("out_f", n0, n1/2+1); + RealView3DType x("x", n0, n1, n2), x_ref("x_ref", n0, n1, n2); + ComplexView3DType out("out", n0, n1, n2/2+1), out1("out1", n0, n1, n2/2+1), out2("out2", n0, n1, n2/2+1), out3("out3", n0, n1, n2/2+1); + ComplexView3DType out_b("out_b", n0, n1, n2/2+1), out_o("out_o", n0, n1, n2/2+1), out_f("out_f", n0, n1, n2/2+1); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, 1); Kokkos::deep_copy(x_ref, x); Kokkos::fence(); - // np.rfft2 is identical to np.fft(np.rfft(x, axis=1), axis=0) - KokkosFFT::rfft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - KokkosFFT::fft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + // np.rfftn for 3D array is identical to np.fft(np.fft(np.rfft(x, axis=2), axis=1), axis=0) + KokkosFFT::rfft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); + KokkosFFT::fft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + KokkosFFT::fft(out2, out3, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::rfftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::rfftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::rfftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + KokkosFFT::rfftn(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); - multiply(out_o, sqrt(static_cast(n0 * n1))); - multiply(out_f, static_cast(n0 * n1)); + multiply(out_o, sqrt(static_cast(n0 * n1 * n2))); + multiply(out_f, 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) ); + + // 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(x, out, axes_type{-3, -2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out_b, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out_o, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(x, out_f, axes_type{-3, -2, -1}, KokkosFFT::FFT_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) ); } -TEST(FFT2D, 2DIRFFT_2DLeftView) { +template +void test_irfftn_2dfft_2dview() { const int n0 = 4, n1 = 6; + using RealView2DType = Kokkos::View; + using ComplexView2DType = Kokkos::View**, LayoutType, execution_space>; - LeftView2D > x("x", n0, n1/2+1), x_ref("x_ref", n0, n1/2+1); - LeftView2D > out1("out1", n0, n1/2+1); - LeftView2D out2("out2", n0, n1), out("out", n0, n1); - LeftView2D out_b("out_b", n0, n1), out_o("out_o", n0, n1), out_f("out_f", n0, n1); + 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); + 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); - Kokkos::fence(); + // np.irfftn for 2D array is identical to np.irfft(np.ifft(x, axis=0), axis=1) + KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + KokkosFFT::irfft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); - // np.irfft2 is identical to np.irfft(np.ifft(x, axis=0), axis=1) - KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, 0); - KokkosFFT::irfft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, 1); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::irfftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::irfftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::irfftn(x, out_f, KokkosFFT::FFT_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>; Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + KokkosFFT::irfftn(x, out, axes_type{-2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out_b, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out_o, axes_type{-2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out_f, axes_type{-2, -1}, KokkosFFT::FFT_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))); + 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) ); @@ -735,42 +1043,204 @@ TEST(FFT2D, 2DIRFFT_2DLeftView) { EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } -TEST(FFT2D, 2DIRFFT_2DRightView) { - const int n0 = 4, n1 = 6; +template +void test_irfftn_3dfft_3dview() { + const int n0 = 4, n1 = 6, n2 = 8; + using RealView3DType = Kokkos::View; + using ComplexView3DType = Kokkos::View***, LayoutType, execution_space>; - RightView2D > x("x", n0, n1/2+1), x_ref("x_ref", n0, n1/2+1); - RightView2D > out1("out1", n0, n1/2+1); - RightView2D out2("out2", n0, n1), out("out", n0, n1); - RightView2D 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); + 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); - Kokkos::fence(); + // np.irfftn for 3D array is identical to np.irfft(np.ifft(np.ifft(x, axis=0), axis=1), axis=2) + KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/0); + KokkosFFT::ifft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/1); + KokkosFFT::irfft(out2, out3, KokkosFFT::FFT_Normalization::BACKWARD, /*axis=*/2); - // np.irfft2 is identical to np.irfft(np.ifft(x, axis=0), axis=1) - KokkosFFT::ifft(x, out1, KokkosFFT::FFT_Normalization::BACKWARD, 0); - KokkosFFT::irfft(out1, out2, KokkosFFT::FFT_Normalization::BACKWARD, 1); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out); // default: KokkosFFT::FFT_Normalization::BACKWARD + KokkosFFT::irfftn(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out_b, KokkosFFT::FFT_Normalization::BACKWARD); + KokkosFFT::irfftn(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfft2(x, out_o, KokkosFFT::FFT_Normalization::ORTHO); + KokkosFFT::irfftn(x, out_f, KokkosFFT::FFT_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.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(x, x_ref); - KokkosFFT::irfft2(x, out_f, KokkosFFT::FFT_Normalization::FORWARD); + KokkosFFT::irfftn(x, out, axes_type{-3, -2, -1}); // default: KokkosFFT::FFT_Normalization::BACKWARD - multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); - multiply(out_f, 1.0/static_cast(n0 * n1)); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out_b, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::BACKWARD); - 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) ); + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out_o, axes_type{-3, -2, -1}, KokkosFFT::FFT_Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(x, out_f, axes_type{-3, -2, -1}, KokkosFFT::FFT_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) ); +} + +// fftn on 2D Views +TEST(FFTN, 2DFFT_2DLeftViewFloat) { + test_fftn_2dfft_2dview(); +} + +TEST(FFTN, 2DFFT_2DLeftViewDouble) { + test_fftn_2dfft_2dview(); +} + +TEST(FFTN, 2DFFT_2DRightViewFloat) { + test_fftn_2dfft_2dview(); +} + +TEST(FFTN, 2DFFT_2DRightViewDouble) { + test_fftn_2dfft_2dview(); +} + +// fftn on 3D Views +TEST(FFTN, 3DFFT_3DLeftViewFloat) { + test_fftn_3dfft_3dview(); +} + +TEST(FFTN, 3DFFT_3DLeftViewDouble) { + test_fftn_3dfft_3dview(); +} + +TEST(FFTN, 3DFFT_3DRightViewFloat) { + test_fftn_3dfft_3dview(); +} + +TEST(FFTN, 3DFFT_3DRightViewDouble) { + test_fftn_3dfft_3dview(); +} + +// ifftn on 2D Views +TEST(FFTN, 2DIFFT_2DLeftViewFloat) { + test_ifftn_2dfft_2dview(); +} + +TEST(FFTN, 2DIFFT_2DLeftViewDouble) { + test_ifftn_2dfft_2dview(); +} + +TEST(FFTN, 2DIFFT_2DRightViewFloat) { + test_ifftn_2dfft_2dview(); +} + +TEST(FFTN, 2DIFFT_2DRightViewDouble) { + test_ifftn_2dfft_2dview(); +} + +// ifftn on 3D Views +TEST(FFTN, 3DIFFT_3DLeftViewFloat) { + test_ifftn_3dfft_3dview(); +} + +TEST(FFTN, 3DIFFT_3DLeftViewDouble) { + test_ifftn_3dfft_3dview(); +} + +TEST(FFTN, 3DIFFT_3DRightViewFloat) { + test_ifftn_3dfft_3dview(); +} + +TEST(FFTN, 3DIFFT_3DRightViewDouble) { + test_ifftn_3dfft_3dview(); +} + +// rfftn on 2D Views +TEST(FFTN, 2DRFFT_2DLeftViewFloat) { + test_rfftn_2dfft_2dview(); +} + +TEST(FFTN, 2DRFFT_2DLeftViewDouble) { + test_rfftn_2dfft_2dview(); +} + +TEST(FFTN, 2DRFFT_2DRightViewFloat) { + test_rfftn_2dfft_2dview(); +} + +TEST(FFTN, 2DRFFT_2DRightViewDouble) { + test_rfftn_2dfft_2dview(); +} + +// rfftn on 3D Views +TEST(FFTN, 3DRFFT_3DLeftViewFloat) { + test_rfftn_3dfft_3dview(); +} + +TEST(FFTN, 3DRFFT_3DLeftViewDouble) { + test_rfftn_3dfft_3dview(); +} + +TEST(FFTN, 3DRFFT_3DRightViewFloat) { + test_rfftn_3dfft_3dview(); +} + +TEST(FFTN, 3DRFFT_3DRightViewDouble) { + test_rfftn_3dfft_3dview(); +} + +// irfftn on 2D Views +TEST(FFTN, 2DIRFFT_2DLeftViewFloat) { + test_irfftn_2dfft_2dview(); +} + +TEST(FFTN, 2DIRFFT_2DLeftViewDouble) { + test_irfftn_2dfft_2dview(); +} + +TEST(FFTN, 2DIRFFT_2DRightViewFloat) { + test_irfftn_2dfft_2dview(); +} + +TEST(FFTN, 2DIRFFT_2DRightViewDouble) { + test_irfftn_2dfft_2dview(); +} + +// irfftn on 3D Views +TEST(FFTN, 3DIRFFT_3DLeftViewFloat) { + test_irfftn_3dfft_3dview(); +} + +TEST(FFTN, 3DIRFFT_3DLeftViewDouble) { + test_irfftn_3dfft_3dview(); +} + +TEST(FFTN, 3DIRFFT_3DRightViewFloat) { + test_irfftn_3dfft_3dview(); +} + +TEST(FFTN, 3DIRFFT_3DRightViewDouble) { + test_irfftn_3dfft_3dview(); } \ No newline at end of file