From 6e1c7cc3c419357559bd60e04e812c1f6fafe6ea Mon Sep 17 00:00:00 2001 From: Weiqun Zhang Date: Mon, 24 Aug 2026 14:05:49 -0700 Subject: [PATCH 1/2] FFT::R2C: add post-forward support for batch size > 1 --- Docs/sphinx_documentation/source/FFT.rst | 24 +++++++- Src/FFT/AMReX_FFT_R2C.H | 76 ++++++++++++++++++------ Tests/FFT/Batch/main.cpp | 26 ++++++++ 3 files changed, 107 insertions(+), 19 deletions(-) diff --git a/Docs/sphinx_documentation/source/FFT.rst b/Docs/sphinx_documentation/source/FFT.rst index 80b34b1b7a0..dbe04763834 100644 --- a/Docs/sphinx_documentation/source/FFT.rst +++ b/Docs/sphinx_documentation/source/FFT.rst @@ -101,10 +101,32 @@ 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. In +that case, the callable can take a :cpp:`CellData` argument instead of a +:cpp:`GpuComplex&`, so that it receives all components of the batch at +a given spectral point at once. This is what one needs when the components +are coupled, such as when projecting a vector field in spectral space. + +.. highlight:: c++ + +:: + + auto scaling = 1. / geom.Domain().d_numPts(); + + r2c.forwardThenBackward(mf, mf2, + [=] AMREX_GPU_DEVICE (int, int, int, CellData> sp) + { + for (int n = 0; n < sp.nComp(); ++n) { + sp[n] *= scaling; + } + }); + +Note that a callable taking a :cpp:`GpuComplex&` sees only component 0 +of the batch. + .. _sec:FFT:c2c: FFT::C2C Class diff --git a/Src/FFT/AMReX_FFT_R2C.H b/Src/FFT/AMReX_FFT_R2C.H index 714fce2e020..8cbcf02ad34 100644 --- a/Src/FFT/AMReX_FFT_R2C.H +++ b/Src/FFT/AMReX_FFT_R2C.H @@ -194,7 +194,14 @@ public: * is `void(int,int,int,GpuComplex&)`, 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, one can instead use the interface + * `void(int,int,int,CellData>)`, which + * receives all components of the batch at that + * spectral point at once. This is required if the + * components are coupled (e.g., projecting a vector + * field in spectral space). Note that the + * `GpuComplex&` interface only sees component 0. * \param incomp component index of input data * \param outcomp component index of output data */ @@ -204,8 +211,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); @@ -1217,11 +1224,36 @@ R2C::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& version, 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 + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + auto call_post_forward (F const& f, int i, int j, int k, Array4 const& a, + int ii, int jj, int kk, int) + noexcept -> decltype(f(0,0,0,a(0,0,0))) + { + f(i,j,k,a(ii,jj,kk)); + } + + template + AMREX_GPU_DEVICE AMREX_FORCE_INLINE + auto call_post_forward (F const& f, int i, int j, int k, Array4 const& a, + int ii, int jj, int kk, long) + noexcept -> decltype(f(0,0,0,a.cellData(0,0,0))) + { + f(i,j,k,a.cellData(ii,jj,kk)); + } +} + template template void R2C::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) { @@ -1229,38 +1261,43 @@ void R2C::post_forward_doit_0 (F const& 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(post_forward(0,0,0,static_cast(sp))) { - post_forward(0, i, 0, sp); + post_forward(0, i, 0, static_cast(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(post_forward(0,0,0,static_cast(sp))) { - post_forward(0, 0, i, sp); + post_forward(0, 0, i, static_cast(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(post_forward(0,0,0,static_cast(sp))) { - post_forward(0, i, 0, sp); + post_forward(0, i, 0, static_cast(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(post_forward(0,0,0,static_cast(sp))) { - post_forward(0, i, j, sp); + post_forward(0, i, j, static_cast(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(post_forward(0,0,0,static_cast(sp))) { - post_forward(i, 0, j, sp); + post_forward(i, 0, j, static_cast(sp)); }); } else { amrex::Abort("R2c::post_forward_doit_0: how did this happen?"); @@ -1276,7 +1313,7 @@ template template void R2C::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?"); @@ -1288,7 +1325,8 @@ void R2C::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()) { @@ -1298,7 +1336,8 @@ void R2C::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 { @@ -1308,7 +1347,8 @@ void R2C::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); }); } } diff --git a/Tests/FFT/Batch/main.cpp b/Tests/FFT/Batch/main.cpp index 2c12a4a8838..88c193313f2 100644 --- a/Tests/FFT/Batch/main.cpp +++ b/Tests/FFT/Batch/main.cpp @@ -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 r2c(geom.Domain(), info); + r2c.forwardThenBackward(mf, mf2, + [=] AMREX_GPU_DEVICE (int, int, int, CellData> 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 r2c(geom.Domain()); for (int icomp = 0; icomp < batch_size; ++icomp) { From e4af5090bb5033a977cdff0bedb5095642956722 Mon Sep 17 00:00:00 2001 From: Weiqun Zhang Date: Mon, 24 Aug 2026 22:43:46 -0700 Subject: [PATCH 2/2] FFT::R2C: call post_forward for every batch component --- Docs/sphinx_documentation/source/FFT.rst | 14 ++++----- Src/FFT/AMReX_FFT_R2C.H | 38 +++++++++++++----------- 2 files changed, 27 insertions(+), 25 deletions(-) diff --git a/Docs/sphinx_documentation/source/FFT.rst b/Docs/sphinx_documentation/source/FFT.rst index dbe04763834..e15c95bced3 100644 --- a/Docs/sphinx_documentation/source/FFT.rst +++ b/Docs/sphinx_documentation/source/FFT.rst @@ -104,11 +104,12 @@ in an :cpp:`FFT::Info` object passed to the constructor of r2c.backward(cmf, mf); -Function :cpp:`forwardThenBackward` is also supported for a batched FFT. In -that case, the callable can take a :cpp:`CellData` argument instead of a -:cpp:`GpuComplex&`, so that it receives all components of the batch at -a given spectral point at once. This is what one needs when the components -are coupled, such as when projecting a vector field in spectral space. +Function :cpp:`forwardThenBackward` is also supported for a batched FFT. A +callable taking a :cpp:`GpuComplex&` 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++ @@ -124,9 +125,6 @@ are coupled, such as when projecting a vector field in spectral space. } }); -Note that a callable taking a :cpp:`GpuComplex&` sees only component 0 -of the batch. - .. _sec:FFT:c2c: FFT::C2C Class diff --git a/Src/FFT/AMReX_FFT_R2C.H b/Src/FFT/AMReX_FFT_R2C.H index 8cbcf02ad34..1793803732c 100644 --- a/Src/FFT/AMReX_FFT_R2C.H +++ b/Src/FFT/AMReX_FFT_R2C.H @@ -195,13 +195,14 @@ public: * are indices in the spectral space, and the reference * to the complex number allows for the modification of * the spectral data at that location. For a batched - * transform, one can instead use the interface + * 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>)`, which * receives all components of the batch at that - * spectral point at once. This is required if the - * components are coupled (e.g., projecting a vector - * field in spectral space). Note that the - * `GpuComplex&` interface only sees component 0. + * spectral point at once. * \param incomp component index of input data * \param outcomp component index of output data */ @@ -1226,24 +1227,27 @@ R2C::make_c2c_plans (cMF& inout, int ndims) const namespace fft_detail { // The trailing int/long parameter orders these two overloads. A call - // passing a literal 0 prefers the GpuComplex& version, 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. + // passing a literal 0 prefers the GpuComplex& 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 AMREX_GPU_DEVICE AMREX_FORCE_INLINE auto call_post_forward (F const& f, int i, int j, int k, Array4 const& a, int ii, int jj, int kk, int) - noexcept -> decltype(f(0,0,0,a(0,0,0))) + noexcept -> decltype(void(f(0,0,0,a(0,0,0)))) { - f(i,j,k,a(ii,jj,kk)); + for (int n = 0; n < a.nComp(); ++n) { + f(i,j,k,a(ii,jj,kk,n)); + } } template AMREX_GPU_DEVICE AMREX_FORCE_INLINE auto call_post_forward (F const& f, int i, int j, int k, Array4 const& a, int ii, int jj, int kk, long) - noexcept -> decltype(f(0,0,0,a.cellData(0,0,0))) + noexcept -> decltype(void(f(0,0,0,a.cellData(0,0,0)))) { f(i,j,k,a.cellData(ii,jj,kk)); } @@ -1262,7 +1266,7 @@ void R2C::post_forward_doit_0 (F const& post_forward) // 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) - -> decltype(post_forward(0,0,0,static_cast(sp))) + -> decltype(void(post_forward(0,0,0,static_cast(sp)))) { post_forward(0, i, 0, static_cast(sp)); }); @@ -1271,7 +1275,7 @@ void R2C::post_forward_doit_0 (F const& post_forward) // 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) - -> decltype(post_forward(0,0,0,static_cast(sp))) + -> decltype(void(post_forward(0,0,0,static_cast(sp)))) { post_forward(0, 0, i, static_cast(sp)); }); @@ -1279,7 +1283,7 @@ void R2C::post_forward_doit_0 (F const& post_forward) // 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) - -> decltype(post_forward(0,0,0,static_cast(sp))) + -> decltype(void(post_forward(0,0,0,static_cast(sp)))) { post_forward(0, i, 0, static_cast(sp)); }); @@ -1287,7 +1291,7 @@ void R2C::post_forward_doit_0 (F const& post_forward) // 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) - -> decltype(post_forward(0,0,0,static_cast(sp))) + -> decltype(void(post_forward(0,0,0,static_cast(sp)))) { post_forward(0, i, j, static_cast(sp)); }); @@ -1295,7 +1299,7 @@ void R2C::post_forward_doit_0 (F const& post_forward) // 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) - -> decltype(post_forward(0,0,0,static_cast(sp))) + -> decltype(void(post_forward(0,0,0,static_cast(sp)))) { post_forward(i, 0, j, static_cast(sp)); });