From f3a3385091d0808e2b67474a4184e38e68da8568 Mon Sep 17 00:00:00 2001 From: Thomas Padioleau Date: Fri, 11 Oct 2024 15:43:01 +0200 Subject: [PATCH] Rework internal allocations (#167) * Rework internal allocations * Update common/src/KokkosFFT_transpose.hpp Co-authored-by: yasahi-hpc <57478230+yasahi-hpc@users.noreply.github.com> * Revert example --------- Co-authored-by: yasahi-hpc <57478230+yasahi-hpc@users.noreply.github.com> --- common/src/KokkosFFT_transpose.hpp | 85 ++---- common/src/KokkosFFT_utils.hpp | 139 +--------- common/unit_test/CMakeLists.txt | 1 - common/unit_test/Test_Transpose.cpp | 84 +++--- common/unit_test/Test_prep_transpose_view.cpp | 258 ------------------ fft/src/KokkosFFT_Transform.hpp | 9 +- 6 files changed, 87 insertions(+), 489 deletions(-) delete mode 100644 common/unit_test/Test_prep_transpose_view.cpp diff --git a/common/src/KokkosFFT_transpose.hpp b/common/src/KokkosFFT_transpose.hpp index e345258c..dd78e142 100644 --- a/common/src/KokkosFFT_transpose.hpp +++ b/common/src/KokkosFFT_transpose.hpp @@ -75,41 +75,24 @@ auto get_map_axes(const ViewType& view, int axis) { return get_map_axes(view, axis_type<1>({axis})); } -template -void prep_transpose_view(InViewType& in, OutViewType& out, - axis_type map) { - constexpr int rank = OutViewType::rank(); - - // Assign a View if not a shallow copy - bool is_out_view_ready = true; - std::array out_extents; - for (int i = 0; i < rank; i++) { - out_extents.at(i) = in.extent(map.at(i)); - if (static_cast(out_extents.at(i)) != out.extent(i)) { - is_out_view_ready = false; - } +template +axis_type compute_transpose_extents( + ViewType const& view, axis_type const& map) { + static_assert(Kokkos::is_view_v, + "compute_transpose_extents: ViewType must be a Kokkos::View."); + constexpr std::size_t rank = ViewType::rank(); + + axis_type out_extents; + for (std::size_t i = 0; i < rank; ++i) { + out_extents.at(i) = view.extent(map.at(i)); } - if constexpr (std::is_const_v) { - KOKKOSFFT_THROW_IF( - !is_out_view_ready, - "prep_transpose_view: OutViewType is const, but does not " - "have the required extents"); - } else { - if (!is_out_view_ready) { - if constexpr (!OutViewType::memory_traits::is_unmanaged) { - KokkosFFT::Impl::create_view(out, "out", out_extents); - } else { - // try to reshape out if it currently has enough memory available - KokkosFFT::Impl::reshape_view(out, out_extents); - } - } - } + return out_extents; } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<2> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<2> /*_map*/) { constexpr std::size_t DIM = 2; using range_type = Kokkos::MDRangePolicy< @@ -125,16 +108,14 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::parallel_for( "KokkosFFT::transpose", range, KOKKOS_LAMBDA(int i0, int i1) { out(i1, i0) = in(i0, i1); }); } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<3> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<3> _map) { constexpr std::size_t DIM = 3; constexpr std::size_t rank = InViewType::rank(); @@ -151,8 +132,6 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, tile_type{{4, 4, 4}} // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::Array map = {_map[0], _map[1], _map[2]}; Kokkos::parallel_for( "KokkosFFT::transpose", range, KOKKOS_LAMBDA(int i0, int i1, int i2) { @@ -166,8 +145,8 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<4> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<4> _map) { constexpr std::size_t DIM = 4; constexpr std::size_t rank = InViewType::rank(); @@ -185,8 +164,6 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::Array map = {_map[0], _map[1], _map[2], _map[3]}; Kokkos::parallel_for( "KokkosFFT::transpose", range, @@ -202,8 +179,8 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<5> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<5> _map) { constexpr std::size_t DIM = 5; constexpr std::size_t rank = InViewType::rank(); @@ -222,8 +199,6 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::Array map = {_map[0], _map[1], _map[2], _map[3], _map[4]}; Kokkos::parallel_for( "KokkosFFT::transpose", range, @@ -240,8 +215,8 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<6> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<6> _map) { constexpr std::size_t DIM = 6; constexpr std::size_t rank = InViewType::rank(); @@ -261,8 +236,6 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::Array map = {_map[0], _map[1], _map[2], _map[3], _map[4], _map[5]}; Kokkos::parallel_for( @@ -281,8 +254,8 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<7> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<7> _map) { constexpr std::size_t DIM = 6; constexpr std::size_t rank = InViewType::rank(); @@ -302,8 +275,6 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::Array map = {_map[0], _map[1], _map[2], _map[3], _map[4], _map[5], _map[6]}; Kokkos::parallel_for( @@ -326,8 +297,8 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, } template -void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type<8> _map) { +void transpose_impl(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type<8> _map) { constexpr std::size_t DIM = 6; constexpr std::size_t rank = InViewType::rank(); @@ -349,8 +320,6 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, // [TO DO] Choose optimal tile sizes for each device ); - prep_transpose_view(in, out, _map); - Kokkos::Array map = {_map[0], _map[1], _map[2], _map[3], _map[4], _map[5], _map[6], _map[7]}; Kokkos::parallel_for( @@ -396,8 +365,8 @@ void transpose_impl(const ExecutionSpace& exec_space, InViewType& in, */ template -void transpose(const ExecutionSpace& exec_space, InViewType& in, - OutViewType& out, axis_type map) { +void transpose(const ExecutionSpace& exec_space, const InViewType& in, + const OutViewType& out, axis_type map) { static_assert( KokkosFFT::Impl::are_operatable_views_v, diff --git a/common/src/KokkosFFT_utils.hpp b/common/src/KokkosFFT_utils.hpp index 6e2ca757..1de7b7b8 100644 --- a/common/src/KokkosFFT_utils.hpp +++ b/common/src/KokkosFFT_utils.hpp @@ -204,136 +204,15 @@ auto extract_extents(const ViewType& view) { return extents; } -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1], extents[2]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1], extents[2], extents[3]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1], extents[2], extents[3], - extents[4]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5], extents[6]); -} - -template -void create_view(ViewType& out, const Label& label, - const std::array& extents) { - out = ViewType(label, extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5], extents[6], extents[7]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF(ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size(extents[0]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF( - ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size(extents[0], extents[1]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0], extents[1]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF(ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size( - extents[0], extents[1], extents[2]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0], extents[1], extents[2]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF(ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size( - extents[0], extents[1], extents[2], extents[3]), - "reshape_view: insufficient memory"); - - out = ViewType(out.data(), extents[0], extents[1], extents[2], extents[3]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF( - ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size(extents[0], extents[1], extents[2], - extents[3], extents[4]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0], extents[1], extents[2], extents[3], - extents[4]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF(ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size( - extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF(ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size( - extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5], extents[6]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5], extents[6]); -} - -template -void reshape_view(ViewType& out, const std::array& extents) { - KOKKOSFFT_THROW_IF(ViewType::required_allocation_size(out.layout()) < - ViewType::required_allocation_size( - extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5], extents[6], extents[7]), - "reshape_view: insufficient memory"); - out = ViewType(out.data(), extents[0], extents[1], extents[2], extents[3], - extents[4], extents[5], extents[6], extents[7]); +template +Layout create_layout(const std::array& extents) { + static_assert(std::is_same_v || + std::is_same_v, + "create_layout: Layout must be either Kokkos::LayoutLeft or " + "Kokkos::LayoutRight."); + Layout layout; + std::copy_n(extents.begin(), N, layout.dimension); + return layout; } } // namespace Impl diff --git a/common/unit_test/CMakeLists.txt b/common/unit_test/CMakeLists.txt index db66d67b..12f2c704 100644 --- a/common/unit_test/CMakeLists.txt +++ b/common/unit_test/CMakeLists.txt @@ -11,7 +11,6 @@ add_executable(unit-tests-kokkos-fft-common Test_Layouts.cpp Test_Padding.cpp Test_Helpers.cpp - Test_prep_transpose_view.cpp ) target_compile_features(unit-tests-kokkos-fft-common PUBLIC cxx_std_17) diff --git a/common/unit_test/Test_Transpose.cpp b/common/unit_test/Test_Transpose.cpp index 0999245c..a59d7de9 100644 --- a/common/unit_test/Test_Transpose.cpp +++ b/common/unit_test/Test_Transpose.cpp @@ -339,7 +339,7 @@ void test_transpose_1d_1dview() { // When transpose is not necessary, we should not call transpose method const int len = 30; View1D x("x", len), ref("ref", len); - View1D xt; + View1D xt("xt", len); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, 1.0); @@ -377,7 +377,7 @@ void test_transpose_1d_2dview() { } auto [_n0, _n1] = out_extents; - RealView2Dtype xt; + RealView2Dtype xt("xt", _n0, _n1); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -401,7 +401,7 @@ void test_transpose_1d_2dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView2Dtype _x; + RealView2Dtype _x("_x", n0, n1); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -432,7 +432,7 @@ void test_transpose_1d_3dview() { } auto [_n0, _n1, _n2] = out_extents; - RealView3Dtype xt; + RealView3Dtype xt("xt", _n0, _n1, _n2); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -462,7 +462,7 @@ void test_transpose_1d_3dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView3Dtype _x; + RealView3Dtype _x("_x", n0, n1, n2); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -493,7 +493,7 @@ void test_transpose_1d_4dview() { } auto [_n0, _n1, _n2, _n3] = out_extents; - RealView4Dtype xt; + RealView4Dtype xt("xt", _n0, _n1, _n2, _n3); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -539,7 +539,7 @@ void test_transpose_1d_4dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView4Dtype _x; + RealView4Dtype _x("_x", n0, n1, n2, n3); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -570,7 +570,7 @@ void test_transpose_1d_5dview() { } auto [_n0, _n1, _n2, _n3, _n4] = out_extents; - RealView5Dtype xt; + RealView5Dtype xt("xt", _n0, _n1, _n2, _n3, _n4); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -627,7 +627,7 @@ void test_transpose_1d_5dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView5Dtype _x; + RealView5Dtype _x("_x", n0, n1, n2, n3, n4); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -659,7 +659,7 @@ void test_transpose_1d_6dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5] = out_extents; - RealView6Dtype xt; + RealView6Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -731,7 +731,7 @@ void test_transpose_1d_6dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView6Dtype _x; + RealView6Dtype _x("_x", n0, n1, n2, n3, n4, n5); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -763,7 +763,7 @@ void test_transpose_1d_7dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = out_extents; - RealView7Dtype xt; + RealView7Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5, _n6); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -851,7 +851,7 @@ void test_transpose_1d_7dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView7Dtype _x; + RealView7Dtype _x("_x", n0, n1, n2, n3, n4, n5, n6); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -883,7 +883,7 @@ void test_transpose_1d_8dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = out_extents; - RealView8Dtype xt; + RealView8Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -989,7 +989,7 @@ void test_transpose_1d_8dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView8Dtype _x; + RealView8Dtype _x("_x", n0, n1, n2, n3, n4, n5, n6, n7); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1049,7 +1049,7 @@ void test_transpose_2d_2dview() { using RealView2Dtype = Kokkos::View; const int n0 = 3, n1 = 5; RealView2Dtype x("x", n0, n1), _x("_x", n0, n1), ref("ref", n1, n0); - RealView2Dtype xt_axis01, xt_axis10; // views are allocated internally + RealView2Dtype xt_axis01("xt_axis01", n0, n1), xt_axis10("xt_axis10", n1, n0); Kokkos::Random_XorShift64_Pool<> random_pool(12345); Kokkos::fill_random(x, random_pool, 1.0); @@ -1110,7 +1110,7 @@ void test_transpose_2d_3dview() { } auto [_n0, _n1, _n2] = out_extents; - RealView3Dtype xt; + RealView3Dtype xt("xt", _n0, _n1, _n2); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1141,7 +1141,7 @@ void test_transpose_2d_3dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView3Dtype _x; + RealView3Dtype _x("_x", n0, n1, n2); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1177,7 +1177,7 @@ void test_transpose_2d_4dview() { } auto [_n0, _n1, _n2, _n3] = out_extents; - RealView4Dtype xt; + RealView4Dtype xt("xt", _n0, _n1, _n2, _n3); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1224,7 +1224,7 @@ void test_transpose_2d_4dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView4Dtype _x; + RealView4Dtype _x("_x", n0, n1, n2, n3); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1260,7 +1260,7 @@ void test_transpose_2d_5dview() { } auto [_n0, _n1, _n2, _n3, _n4] = out_extents; - RealView5Dtype xt; + RealView5Dtype xt("xt", _n0, _n1, _n2, _n3, _n4); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1318,7 +1318,7 @@ void test_transpose_2d_5dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView5Dtype _x; + RealView5Dtype _x("_x", n0, n1, n2, n3, n4); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1355,7 +1355,7 @@ void test_transpose_2d_6dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5] = out_extents; - RealView6Dtype xt; + RealView6Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1428,7 +1428,7 @@ void test_transpose_2d_6dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView6Dtype _x; + RealView6Dtype _x("_x", n0, n1, n2, n3, n4, n5); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1465,7 +1465,7 @@ void test_transpose_2d_7dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = out_extents; - RealView7Dtype xt; + RealView7Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5, _n6); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1554,7 +1554,7 @@ void test_transpose_2d_7dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView7Dtype _x; + RealView7Dtype _x("_x", n0, n1, n2, n3, n4, n5, n6); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1591,7 +1591,7 @@ void test_transpose_2d_8dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = out_extents; - RealView8Dtype xt; + RealView8Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1698,7 +1698,7 @@ void test_transpose_2d_8dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView8Dtype _x; + RealView8Dtype _x("_x", n0, n1, n2, n3, n4, n5, n6, n7); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1753,8 +1753,12 @@ void test_transpose_3d_3dview() { using RealView3Dtype = Kokkos::View; const int n0 = 3, n1 = 5, n2 = 8; RealView3Dtype x("x", n0, n1, n2); - RealView3Dtype xt_axis012, xt_axis021, xt_axis102, xt_axis120, xt_axis201, - xt_axis210; // views are allocated internally + RealView3Dtype xt_axis012("xt_axis012", n0, n1, n2), + xt_axis021("xt_axis021", n0, n2, n1), + xt_axis102("xt_axis102", n1, n0, n2), + xt_axis120("xt_axis120", n1, n2, n0), + xt_axis201("xt_axis201", n2, n0, n1), + xt_axis210("xt_axis210", n2, n1, n0); // views are allocated internally RealView3Dtype ref_axis021("ref_axis021", n0, n2, n1), ref_axis102("ref_axis102", n1, n0, n2); RealView3Dtype ref_axis120("ref_axis120", n1, n2, n0), @@ -1854,7 +1858,7 @@ void test_transpose_3d_4dview() { } auto [_n0, _n1, _n2, _n3] = out_extents; - RealView4Dtype xt; + RealView4Dtype xt("xt", _n0, _n1, _n2, _n3); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -1902,7 +1906,7 @@ void test_transpose_3d_4dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView4Dtype _x; + RealView4Dtype _x("_x", n0, n1, n2, n3); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -1941,7 +1945,7 @@ void test_transpose_3d_5dview() { } auto [_n0, _n1, _n2, _n3, _n4] = out_extents; - RealView5Dtype xt; + RealView5Dtype xt("xt", _n0, _n1, _n2, _n3, _n4); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -2000,7 +2004,7 @@ void test_transpose_3d_5dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView5Dtype _x; + RealView5Dtype _x("_x", n0, n1, n2, n3, n4); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -2040,7 +2044,7 @@ void test_transpose_3d_6dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5] = out_extents; - RealView6Dtype xt; + RealView6Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -2114,7 +2118,7 @@ void test_transpose_3d_6dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView6Dtype _x; + RealView6Dtype _x("_x", n0, n1, n2, n3, n4, n5); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -2154,7 +2158,7 @@ void test_transpose_3d_7dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6] = out_extents; - RealView7Dtype xt; + RealView7Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5, _n6); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -2244,7 +2248,7 @@ void test_transpose_3d_7dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView7Dtype _x; + RealView7Dtype _x("_x", n0, n1, n2, n3, n4, n5, n6); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } @@ -2284,7 +2288,7 @@ void test_transpose_3d_8dview() { } auto [_n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7] = out_extents; - RealView8Dtype xt; + RealView8Dtype xt("xt", _n0, _n1, _n2, _n3, _n4, _n5, _n6, _n7); if (map == default_axes) { EXPECT_THROW(KokkosFFT::Impl::transpose(execution_space(), x, xt, map), // xt is identical to x @@ -2393,7 +2397,7 @@ void test_transpose_3d_8dview() { EXPECT_TRUE(allclose(xt, ref, 1.e-5, 1.e-12)); // Inverse (transpose of transpose is identical to the original) - RealView8Dtype _x; + RealView8Dtype _x("_x", n0, n1, n2, n3, n4, n5, n6, n7); KokkosFFT::Impl::transpose(execution_space(), xt, _x, map_inv); EXPECT_TRUE(allclose(_x, x, 1.e-5, 1.e-12)); } diff --git a/common/unit_test/Test_prep_transpose_view.cpp b/common/unit_test/Test_prep_transpose_view.cpp deleted file mode 100644 index 68a7626c..00000000 --- a/common/unit_test/Test_prep_transpose_view.cpp +++ /dev/null @@ -1,258 +0,0 @@ -// SPDX-FileCopyrightText: (C) The Kokkos-FFT development team, see COPYRIGHT.md file -// -// SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception - -#include -#include -#include -#include -#include "KokkosFFT_transpose.hpp" -#include "Test_Types.hpp" -#include "Test_Utils.hpp" - -template -void test_managed_prep_transpose_view() { - constexpr int DIMS = InViewType::rank(); - static_assert(InViewType::rank() == OutViewType::rank(), - "input and output have different ranks"); - - using InManagedViewType = - typename KokkosFFT::Impl::manageable_view_type::type; - using OutManagedViewType = - typename KokkosFFT::Impl::manageable_view_type::type; - static_assert(!InManagedViewType::memory_traits::is_unmanaged, - "Unable to get managed input view type"); - static_assert(!OutManagedViewType::memory_traits::is_unmanaged, - "Unable to get managed output view type"); - - static_assert(!InViewType::memory_traits::is_unmanaged, - "Unable to get managed input view type"); - static_assert(!OutViewType::memory_traits::is_unmanaged, - "Unable to get managed output view type"); - using LayoutType = typename InViewType::array_layout; - - // no need to allocate - { - LayoutType layout; - KokkosFFT::axis_type map; - for (int i = 0; i < DIMS; ++i) { - layout.dimension[i] = 5; - map[i] = i; - } - InViewType in("in", layout); - OutViewType out("out", layout); - auto data_prev = out.data(); - KokkosFFT::Impl::prep_transpose_view(in, out, map); - // ensure no allocation - EXPECT_EQ(data_prev, out.data()); - // check shape - for (int i = 0; i < DIMS; ++i) { - EXPECT_EQ(out.extent(i), 5); - } - } - // allocate - { - LayoutType layout; - KokkosFFT::axis_type map; - for (int i = 0; i < DIMS; ++i) { - layout.dimension[i] = 5; - map[i] = i; - } - InViewType in("in", layout); - OutViewType out; - KokkosFFT::Impl::prep_transpose_view(in, out, map); - // check shape - for (int i = 0; i < DIMS; ++i) { - EXPECT_EQ(out.extent(i), 5); - } - } -} - -template -void test_unmanaged_prep_transpose_view() { - constexpr int DIMS = InViewType::rank(); - static_assert(InViewType::rank() == OutViewType::rank(), - "input and output have different ranks"); - - using InManagedViewType = - typename KokkosFFT::Impl::manageable_view_type::type; - using OutManagedViewType = - typename KokkosFFT::Impl::manageable_view_type::type; - static_assert(!InManagedViewType::memory_traits::is_unmanaged, - "Unable to get managed input view type"); - static_assert(!OutManagedViewType::memory_traits::is_unmanaged, - "Unable to get managed output view type"); - using LayoutType = typename InViewType::array_layout; - - // no need to reshape - { - LayoutType layout; - KokkosFFT::axis_type map; - for (int i = 0; i < DIMS; ++i) { - layout.dimension[i] = 5; - map[i] = i; - } - InManagedViewType in("in", layout); - OutManagedViewType out("out", layout); - OutViewType u_out(out.data(), layout); - auto data_prev = out.data(); - KokkosFFT::Impl::prep_transpose_view(in, u_out, map); - EXPECT_EQ(data_prev, u_out.data()); - // check shape - for (int i = 0; i < DIMS; ++i) { - EXPECT_EQ(u_out.extent(i), 5); - } - } - // reshape success - { - LayoutType layout; - KokkosFFT::axis_type map; - for (int i = 0; i < DIMS; ++i) { - layout.dimension[i] = 5; - map[i] = i; - } - - LayoutType layout_orig; - layout_orig.dimension[0] = 5; - if (DIMS == 0) { - // give the 1D version a larger original shape so it will reshape down to - // a portion of the allocation - layout_orig.dimension[0] = 10; - } - for (int i = 1; i < DIMS; ++i) { - layout_orig.dimension[0] *= 5; - layout_orig.dimension[i] = 1; - } - InManagedViewType in("in", layout); - OutManagedViewType out("out", layout_orig); - OutViewType u_out(out.data(), layout_orig); - KokkosFFT::Impl::prep_transpose_view(in, u_out, map); - // check shape - for (int i = 0; i < DIMS; ++i) { - EXPECT_EQ(u_out.extent(i), 5); - } - } - // reshape failure - { - LayoutType layout; - KokkosFFT::axis_type map; - for (int i = 0; i < DIMS; ++i) { - layout.dimension[i] = 5; - map[i] = i; - } - - LayoutType layout_orig; - for (int i = 0; i < DIMS; ++i) { - layout_orig.dimension[i] = 1; - } - InManagedViewType in("in", layout); - OutManagedViewType out("out", layout_orig); - OutViewType u_out(out.data(), layout_orig); - EXPECT_THROW(KokkosFFT::Impl::prep_transpose_view(in, u_out, map), - std::runtime_error); - } -} - -TEST(prep_transpose_view, 1DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 1DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 2DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 2DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 3DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 3DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 4DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 4DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 5DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 5DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 6DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 6DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 7DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 7DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} - -TEST(prep_transpose_view, 8DManaged) { - test_managed_prep_transpose_view< - Kokkos::View, - Kokkos::View>(); -} - -TEST(prep_transpose_view, 8DUnmanaged) { - test_unmanaged_prep_transpose_view< - Kokkos::View, - Kokkos::View>>(); -} diff --git a/fft/src/KokkosFFT_Transform.hpp b/fft/src/KokkosFFT_Transform.hpp index b4e6e52c..7a2a1350 100644 --- a/fft/src/KokkosFFT_Transform.hpp +++ b/fft/src/KokkosFFT_Transform.hpp @@ -99,8 +99,13 @@ void fft_exec_impl( } if (plan.is_transpose_needed()) { - ManagableInViewType in_T; - ManagableOutViewType out_T; + using LayoutType = typename ManagableInViewType::array_layout; + ManagableInViewType const in_T( + "in_T", + create_layout(compute_transpose_extents(_in, plan.map()))); + ManagableOutViewType const out_T( + "out_T", + create_layout(compute_transpose_extents(out, plan.map()))); KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map());