Add FFT fast-length helper and pad OpenBC FFT domains (#5494)

Add FFT::nextFastLen for choosing 2/3/5/7/11-smooth padded FFT sizes.

Use the helper in OpenBCSolver to pad internal doubled convolution
domains to fast lengths by default. Add FFT::Info controls to disable or
tune OpenBC padding, pass Info through PoissonOpenBC, update docs, and
test padded vs. unpadded OpenBC solves.

Author: WeiqunZhang

Docs/sphinx_documentation/source/FFT.rst

@@ -66,6 +66,18 @@ 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.
 
+Choosing FFT lengths
+--------------------
+
+:cpp:`FFT::nextFastLen(target, nfactors)` returns the smallest FFT length
+greater than or equal to :cpp:`target` whose prime factors are limited to the
+first :cpp:`nfactors` values from :cpp:`{2, 3, 5, 7, 11, 13}`. If
+:cpp:`nfactors` is omitted, :cpp:`FFT::FastNumPrimeFactors()` provides the
+platform-dependent default. It currently returns 5 for CUDA and 6 for other
+backends. This default is a performance tuning policy and may change in the
+future. This helper can be used to choose a padded FFT domain size that is
+expected to perform well with common FFT backends.
+
 
 Class template `FFT::R2C` also supports batched FFTs. The batch size is set
 in an :cpp:`FFT::Info` object passed to the constructor of
@@ -177,7 +189,14 @@ boundaries.
 
 :cpp:`FFT::OpenBCSolver` currently supports one right-hand-side component per
 solve. It does not support :cpp:`FFT::Info::setBatchSize` values greater than
-one.
+one. The solver uses an internal doubled convolution domain in each transformed
+direction. By default, the one-sided length is rounded up with
+:cpp:`FFT::nextFastLen(n)` before doubling, using the platform-dependent
+:cpp:`FFT::FastNumPrimeFactors()` default described above. This extra padding
+changes only the internal FFT work arrays, not the user-provided
+:cpp:`MultiFab` domains. Users can disable it with
+:cpp:`FFT::Info::setOpenBCPadding(false)` or tune the factor count with
+:cpp:`FFT::Info::setOpenBCPaddingNumPrimeFactors(nfactors)`.
 
 .. highlight:: c++
 

Src/FFT/AMReX_FFT_Helper.H

@@ -60,6 +60,24 @@ AMREX_ENUM( Boundary, periodic, even, odd );
 enum struct Kind { none, r2c_f, r2c_b, c2c_f, c2c_b, r2r_ee_f, r2r_ee_b,
                    r2r_oo_f, r2r_oo_b, r2r_eo, r2r_oe };
 
+/**
+ * \brief Return the default factor count used for selecting fast FFT lengths.
+ *
+ * The default is platform dependent. CUDA currently uses 5; other backends
+ * currently use 6. This tuning default may change in the future.
+ *
+ * \return Default number of prime factors used by FFT::nextFastLen.
+ */
+[[nodiscard]] constexpr int
+FastNumPrimeFactors () noexcept
+{
+#if defined(AMREX_USE_CUDA)
+    return 5;
+#else
+    return 6;
+#endif
+}
+
 struct Info
 {
     //! Domain composition strategy.
@@ -89,6 +107,13 @@ struct Info
     //! Max number of processes to use
     int nprocs = std::numeric_limits<int>::max();
 
+    //! Whether OpenBCSolver pads internal FFT lengths for better performance.
+    bool openbc_padding = true;
+
+    //! Number of prime factors used by FFT::nextFastLen for OpenBCSolver padding.
+    //! Defaults to FFT::FastNumPrimeFactors().
+    int openbc_padding_nfactors = FastNumPrimeFactors();
+
     /**
      * \brief Select how the domain is decomposed across MPI ranks.
      *
@@ -135,8 +160,112 @@ struct Info
      * \return Reference to this Info.
      */
     Info& setNumProcs (int n) { nprocs = n; return *this; }
+    /**
+     * \brief Enable or disable OpenBCSolver internal FFT padding.
+     *
+     * \param x True to pad OpenBCSolver FFT lengths to fast sizes.
+     * \return Reference to this Info.
+     */
+    Info& setOpenBCPadding (bool x) { openbc_padding = x; return *this; }
+    /**
+     * \brief Set the factor count used for OpenBCSolver FFT padding.
+     *
+     * \param n Number of prime factors passed to FFT::nextFastLen. It must be between 3 and 6.
+     * \return Reference to this Info.
+     */
+    Info& setOpenBCPaddingNumPrimeFactors (int n)
+    {
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+            3 <= n && 6 >= n,
+            "amrex::FFT::Info::setOpenBCPaddingNumPrimeFactors: nfactors must be between 3 and 6");
+        openbc_padding_nfactors = n;
+        return *this;
+    }
 };
 
