Skip to content

FFT::R2C: add post-forward support for batch size > 1 #5642

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

Merged
WeiqunZhang merged 2 commits into from
Aug 27, 2026
Merged

FFT::R2C: add post-forward support for batch size > 1 #5642

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6e1c7cc
FFT::R2C: add post-forward support for batch size > 1
WeiqunZhang Aug 24, 2026
e4af509
FFT::R2C: call post_forward for every batch component
WeiqunZhang Aug 25, 2026
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
22 changes: 21 additions & 1 deletion Docs/sphinx_documentation/source/FFT.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,10 +101,30 @@ in an :cpp:`FFT::Info` object passed to the constructor of
r2c.forward(mf, cmf);

// Do work on cmf.
// Function forwardThenBackward is not yet supported for a batched FFT.

r2c.backward(cmf, mf);

Function :cpp:`forwardThenBackward` is also supported for a batched FFT. A
callable taking a :cpp:`GpuComplex<Real>&` is called once for each component
of the batch at each spectral point. When the components are coupled, such as
when projecting a vector field in spectral space, the callable can take a
:cpp:`CellData` argument instead, so that it receives all components of the
batch at a given spectral point at once.

.. highlight:: c++

::

auto scaling = 1. / geom.Domain().d_numPts();

r2c.forwardThenBackward(mf, mf2,
[=] AMREX_GPU_DEVICE (int, int, int, CellData<GpuComplex<Real>> sp)
{
for (int n = 0; n < sp.nComp(); ++n) {
sp[n] *= scaling;
}
});

.. _sec:FFT:c2c:

FFT::C2C Class
Expand Down
80 changes: 62 additions & 18 deletions Src/FFT/AMReX_FFT_R2C.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,15 @@ public:
* is `void(int,int,int,GpuComplex<T>&)`, where the integers
* are indices in the spectral space, and the reference
* to the complex number allows for the modification of
* the spectral data at that location.
* the spectral data at that location. For a batched
* transform, this callable is called once for each
* component of the batch at each spectral point. If
* the components are coupled (e.g., projecting a
* vector field in spectral space), one can instead
* use the interface
* `void(int,int,int,CellData<GpuComplex<T>>)`, which
* receives all components of the batch at that
* spectral point at once.
* \param incomp component index of input data
* \param outcomp component index of output data
*/
Expand All @@ -204,8 +212,8 @@ public:
int incomp = 0, int outcomp = 0)
{
AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
!m_info.twod_mode && m_info.batch_size == 1,
"FFT::R2C::forwardThenBackward(post_forward) currently supports only !twod_mode and batch_size==1");
!m_info.twod_mode,
"FFT::R2C::forwardThenBackward(post_forward) currently supports only !twod_mode");
BL_PROFILE("FFT::R2C::forwardbackward");
this->forward(inmf, incomp);
this->post_forward_doit_0(post_forward);
Expand Down Expand Up @@ -1217,50 +1225,83 @@ R2C<T,D,C>::make_c2c_plans (cMF& inout, int ndims) const
return {fwd, bwd};
}

namespace fft_detail {
// The trailing int/long parameter orders these two overloads. A call
// passing a literal 0 prefers the GpuComplex<T>& version, which is
// called once per component of the batch, and falls back to the CellData
// version only when the former is not viable. Without it, a functor
// taking a generic parameter by value or by const reference would make
// the two overloads ambiguous.
template <typename F, typename T>
AMREX_GPU_DEVICE AMREX_FORCE_INLINE
auto call_post_forward (F const& f, int i, int j, int k, Array4<T> const& a,
int ii, int jj, int kk, int)
noexcept -> decltype(void(f(0,0,0,a(0,0,0))))
{
for (int n = 0; n < a.nComp(); ++n) {
f(i,j,k,a(ii,jj,kk,n));
}
}

template <typename F, typename T>
AMREX_GPU_DEVICE AMREX_FORCE_INLINE
auto call_post_forward (F const& f, int i, int j, int k, Array4<T> const& a,
int ii, int jj, int kk, long)
noexcept -> decltype(void(f(0,0,0,a.cellData(0,0,0))))
{
f(i,j,k,a.cellData(ii,jj,kk));
}
}

