Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cleanup plan class #91

Merged
merged 4 commits into from
Mar 20, 2024
Merged

Cleanup plan class #91

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d6878b7
disallow non-default constructors of Plan
Mar 15, 2024
0b39653
Add tests for shape optional arguments
Mar 15, 2024
0ba0cb7
remove unused changes
Mar 15, 2024
a495b0c
protect input data for hipfft backend
Mar 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 26 additions & 5 deletions fft/src/KokkosFFT_Plans.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,9 @@ class Plan {
//! The type of extents of input/output views
using extents_type = shape_type<InViewType::rank()>;

//! Execution space
execSpace m_exec_space;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do you hold on to an exec space?
When one of the API function is called presumably you take an exec space as argument and you pass the plan. Then which exec space do you pick? Do they have to match? etc..

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the moment, this internal exec_space is unused. I will suppress this change in this PR.
Originally, I was inclined to use internal exec_space for all the execution of FFTs and Kokkos functions, but it turns out that the changes are too big to be consistent.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for your review. I will merge this. Your reviews and/or suggestions are always appreciated!


//! Dynamically allocatable fft plan.
std::unique_ptr<fft_plan_type> m_plan;

Expand All @@ -115,6 +118,9 @@ class Plan {
//! axes for fft
axis_type<DIM> m_axes;

//! Shape of the transformed axis of the output
shape_type<DIM> m_shape;

//! directions of fft
KokkosFFT::Direction m_direction;

Expand All @@ -140,10 +146,16 @@ class Plan {
/// \param out [in] Ouput data
/// \param direction [in] Direction of FFT (forward/backward)
/// \param axis [in] Axis over which FFT is performed
/// \param n [in] Length of the transformed axis of the output (optional)
//
explicit Plan(const ExecutionSpace& exec_space, InViewType& in,
OutViewType& out, KokkosFFT::Direction direction, int axis)
: m_fft_size(1), m_is_transpose_needed(false), m_direction(direction) {
OutViewType& out, KokkosFFT::Direction direction, int axis,
std::optional<std::size_t> n = std::nullopt)
: m_exec_space(exec_space),
m_fft_size(1),
m_is_transpose_needed(false),
m_direction(direction),
m_axes({axis}) {
static_assert(Kokkos::is_view<InViewType>::value,
"Plan::Plan: InViewType is not a Kokkos::View.");
static_assert(Kokkos::is_view<OutViewType>::value,
Expand Down Expand Up @@ -172,7 +184,6 @@ class Plan {
ExecutionSpace, typename OutViewType::memory_space>::accessible,
"Plan::Plan: execution_space cannot access data in OutViewType");

m_axes = {axis};
m_in_extents = KokkosFFT::Impl::extract_extents(in);
m_out_extents = KokkosFFT::Impl::extract_extents(out);
std::tie(m_map, m_map_inv) = KokkosFFT::Impl::get_map_axes(in, axis);
Expand All @@ -188,11 +199,13 @@ class Plan {
/// \param out [in] Ouput data
/// \param direction [in] Direction of FFT (forward/backward)
/// \param axes [in] Axes over which FFT is performed
/// \param s [in] Shape of the transformed axis of the output (optional)
//
explicit Plan(const ExecutionSpace& exec_space, InViewType& in,
OutViewType& out, KokkosFFT::Direction direction,
axis_type<DIM> axes)
: m_fft_size(1),
axis_type<DIM> axes, shape_type<DIM> s = {0})
: m_exec_space(exec_space),
m_fft_size(1),
m_is_transpose_needed(false),
m_direction(direction),
m_axes(axes) {
Expand Down Expand Up @@ -238,6 +251,11 @@ class Plan {
_destroy_plan<ExecutionSpace, fft_plan_type>(m_plan);
}

Plan(const Plan&) = delete;
Plan& operator=(const Plan&) = delete;
Plan& operator=(Plan&&) = delete;
Plan(Plan&&) = delete;

/// \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
Expand Down Expand Up @@ -298,6 +316,9 @@ class Plan {
}
}

/// \brief Return the internal 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; }

Expand Down
12 changes: 8 additions & 4 deletions fft/src/KokkosFFT_Transform.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,20 +319,23 @@ void ifft(const ExecutionSpace& exec_space, const InViewType& in,
"ifft: execution_space cannot access data in OutViewType");

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}));

/* [FIX THIS] Shallow copy should be sufficient
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;
}
*/
KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape);
} else {
_in = in;
}
std::cout << "_in.extent(0): " << _in.extent(0) << std::endl;