+/// \cond DOXYGEN_IGNORE
+namespace detail
+{
+inline void next_fast_len_doit (Long target, Long n, int const* factors,
+                                int nfactors, int ifactor, Long& best)
+{
+    if (n >= target) {
+        best = std::min(best, n);
+        return;
+    }
+
+    if (ifactor == nfactors) {
+        // Factors 2 and 3 are handled here: try each n * 3^b, multiply it
+        // by the smallest power of 2 that reaches target, and keep the best.
+        Long const max_m3 = std::numeric_limits<int>::max() / 3;
+        Long const max_m2 = std::numeric_limits<int>::max() / 2;
+        for (Long m3 = n; m3 < best; ) {
+            Long k = m3;
+            bool overflow = false;
+            while (k < target) {
+                if (k > max_m2) {
+                    overflow = true;
+                    break;
+                }
+                k *= 2;
+            }
+            if (!overflow && k < best) {
+                best = k;
+            }
+            if (m3 > max_m3) {
+                break;
+            }
+            m3 *= 3;
+        }
+        return;
+    }
+
+    int const factor = factors[ifactor];
+    Long const max_len = std::numeric_limits<int>::max();
+    Long const max_m = max_len / factor;
+    for (Long m = n; m < best; ) {
+        next_fast_len_doit(target, m, factors, nfactors, ifactor+1, best);
+        if (m > max_m) {
+            break;
+        }
+        m *= factor;
+    }
+}
+}
+/// \endcond
+
+/**
+ * \brief Return the smallest fast FFT length greater than or equal to \p target.
+ *
+ * The allowed factors are the first \p nfactors values from {2, 3, 5, 7, 11, 13}.
+ * When \p nfactors is omitted, FFT::FastNumPrimeFactors() provides a
+ * platform-dependent default.
+ *
+ * \param target Minimum requested FFT length. It must be positive.
+ * \param nfactors Number of supported prime factors. It must be between 3 and 6.
+ * \return A fast FFT length no smaller than \p target.
+ */
+[[nodiscard]] inline int
+nextFastLen (int target, int nfactors = FastNumPrimeFactors())
+{
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(target > 0,
+                                     "amrex::FFT::nextFastLen: target must be positive");
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(3 <= nfactors && 6 >= nfactors,
+                                     "amrex::FFT::nextFastLen: nfactors must be between 3 and 6");
+
+    int constexpr factors[] = {5, 7, 11, 13, 17}; // The extra '17' is to make the buggy GCC -Warray-bounds quiet.
+
+    Long const max_len = std::numeric_limits<int>::max();
+    Long best = max_len + Long(1);
+
+    // Factors 2 and 3 are handled directly in the base case.
+    detail::next_fast_len_doit(target, Long(1), factors, nfactors-2, 0, best);
+
+    AMREX_ALWAYS_ASSERT_WITH_MESSAGE(best <= max_len,
+                                     "amrex::FFT::nextFastLen: result exceeds int range");
+    return static_cast<int>(best);
+}
+
 #ifdef AMREX_USE_HIP
 /// \cond DOXYGEN_IGNORE
 namespace detail { void hip_execute (rocfft_plan plan, void **in, void **out); }

Src/FFT/AMReX_FFT_OpenBCSolver.H

@@ -60,41 +60,78 @@ public:
      */
     [[nodiscard]] Box const& Domain () const { return m_domain; }
 
+    /**
+     * \brief Access the one-sided padded length used to build the internal FFT domain.
+     *
+     * \return Padded length before the open-boundary convolution domain is doubled.
+     */
+    [[nodiscard]] IntVect const& PaddedLength () const { return m_padded_length; }
+
 private:
-    static Box make_grown_domain (Box const& domain, Info const& info);
+    static IntVect make_padded_length (Box const& domain, Info const& info);
+    static Box make_grown_domain (Box const& domain, IntVect const& padded_len,
+                                  Info const& info);
 
     Box m_domain;
     Info m_info;
+    IntVect m_padded_length;
     R2C<T> m_r2c;
     cMF m_G_fft;
     std::unique_ptr<R2C<T>> m_r2c_green;
 };
 
 template <typename T>
