Documentation

Author: WeiqunZhang

Docs/sphinx_documentation/source/FFT.rst

@@ -0,0 +1,68 @@
+.. role:: cpp(code)
+   :language: c++
+
+.. _sec:FFT:r2c:
+
+FFT::R2C Class
+==============
+
+Class template `FFT::R2C` supports discrete Fourier transforms between real
+and complex data. The name R2C indicates that the forward transform converts
+real data to complex data, while the backward transform converts complex
+data to real data. It should be noted that both directions of transformation
+are supported, not just from real to complex.
+
+The implementation utilizes cuFFT, rocFFT, oneMKL and FFTW, for CUDA, HIP,
+SYCL and CPU builds, respectively. Because the parallel communication is
+handled by AMReX, it does not need the parallel version of
+FFTW. Furthermore, there is no constraint on the domain decomposition such
+as one Box per process. This class performs parallel FFT on AMReX's parallel
+data containers (e.g., :cpp:`MultiFab` and
+:cpp:`FabArray<BaseFab<ComplexData<Real>>>`. For local FFT, the users can
+use FFTW, cuFFT, rocFFT, or oneMKL directly.
+
+The scaling follows the FFTW convention, where applying the forward
+transform followed by the backward transform scales the original data by the
+size of the input array. The layout of the complex data also follows the
+FFTW convention, where the complex Hermitian output array has
+`(nx/2+1,ny,nz)` elements. Here `nx`, `ny` and `nz` are the sizes of the
+real array and the division is rounded down.
+
+Below are examples of using :cpp:`FFT:R2C`.
+
+.. highlight:: c++
+
+::
+
+    Geometry geom(...);
+    MultiFab mfin(...);
+    MultiFab mfout(...);
+
+    auto scaling = 1. / geom.Domain().d_numPts();
+
+    FFT::R2C r2c(geom.Domain());
+    r2c.forwardThenBackward(mfin, mfout,
+        [=] AMREX_GPU_DEVICE (int, int, int, auto& sp)
+        {
+            sp *= scaling;
+        });
+
+    cMultiFab cmf(...);
+    FFT::R2C<Real,FFT::Direction::forward> r2c_forward(geom.Domain());
+    r2c_forward(mfin, cmf);
+
+    FFT::R2C<Real,FFT::Direction::backward> r2c_backward(geom.Domain());
+    r2c_backward(cmf, mfout);
+
+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.
+
+
+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.

Docs/sphinx_documentation/source/FFT_Chapter.rst

@@ -0,0 +1,16 @@
+.. _Chap:FFT:
+
+.. _sec:FFT:FFTOverview:
+
+Discrete Fourier Transform
+==========================
+
+AMReX provides support for parallel discrete Fourier transform. The
+implementation utilizes cuFFT, rocFFT, oneMKL and FFTW, for CUDA, HIP, SYCL
+and CPU builds, respectively. It also provides FFT based Poisson
+solvers.
+
+.. toctree::
+   :maxdepth: 1
+
+   FFT

Docs/sphinx_documentation/source/index.rst

@@ -52,6 +52,7 @@ Documentation on migration from BoxLib is available in the AMReX repository at D
    Fortran_Chapter
    Python_Chapter
    EB_Chapter
+   FFT_Chapter
    TimeIntegration_Chapter
    GPU_Chapter
    Visualization_Chapter

Src/FFT/AMReX_FFT.H

@@ -49,6 +49,12 @@ public:
                                   MultiFab, FabArray<BaseFab<T> > >;
     using cMF = FabArray<BaseFab<GpuComplex<T> > >;
 
+    /**
+     * \brief Constructor
+     *
+     * \param domain the forward domain (i.e., the domain of the real data)
+     * \param info optional information
+     */
     explicit R2C (Box const& domain, Info const& info = Info{});
 
     ~R2C ();
@@ -58,6 +64,25 @@ public:
     R2C& operator= (R2C const&) = delete;
     R2C& operator= (R2C &&) = delete;
 
+    /**
+     * \brief Forward and then backward transform
+     *
+     * This function is available only when this class template is
+     * instantiated for transforms in both directions. It's more efficient
+     * than calling the forward function that stores the spectral data in a
+     * caller provided container followed by the backward function, because
+     * this can avoid parallel communication between the internal data and
+     * the caller's data container.
+     *
+     * \param inmf         input data in MultiFab or FabArray<BaseFab<float>>
+     * \param outmf        output data in MultiFab or FabArray<BaseFab<float>>
+     * \param post_forward a callable object for processing the post-forward
+     *                     data before the backward transform. Its interface
+     *                     is `(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.
+     */
     template <typename F, Direction DIR=D,
               std::enable_if_t<DIR == Direction::both, int> = 0>
     void forwardThenBackward (MF const& inmf, MF& outmf, F const& post_forward)
@@ -67,21 +92,65 @@ public:
         this->backward(outmf);
     }
 