KokkosFFT::Impl::Plan plan(exec_space, _in, out,
KokkosFFT::Direction::backward, axis);
Expand Down Expand Up @@ -393,17 +396,18 @@ void ifft(const ExecutionSpace& exec_space, const InViewType& in,
"ifft: execution_space cannot access data in OutViewType");

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}));
/* [FIX THIS] Shallow copy should be sufficient
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;
}
*/
KokkosFFT::Impl::crop_or_pad(exec_space, in, _in, modified_shape);
} else {
_in = in;
}
Expand Down
68 changes: 68 additions & 0 deletions fft/unit_test/Test_Transform.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -580,6 +580,65 @@ void test_fft1_1dihfft_1dview() {
EXPECT_TRUE(allclose(out2_f, x_herm_ref, 1.e-5, 1.e-6));
}

template <typename T, typename LayoutType>
void test_fft1_shape(T atol = 1.0e-12) {
const int n = 32;
using RealView1DType = Kokkos::View<T*, LayoutType, execution_space>;
using ComplexView1DType =
Kokkos::View<Kokkos::complex<T>*, LayoutType, execution_space>;

RealView1DType xr("xr", n);
ComplexView1DType x("x", n / 2 + 1);

const Kokkos::complex<T> I(1.0, 1.0);
Kokkos::Random_XorShift64_Pool<> random_pool(/*seed=*/12345);
Kokkos::fill_random(xr, random_pool, 1.0);
Kokkos::fill_random(x, random_pool, I);
Kokkos::fence();

std::vector<int> shapes = {n / 2, n, n * 2};
for (auto&& shape : shapes) {
// Real to comple
ComplexView1DType outr("outr", shape / 2 + 1),
outr_b("outr_b", shape / 2 + 1), outr_o("outr_o", shape / 2 + 1),
outr_f("outr_f", shape / 2 + 1);
KokkosFFT::rfft(execution_space(), xr, outr, KokkosFFT::Normalization::none,
-1, shape);
KokkosFFT::rfft(execution_space(), xr, outr_b,
KokkosFFT::Normalization::backward, -1, shape);
KokkosFFT::rfft(execution_space(), xr, outr_o,
KokkosFFT::Normalization::ortho, -1, shape);
KokkosFFT::rfft(execution_space(), xr, outr_f,
KokkosFFT::Normalization::forward, -1, shape);

multiply(outr_o, sqrt(static_cast<T>(shape)));
multiply(outr_f, static_cast<T>(shape));

EXPECT_TRUE(allclose(outr_b, outr, 1.e-5, atol));
EXPECT_TRUE(allclose(outr_o, outr, 1.e-5, atol));
EXPECT_TRUE(allclose(outr_f, outr, 1.e-5, atol));

// Complex to real
RealView1DType out("out", shape), out_b("out_b", shape),
out_o("out_o", shape), out_f("out_f", shape);
KokkosFFT::irfft(execution_space(), x, out, KokkosFFT::Normalization::none,
-1, shape);
KokkosFFT::irfft(execution_space(), x, out_b,
KokkosFFT::Normalization::backward, -1, shape);
KokkosFFT::irfft(execution_space(), x, out_o,
KokkosFFT::Normalization::ortho, -1, shape);
KokkosFFT::irfft(execution_space(), x, out_f,
KokkosFFT::Normalization::forward, -1, shape);

multiply(out_o, sqrt(static_cast<T>(shape)));
multiply(out_b, static_cast<T>(shape));

EXPECT_TRUE(allclose(out_b, out, 1.e-5, atol));
EXPECT_TRUE(allclose(out_o, out, 1.e-5, atol));
EXPECT_TRUE(allclose(out_f, out, 1.e-5, atol));
}
}

template <typename T, typename LayoutType>
void test_fft1_1dfft_2dview(T atol = 1.e-12) {
const int n0 = 10, n1 = 12;
Expand Down Expand Up @@ -1218,6 +1277,15 @@ TYPED_TEST(FFT1D, IHFFT_1DView) {
test_fft1_1dihfft_1dview<float_type, layout_type>();
}

// fft1 on 1D Views with shape argument
TYPED_TEST(FFT1D, FFT_1DView_shape) {
using float_type = typename TestFixture::float_type;
using layout_type = typename TestFixture::layout_type;

float_type atol = std::is_same_v<float_type, float> ? 1.0e-6 : 1.0e-12;
test_fft1_shape<float_type, layout_type>(atol);
}

// batced fft1 on 2D Views
TYPED_TEST(FFT1D, FFT_batched_2DView) {
using float_type = typename TestFixture::float_type;
Expand Down
Loading