Skip to content

Commit

Permalink
Make Plan class as a part of API (#165)
Browse files Browse the repository at this point in the history 
* move Plan class outside of Impl namespace

* Update docs based on changes

* Make execute a method of Plan class

* Make arguments const for Plan constructors

* mfix: ake internal tyypes of Plan class private

* improve docs

---------

Co-authored-by: Yuuichi Asahi <[email protected]>
  • Loading branch information
@yasahi-hpc
yasahi-hpc and Yuuichi Asahi authored Oct 16, 2024
1 parent 1ce1313 commit 06d60a7
Show file tree
Hide file tree
Showing 8 changed files with 497 additions and 604 deletions.
6 changes: 3 additions & 3 deletions docs/api/plan.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
..
.. SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception
KokkosFFT::Impl::Plan
---------------------
KokkosFFT::Plan
---------------

.. doxygenclass:: KokkosFFT::Impl::Plan
.. doxygenclass:: KokkosFFT::Plan
:members:
4 changes: 2 additions & 2 deletions docs/api_reference.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,8 @@ API Reference

This section documents the public user interface of ``Kokkos-fft``.
APIs are defined in ``KokkosFFT`` namespace and implementation details are defined in ``KokkosFFT::Impl`` namespace.
Thus, it is highly discouraged for users to access functions in ``KokkosFFT::Impl`` namespace except for ``Plan``.
Except for ``KokkosFFT::Impl::Plan``, there are corresponding functions in ``numpy.fft`` as shown below.
Thus, it is highly discouraged for users to access functions in ``KokkosFFT::Impl`` namespace.
Except for ``KokkosFFT::Plan``, there are corresponding functions in ``numpy.fft`` as shown below.

FFT Plan
--------
Expand Down
4 changes: 2 additions & 2 deletions docs/intro/using.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,8 +122,8 @@ The following listing shows examples with and without transpose operation.
Reuse FFT plan
--------------

Apart from the basic APIs, Kokkos-fft APIs include overloaded APIs which can take a FFT plan as an argument.
Using these overloaded APIs, we can reuse the FFT plan created before.
Apart from the basic APIs, Kokkos-fft offers the capability to create a FFT plan wrapping the FFT plans of backend libraries.
We can reuse the FFT plan created once to perform FFTs multiple times on different data given that they have the same properties.
In some backend, FFT plan creation leads to some overhead, wherein we need this functionality.
(see :doc:`minimum working example<../samples/06_1DFFT_reuse_plans>`)

Expand Down
24 changes: 12 additions & 12 deletions examples/06_1DFFT_reuse_plans/06_1DFFT_reuse_plans.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,31 +27,31 @@ int main(int argc, char* argv[]) {
Kokkos::fill_random(exec, xc2c, random_pool, z);

int axis = -1;
KokkosFFT::Impl::Plan fft_plan(exec, xc2c, xc2c_hat,
KokkosFFT::Direction::forward, axis);
KokkosFFT::Impl::fft_exec_impl(fft_plan, xc2c, xc2c_hat);
KokkosFFT::Plan fft_plan(exec, xc2c, xc2c_hat,
KokkosFFT::Direction::forward, axis);
fft_plan.execute(xc2c, xc2c_hat);

KokkosFFT::Impl::Plan ifft_plan(exec, xc2c_hat, xc2c_inv,
KokkosFFT::Direction::backward, axis);
KokkosFFT::Impl::fft_exec_impl(ifft_plan, xc2c_hat, xc2c_inv);
KokkosFFT::Plan ifft_plan(exec, xc2c_hat, xc2c_inv,
KokkosFFT::Direction::backward, axis);
ifft_plan.execute(xc2c_hat, xc2c_inv);

// 1D R2C FFT
View1D<double> xr2c("xr2c", n0);
View1D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0 / 2 + 1);
Kokkos::fill_random(exec, xr2c, random_pool, 1);

KokkosFFT::Impl::Plan rfft_plan(exec, xr2c, xr2c_hat,
KokkosFFT::Direction::forward, axis);
KokkosFFT::Impl::fft_exec_impl(rfft_plan, xr2c, xr2c_hat);
KokkosFFT::Plan rfft_plan(exec, xr2c, xr2c_hat,
KokkosFFT::Direction::forward, axis);
rfft_plan.execute(xr2c, xr2c_hat);

// 1D C2R FFT
View1D<Kokkos::complex<double> > xc2r("xc2r_hat", n0 / 2 + 1);
View1D<double> xc2r_hat("xc2r", n0);
Kokkos::fill_random(exec, xc2r, random_pool, z);

KokkosFFT::Impl::Plan irfft_plan(exec, xc2r, xc2r_hat,
KokkosFFT::Direction::backward, axis);
KokkosFFT::Impl::fft_exec_impl(irfft_plan, xc2r, xc2r_hat);
KokkosFFT::Plan irfft_plan(exec, xc2r, xc2r_hat,
KokkosFFT::Direction::backward, axis);
irfft_plan.execute(xc2r, xc2r_hat);
exec.fence();
}
Kokkos::finalize();
Expand Down
142 changes: 91 additions & 51 deletions fft/src/KokkosFFT_Plans.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,38 +15,48 @@
#include "KokkosFFT_default_types.hpp"
#include "KokkosFFT_traits.hpp"
#include "KokkosFFT_transpose.hpp"
#include "KokkosFFT_normalization.hpp"
#include "KokkosFFT_padding.hpp"
#include "KokkosFFT_utils.hpp"