template <typename T, Direction D, bool C>
template <typename F>
void R2C<T,D,C>::post_forward_doit_0 (F const& post_forward)
{
if (m_info.twod_mode || m_info.batch_size > 1) {
if (m_info.twod_mode) {
amrex::Abort("xxxxx todo: post_forward");
#if (AMREX_SPACEDIM > 1)
} else if (m_r2c_sub) {
// We need to pass the originally ordered indices to post_forward.
#if (AMREX_SPACEDIM == 2)
// The original domain is (1,ny). The sub domain is (ny,1).
m_r2c_sub->post_forward_doit_1
([=] AMREX_GPU_DEVICE (int i, int, int, auto& sp)
([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
-> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
{
post_forward(0, i, 0, sp);
post_forward(0, i, 0, static_cast<decltype(sp)>(sp));
});
#else
if (m_real_domain.length(0) == 1 && m_real_domain.length(1) == 1) {
// Original domain: (1, 1, nz). Sub domain: (nz, 1, 1)
m_r2c_sub->post_forward_doit_1
([=] AMREX_GPU_DEVICE (int i, int, int, auto& sp)
([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
-> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
{
post_forward(0, 0, i, sp);
post_forward(0, 0, i, static_cast<decltype(sp)>(sp));
});
} else if (m_real_domain.length(0) == 1 && m_real_domain.length(2) == 1) {
// Original domain: (1, ny, 1). Sub domain: (ny, 1, 1)
m_r2c_sub->post_forward_doit_1
([=] AMREX_GPU_DEVICE (int i, int, int, auto& sp)
([=] AMREX_GPU_DEVICE (int i, int, int, auto&& sp)
-> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
{
post_forward(0, i, 0, sp);
post_forward(0, i, 0, static_cast<decltype(sp)>(sp));
});
} else if (m_real_domain.length(0) == 1) {
// Original domain: (1, ny, nz). Sub domain: (ny, nz, 1)
m_r2c_sub->post_forward_doit_1
([=] AMREX_GPU_DEVICE (int i, int j, int, auto& sp)
([=] AMREX_GPU_DEVICE (int i, int j, int, auto&& sp)
-> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
{
post_forward(0, i, j, sp);
post_forward(0, i, j, static_cast<decltype(sp)>(sp));
});
} else if (m_real_domain.length(1) == 1) {
// Original domain: (nx, 1, nz). Sub domain: (nx, nz, 1)
m_r2c_sub->post_forward_doit_1
([=] AMREX_GPU_DEVICE (int i, int j, int, auto& sp)
([=] AMREX_GPU_DEVICE (int i, int j, int, auto&& sp)
-> decltype(void(post_forward(0,0,0,static_cast<decltype(sp)>(sp))))
{
post_forward(i, 0, j, sp);
post_forward(i, 0, j, static_cast<decltype(sp)>(sp));
});
} else {
amrex::Abort("R2c::post_forward_doit_0: how did this happen?");
Expand All @@ -1276,7 +1317,7 @@ template <typename T, Direction D, bool C>
template <typename F>
void R2C<T,D,C>::post_forward_doit_1 (F const& post_forward)
{
if (m_info.twod_mode || m_info.batch_size > 1) {
if (m_info.twod_mode) {
amrex::Abort("xxxxx todo: post_forward");
} else if (m_r2c_sub) {
amrex::Abort("R2C::post_forward_doit_1: How did this happen?");
Expand All @@ -1288,7 +1329,8 @@ void R2C<T,D,C>::post_forward_doit_1 (F const& post_forward)
ParallelForOMP(spectral_fab->box(),
[=] AMREX_GPU_DEVICE (int iz, int jx, int ky)
{
post_forward(jx,ky,iz,a(iz,jx,ky));
fft_detail::call_post_forward(post_forward,
jx,ky,iz,a,iz,jx,ky,0);
});
}
} else if ( ! m_cy.empty()) {
Expand All @@ -1298,7 +1340,8 @@ void R2C<T,D,C>::post_forward_doit_1 (F const& post_forward)
ParallelForOMP(spectral_fab->box(),
[=] AMREX_GPU_DEVICE (int iy, int jx, int k)
{
post_forward(jx,iy,k,a(iy,jx,k));
fft_detail::call_post_forward(post_forward,
jx,iy,k,a,iy,jx,k,0);
});
}
} else {
Expand All @@ -1308,7 +1351,8 @@ void R2C<T,D,C>::post_forward_doit_1 (F const& post_forward)
ParallelForOMP(spectral_fab->box(),
[=] AMREX_GPU_DEVICE (int i, int j, int k)
{
post_forward(i,j,k,a(i,j,k));
fft_detail::call_post_forward(post_forward,
i,j,k,a,i,j,k,0);
});
}
}
Expand Down
26 changes: 26 additions & 0 deletions Tests/FFT/Batch/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,32 @@ int main (int argc, char* argv[])
AMREX_ALWAYS_ASSERT(error < eps);
}

// Batched forwardThenBackward. The callable takes a CellData so that
// it sees all components of the batch at a spectral point at once.
{
FFT::Info info{};
info.setBatchSize(batch_size);
FFT::R2C<Real,FFT::Direction::both> r2c(geom.Domain(), info);
r2c.forwardThenBackward(mf, mf2,
[=] AMREX_GPU_DEVICE (int, int, int, CellData<GpuComplex<Real>> sp)
{
for (int n = 0; n < sp.nComp(); ++n) {
sp[n] *= scaling;
}
});

MultiFab::Subtract(mf2, mf, 0, 0, batch_size, 0);

auto error = mf2.norminf(0, batch_size, IntVect(0));
amrex::Print() << " Expected to be close to zero: " << error << "\n";
#ifdef AMREX_USE_FLOAT
auto eps = 3.e-6F;
#else
auto eps = 1.e-13;
#endif
AMREX_ALWAYS_ASSERT(error < eps);
}

{
FFT::R2C<Real,FFT::Direction::backward> r2c(geom.Domain());
for (int icomp = 0; icomp < batch_size; ++icomp) {
Expand Down
Loading