+    /**
+     * \brief Forward transform
+     *
+     * The output is stored in this object's internal data. This function is
+     * not available when this class template is instantiated for
+     * backward-only transform.
+     *
+     * \param inmf input data in MultiFab or FabArray<BaseFab<float>>
+     */
     template <Direction DIR=D, std::enable_if_t<DIR == Direction::forward ||
                                                 DIR == Direction::both, int> = 0>
     void forward (MF const& inmf);
 
+    /**
+     * \brief Forward transform
+     *
+     * This function is not available when this class template is
+     * instantiated for backward-only transform.
+     *
+     * \param inmf input data in MultiFab or FabArray<BaseFab<float>>
+     * \param outmf output data in FabArray<BaseFab<GpuComplex<T>>>
+     */
     template <Direction DIR=D, std::enable_if_t<DIR == Direction::forward ||
                                                 DIR == Direction::both, int> = 0>
     void forward (MF const& inmf, cMF& outmf);
 
+    /**
+     * \brief Backward transform
+     *
+     * This function is available only when this class template is
+     * instantiated for transforms in both directions.
+     *
+     * \param outmf output data in MultiFab or FabArray<BaseFab<float>>
+     */
     template <Direction DIR=D, std::enable_if_t<DIR == Direction::both, int> = 0>
     void backward (MF& outmf);
 
+    /**
+     * \brief Backward transform
+     *
+     * This function is not available when this class template is
+     * instantiated for forward-only transform.
+     *
+     * \param inmf input data in FabArray<BaseFab<GpuComplex<T>>>
+     * \param outmf output data in MultiFab or FabArray<BaseFab<float>>
+     */
     template <Direction DIR=D, std::enable_if_t<DIR == Direction::backward ||
                                                 DIR == Direction::both, int> = 0>
     void backward (cMF const& inmf, MF& outmf);
 
+    /**
+     * \brief Get the internal spectral data
+     *
+     * This function is not available when this class template is
+     * instantiated for backward-only transform. For performance reasons,
+     * the returned data array does not have the usual ordering of
+     * `(x,y,z)`. The order is specified in the second part of the return
+     * value.
+     */
     template <Direction DIR=D, std::enable_if_t<DIR == Direction::forward ||
                                                 DIR == Direction::both, int> = 0>
     std::pair<cMF*,IntVect> getSpectralData ();

Src/FFT/AMReX_FFT_Poisson.H

@@ -7,6 +7,9 @@
 namespace amrex::FFT
 {
 
+/**
+ * \brief Poisson solver for all periodic boundaries using FFT
+ */
 template <typename MF>
 class Poisson
 {
@@ -26,6 +29,10 @@ private:
     R2C<typename MF::value_type, Direction::both> m_r2c;
 };
 
+/**
+ * \brief 3D Poisson solver for periodic boundaries in the first two
+ * dimensions and Neumann in the last dimension.
+ */
 template <typename MF>
 class PoissonHybrid
 {