#if defined(KOKKOS_ENABLE_CUDA)
#include "KokkosFFT_Cuda_plans.hpp"
#include "KokkosFFT_Cuda_transform.hpp"
#ifdef ENABLE_HOST_AND_DEVICE
#include "KokkosFFT_Host_plans.hpp"
#include "KokkosFFT_Host_transform.hpp"
#endif
#elif defined(KOKKOS_ENABLE_HIP)
#if defined(KOKKOSFFT_ENABLE_TPL_ROCFFT)
#include "KokkosFFT_ROCM_plans.hpp"
#include "KokkosFFT_ROCM_transform.hpp"
#else
#include "KokkosFFT_HIP_plans.hpp"
#include "KokkosFFT_HIP_transform.hpp"
#endif
#ifdef ENABLE_HOST_AND_DEVICE
#include "KokkosFFT_Host_plans.hpp"
#include "KokkosFFT_Host_transform.hpp"
#endif
#elif defined(KOKKOS_ENABLE_SYCL)
#include "KokkosFFT_SYCL_plans.hpp"
#include "KokkosFFT_SYCL_transform.hpp"
#ifdef ENABLE_HOST_AND_DEVICE
#include "KokkosFFT_Host_plans.hpp"
#include "KokkosFFT_Host_transform.hpp"
#endif
#elif defined(KOKKOS_ENABLE_OPENMP)
#include "KokkosFFT_Host_plans.hpp"
#include "KokkosFFT_Host_transform.hpp"
#elif defined(KOKKOS_ENABLE_THREADS)
#include "KokkosFFT_Host_plans.hpp"
#include "KokkosFFT_Host_transform.hpp"
#else
#include "KokkosFFT_Host_plans.hpp"
#include "KokkosFFT_Host_transform.hpp"
#endif

