From 126510fc249ba7764ca2d6f03f01c45d4b71de1e Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 8 Jan 2024 16:30:20 +0900 Subject: [PATCH 1/3] Reuse plan capability --- fft/src/KokkosFFT_Plans.hpp | 33 +- fft/src/KokkosFFT_Transform.hpp | 496 +++++++++++++++++++++++++++++++ fft/unit_test/Test_Transform.cpp | 344 ++++++++++++++++++++- 3 files changed, 854 insertions(+), 19 deletions(-) diff --git a/fft/src/KokkosFFT_Plans.hpp b/fft/src/KokkosFFT_Plans.hpp index 75fb5ceb..d303cd08 100644 --- a/fft/src/KokkosFFT_Plans.hpp +++ b/fft/src/KokkosFFT_Plans.hpp @@ -33,6 +33,7 @@ namespace KokkosFFT { namespace Impl { template class Plan { + using execSpace = ExecutionSpace; using in_value_type = typename InViewType::non_const_value_type; using out_value_type = typename OutViewType::non_const_value_type; using float_type = KokkosFFT::Impl::real_type_t; @@ -46,24 +47,27 @@ namespace Impl { fft_size_type m_fft_size; map_type m_map, m_map_inv; bool m_is_transpose_needed; + axis_type m_axes; + KokkosFFT::Impl::Direction m_direction; // Only used when transpose needed nonConstInViewType m_in_T; nonConstOutViewType m_out_T; public: - explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, int axis) : m_fft_size(1), m_is_transpose_needed(false) { + explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, int axis) : m_fft_size(1), m_is_transpose_needed(false), m_direction(direction) { static_assert(Kokkos::is_view::value, "KokkosFFT::Plan: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, "KokkosFFT::Plan: OutViewType is not a Kokkos::View."); + m_axes = {axis}; std::tie(m_map, m_map_inv) = KokkosFFT::Impl::get_map_axes(in, axis); m_is_transpose_needed = KokkosFFT::Impl::is_transpose_needed(m_map); - m_fft_size = KokkosFFT::Impl::_create(exec_space, m_plan, in, out, direction, axis_type<1>{axis}); + m_fft_size = KokkosFFT::Impl::_create(exec_space, m_plan, in, out, direction, m_axes); } - explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, axis_type axes) : m_fft_size(1), m_is_transpose_needed(false) { + explicit Plan(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, KokkosFFT::Impl::Direction direction, axis_type axes) : m_fft_size(1), m_is_transpose_needed(false), m_direction(direction), m_axes(axes) { static_assert(Kokkos::is_view::value, "KokkosFFT::Plan: InViewType is not a Kokkos::View."); static_assert(Kokkos::is_view::value, @@ -78,6 +82,29 @@ namespace Impl { _destroy(m_plan); } + template + void good(KokkosFFT::Impl::Direction direction, axis_type axes) const { + static_assert(std::is_same_v, + "KokkosFFT::Plan: is_good: Execution spaces for plan and execution are inconsistent."); + + using nonConstInViewType2 = std::remove_cv_t; + using nonConstOutViewType2 = std::remove_cv_t; + static_assert(std::is_same_v, + "KokkosFFT::Plan: is_good: InViewType for plan and execution are inconsistent."); + static_assert(std::is_same_v, + "KokkosFFT::Plan: is_good: OutViewType for plan and execution are inconsistent."); + + if(direction != m_direction) { + throw std::runtime_error("KokkosFFT::Impl::Plan::good: direction for plan and execution are inconsistent."); + } + + if(axes != m_axes) { + throw std::runtime_error("KokkosFFT::Impl::Plan::good: axes for plan and execution are inconsistent."); + } + + // [TO DO] Check view extents + } + fft_plan_type plan() const { return m_plan; } fft_size_type fft_size() const { return m_fft_size; } bool is_transpose_needed() const { return m_is_transpose_needed; } diff --git a/fft/src/KokkosFFT_Transform.hpp b/fft/src/KokkosFFT_Transform.hpp index 604ce298..d57ff1ac 100644 --- a/fft/src/KokkosFFT_Transform.hpp +++ b/fft/src/KokkosFFT_Transform.hpp @@ -74,6 +74,7 @@ namespace Impl { } // namespace KokkosFFT namespace KokkosFFT { + // 1D FFT template void fft(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, @@ -127,6 +128,62 @@ namespace KokkosFFT { } } + template + void fft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::fft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::fft: OutViewType is not a Kokkos::View."); + + static_assert( + Kokkos::SpaceAccessibility::accessible, + "KokkosFFT::fft: execution_space cannot access data in InViewType" + ); + + static_assert( + Kokkos::SpaceAccessibility::accessible, + "KokkosFFT::fft: execution_space cannot access data in OutViewType" + ); + + InViewType _in; + if(n) { + std::size_t _n = n.value(); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + plan.template good(KokkosFFT::Impl::Direction::Forward, axis_type<1>{axis}); + + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); + + KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); + + KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); + + } else { + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); + } + } + template void ifft(const ExecutionSpace& exec_space, const InViewType& in, @@ -171,6 +228,52 @@ namespace KokkosFFT { } } + template + void ifft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifft: OutViewType is not a Kokkos::View."); + + InViewType _in; + // [TO DO] Modify crop_or_pad to perform the following lines + // KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, n); + if(n) { + std::size_t _n = n.value(); + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, shape_type<1>({_n})); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + plan.template good(KokkosFFT::Impl::Direction::Backward, axis_type<1>{axis}); + + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); + + KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); + + KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); + + } else { + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); + } + } + template void rfft(const ExecutionSpace& exec_space, const InViewType& in, @@ -194,6 +297,30 @@ namespace KokkosFFT { fft(exec_space, in, out, norm, axis, n); } + template + void rfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfft: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + + static_assert(std::is_floating_point::value, + "KokkosFFT::rfft: InViewType must be real"); + static_assert(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::rfft: OutViewType must be complex"); + + fft(exec_space, in, out, plan, norm, axis, n); + } + template void irfft(const ExecutionSpace& exec_space, const InViewType& in, @@ -221,6 +348,34 @@ namespace KokkosFFT { } } + template + void irfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfft: 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(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::irfft: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::irfft: OutViewType must be real"); + if(n) { + std::size_t _n = n.value() / 2 + 1; + ifft(exec_space, in, out, plan, norm, axis, _n); + } else { + ifft(exec_space, in, out, plan, norm, axis); + } + } + template void hfft(const ExecutionSpace& exec_space, const InViewType& in, @@ -249,6 +404,35 @@ namespace KokkosFFT { irfft(exec_space, in_conj, out, new_norm, axis, n); } + template + void hfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::hfft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::hfft: OutViewType is not a Kokkos::View."); + + // [TO DO] + // allow real type as input, need to obtain complex view type from in view type + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + static_assert(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::hfft: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::hfft: OutViewType must be real"); + auto new_norm = KokkosFFT::Impl::swap_direction(norm); + //using ComplexViewType = typename KokkosFFT::Impl::complex_view_type::type; + //ComplexViewType in_conj; + InViewType in_conj; + KokkosFFT::Impl::conjugate(exec_space, in, in_conj); + irfft(exec_space, in_conj, out, plan, new_norm, axis, n); + } + template void ihfft(const ExecutionSpace& exec_space, const InViewType& in, @@ -275,6 +459,33 @@ namespace KokkosFFT { out = out_conj; } + template + void ihfft(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + int axis=-1, + std::optional n = std::nullopt) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ihfft: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ihfft: OutViewType is not a Kokkos::View."); + + using in_value_type = typename InViewType::non_const_value_type; + using out_value_type = typename OutViewType::non_const_value_type; + static_assert(std::is_floating_point::value, + "KokkosFFT::rfft: InViewType must be real"); + static_assert(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::ihfft: OutViewType must be complex"); + + auto new_norm = KokkosFFT::Impl::swap_direction(norm); + OutViewType out_conj; + rfft(exec_space, in, out, plan, new_norm, axis, n); + KokkosFFT::Impl::conjugate(exec_space, out, out_conj); + out = out_conj; + } + template void fft2(const ExecutionSpace& exec_space, const InViewType& in, OutViewType& out, @@ -315,6 +526,49 @@ namespace KokkosFFT { } } + // 2D FFT + template + void fft2(const ExecutionSpace& exec_space, + const InViewType& in, OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::fft2: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::fft2: OutViewType is not a Kokkos::View."); + + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + plan.template good(KokkosFFT::Impl::Direction::Forward, axes); + + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); + + KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); + + KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); + } else { + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); + } + } + template void ifft2(const ExecutionSpace& exec_space, const InViewType& in, @@ -356,6 +610,49 @@ namespace KokkosFFT { } } + template + void ifft2(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifft2: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifft2: OutViewType is not a Kokkos::View."); + + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + plan.template good(KokkosFFT::Impl::Direction::Backward, axes); + + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); + + KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); + + KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); + } else { + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); + } + } + template void rfft2(const ExecutionSpace& exec_space, const InViewType& in, @@ -379,6 +676,30 @@ namespace KokkosFFT { fft2(exec_space, in, out, norm, axes, s); } + template + void rfft2(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfft2: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfft2: 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::rfft2: InViewType must be real"); + static_assert(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::rfft2: OutViewType must be complex"); + + fft2(exec_space, in, out, plan, norm, axes, s); + } + template void irfft2(const ExecutionSpace& exec_space, const InViewType& in, @@ -410,6 +731,39 @@ namespace KokkosFFT { ifft2(exec_space, in, out, norm, axes, _s); } + template + void irfft2(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + axis_type<2> axes={-2, -1}, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfft2: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfft2: 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(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::irfft2: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::irfft2: OutViewType must be real"); + + shape_type zeros = {0}; // default shape means no crop or pad + shape_type _s = {0}; + if( s != zeros ) { + for(int i=0; i void fftn(const ExecutionSpace& exec_space, const InViewType& in, @@ -496,6 +850,49 @@ namespace KokkosFFT { } } + template + void fftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::fftn: OutViewType is not a Kokkos::View."); + + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + plan.template good(KokkosFFT::Impl::Direction::Forward, axes); + + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); + + KokkosFFT::Impl::_fft(exec_space, plan, in_T, out_T, norm); + + KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); + } else { + KokkosFFT::Impl::_fft(exec_space, plan, _in, out, norm); + } + } + template void ifftn(const ExecutionSpace& exec_space, const InViewType& in, @@ -582,6 +979,49 @@ namespace KokkosFFT { } } + template + void ifftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "KokkosFFT::ifftn: OutViewType is not a Kokkos::View."); + + InViewType _in; + shape_type zeros = {0}; // default shape means no crop or pad + if( s != zeros ) { + auto modified_shape = KokkosFFT::Impl::get_modified_shape(in, s); + if( KokkosFFT::Impl::is_crop_or_pad_needed(in, modified_shape) ) { + KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape); + } else { + _in = in; + } + } else { + _in = in; + } + + plan.template good(KokkosFFT::Impl::Direction::Backward, axes); + + if(plan.is_transpose_needed()) { + InViewType in_T; + OutViewType out_T; + + KokkosFFT::Impl::transpose(exec_space, _in, in_T, plan.map()); + KokkosFFT::Impl::transpose(exec_space, out, out_T, plan.map()); + + KokkosFFT::Impl::_ifft(exec_space, plan, in_T, out_T, norm); + + KokkosFFT::Impl::transpose(exec_space, out_T, out, plan.map_inv()); + } else { + KokkosFFT::Impl::_ifft(exec_space, plan, _in, out, norm); + } + } + template void rfftn(const ExecutionSpace& exec_space, const InViewType& in, @@ -604,6 +1044,30 @@ namespace KokkosFFT { fftn(exec_space, in, out, norm, s); } + template + void rfftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::rfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "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(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::rfftn: OutViewType must be complex"); + + fftn(exec_space, in, out, plan, axes, norm, s); + } + template void rfftn(const ExecutionSpace& exec_space, const InViewType& in, @@ -687,6 +1151,38 @@ namespace KokkosFFT { ifftn(exec_space, in, out, axes, norm, _s); } + + template + void irfftn(const ExecutionSpace& exec_space, + const InViewType& in, + OutViewType& out, + const PlanType& plan, + axis_type axes, + KokkosFFT::Normalization norm=KokkosFFT::Normalization::BACKWARD, + shape_type s={0}) { + static_assert(Kokkos::is_view::value, + "KokkosFFT::irfftn: InViewType is not a Kokkos::View."); + static_assert(Kokkos::is_view::value, + "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(KokkosFFT::Impl::is_complex::value, + "KokkosFFT::irfftn: InViewType must be complex"); + static_assert(std::is_floating_point::value, + "KokkosFFT::irfftn: OutViewType must be real"); + + shape_type zeros = {0}; // default shape means no crop or pad + shape_type _s = {0}; + if( s != zeros ) { + for(int i=0; i +void test_fft1_identity_reuse_plan(T atol=1.0e-12) { + const int maxlen = 30; + using RealView1DType = Kokkos::View; + using ComplexView1DType = Kokkos::View*, LayoutType, execution_space>; + + 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); + 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(); + + int axis = -1; + KokkosFFT::Impl::Plan fft_plan(execution_space(), a, out, KokkosFFT::Impl::Direction::Forward, axis); + KokkosFFT::fft(execution_space(), a, out, fft_plan); + + KokkosFFT::Impl::Plan ifft_plan(execution_space(), out, _a, KokkosFFT::Impl::Direction::Backward, axis); + KokkosFFT::ifft(execution_space(), out, _a, ifft_plan); + + KokkosFFT::Impl::Plan rfft_plan(execution_space(), ar, outr, KokkosFFT::Impl::Direction::Forward, axis); + KokkosFFT::rfft(execution_space(), ar, outr, rfft_plan); + + KokkosFFT::Impl::Plan irfft_plan(execution_space(), outr, _ar, KokkosFFT::Impl::Direction::Backward, axis); + KokkosFFT::irfft(execution_space(), outr, _ar, irfft_plan); + + EXPECT_TRUE( allclose(_a, a_ref, 1.e-5, atol) ); + EXPECT_TRUE( allclose(_ar, ar_ref, 1.e-5, atol) ); + + // Check if errors are correctly raised aginst wrong axis + int wrong_axis = 0; + EXPECT_THROW( + KokkosFFT::fft(execution_space(), a, out, fft_plan, KokkosFFT::Normalization::BACKWARD, wrong_axis), + std::runtime_error + ); + + EXPECT_THROW( + KokkosFFT::ifft(execution_space(), out, _a, ifft_plan, KokkosFFT::Normalization::BACKWARD, wrong_axis), + std::runtime_error + ); + + EXPECT_THROW( + KokkosFFT::rfft(execution_space(), ar, outr, rfft_plan, KokkosFFT::Normalization::BACKWARD, wrong_axis), + std::runtime_error + ); + + EXPECT_THROW( + KokkosFFT::irfft(execution_space(), outr, _ar, irfft_plan, KokkosFFT::Normalization::BACKWARD, wrong_axis), + std::runtime_error + ); + + // Check if errors are correctly raised aginst wrong dirction + KokkosFFT::Impl::Direction wrong_fft_direction = KokkosFFT::Impl::Direction::Backward; + KokkosFFT::Impl::Plan wrong_fft_plan(execution_space(), a, out, wrong_fft_direction, axis); + + KokkosFFT::Impl::Direction wrong_ifft_direction = KokkosFFT::Impl::Direction::Forward; + KokkosFFT::Impl::Plan wrong_ifft_plan(execution_space(), out, _a, wrong_ifft_direction, axis); + + KokkosFFT::Impl::Plan wrong_rfft_plan(execution_space(), ar, outr, wrong_fft_direction, axis); + KokkosFFT::Impl::Plan wrong_irfft_plan(execution_space(), outr, _ar, wrong_ifft_direction, axis); + + EXPECT_THROW( + KokkosFFT::fft(execution_space(), a, out, wrong_fft_plan, KokkosFFT::Normalization::BACKWARD, axis), + std::runtime_error + ); + + EXPECT_THROW( + KokkosFFT::ifft(execution_space(), out, _a, wrong_ifft_plan, KokkosFFT::Normalization::BACKWARD, axis), + std::runtime_error + ); + + EXPECT_THROW( + KokkosFFT::rfft(execution_space(), ar, outr, wrong_rfft_plan, KokkosFFT::Normalization::BACKWARD, axis), + std::runtime_error + ); + + EXPECT_THROW( + KokkosFFT::irfft(execution_space(), outr, _ar, wrong_irfft_plan, KokkosFFT::Normalization::BACKWARD, axis), + std::runtime_error + ); +} + template void test_fft1_1dfft_1dview() { const int len = 30; @@ -262,6 +350,30 @@ void test_fft1_1dhfft_1dview() { EXPECT_TRUE( allclose(out_b, out, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_o, out, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_f, out, 1.e-5, 1.e-6) ); + + // Reuse plans + int axis = -1; + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::Impl::Plan hfft_plan(execution_space(), x_herm, out, KokkosFFT::Impl::Direction::Backward, axis); + KokkosFFT::hfft(execution_space(), x_herm, out, hfft_plan); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out_b, hfft_plan, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out_o, hfft_plan, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out_f, hfft_plan, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, sqrt(static_cast(len))); + multiply(out_f, static_cast(len)); + + EXPECT_TRUE( allclose(out, ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out, 1.e-5, 1.e-6) ); } template @@ -310,6 +422,32 @@ void test_fft1_1dihfft_1dview() { EXPECT_TRUE( allclose(out2_b, x_herm_ref, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out2_o, x_herm_ref, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out2_f, x_herm_ref, 1.e-5, 1.e-6) ); + + // Reuse plans + int axis = -1; + KokkosFFT::Impl::Plan hfft_plan(execution_space(), x_herm, out1, KokkosFFT::Impl::Direction::Backward, axis); + KokkosFFT::Impl::Plan ihfft_plan(execution_space(), out1, out2, KokkosFFT::Impl::Direction::Forward, axis); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1, hfft_plan); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::ihfft(execution_space(), out1, out2, ihfft_plan); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1_b, hfft_plan, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::ihfft(execution_space(), out1_b, out2_b, ihfft_plan, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1_o, hfft_plan, KokkosFFT::Normalization::ORTHO); + KokkosFFT::ihfft(execution_space(), out1_o, out2_o, ihfft_plan, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x_herm, x_herm_ref); + KokkosFFT::hfft(execution_space(), x_herm, out1_f, hfft_plan, KokkosFFT::Normalization::FORWARD); + KokkosFFT::ihfft(execution_space(), out1_f, out2_f, ihfft_plan, KokkosFFT::Normalization::FORWARD); + + EXPECT_TRUE( allclose(out2, x_herm_ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out2_b, x_herm_ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out2_o, x_herm_ref, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out2_f, x_herm_ref, 1.e-5, 1.e-6) ); } template @@ -494,6 +632,15 @@ TYPED_TEST(FFT1D, Identity_1DView) { test_fft1_identity(atol); } +// Identity tests on 1D Views with plan reuse +TYPED_TEST(FFT1D, Identity_1DView_reuse_plans) { + using float_type = typename TestFixture::float_type; + using layout_type = typename TestFixture::layout_type; + + float_type atol = std::is_same_v ? 1.0e-6 : 1.0e-12; + test_fft1_identity_reuse_plan(atol); +} + // fft on 1D Views TYPED_TEST(FFT1D, FFT_1DView) { using float_type = typename TestFixture::float_type; @@ -576,6 +723,23 @@ void test_fft2_2dfft_2dview() { EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); + + // Reuse plans + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes={-2, -1}; + KokkosFFT::Impl::Plan fft2_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Forward, axes); + KokkosFFT::fft2(execution_space(), x, out, fft2_plan); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::fft2(execution_space(), x, out_b, fft2_plan, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::fft2(execution_space(), x, out_o, fft2_plan, KokkosFFT::Normalization::ORTHO); + KokkosFFT::fft2(execution_space(), x, out_f, fft2_plan, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, sqrt(static_cast(n0 * n1))); + multiply(out_f, static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } template @@ -609,6 +773,24 @@ void test_fft2_2difft_2dview() { EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); + + // Reuse plans + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes={-2, -1}; + KokkosFFT::Impl::Plan ifft2_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Backward, axes); + + KokkosFFT::ifft2(execution_space(), x, out, ifft2_plan); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::ifft2(execution_space(), x, out_b, ifft2_plan, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::ifft2(execution_space(), x, out_o, ifft2_plan, KokkosFFT::Normalization::ORTHO); + KokkosFFT::ifft2(execution_space(), x, out_f, ifft2_plan, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); + multiply(out_f, 1.0/static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } template @@ -649,6 +831,31 @@ void test_fft2_2drfft_2dview() { EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); + + // Reuse plans + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes={-2, -1}; + KokkosFFT::Impl::Plan rfft2_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Forward, axes); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(execution_space(), x, out, rfft2_plan); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(execution_space(), x, out_b, rfft2_plan, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(execution_space(), x, out_o, rfft2_plan, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfft2(execution_space(), x, out_f, rfft2_plan, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, sqrt(static_cast(n0 * n1))); + multiply(out_f, static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } template @@ -690,6 +897,31 @@ void test_fft2_2dirfft_2dview() { EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); + + // Reuse plans + using axes_type = KokkosFFT::axis_type<2>; + axes_type axes={-2, -1}; + KokkosFFT::Impl::Plan irfft2_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Backward, axes); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out, irfft2_plan); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out_b, irfft2_plan, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out_o, irfft2_plan, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfft2(execution_space(), x, out_f, irfft2_plan, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); + multiply(out_f, 1.0/static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); } // fft2 on 2D Views @@ -760,11 +992,28 @@ void test_fftn_2dfft_2dview() { // Same tests with specifying axes // np.fftn for 2D array is identical to np.fft(np.fft(x, axis=1), axis=0) using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; - KokkosFFT::fftn(execution_space(), x, out, axes_type{-2, -1}); // default: KokkosFFT::Normalization::BACKWARD - KokkosFFT::fftn(execution_space(), x, out_b, axes_type{-2, -1}, KokkosFFT::Normalization::BACKWARD); - KokkosFFT::fftn(execution_space(), x, out_o, axes_type{-2, -1}, KokkosFFT::Normalization::ORTHO); - KokkosFFT::fftn(execution_space(), x, out_f, axes_type{-2, -1}, KokkosFFT::Normalization::FORWARD); + KokkosFFT::fftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::fftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::fftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ORTHO); + KokkosFFT::fftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, sqrt(static_cast(n0 * n1))); + 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) ); + + // Reuse plans + KokkosFFT::Impl::Plan fftn_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Forward, axes); + + KokkosFFT::fftn(execution_space(), x, out, fftn_plan, axes); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::fftn(execution_space(), x, out_b, fftn_plan, axes, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::fftn(execution_space(), x, out_o, fftn_plan, axes, KokkosFFT::Normalization::ORTHO); + KokkosFFT::fftn(execution_space(), x, out_f, fftn_plan, axes, KokkosFFT::Normalization::FORWARD); multiply(out_o, sqrt(static_cast(n0 * n1))); multiply(out_f, static_cast(n0 * n1)); @@ -861,11 +1110,27 @@ void test_ifftn_2dfft_2dview() { // Same tests with specifying axes // np.fftn for 2D array is identical to np.fft(np.fft(x, axis=1), axis=0) using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; + + KokkosFFT::ifftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::ifftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::ifftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ORTHO); + KokkosFFT::ifftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); + multiply(out_f, 1.0/static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); - KokkosFFT::ifftn(execution_space(), x, out, axes_type{-2, -1}); // default: KokkosFFT::Normalization::BACKWARD - KokkosFFT::ifftn(execution_space(), x, out_b, axes_type{-2, -1}, KokkosFFT::Normalization::BACKWARD); - KokkosFFT::ifftn(execution_space(), x, out_o, axes_type{-2, -1}, KokkosFFT::Normalization::ORTHO); - KokkosFFT::ifftn(execution_space(), x, out_f, axes_type{-2, -1}, KokkosFFT::Normalization::FORWARD); + // Reuse plans + KokkosFFT::Impl::Plan ifftn_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Backward, axes); + KokkosFFT::ifftn(execution_space(), x, out, ifftn_plan, axes); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::ifftn(execution_space(), x, out_b, ifftn_plan, axes, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::ifftn(execution_space(), x, out_o, ifftn_plan, axes, KokkosFFT::Normalization::ORTHO); + KokkosFFT::ifftn(execution_space(), x, out_f, ifftn_plan, axes, KokkosFFT::Normalization::FORWARD); multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); multiply(out_f, 1.0/static_cast(n0 * n1)); @@ -969,18 +1234,42 @@ void test_rfftn_2dfft_2dview() { // Same tests with specifying axes // np.rfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out, axes_type{-2, -1}); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::rfftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::BACKWARD Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_b, axes_type{-2, -1}, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::rfftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::BACKWARD); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_o, axes_type{-2, -1}, KokkosFFT::Normalization::ORTHO); + KokkosFFT::rfftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ORTHO); Kokkos::deep_copy(x, x_ref); - KokkosFFT::rfftn(execution_space(), x, out_f, axes_type{-2, -1}, KokkosFFT::Normalization::FORWARD); + KokkosFFT::rfftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, sqrt(static_cast(n0 * n1))); + multiply(out_f, static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); + + // Reuse plans + KokkosFFT::Impl::Plan rfftn_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Forward, axes); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(execution_space(), x, out, rfftn_plan, axes); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(execution_space(), x, out_b, rfftn_plan, axes, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(execution_space(), x, out_o, rfftn_plan, axes, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::rfftn(execution_space(), x, out_f, rfftn_plan, axes, KokkosFFT::Normalization::FORWARD); multiply(out_o, sqrt(static_cast(n0 * n1))); multiply(out_f, static_cast(n0 * n1)); @@ -1099,18 +1388,41 @@ void test_irfftn_2dfft_2dview() { // Same tests with specifying axes // np.irfftn for 2D array is identical to np.fft(np.rfft(x, axis=1), axis=0) using axes_type = KokkosFFT::axis_type<2>; + axes_type axes = {-2, -1}; + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out, axes); // default: KokkosFFT::Normalization::BACKWARD + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_b, axes, KokkosFFT::Normalization::BACKWARD); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_o, axes, KokkosFFT::Normalization::ORTHO); + + Kokkos::deep_copy(x, x_ref); + KokkosFFT::irfftn(execution_space(), x, out_f, axes, KokkosFFT::Normalization::FORWARD); + + multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); + multiply(out_f, 1.0/static_cast(n0 * n1)); + + EXPECT_TRUE( allclose(out, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_b, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_o, out2, 1.e-5, 1.e-6) ); + EXPECT_TRUE( allclose(out_f, out2, 1.e-5, 1.e-6) ); + // Reuse plans + KokkosFFT::Impl::Plan irfftn_plan(execution_space(), x, out, KokkosFFT::Impl::Direction::Backward, axes); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out, axes_type{-2, -1}); // default: KokkosFFT::Normalization::BACKWARD + KokkosFFT::irfftn(execution_space(), x, out, irfftn_plan, axes); // default: KokkosFFT::Normalization::BACKWARD Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_b, axes_type{-2, -1}, KokkosFFT::Normalization::BACKWARD); + KokkosFFT::irfftn(execution_space(), x, out_b, irfftn_plan, axes, KokkosFFT::Normalization::BACKWARD); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_o, axes_type{-2, -1}, KokkosFFT::Normalization::ORTHO); + KokkosFFT::irfftn(execution_space(), x, out_o, irfftn_plan, axes, KokkosFFT::Normalization::ORTHO); Kokkos::deep_copy(x, x_ref); - KokkosFFT::irfftn(execution_space(), x, out_f, axes_type{-2, -1}, KokkosFFT::Normalization::FORWARD); + KokkosFFT::irfftn(execution_space(), x, out_f, irfftn_plan, axes, KokkosFFT::Normalization::FORWARD); multiply(out_o, 1.0/sqrt(static_cast(n0 * n1))); multiply(out_f, 1.0/static_cast(n0 * n1)); From 6b43110b9b6a1e33246f9d03400b2393eaed3a30 Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 8 Jan 2024 16:30:40 +0900 Subject: [PATCH 2/3] Add an example to reuse a plan --- .../06_1DFFT_reuse_plans.cpp | 49 +++++++++++++++++++ examples/06_1DFFT_reuse_plans/CMakeLists.txt | 2 + examples/06_1DFFT_reuse_plans/numpy_1DFFT.py | 17 +++++++ examples/CMakeLists.txt | 5 +- 4 files changed, 70 insertions(+), 3 deletions(-) create mode 100644 examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp create mode 100644 examples/06_1DFFT_reuse_plans/CMakeLists.txt create mode 100644 examples/06_1DFFT_reuse_plans/numpy_1DFFT.py diff --git a/examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp b/examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp new file mode 100644 index 00000000..801e61df --- /dev/null +++ b/examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp @@ -0,0 +1,49 @@ +#include +#include +#include +#include + +using execution_space = Kokkos::DefaultExecutionSpace; +template using View1D = Kokkos::View; + +int main( int argc, char* argv[] ) { + Kokkos::initialize( argc, argv ); + { + constexpr int n0 = 128; + const Kokkos::complex I(1.0, 1.0); + + // 1D C2C FFT (Forward and Backward) + View1D > xc2c("xc2c", n0); + View1D > xc2c_hat("xc2c_hat", n0); + View1D > xc2c_inv("xc2c_inv", n0); + + Kokkos::Random_XorShift64_Pool<> random_pool(12345); + Kokkos::fill_random(xc2c, random_pool, I); + + int axis = -1; + KokkosFFT::Impl::Plan fft_plan(execution_space(), xc2c, xc2c_hat, KokkosFFT::Impl::Direction::Forward, axis); + KokkosFFT::fft(execution_space(), xc2c, xc2c_hat, fft_plan); + + KokkosFFT::Impl::Plan ifft_plan(execution_space(), xc2c_hat, xc2c_inv, KokkosFFT::Impl::Direction::Backward, axis); + KokkosFFT::ifft(execution_space(), xc2c_hat, xc2c_inv, ifft_plan); + + // 1D R2C FFT + View1D xr2c("xr2c", n0); + View1D > xr2c_hat("xr2c_hat", n0/2+1); + Kokkos::fill_random(xr2c, random_pool, 1); + + KokkosFFT::Impl::Plan rfft_plan(execution_space(), xr2c, xr2c_hat, KokkosFFT::Impl::Direction::Forward, axis); + KokkosFFT::rfft(execution_space(), xr2c, xr2c_hat, rfft_plan); + + // 1D C2R FFT + View1D > xc2r("xc2r_hat", n0/2+1); + View1D xc2r_hat("xc2r", n0); + Kokkos::fill_random(xc2r, random_pool, I); + + KokkosFFT::Impl::Plan irfft_plan(execution_space(), xc2r, xc2r_hat, KokkosFFT::Impl::Direction::Backward, axis); + KokkosFFT::irfft(execution_space(), xc2r, xc2r_hat, irfft_plan); + } + Kokkos::finalize(); + + return 0; +} \ No newline at end of file diff --git a/examples/06_1DFFT_reuse_plans/CMakeLists.txt b/examples/06_1DFFT_reuse_plans/CMakeLists.txt new file mode 100644 index 00000000..61ec93eb --- /dev/null +++ b/examples/06_1DFFT_reuse_plans/CMakeLists.txt @@ -0,0 +1,2 @@ +add_executable(06_1DFFT_reuse_plans 06_1DFFT_reuse_plans.cpp) +target_link_libraries(06_1DFFT_reuse_plans PUBLIC Kokkos::fft) \ No newline at end of file diff --git a/examples/06_1DFFT_reuse_plans/numpy_1DFFT.py b/examples/06_1DFFT_reuse_plans/numpy_1DFFT.py new file mode 100644 index 00000000..274978f5 --- /dev/null +++ b/examples/06_1DFFT_reuse_plans/numpy_1DFFT.py @@ -0,0 +1,17 @@ +import numpy as np + +if __name__ == '__main__': + n0 = 128 + + # 1D C2C FFT (Forward and Backward) + xc2c = np.random.rand(n0) + 1j * np.random.rand(n0) + xc2c_hat = np.fft.fft(xc2c) + xc2c_inv = np.fft.ifft(xc2c_hat) + + # 1D R2C FFT + xr2c = np.random.rand(n0) + xr2c_hat = np.fft.rfft(xr2c) + + # 1D C2R FFT + xc2r = np.random.rand(n0//2+1) + xc2r_hat = np.fft.irfft(xc2r) \ No newline at end of file diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index a4536b1c..47c3d9dd 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -2,6 +2,5 @@ add_subdirectory(01_1DFFT) add_subdirectory(02_2DFFT) add_subdirectory(03_NDFFT) add_subdirectory(04_batchedFFT) -if(KokkosFFT_ENABLE_HOST_AND_DEVICE) - add_subdirectory(05_1DFFT_HOST_DEVICE) -endif() \ No newline at end of file +add_subdirectory(05_1DFFT_HOST_DEVICE) +add_subdirectory(06_1DFFT_reuse_plans) \ No newline at end of file From 7270941ff69b9c9c80b2243e49e8347e8364cc86 Mon Sep 17 00:00:00 2001 From: Yuuichi Asahi Date: Mon, 8 Jan 2024 16:30:55 +0900 Subject: [PATCH 3/3] matrix build on AMD --- .github/workflows/cmake.yml | 35 +++++++++++++++++++++-------------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/.github/workflows/cmake.yml b/.github/workflows/cmake.yml index 40d229c7..708a87ca 100644 --- a/.github/workflows/cmake.yml +++ b/.github/workflows/cmake.yml @@ -82,11 +82,15 @@ jobs: runs-on: ubuntu-latest env: - backends: HIP HIP_HOST_DEVICE + #backends: HIP HIP_HOST_DEVICE architecture: VEGA90A CMAKE_CXX_COMPILER: hipcc container: amd_env + strategy: + matrix: + backend: [ {name: HIP, option: ""}, {name: HIP_HOST_DEVICE, option: "-DKokkosFFT_ENABLE_HOST_AND_DEVICE=ON"} ] + steps: - name: Free Disk Space (Ubuntu) uses: jlumbroso/free-disk-space@v1.2.0 @@ -107,22 +111,25 @@ jobs: # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make. # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type run: | - docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_HIP \ + docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_${{matrix.backend.name}} \ -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} -DCMAKE_CXX_COMPILER=${{env.CMAKE_CXX_COMPILER}} \ - -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON + -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON ${{matrix.backend.option}} - - name: Configure CMake for HIP backend with HOST and DEVICE option - # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make. - # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type - run: | - docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_HIP_HOST_DEVICE \ - -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} -DCMAKE_CXX_COMPILER=${{env.CMAKE_CXX_COMPILER}} \ - -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON \ - -DKokkosFFT_ENABLE_HOST_AND_DEVICE=ON + #- name: Configure CMake for HIP backend with HOST and DEVICE option + # # Configure CMake in a 'build' subdirectory. `CMAKE_BUILD_TYPE` is only required if you are using a single-configuration generator such as make. + # # See https://cmake.org/cmake/help/latest/variable/CMAKE_BUILD_TYPE.html?highlight=cmake_build_type + # run: | + # docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake -B build_HIP_HOST_DEVICE \ + # -DCMAKE_BUILD_TYPE=${{env.BUILD_TYPE}} -DCMAKE_CXX_COMPILER=${{env.CMAKE_CXX_COMPILER}} \ + # -DCMAKE_CXX_STANDARD=17 -DKokkos_ENABLE_HIP=ON -DKokkos_ARCH_${{env.architecture}}=ON -DBUILD_TESTING=ON \ + # -DKokkosFFT_ENABLE_HOST_AND_DEVICE=ON - name: Build # Build your program with the given configuration run: | - for backend in ${{ env.backends }}; do - docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake --build build_${backend} --config ${{env.BUILD_TYPE}} -j 2 - done \ No newline at end of file + docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake --build build_${{matrix.backend.name}} --config ${{env.BUILD_TYPE}} -j 2 + + #run: | + # for backend in ${{ env.backends }}; do + # docker run -v ${{github.workspace}}:/work ${{ env.container }} cmake --build build_${backend} --config ${{env.BUILD_TYPE}} -j 2 + # done \ No newline at end of file