From 9fd7cc85f6ace54245f3e8f02e0761d1765a4dda Mon Sep 17 00:00:00 2001
From: Weiqun Zhang <weiqunzhang@lbl.gov>
Date: Mon, 21 Oct 2024 16:22:41 -0700
Subject: [PATCH] Use FFT for MAC projection cases

For level 0 MAC projection with all periodic boundaries or periodic in the
first two dimensions and Neumann in the last dimension, there is now the
option of using FFT. To compile with FFT support, use `USE_FFT=TRUE` for GNU
Make and `HYDRO_FFT=ON` (which will turn on AMReX_FFT) for CMake. At
runtime, one needs to specify `mac_proj.use_fft = 1`. If the conditions for
FFT support are not satisfied, a runtime error may occur. So don't set
`mac_proj.use_fft=1` unless FFT is supported for the problem.

Note that the current implementation is still experimental.
---
 CMakeLists.txt                     |  4 +++
 Projections/hydro_MacProjector.H   | 10 ++++++
 Projections/hydro_MacProjector.cpp | 52 +++++++++++++++++++++++++++++-
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4acb91dd9..cb0f4dc60 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -28,6 +28,7 @@ option( HYDRO_PROJECTIONS  "Enable projections" YES)
 option( HYDRO_EB     "Enable Embedded-Boundary support" YES)
 option( HYDRO_OMP    "Enable OpenMP" NO )
 option( HYDRO_MPI    "Enable MPI"   YES )
+option( HYDRO_FFT    "Enable FFT"   NO )
 
 
 set(HYDRO_GPU_BACKEND_VALUES NONE SYCL CUDA HIP)
@@ -71,6 +72,9 @@ if (NOT TARGET AMReX::amrex)
    if (HYDRO_EB)
       list(APPEND AMREX_REQUIRED_COMPONENTS EB)
    endif ()
+   if (HYDRO_FFT)
+      list(APPEND AMREX_REQUIRED_COMPONENTS FFT)
+   endif ()
    find_package(AMReX CONFIG REQUIRED ${AMREX_REQUIRED_COMPONENTS} )
 endif ()
 
diff --git a/Projections/hydro_MacProjector.H b/Projections/hydro_MacProjector.H
index 85301c6e5..d9c789e7d 100644
--- a/Projections/hydro_MacProjector.H
+++ b/Projections/hydro_MacProjector.H
@@ -10,6 +10,10 @@
 #include <AMReX_MLEBABecLap.H>
 #endif
 
+#ifdef AMREX_USE_FFT
+#include <AMReX_FFT_Poisson.H>
+#endif
+
 #ifdef AMREX_USE_HYPRE
 #include <AMReX_HypreMLABecLap.H>
 #endif
@@ -212,6 +216,12 @@ private:
 #ifdef AMREX_USE_HYPRE
     std::unique_ptr<amrex::HypreMLABecLap> m_hypremlabeclap;
 #endif
+
+    bool m_use_fft = false;
+#ifdef AMREX_USE_FFT
+    std::unique_ptr<amrex::FFT::Poisson<amrex::MultiFab>> m_fft_poisson;
+    std::unique_ptr<amrex::FFT::PoissonHybrid<amrex::MultiFab>> m_fft_poisson_hybrid;
+#endif
 };
 
 }
diff --git a/Projections/hydro_MacProjector.cpp b/Projections/hydro_MacProjector.cpp
index d7da86c69..1805ae0b3 100644
--- a/Projections/hydro_MacProjector.cpp
+++ b/Projections/hydro_MacProjector.cpp
@@ -259,6 +259,13 @@ MacProjector::setDomainBC (const Array<LinOpBCType,AMREX_SPACEDIM>& lobc,
 #ifdef AMREX_USE_HYPRE
     m_hypremlabeclap.reset();
 #endif
+#ifdef AMREX_USE_FFT
+    if (m_fft_poisson_hybrid) {
+        AMREX_ALWAYS_ASSERT_WITH_MESSAGE(lobc.back() == LinOpBCType::Neumann &&
+                                         hibc.back() == LinOpBCType::Neumann,
+                                         "FFT::PoissonHybrid supports Neumann BC in z-direction only");
+    }
+#endif
 }
 
 
@@ -387,6 +394,18 @@ MacProjector::project (Real reltol, Real atol)
             m_mlmg->prepareForFluxes(amrex::GetVecOfConstPtrs(m_phi));
         }
     } else
+#endif
+#ifdef AMREX_USE_FFT
+    if (m_fft_poisson || m_fft_poisson_hybrid) {
+        m_mlmg->makeSolvable(0,0,m_rhs[0]);
+        if (m_fft_poisson) {
+            m_fft_poisson->solve(m_phi[0], m_rhs[0]);
+        } else {
+            m_fft_poisson_hybrid->solve(m_phi[0], m_rhs[0]);
+        }
+        // The solve below should not do any iterations.
+        m_mlmg->solve(amrex::GetVecOfPtrs(m_phi), amrex::GetVecOfConstPtrs(m_rhs), reltol, atol);
+    } else
 #endif
     {
         m_mlmg->solve(amrex::GetVecOfPtrs(m_phi), amrex::GetVecOfConstPtrs(m_rhs), reltol, atol);
@@ -565,6 +584,27 @@ void MacProjector::initProjector (Vector<BoxArray> const& a_grids,
     }
     auto const& dm = a_dmap;
 
+#ifdef AMREX_USE_FFT
+    {
+        ParmParse pp("mac_proj");
+        pp.query("use_fft", m_use_fft);
+    }
+    if (m_use_fft) {
+        bool fft_supported = (nlevs == 1) && (a_overset_mask.empty()) &&
+            m_geom[0].Domain().numPts() == ba[0].numPts();
+#if (AMREX_SPACEDIM == 3)
+        fft_supported = fft_supported && (m_geom[0].isPeriodic(0) &&
+                                          m_geom[0].isPeriodic(1));
+#else
+        fft_supported = fft_supported && m_geom[0].isAllPeriodic();
+#endif
+        if (!fft_supported) {
+            m_use_fft = false;
+            amrex::Warning("MacProjector: FFT not support");
+        }
+    }
+#endif
+
     m_rhs.resize(nlevs);
     m_phi.resize(nlevs);
     m_fluxes.resize(nlevs);
@@ -582,7 +622,7 @@ void MacProjector::initProjector (Vector<BoxArray> const& a_grids,
         }
     }
 
-    if (m_use_mlhypre) { a_lpinfo.setMaxCoarseningLevel(0); }
+    if (m_use_mlhypre || m_use_fft) { a_lpinfo.setMaxCoarseningLevel(0); }
 
     if (a_overset_mask.empty()) {
         m_poisson = std::make_unique<MLPoisson>(m_geom, ba, dm, a_lpinfo);
@@ -608,6 +648,16 @@ void MacProjector::initProjector (Vector<BoxArray> const& a_grids,
     }
 #endif
 
+#ifdef AMREX_USE_FFT
+    if (m_use_fft) {
+        if (m_geom[0].isAllPeriodic()) {
+            m_fft_poisson = std::make_unique<FFT::Poisson<MultiFab>>(m_geom[0]);
+        } else {
+            m_fft_poisson_hybrid = std::make_unique<FFT::PoissonHybrid<MultiFab>>(m_geom[0]);
+        }
+    }
+#endif
+
     setOptions();
 
     m_needs_init = false;