Skip to content

Commit

Permalink
Add inplace transform example (#190)
Browse files Browse the repository at this point in the history
* Add inplace trasnform example

* fix: in_inplace is unsed for C2C

* fix: inplace example based on reviews

* remove unused types

---------

Co-authored-by: Yuuichi Asahi <[email protected]>
  • Loading branch information
yasahi-hpc and Yuuichi Asahi authored Nov 5, 2024
1 parent bbd4ec6 commit 29389f2
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 1 deletion.
2 changes: 1 addition & 1 deletion common/src/KokkosFFT_Extents.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ namespace Impl {
template <typename InViewType, typename OutViewType, std::size_t DIM = 1>
auto get_extents(const InViewType& in, const OutViewType& out,
axis_type<DIM> axes, shape_type<DIM> shape = {},
bool is_inplace = false) {
[[maybe_unused]] bool is_inplace = false) {
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 Down
75 changes: 75 additions & 0 deletions examples/08_inplace_FFT/08_inplace_FFT.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
// SPDX-FileCopyrightText: (C) The kokkos-fft development team, see COPYRIGHT.md file
//
// SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception

#include <Kokkos_Core.hpp>
#include <Kokkos_Complex.hpp>
#include <Kokkos_Random.hpp>
#include <KokkosFFT.hpp>

using execution_space = Kokkos::DefaultExecutionSpace;

template <typename T>
using RightView2D = Kokkos::View<T **, Kokkos::LayoutRight, execution_space>;

int main(int argc, char *argv[]) {
Kokkos::initialize(argc, argv);
{
const int n0 = 128, n1 = 128;
const Kokkos::complex<double> z(1.0, 1.0);

// Forward and backward complex to complex transform
// Define a 2D complex view to handle data
RightView2D<Kokkos::complex<double>> xc2c("xc2c", n0, n1);

// Fill the input view with random data
Kokkos::Random_XorShift64_Pool<> random_pool(12345);
execution_space exec;
Kokkos::fill_random(exec, xc2c, random_pool, z);

KokkosFFT::fft2(exec, xc2c, xc2c);
KokkosFFT::ifft2(exec, xc2c, xc2c);

// Real to complex transform
// Define a 2D complex view to handle data
RightView2D<Kokkos::complex<double>> xr2c_hat("xr2c", n0, n1 / 2 + 1);

// Create unmanaged views on the same data with the FFT shape,
// that is (n0, n1) -> (n0, n1/2+1) R2C transform
// The shape is incorrect from the view point of casting to real
// For casting, the shape should be (n0, (n0/2+1) * 2)
RightView2D<double> xr2c(reinterpret_cast<double *>(xr2c_hat.data()), n0,
n1),
xr2c_padded(reinterpret_cast<double *>(xr2c_hat.data()), n0,
(n0 / 2 + 1) * 2);

// Fill the input view with random data in real space through xr2c_padded
auto sub_xr2c_padded =
Kokkos::subview(xr2c_padded, Kokkos::ALL, Kokkos::make_pair(0, n1));
Kokkos::fill_random(exec, sub_xr2c_padded, random_pool, 1.0);

// Perform the real to complex transform
// [Important] You must use xr2c to define the FFT shape correctly
KokkosFFT::rfft2(exec, xr2c, xr2c_hat);

// Complex to real transform
// Define a 2D complex view to handle data
RightView2D<Kokkos::complex<double>> xc2r("xc2r", n0, n1 / 2 + 1);

// Create an unmanaged view on the same data with the FFT shape
RightView2D<double> xc2r_hat(reinterpret_cast<double *>(xc2r.data()), n0,
n1);

// Fill the input view with random data in complex space
Kokkos::fill_random(exec, xc2r, random_pool, z);

// Perform the complex to real transform
// [Important] You must use xc2r_hat to define the FFT shape correctly
KokkosFFT::irfft2(exec, xc2r, xc2r_hat);

exec.fence();
}
Kokkos::finalize();

return 0;
}
6 changes: 6 additions & 0 deletions examples/08_inplace_FFT/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
# SPDX-FileCopyrightText: (C) The kokkos-fft development team, see COPYRIGHT.md file
#
# SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception

add_executable(08_inplace_FFT 08_inplace_FFT.cpp)
target_link_libraries(08_inplace_FFT PUBLIC KokkosFFT::fft)
1 change: 1 addition & 0 deletions examples/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,4 @@ add_subdirectory(04_batchedFFT)
add_subdirectory(05_1DFFT_HOST_DEVICE)
add_subdirectory(06_1DFFT_reuse_plans)
add_subdirectory(07_unmanaged_views)
add_subdirectory(08_inplace_FFT)

0 comments on commit 29389f2

Please sign in to comment.