-Box OpenBCSolver<T>::make_grown_domain (Box const& domain, Info const& info)
+IntVect OpenBCSolver<T>::make_padded_length (Box const& domain, Info const& info)
 {
     IntVect len = domain.length();
+    int ndims = AMREX_SPACEDIM;
+#if (AMREX_SPACEDIM == 3)
+    if (info.twod_mode) { ndims = 2; }
+#else
+    amrex::ignore_unused(info);
+#endif
+    if (info.openbc_padding) {
+        for (int idim = 0; idim < ndims; ++idim) {
+            len[idim] = FFT::nextFastLen(len[idim], info.openbc_padding_nfactors);
+        }
+    }
+    return len;
+}
+
+template <typename T>
+Box OpenBCSolver<T>::make_grown_domain (Box const& domain, IntVect const& padded_len,
+                                        Info const& info)
+{
+    IntVect len = padded_len;
+    int ndims = AMREX_SPACEDIM;
 #if (AMREX_SPACEDIM == 3)
-    if (info.twod_mode) { len[2] = 0; }
+    if (info.twod_mode) { ndims = 2; }
 #else
     amrex::ignore_unused(info);
 #endif
-    return Box(domain.smallEnd(), domain.bigEnd()+len, domain.ixType());
+    for (int idim = 0; idim < ndims; ++idim) {
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(
+            len[idim] <= std::numeric_limits<int>::max()/2,
+            "FFT::OpenBCSolver: padded domain length exceeds int range");
+        len[idim] *= 2;
+    }
+    return Box(domain.smallEnd(), domain.smallEnd()+len-IntVect(1), domain.ixType());
 }
 
 template <typename T>
 OpenBCSolver<T>::OpenBCSolver (Box const& domain, Info const& info)
     : m_domain(domain),
       m_info(info),
-      m_r2c(OpenBCSolver<T>::make_grown_domain(domain,info),
+      m_padded_length(OpenBCSolver<T>::make_padded_length(domain, info)),
+      m_r2c(OpenBCSolver<T>::make_grown_domain(domain, m_padded_length, info),
             m_info.setDomainStrategy(FFT::DomainStrategy::slab))
 {
     AMREX_ALWAYS_ASSERT_WITH_MESSAGE(m_info.batch_size == 1,
         "FFT::OpenBCSolver does not support FFT::Info::batch_size > 1");
 
 #if (AMREX_SPACEDIM == 3)
     if (m_info.twod_mode) {
-        auto gdom = make_grown_domain(domain,m_info);
+        auto gdom = make_grown_domain(domain, m_padded_length, m_info);
         gdom.enclosedCells(2);
         gdom.setSmall(2, 0);
         int nprocs = std::min({ParallelContext::NProcsSub(),
@@ -124,7 +161,8 @@ void OpenBCSolver<T>::setGreensFunction (F const& greens_function)
         :                             detail::get_fab(m_r2c.m_rx);
     auto const& lo = m_domain.smallEnd();
     auto const& lo3 = lo.dim3();
-    auto const& len = m_domain.length3d();
+    auto const len3d = m_padded_length.dim3();
+    GpuArray<int,3> len{len3d.x, len3d.y, len3d.z};
     if (infab) {
         auto const& a = infab->array();
         auto box = infab->box();

Src/FFT/AMReX_FFT_Poisson.H

@@ -149,12 +149,14 @@ public:
      * \param geom    Geometry describing the problem domain.
      * \param ixtype  Index type (cell vs. node) for the working arrays.
      * \param ngrow   Number of grow cells applied prior to solving.
+     * \param info    OpenBC solver options.
      */
     template <typename FA=MF>
     requires (IsFabArray_v<FA>)
     explicit PoissonOpenBC (Geometry const& geom,
                             IndexType ixtype = IndexType::TheCellType(),
-                            IntVect const& ngrow = IntVect(0));
+                            IntVect const& ngrow = IntVect(0),
+                            Info const& info = Info{});
 
     /**
      * \brief Solve the open-boundary Poisson problem.
@@ -169,7 +171,16 @@ public:
      */
     void define_doit (); // has to be public for cuda
 
+    /**
+     * \brief Access the one-sided padded length used by the internal OpenBC solver.
+     *
+     * \return Padded length before the open-boundary convolution domain is doubled.
+     */
+    [[nodiscard]] IntVect const& PaddedLength () const { return m_solver.PaddedLength(); }
+
 private:
+    static Info make_solver_info (Info const& info);
+
     Geometry m_geom;
     Box m_grown_domain;
     IntVect m_ngrow;
@@ -398,15 +409,24 @@ template <typename MF>
 template <typename FA>
 requires (IsFabArray_v<FA>)
 PoissonOpenBC<MF>::PoissonOpenBC (Geometry const& geom, IndexType ixtype,
-                                  IntVect const& ngrow)
+                                  IntVect const& ngrow, Info const& info)
     : m_geom(geom),
       m_grown_domain(amrex::grow(amrex::convert(geom.Domain(),ixtype),ngrow)),
       m_ngrow(ngrow),
-      m_solver(m_grown_domain)
+      m_solver(m_grown_domain, make_solver_info(info))
 {
     define_doit();
 }
 
+template <typename MF>
+Info PoissonOpenBC<MF>::make_solver_info (Info const& info)
+{
+    Info solver_info;
+    solver_info.openbc_padding = info.openbc_padding;
+    solver_info.openbc_padding_nfactors = info.openbc_padding_nfactors;
+    return solver_info;
+}
+
 template <typename MF>
 void PoissonOpenBC<MF>::define_doit ()
 {