namespace KokkosFFT {
namespace Impl {
/// \brief A class that manages a FFT plan of backend FFT library.
///
/// This class is used to manage the FFT plan of backend FFT library.
Expand All @@ -57,7 +67,7 @@ namespace Impl {
template <typename ExecutionSpace, typename InViewType, typename OutViewType,
std::size_t DIM = 1>
class Plan {
public:
private:
//! The type of Kokkos execution pace
using execSpace = ExecutionSpace;

Expand Down Expand Up @@ -87,12 +97,6 @@ class Plan {
//! The type of map for transpose
using map_type = axis_type<InViewType::rank()>;

//! The non-const View type of input view
using nonConstInViewType = std::remove_cv_t<InViewType>;

//! The non-const View type of output view
using nonConstOutViewType = std::remove_cv_t<OutViewType>;

//! Naive 1D View for work buffer
using BufferViewType =
Kokkos::View<Kokkos::complex<float_type>*, layout_type, execSpace>;
Expand Down Expand Up @@ -136,12 +140,6 @@ class Plan {
extents_type m_in_extents, m_out_extents;
///@}

//! @{
//! Internal buffers used for transpose
nonConstInViewType m_in_T;
nonConstOutViewType m_out_T;
//! @}

//! Internal work buffer (for rocfft)
BufferViewType m_buffer;

Expand All @@ -156,9 +154,9 @@ class Plan {
/// \param n [in] Length of the transformed axis of the output (default,
/// nullopt)
//
explicit Plan(const ExecutionSpace& exec_space, InViewType& in,
OutViewType& out, KokkosFFT::Direction direction, int axis,
std::optional<std::size_t> n = std::nullopt)
explicit Plan(const ExecutionSpace& exec_space, const InViewType& in,
const OutViewType& out, KokkosFFT::Direction direction,
int axis, std::optional<std::size_t> n = std::nullopt)
: m_exec_space(exec_space), m_axes({axis}), m_direction(direction) {
static_assert(KokkosFFT::Impl::is_AllowedSpace_v<ExecutionSpace>,
"Plan::Plan: ExecutionSpace is not allowed ");
Expand Down Expand Up @@ -214,8 +212,8 @@ class Plan {
/// \param axes [in] Axes over which FFT is performed
/// \param s [in] Shape of the transformed axis of the output (default, {})
//
explicit Plan(const ExecutionSpace& exec_space, InViewType& in,
OutViewType& out, KokkosFFT::Direction direction,
explicit Plan(const ExecutionSpace& exec_space, const InViewType& in,
const OutViewType& out, KokkosFFT::Direction direction,
axis_type<DIM> axes, shape_type<DIM> s = {})
: m_exec_space(exec_space), m_axes(axes), m_direction(direction) {
static_assert(KokkosFFT::Impl::is_AllowedSpace_v<ExecutionSpace>,
Expand Down Expand Up @@ -261,8 +259,8 @@ class Plan {
}

~Plan() {
destroy_plan_and_info<ExecutionSpace, fft_plan_type, fft_info_type>(m_plan,
m_info);
KokkosFFT::Impl::destroy_plan_and_info<ExecutionSpace, fft_plan_type,
fft_info_type>(m_plan, m_info);
}

Plan() = delete;
Expand All @@ -271,23 +269,84 @@ class Plan {
Plan& operator=(Plan&&) = delete;
Plan(Plan&&) = delete;

/// \brief Execute FFT on input and output Views with normalization
///
/// \param in [in] Input data
/// \param out [out] Ouput data
/// \param norm [in] How the normalization is applied (default, backward)
void execute(const InViewType& in, const OutViewType& out,
KokkosFFT::Normalization norm =
KokkosFFT::Normalization::backward) const {
static_assert(
KokkosFFT::Impl::are_operatable_views_v<execSpace, InViewType,
OutViewType>,
"Plan::execute: InViewType and OutViewType must have the same base "
"floating point "
"type (float/double), the same layout (LayoutLeft/LayoutRight), and "
"the "
"same rank. ExecutionSpace must be accessible to the data in "
"InViewType "
"and OutViewType.");

// sanity check that the plan is consistent with the input/output views
good(in, out);

using ManagableInViewType =
typename KokkosFFT::Impl::manageable_view_type<InViewType>::type;
using ManagableOutViewType =
typename KokkosFFT::Impl::manageable_view_type<OutViewType>::type;
ManagableInViewType in_s;
InViewType in_tmp;
if (m_is_crop_or_pad_needed) {
KokkosFFT::Impl::crop_or_pad(m_exec_space, in, in_s, m_shape);
in_tmp = in_s;
} else {
in_tmp = in;
}

if (m_is_transpose_needed) {
using LayoutType = typename ManagableInViewType::array_layout;
ManagableInViewType const in_T(
"in_T",
KokkosFFT::Impl::create_layout<LayoutType>(
KokkosFFT::Impl::compute_transpose_extents(in_tmp, m_map)));
ManagableOutViewType const out_T(
"out_T", KokkosFFT::Impl::create_layout<LayoutType>(
KokkosFFT::Impl::compute_transpose_extents(out, m_map)));

KokkosFFT::Impl::transpose(m_exec_space, in_tmp, in_T, m_map);
KokkosFFT::Impl::transpose(m_exec_space, out, out_T, m_map);

execute_fft(in_T, out_T, norm);

KokkosFFT::Impl::transpose(m_exec_space, out_T, out, m_map_inv);
} else {
execute_fft(in_tmp, out, norm);
}
}

private:
void execute_fft(const InViewType& in, const OutViewType& out,
KokkosFFT::Normalization norm) const {
auto* idata = reinterpret_cast<typename KokkosFFT::Impl::fft_data_type<
execSpace, in_value_type>::type*>(in.data());
auto* odata = reinterpret_cast<typename KokkosFFT::Impl::fft_data_type<
execSpace, out_value_type>::type*>(out.data());

auto const direction =
KokkosFFT::Impl::direction_type<execSpace>(m_direction);
KokkosFFT::Impl::exec_plan(*m_plan, idata, odata, direction, m_info);
KokkosFFT::Impl::normalize(m_exec_space, out, m_direction, norm,
m_fft_size);
}

/// \brief Sanity check of the plan used to call FFT interface with
/// pre-defined FFT plan. This raises an error if there is an
/// incosistency between FFT function and plan
///
/// \param in [in] Input data
/// \param out [in] Ouput data
template <typename InViewType2, typename OutViewType2>
void good(const InViewType2& in, const OutViewType2& out) const {
using nonConstInViewType2 = std::remove_cv_t<InViewType2>;
using nonConstOutViewType2 = std::remove_cv_t<OutViewType2>;
static_assert(std::is_same_v<nonConstInViewType2, nonConstInViewType>,
"Plan::good: InViewType for plan and execution "
"are not identical.");
static_assert(std::is_same_v<nonConstOutViewType2, nonConstOutViewType>,
"Plan::good: OutViewType for plan and "
"execution are not identical.");

void good(const InViewType& in, const OutViewType& out) const {
auto in_extents = KokkosFFT::Impl::extract_extents(in);
auto out_extents = KokkosFFT::Impl::extract_extents(out);

Expand All @@ -299,26 +358,7 @@ class Plan {
out_extents != m_out_extents,
"extents of output View for plan and execution are not identical.");
}

/// \brief Return the execution space
execSpace const& exec_space() const noexcept { return m_exec_space; }

/// \brief Return the FFT plan
fft_plan_type& plan() const { return *m_plan; }

/// \brief Return the FFT info
fft_info_type const& info() const { return m_info; }

/// \brief Return the FFT size
fft_size_type fft_size() const { return m_fft_size; }
KokkosFFT::Direction direction() const { return m_direction; }
bool is_transpose_needed() const { return m_is_transpose_needed; }
bool is_crop_or_pad_needed() const { return m_is_crop_or_pad_needed; }
extents_type shape() const { return m_shape; }
map_type map() const { return m_map; }
map_type map_inv() const { return m_map_inv; }
};
} // namespace Impl
} // namespace KokkosFFT

#endif
Loading

0 comments on commit 06d60a7

Please sign in to comment.