Skip to content

Commit

Permalink
Add inplace trasnform example
Browse files Browse the repository at this point in the history
  • Loading branch information
Yuuichi Asahi committed Nov 4, 2024
1 parent 0bff334 commit cd8e05e
Show file tree
Hide file tree
Showing 3 changed files with 95 additions and 0 deletions.
88 changes: 88 additions & 0 deletions examples/08_inplace_FFT/08_inplace_FFT.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
// 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>;

template <std::size_t DIM>
using axis_type = KokkosFFT::axis_type<DIM>;
template <std::size_t DIM>
using shape_type = KokkosFFT::shape_type<DIM>;

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

shape_type<2> shape;
shape[0] = n0;
shape[1] = n1;

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

// Create an unmanaged view on the same data
RightView2D<Kokkos::complex<double>> xc2c_hat(xc2c.data(), n0, n1),
xc2c_inv(xc2c.data(), 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_hat);
KokkosFFT::ifft2(exec, xc2c_hat, xc2c_inv);

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

// Create an unmanaged view 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);

Kokkos::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 cd8e05e

Please sign in to comment.