FFT Poisson Solver: Neumann and Dirichlet Boundaries (#4202)

Add support for Neumann and Dirichlet boundaries in the FFT based
Poisson solver. This requires cosine and sine transforms. For CPU
builds, we use FFTW for these transforms. For GPU builds, we have
implemented cosine and sine transforms using the real-to-complex
transform provided by cuFFT, rocFFT and oneMKL.

Author: WeiqunZhang

Docs/sphinx_documentation/source/FFT.rst

@@ -47,7 +47,9 @@ Below are examples of using :cpp:`FFT:R2C`.
             sp *= scaling;
         });
 
-    cMultiFab cmf(...);
+    // Use R2C provided spectral data layout.
+    auto const& [cba, cdm] = r2c.getSpectralDataLayout();
+    cMultiFab cmf(cba, cdm, 1, 0);
     FFT::R2C<Real,FFT::Direction::forward> r2c_forward(geom.Domain());
     r2c_forward(mfin, cmf);
 
@@ -56,16 +58,57 @@ Below are examples of using :cpp:`FFT:R2C`.
 
 Note that using :cpp:`forwardThenBackward` is expected to be more efficient
 than separate calls to :cpp:`forward` and :cpp:`backward` because some
-parallel communication can be avoided. It should also be noted that a lot of
+parallel communication can be avoided. For the spectral data, the example
+above builds :cpp:`cMultiFab` using :cpp:`FFT::R2C` provided layout. You can
+also use your own :cpp:`BoxArray` and :cpp:`DistributionMapping`, but it
+might result in extra communication. It should also be noted that a lot of
 preparation works are done in the construction of an :cpp:`FFT::R2C`
-object. Therefore, one should cache it for reuse if possible.
+object. Therefore, one should cache it for reuse if possible. Although
+:cpp:`FFT::R2C` does not have a default constructor, one could always use
+:cpp:`std::unique_ptr<FFT::R2C<Real>>` to store an object in one's class.
 
 
 Poisson Solver
 ==============
 
-AMReX provides FFT based Poisson solvers. :cpp:`FFT::Poisson` supports all
-periodic boundaries using purely FFT. :cpp:`FFT::PoissonHybrid` is a 3D only
-solver that supports periodic boundaries in the first two dimensions and
-Neumann boundary in the last dimension. Similar to :cpp:`FFT::R2C`, the
-Poisson solvers should be cached for reuse.
+AMReX provides FFT based Poisson solvers. :cpp:`FFT::Poisson` supports
+periodic (:cpp:`FFT::Boundary::periodic`), homogeneous Neumann
+(:cpp:`FFT::Boundary::even`), and homogeneous Dirichlet
+(:cpp:`FFT::Boundary::odd`) boundaries using FFT. Below is an example of
+using the solver.
+
+.. highlight:: c++
+
+::
+
+    Geometry geom(...);
+    MultiFab soln(...);
+    MultiFab rhs(...);
+
+    Array<std::pair<FFT::Boundary,FFT::Boundary>,AMREX_SPACEDIM>
+            fft_bc{...};
+
+    bool has_dirichlet = false;
+    for (int idim = 0; idim < AMREX_SPACEDIM; ++idim) {
+        has_dirichlet = has_dirichlet ||
+            fft_bc[idim].first == FFT::Boundary::odd ||
+            fft_bc[idim].second == FFT::Boundary::odd;
+    }
+    if (! has_dirichlet) {
+        // Shift rhs so that its sum is zero.
+        auto rhosum = rhs.sum(0);
+        rhs.plus(-rhosum/geom.Domain().d_numPts(), 0, 1);
+    }
+
+    FFT::Poisson fft_poisson(geom, fft_bc);
+    fft_poisson.solve(soln, rhs);
+
+:cpp:`FFT::PoissonHybrid` is a 3D only solver that supports periodic
+boundaries in the first two dimensions and Neumann boundary in the last
+dimension. The last dimension is solved with a tridiagonal solver that can
+support non-uniform cell size in the z-direction. For most applications,
+:cpp:`FFT::Poisson` should be used.
+
+Similar to :cpp:`FFT::R2C`, the Poisson solvers should be cached for reuse,
+and one might need to use :cpp:`std::unique_ptr<FFT::Poisson<MultiFab>>`
+because there is no default constructor.

Src/Base/AMReX_GpuDevice.cpp

@@ -35,6 +35,13 @@
 #include <roctracer/roctx.h>
 #endif
 #endif
+#if defined(AMREX_USE_FFT)
+#  if __has_include(<rocfft/rocfft.h>)  // ROCm 5.3+
+#    include <rocfft/rocfft.h>
+#  else
+#    include <rocfft.h>
+#  endif
+#endif
 #endif
 
 #ifdef AMREX_USE_ACC
@@ -310,6 +317,10 @@ Device::Initialize ()
     }
 #endif /* AMREX_USE_MPI */
 
+#if defined(AMREX_USE_HIP) && defined(AMREX_USE_FFT)
+    AMREX_ROCFFT_SAFE_CALL(rocfft_setup());
+#endif
+
     if (amrex::Verbose()) {
 #if defined(AMREX_USE_CUDA)
         amrex::Print() << "CUDA"
@@ -349,6 +360,10 @@ Device::Finalize ()
 #ifdef AMREX_USE_GPU
     Device::profilerStop();
 
+#if defined(AMREX_USE_HIP) && defined(AMREX_USE_FFT)
+    AMREX_ROCFFT_SAFE_CALL(rocfft_cleanup());
+#endif
+
 #ifdef AMREX_USE_SYCL
     for (auto& s : gpu_stream_pool) {
         delete s.queue;