diff --git a/Docs/sphinx_documentation/source/FFT.rst b/Docs/sphinx_documentation/source/FFT.rst index 80b34b1b7a0..e15c95bced3 100644 --- a/Docs/sphinx_documentation/source/FFT.rst +++ b/Docs/sphinx_documentation/source/FFT.rst @@ -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&` 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> sp) + { + for (int n = 0; n < sp.nComp(); ++n) { + sp[n] *= scaling; + } + }); + .. _sec:FFT:c2c: FFT::C2C Class diff --git a/Src/FFT/AMReX_FFT_R2C.H b/Src/FFT/AMReX_FFT_R2C.H index 714fce2e020..1793803732c 100644 --- a/Src/FFT/AMReX_FFT_R2C.H +++ b/Src/FFT/AMReX_FFT_R2C.H @@ -194,7 +194,15 @@ 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, 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. * \param incomp component index of input data * \param outcomp component index of output data */ @@ -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); @@ -1217,11 +1225,39 @@ 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, 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(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 + 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(void(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 +1265,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(void(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(void(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(void(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(void(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(void(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 +1317,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 +1329,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 +1340,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 +1351,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) {