From 9fca6b1df0c7194385d79b5effbf9fdb755a7ff5 Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Tue, 27 Feb 2024 21:36:41 +0900 Subject: [PATCH] tests of 3D FFT over 8D --- fft/unit_test/Test_Transform.cpp | 414 +++++++++++++++++++++++++++++++ 1 file changed, 414 insertions(+) diff --git a/fft/unit_test/Test_Transform.cpp b/fft/unit_test/Test_Transform.cpp index acf2f435..05cc2001 100644 --- a/fft/unit_test/Test_Transform.cpp +++ b/fft/unit_test/Test_Transform.cpp @@ -2798,6 +2798,376 @@ void test_irfftn_3dfft_3dview() { EXPECT_TRUE(allclose(out_f, out3, 1.e-5, 1.e-6)); } +template +void test_fftn_3dfft_4dview(T atol = 1.e-12) { + const int n0 = 10, n1 = 6, n2 = 8, n3 = 5; + using RealView4DType = Kokkos::View; + 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); + + using axes_type = KokkosFFT::axis_type<3>; + + 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 axes = {axis0, axis1, axis2}; + + std::array shape_c2r = shape; + shape_c2r.at(axis2) = shape_c2r.at(axis2) / 2 + 1; + + 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); + + 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); + + Kokkos::deep_copy(ref_x, x); + Kokkos::deep_copy(ref_xr, xr); + + Kokkos::fence(); + + // Along one axis + // Simple identity tests + KokkosFFT::fftn(execution_space(), x, out, axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::ifftn(execution_space(), out, _x, axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfftn(execution_space(), xr, outr, axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::irfftn(execution_space(), outr, _xr, axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } + } + } +} + +template +void test_fftn_3dfft_5dview(T atol = 1.e-12) { + const int n0 = 10, n1 = 6, n2 = 8, n3 = 5, n4 = 4; + using RealView5DType = Kokkos::View; + 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); + + using axes_type = KokkosFFT::axis_type<3>; + KokkosFFT::axis_type default_axes({0, 1, 2, 3, 4}); + + // Too many combinations, choose axes randomly + std::vector list_of_tested_axes; + + constexpr int nb_trials = 32; + auto rng = std::default_random_engine{}; + + for (int i = 0; i < nb_trials; i++) { + auto tmp_axes = default_axes; + std::shuffle(std::begin(tmp_axes), std::end(tmp_axes), rng); + + // pickup 3 elements only + axes_type trimed_axes; + std::copy(std::begin(tmp_axes) + DIM - 3, std::end(tmp_axes), + std::begin(trimed_axes)); + list_of_tested_axes.push_back(trimed_axes); + } + + for (auto& tested_axes : list_of_tested_axes) { + int last_axis = tested_axes.at(2); + std::array shape_c2r = shape; + shape_c2r.at(last_axis) = shape_c2r.at(last_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); + + 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); + + Kokkos::deep_copy(ref_x, x); + Kokkos::deep_copy(ref_xr, xr); + + Kokkos::fence(); + + // Along one axis + // Simple identity tests + KokkosFFT::fftn(execution_space(), x, out, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::ifftn(execution_space(), out, _x, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfftn(execution_space(), xr, outr, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::irfftn(execution_space(), outr, _xr, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } +} + +template +void test_fftn_3dfft_6dview(T atol = 1.e-12) { + const int n0 = 2, n1 = 3, n2 = 4, n3 = 5, n4 = 6, n5 = 7; + using RealView6DType = Kokkos::View; + 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); + + using axes_type = KokkosFFT::axis_type<3>; + KokkosFFT::axis_type default_axes({0, 1, 2, 3, 4, 5}); + + // Too many combinations, choose axes randomly + std::vector list_of_tested_axes; + + constexpr int nb_trials = 32; + auto rng = std::default_random_engine{}; + + for (int i = 0; i < nb_trials; i++) { + auto tmp_axes = default_axes; + std::shuffle(std::begin(tmp_axes), std::end(tmp_axes), rng); + + // pickup 3 elements only + axes_type trimed_axes; + std::copy(std::begin(tmp_axes) + DIM - 3, std::end(tmp_axes), + std::begin(trimed_axes)); + list_of_tested_axes.push_back(trimed_axes); + } + + for (auto& tested_axes : list_of_tested_axes) { + int last_axis = tested_axes.at(2); + std::array shape_c2r = shape; + shape_c2r.at(last_axis) = shape_c2r.at(last_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); + + 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); + + Kokkos::deep_copy(ref_x, x); + Kokkos::deep_copy(ref_xr, xr); + + Kokkos::fence(); + + // Along one axis + // Simple identity tests + KokkosFFT::fftn(execution_space(), x, out, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::ifftn(execution_space(), out, _x, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfftn(execution_space(), xr, outr, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::irfftn(execution_space(), outr, _xr, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } +} + +template +void test_fftn_3dfft_7dview(T atol = 1.e-12) { + const int n0 = 2, n1 = 3, n2 = 4, n3 = 5, n4 = 6, n5 = 7, n6 = 8; + using RealView7DType = Kokkos::View; + 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); + + using axes_type = KokkosFFT::axis_type<3>; + KokkosFFT::axis_type default_axes({0, 1, 2, 3, 4, 5, 6}); + + // Too many combinations, choose axes randomly + std::vector list_of_tested_axes; + + constexpr int nb_trials = 32; + auto rng = std::default_random_engine{}; + + for (int i = 0; i < nb_trials; i++) { + auto tmp_axes = default_axes; + std::shuffle(std::begin(tmp_axes), std::end(tmp_axes), rng); + + // pickup 3 elements only + axes_type trimed_axes; + std::copy(std::begin(tmp_axes) + DIM - 3, std::end(tmp_axes), + std::begin(trimed_axes)); + list_of_tested_axes.push_back(trimed_axes); + } + + for (auto& tested_axes : list_of_tested_axes) { + int last_axis = tested_axes.at(2); + std::array shape_c2r = shape; + shape_c2r.at(last_axis) = shape_c2r.at(last_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); + + 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); + + Kokkos::deep_copy(ref_x, x); + Kokkos::deep_copy(ref_xr, xr); + + Kokkos::fence(); + + // Along one axis + // Simple identity tests + KokkosFFT::fftn(execution_space(), x, out, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::ifftn(execution_space(), out, _x, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfftn(execution_space(), xr, outr, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::irfftn(execution_space(), outr, _xr, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } +} + +template +void test_fftn_3dfft_8dview(T atol = 1.e-12) { + const int n0 = 2, n1 = 3, n2 = 4, n3 = 5, n4 = 6, n5 = 7, n6 = 8, n7 = 9; + 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); + + using axes_type = KokkosFFT::axis_type<3>; + KokkosFFT::axis_type default_axes({0, 1, 2, 3, 4, 5, 6, 7}); + + // Too many combinations, choose axes randomly + std::vector list_of_tested_axes; + + constexpr int nb_trials = 32; + auto rng = std::default_random_engine{}; + + for (int i = 0; i < nb_trials; i++) { + auto tmp_axes = default_axes; + std::shuffle(std::begin(tmp_axes), std::end(tmp_axes), rng); + + // pickup 3 elements only + axes_type trimed_axes; + std::copy(std::begin(tmp_axes) + DIM - 3, std::end(tmp_axes), + std::begin(trimed_axes)); + list_of_tested_axes.push_back(trimed_axes); + } + + for (auto& tested_axes : list_of_tested_axes) { + int last_axis = tested_axes.at(2); + std::array shape_c2r = shape; + shape_c2r.at(last_axis) = shape_c2r.at(last_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); + + 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); + + Kokkos::deep_copy(ref_x, x); + Kokkos::deep_copy(ref_xr, xr); + + Kokkos::fence(); + + // Along one axis + // Simple identity tests + KokkosFFT::fftn(execution_space(), x, out, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::ifftn(execution_space(), out, _x, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_x, ref_x, 1.e-5, atol)); + + // Simple identity tests for r2c and c2r transforms + KokkosFFT::rfftn(execution_space(), xr, outr, tested_axes, + KokkosFFT::Normalization::backward); + + KokkosFFT::irfftn(execution_space(), outr, _xr, tested_axes, + KokkosFFT::Normalization::backward); + + EXPECT_TRUE(allclose(_xr, ref_xr, 1.e-5, atol)); + } +} + // fftn on 2D Views TYPED_TEST(FFTND, 2DFFT_2DView) { using float_type = typename TestFixture::float_type; @@ -2861,4 +3231,48 @@ TYPED_TEST(FFTND, 3DIRFFT_3DView) { using layout_type = typename TestFixture::layout_type; test_irfftn_3dfft_3dview(); +} +// batched fftn on 4D Views +TYPED_TEST(FFTND, 3DFFT_batched_4DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + float_type atol = std::is_same_v ? 1.0e-5 : 1.0e-10; + test_fftn_3dfft_4dview(atol); +} + +// batched fftn on 5D Views +TYPED_TEST(FFTND, 3DFFT_batched_5DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + float_type atol = std::is_same_v ? 1.0e-5 : 1.0e-10; + test_fftn_3dfft_5dview(atol); +} + +// batched fftn on 6D Views +TYPED_TEST(FFTND, 3DFFT_batched_6DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + float_type atol = std::is_same_v ? 1.0e-5 : 1.0e-10; + test_fftn_3dfft_6dview(atol); +} + +// batched fftn on 7D Views +TYPED_TEST(FFTND, 3DFFT_batched_7DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + float_type atol = std::is_same_v ? 1.0e-5 : 1.0e-10; + test_fftn_3dfft_7dview(atol); +} + +// batched fftn on 8D Views +TYPED_TEST(FFTND, 3DFFT_batched_8DView) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + float_type atol = std::is_same_v ? 1.0e-5 : 1.0e-10; + test_fftn_3dfft_8dview(atol); } \ No newline at end of file