Skip to content

Commit

Permalink
Merge pull request #12 from CExA-project/exec-space-as-arg
Browse files Browse the repository at this point in the history 
Exec space as arg
  • Loading branch information
@yasahi-hpc
yasahi-hpc authored Dec 19, 2023
2 parents 2a25891 + b453a15 commit 97f747b
Show file tree
Hide file tree
Showing 15 changed files with 447 additions and 444 deletions.
12 changes: 6 additions & 6 deletions common/src/KokkosFFT_normalization.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,13 @@ namespace KokkosFFT {
ORTHO
};

template <typename ViewType, typename T>
void _normalize(ViewType& inout, const T coef) {
template <typename ExecutionSpace, typename ViewType, typename T>
void _normalize(const ExecutionSpace& exec_space, ViewType& inout, const T coef) {
std::size_t size = inout.size();
auto* data = inout.data();

Kokkos::parallel_for(
Kokkos::RangePolicy<Kokkos::IndexType<std::size_t>>{0, size},
Kokkos::RangePolicy<ExecutionSpace, Kokkos::IndexType<std::size_t>>(exec_space, 0, size),
KOKKOS_LAMBDA(const int& i) { data[i] *= coef; }
);
}
Expand Down Expand Up @@ -53,10 +53,10 @@ namespace KokkosFFT {
return std::tuple<value_type, bool> ({coef, to_normalize});
}

template <typename ViewType>
void normalize(ViewType& inout, FFTDirectionType direction, Normalization normalization, std::size_t fft_size) {
template <typename ExecutionSpace, typename ViewType>
void normalize(const ExecutionSpace& exec_space, ViewType& inout, FFTDirectionType direction, Normalization normalization, std::size_t fft_size) {
auto [coef, to_normalize] = _coefficients(inout, direction, normalization, fft_size);
if(to_normalize) _normalize(inout, coef);
if(to_normalize) _normalize(exec_space, inout, coef);
}
};

Expand Down
22 changes: 11 additions & 11 deletions common/src/KokkosFFT_transpose.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,21 +79,21 @@ namespace KokkosFFT {
return get_map_axes(view, axis_type<1>({axis}));
}

template <typename InViewType, typename OutViewType,
template <typename ExecutionSpace, typename InViewType, typename OutViewType,
std::enable_if_t<InViewType::rank()==1, std::nullptr_t> = nullptr>
void _transpose(InViewType& in, OutViewType& out, [[maybe_unused]] axis_type<1> _map) {
void _transpose(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, [[maybe_unused]] axis_type<1> _map) {
}

template <typename InViewType, typename OutViewType>
void _transpose(InViewType& in, OutViewType& out, axis_type<2> _map) {
template <typename ExecutionSpace, typename InViewType, typename OutViewType>
void _transpose(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, axis_type<2> _map) {
constexpr std::size_t DIM = 2;
static_assert(InViewType::rank() == DIM,
"KokkosFFT::_transpose: Rank of View must be equal to Rank of transpose axes.");

constexpr std::size_t rank = InViewType::rank();
using array_layout_type = typename InViewType::array_layout;

using range_type = Kokkos::MDRangePolicy< Kokkos::Rank<DIM, Kokkos::Iterate::Default, Kokkos::Iterate::Default> >;
using range_type = Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<DIM, Kokkos::Iterate::Default, Kokkos::Iterate::Default> >;
using tile_type = typename range_type::tile_type;
using point_type = typename range_type::point_type;

Expand Down Expand Up @@ -126,15 +126,15 @@ namespace KokkosFFT {
);
}

template <typename InViewType, typename OutViewType>
void _transpose(InViewType& in, OutViewType& out, axis_type<3> _map) {
template <typename ExecutionSpace, typename InViewType, typename OutViewType>
void _transpose(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, axis_type<3> _map) {
constexpr std::size_t DIM = 3;
static_assert(InViewType::rank() == DIM,
"KokkosFFT::_transpose: Rank of View must be equal to Rank of transpose axes.");
constexpr std::size_t rank = InViewType::rank();
using array_layout_type = typename InViewType::array_layout;

using range_type = Kokkos::MDRangePolicy< Kokkos::Rank<DIM, Kokkos::Iterate::Default, Kokkos::Iterate::Default> >;
using range_type = Kokkos::MDRangePolicy<ExecutionSpace, Kokkos::Rank<DIM, Kokkos::Iterate::Default, Kokkos::Iterate::Default> >;
using tile_type = typename range_type::tile_type;
using point_type = typename range_type::point_type;

Expand Down Expand Up @@ -197,8 +197,8 @@ namespace KokkosFFT {
D (5, 7) and axis = -1 -> D' (5, 7)
*
*/
template <typename InViewType, typename OutViewType, std::size_t DIM = 1>
void transpose(InViewType& in, OutViewType& out, axis_type<DIM> _map) {
template <typename ExecutionSpace, typename InViewType, typename OutViewType, std::size_t DIM = 1>
void transpose(const ExecutionSpace& exec_space, InViewType& in, OutViewType& out, axis_type<DIM> _map) {
using in_value_type = typename InViewType::non_const_value_type;
using out_value_type = typename OutViewType::non_const_value_type;
using array_layout_type = typename InViewType::array_layout;
Expand All @@ -215,7 +215,7 @@ namespace KokkosFFT {
throw std::runtime_error("KokkosFFT::transpose: transpose not necessary");
}

_transpose(in, out, _map);
_transpose(exec_space, in, out, _map);
}
};

Expand Down
13 changes: 6 additions & 7 deletions common/unit_test/Test_Normalization.cpp
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

#include <gtest/gtest.h>
#include <Kokkos_Random.hpp>
#include "KokkosFFT_normalization.hpp"
Expand All @@ -21,11 +20,11 @@ TEST(Normalization, Forward) {
Kokkos::fence();

// Backward FFT with Forward Normalization -> Do nothing
KokkosFFT::normalize(x, KOKKOS_FFT_BACKWARD, KokkosFFT::Normalization::FORWARD, len);
KokkosFFT::normalize(execution_space(), x, KOKKOS_FFT_BACKWARD, KokkosFFT::Normalization::FORWARD, len);
EXPECT_TRUE( allclose(x, ref_b, 1.e-5, 1.e-12) );

// Forward FFT with Forward Normalization -> 1/N normalization
KokkosFFT::normalize(x, KOKKOS_FFT_FORWARD, KokkosFFT::Normalization::FORWARD, len);
KokkosFFT::normalize(execution_space(), x, KOKKOS_FFT_FORWARD, KokkosFFT::Normalization::FORWARD, len);
EXPECT_TRUE( allclose(x, ref_f, 1.e-5, 1.e-12) );
}

Expand All @@ -45,11 +44,11 @@ TEST(Normalization, Backward) {
Kokkos::fence();

// Forward FFT with Backward Normalization -> Do nothing
KokkosFFT::normalize(x, KOKKOS_FFT_FORWARD, KokkosFFT::Normalization::BACKWARD, len);
KokkosFFT::normalize(execution_space(), x, KOKKOS_FFT_FORWARD, KokkosFFT::Normalization::BACKWARD, len);
EXPECT_TRUE( allclose(x, ref_f, 1.e-5, 1.e-12) );

// Backward FFT with Backward Normalization -> 1/N normalization
KokkosFFT::normalize(x, KOKKOS_FFT_BACKWARD, KokkosFFT::Normalization::BACKWARD, len);
KokkosFFT::normalize(execution_space(), x, KOKKOS_FFT_BACKWARD, KokkosFFT::Normalization::BACKWARD, len);
EXPECT_TRUE( allclose(x, ref_b, 1.e-5, 1.e-12) );
}

Expand All @@ -72,10 +71,10 @@ TEST(Normalization, Ortho) {
Kokkos::fence();

// Forward FFT with Ortho Normalization -> 1 / sqrt(N) normalization
KokkosFFT::normalize(x_f, KOKKOS_FFT_FORWARD, KokkosFFT::Normalization::ORTHO, len);
KokkosFFT::normalize(execution_space(), x_f, KOKKOS_FFT_FORWARD, KokkosFFT::Normalization::ORTHO, len);
EXPECT_TRUE( allclose(x_f, ref_f, 1.e-5, 1.e-12) );

// Backward FFT with Ortho Normalization -> 1 / sqrt(N) normalization
KokkosFFT::normalize(x_b, KOKKOS_FFT_BACKWARD, KokkosFFT::Normalization::ORTHO, len);
KokkosFFT::normalize(execution_space(), x_b, KOKKOS_FFT_BACKWARD, KokkosFFT::Normalization::ORTHO, len);
EXPECT_TRUE( allclose(x_b, ref_b, 1.e-5, 1.e-12) );
}
34 changes: 17 additions & 17 deletions common/unit_test/Test_Transpose.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,7 @@ void test_transpose_1d() {
Kokkos::fence();

EXPECT_THROW(
KokkosFFT::transpose(x, xt, axes_type<1>({0})),
KokkosFFT::transpose(execution_space(), x, xt, axes_type<1>({0})),
std::runtime_error
);
}
Expand Down Expand Up @@ -325,11 +325,11 @@ TEST(Transpose, 2DLeftView) {
Kokkos::fence();

EXPECT_THROW(
KokkosFFT::transpose(x, xt_axis0, axes_type<2>({0, 1})), // xt is identical to x
KokkosFFT::transpose(execution_space(), x, xt_axis0, axes_type<2>({0, 1})), // xt is identical to x
std::runtime_error
);

KokkosFFT::transpose(x, xt_axis1, axes_type<2>({1, 0})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis1, axes_type<2>({1, 0})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis1, ref_axis1, 1.e-5, 1.e-12) );
}

Expand All @@ -354,11 +354,11 @@ TEST(Transpose, 2DRightView) {
Kokkos::deep_copy(ref_axis0, h_ref_axis0);
Kokkos::fence();

KokkosFFT::transpose(x, xt_axis0, axes_type<2>({1, 0})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis0, axes_type<2>({1, 0})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis0, ref_axis0, 1.e-5, 1.e-12) );

EXPECT_THROW(
KokkosFFT::transpose(x, xt_axis1, axes_type<2>({0, 1})), // xt is identical to x
KokkosFFT::transpose(execution_space(), x, xt_axis1, axes_type<2>({0, 1})), // xt is identical to x
std::runtime_error
);
}
Expand Down Expand Up @@ -404,23 +404,23 @@ TEST(Transpose, 3DLeftView) {
Kokkos::fence();

EXPECT_THROW(
KokkosFFT::transpose(x, xt_axis012, axes_type<3>({0, 1, 2})), // xt is identical to x
KokkosFFT::transpose(execution_space(), x, xt_axis012, axes_type<3>({0, 1, 2})), // xt is identical to x
std::runtime_error
);

KokkosFFT::transpose(x, xt_axis021, axes_type<3>({0, 2, 1})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis021, axes_type<3>({0, 2, 1})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis021, ref_axis021, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis102, axes_type<3>({1, 0, 2})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis102, axes_type<3>({1, 0, 2})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis102, ref_axis102, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis120, axes_type<3>({1, 2, 0})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis120, axes_type<3>({1, 2, 0})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis120, ref_axis120, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis201, axes_type<3>({2, 0, 1})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis201, axes_type<3>({2, 0, 1})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis201, ref_axis201, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis210, axes_type<3>({2, 1, 0})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis210, axes_type<3>({2, 1, 0})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis210, ref_axis210, 1.e-5, 1.e-12) );
}

Expand Down Expand Up @@ -465,22 +465,22 @@ TEST(Transpose, 3DRightView) {
Kokkos::fence();

EXPECT_THROW(
KokkosFFT::transpose(x, xt_axis012, axes_type<3>({0, 1, 2})), // xt is identical to x
KokkosFFT::transpose(execution_space(), x, xt_axis012, axes_type<3>({0, 1, 2})), // xt is identical to x
std::runtime_error
);

KokkosFFT::transpose(x, xt_axis021, axes_type<3>({0, 2, 1})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis021, axes_type<3>({0, 2, 1})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis021, ref_axis021, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis102, axes_type<3>({1, 0, 2})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis102, axes_type<3>({1, 0, 2})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis102, ref_axis102, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis120, axes_type<3>({1, 2, 0})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis120, axes_type<3>({1, 2, 0})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis120, ref_axis120, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis201, axes_type<3>({2, 0, 1})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis201, axes_type<3>({2, 0, 1})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis201, ref_axis201, 1.e-5, 1.e-12) );

KokkosFFT::transpose(x, xt_axis210, axes_type<3>({2, 1, 0})); // xt is the transpose of x
KokkosFFT::transpose(execution_space(), x, xt_axis210, axes_type<3>({2, 1, 0})); // xt is the transpose of x
EXPECT_TRUE( allclose(xt_axis210, ref_axis210, 1.e-5, 1.e-12) );
}
8 changes: 4 additions & 4 deletions examples/01_1DFFT/01_1DFFT.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,22 @@ int main( int argc, char* argv[] ) {
Kokkos::Random_XorShift64_Pool<> random_pool(12345);
Kokkos::fill_random(xc2c, random_pool, I);

KokkosFFT::fft(xc2c, xc2c_hat);
KokkosFFT::ifft(xc2c_hat, xc2c_inv);
KokkosFFT::fft(execution_space(), xc2c, xc2c_hat);
KokkosFFT::ifft(execution_space(), 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(xr2c, random_pool, 1);

KokkosFFT::rfft(xr2c, xr2c_hat);
KokkosFFT::rfft(execution_space(), xr2c, xr2c_hat);

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

KokkosFFT::irfft(xc2r, xc2r_hat);
KokkosFFT::irfft(execution_space(), xc2r, xc2r_hat);
}
Kokkos::finalize();

Expand Down
8 changes: 4 additions & 4 deletions examples/02_2DFFT/02_2DFFT.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,22 @@ int main( int argc, char* argv[] ) {
Kokkos::Random_XorShift64_Pool<> random_pool(12345);
Kokkos::fill_random(xc2c, random_pool, I);

KokkosFFT::fft2(xc2c, xc2c_hat);
KokkosFFT::ifft2(xc2c_hat, xc2c_inv);
KokkosFFT::fft2(execution_space(), xc2c, xc2c_hat);
KokkosFFT::ifft2(execution_space(), xc2c_hat, xc2c_inv);

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

KokkosFFT::rfft2(xr2c, xr2c_hat);
KokkosFFT::rfft2(execution_space(), xr2c, xr2c_hat);

// 2D C2R FFT
View2D<Kokkos::complex<double> > xc2r("xr2c_hat", n0, n1/2+1);
View2D<double> xc2r_hat("xc2r", n0, n1);
Kokkos::fill_random(xc2r, random_pool, I);

KokkosFFT::irfft2(xc2r, xc2r_hat);
KokkosFFT::irfft2(execution_space(), xc2r, xc2r_hat);
}
Kokkos::finalize();

Expand Down
8 changes: 4 additions & 4 deletions examples/03_NDFFT/03_NDFFT.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,22 @@ int main( int argc, char* argv[] ) {
Kokkos::Random_XorShift64_Pool<> random_pool(12345);
Kokkos::fill_random(xc2c, random_pool, I);

KokkosFFT::fftn(xc2c, xc2c_hat);
KokkosFFT::ifftn(xc2c_hat, xc2c_inv);
KokkosFFT::fftn(execution_space(), xc2c, xc2c_hat);
KokkosFFT::ifftn(execution_space(), xc2c_hat, xc2c_inv);

// 3D R2C FFT
View3D<double> xr2c("xr2c", n0, n1, n2);
View3D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0, n1, n2/2+1);
Kokkos::fill_random(xr2c, random_pool, 1);

KokkosFFT::rfftn(xr2c, xr2c_hat);
KokkosFFT::rfftn(execution_space(), xr2c, xr2c_hat);

// 3D C2R FFT
View3D<Kokkos::complex<double> > xc2r("xr2c_hat", n0, n1, n2/2+1);
View3D<double> xc2r_hat("xc2r", n0, n1, n2);
Kokkos::fill_random(xc2r, random_pool, I);

KokkosFFT::irfftn(xc2r, xc2r_hat);
KokkosFFT::irfftn(execution_space(), xc2r, xc2r_hat);
}
Kokkos::finalize();

Expand Down
8 changes: 4 additions & 4 deletions examples/04_batchedFFT/04_batchedFFT.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,22 +20,22 @@ int main( int argc, char* argv[] ) {
Kokkos::Random_XorShift64_Pool<> random_pool(12345);
Kokkos::fill_random(xc2c, random_pool, I);

KokkosFFT::fft(xc2c, xc2c_hat, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);
KokkosFFT::ifft(xc2c_hat, xc2c_inv, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);
KokkosFFT::fft(execution_space(), xc2c, xc2c_hat, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);
KokkosFFT::ifft(execution_space(), xc2c_hat, xc2c_inv, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);

// 1D batched R2C FFT
View3D<double> xr2c("xr2c", n0, n1, n2);
View3D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0, n1, n2/2+1);
Kokkos::fill_random(xr2c, random_pool, 1);

KokkosFFT::rfft(xr2c, xr2c_hat, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);
KokkosFFT::rfft(execution_space(), xr2c, xr2c_hat, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);

// 1D batched C2R FFT
View3D<Kokkos::complex<double> > xc2r("xr2c_hat", n0, n1, n2/2+1);
View3D<double> xc2r_hat("xc2r", n0, n1, n2);
Kokkos::fill_random(xc2r, random_pool, I);

KokkosFFT::irfft(xc2r, xc2r_hat, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);
KokkosFFT::irfft(execution_space(), xc2r, xc2r_hat, KokkosFFT::Normalization::BACKWARD, /*axis=*/-1);
}
Kokkos::finalize();

Expand Down
Loading

0 comments on commit 97f747b

Please sign in to comment.