search

kokkos / kokkos-fft

A shared-memory FFT for the Kokkos ecosystem
https://kokkosfft.readthedocs.io/
Other
22 stars 5 forks source link

Support in-place transform #61

Open yasahi-hpc opened 8 months ago

yasahi-hpc commented 1 month ago

Interface should be something like this

Basic API

// Out-place
KokkosFFT::rfft(execution_space(), x, x_hat, KokkosFFT::Normalization::backward, axis); // FFT along -1 axis and batched along 0th axis

// In-place (x2 should have a size of n+2)
KokkosFFT::rfft(execution_space(), x2, KokkosFFT::Normalization::backward, axis); // FFT along -1 axis and batched along 0th axis

To reuse plans

// Out-place
View1D<double> xr2c("xr2c", n0);
View1D<Kokkos::complex<double> > xr2c_hat("xr2c_hat", n0 / 2 + 1);
KokkosFFT::Impl::Plan rfft_plan(exec, xr2c, xr2c_hat,
                                    KokkosFFT::Direction::forward, axis);
KokkosFFT::Impl::fft_exec_impl(rfft_plan, xr2c, xr2c_hat);

// In-place (x2 should have a size of n+2)
template <typename T>
using UView1D = Kokkos::View<T*, execution_space,
                             Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
View1D<double> xr2("xr2", n0+2);
UView1D <Kokkos::complex<double>> xc2(reinterpret_cast<Kokkos::complex<double>*>(xr2.data()), n0/2+1);
KokkosFFT::Impl::Plan rfft_plan(exec, xr2, xc2,
                                    KokkosFFT::Direction::forward, axis);
KokkosFFT::Impl::fft_exec_impl(rfft_plan, xr2, xc2);

Internally, check the address of input and output views. If it is identical, then construct an in-place plan.