Open Boundary Poisson Solver (AMReX-Codes#2912)

This adds an open boundary Poisson solver based on the James's algorithm.
To use it, the user builds an amrex:OpenBCSolver object, which can be reused
until the grids change, and then call OpenBCSolver::solver.

Currently, this is for 3D cell-centered data only. The solver works on CPU,
Nvidia GPUS, and AMD GPUs.  The SYCL version of a couple of kernels for
Intel GPUs are to be implemented.

GNUmakefile.in

@@ -19,6 +19,9 @@ ifeq ($(USE_FORTRAN_INTERFACE),TRUE)
 endif
 ifeq ($(USE_LINEAR_SOLVERS),TRUE)
    Pdirs += LinearSolvers/MLMG
+   ifeq ($(DIM),3)
+     Pdirs += LinearSolvers/OpenBC
+   endif
    ifeq ($(USE_FORTRAN_INTERFACE),TRUE)
      Pdirs += F_Interfaces/LinearSolvers
    endif

Src/Base/AMReX_DistributionMapping.cpp

@@ -1300,7 +1300,7 @@ DistributionMapping::SFCProcessorMap (const BoxArray& boxes,
 
         for (int i = 0, N = boxes.size(); i < N; ++i)
         {
-            wgts.push_back(boxes[i].volume());
+            wgts.push_back(boxes[i].numPts());
         }
 
         SFCProcessorMapDoIt(boxes,wgts,nprocs);
@@ -1769,7 +1769,7 @@ DistributionMapping::makeSFC (const BoxArray& ba, bool use_box_vol, const int np
     {
         const Box& bx = ba[i];
         tokens.push_back(makeSFCToken(i, bx.smallEnd()));
-        const Long v = use_box_vol ? bx.volume() : Long(1);
+        const Long v = use_box_vol ? bx.numPts() : Long(1);
         vol_sum += v;
         wgts.push_back(v);
     }

Src/Boundary/AMReX_LOUtil_K.H

@@ -34,6 +34,22 @@ void poly_interp_coeff (Real xInt, Real const* AMREX_RESTRICT x, int N, Real* AM
     }
 }
 
+template <int N>
+AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+void poly_interp_coeff (Real xInt, Real const* AMREX_RESTRICT x, Real* AMREX_RESTRICT c) noexcept
+{
+    for (int j = 0; j < N; ++j) {
+        Real num = 1.0, den = 1.0;
+        for (int i = 0; i < N; ++i) {
+            if (i != j) {
+                num *= xInt-x[i];
+                den *= x[j]-x[i];
+            }
+        }
+        c[j] = num / den;
+    }
+}
+
 }
 
 #endif

Src/LinearSolvers/CMakeLists.txt

@@ -98,3 +98,15 @@ if (AMReX_HYPRE)
       MLMG/AMReX_MLNodeLaplacian_hypre.cpp
       )
 endif ()
+
+if (AMReX_SPACEDIM EQUAL 3)
+
+   target_include_directories(amrex PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/OpenBC>)
+
+   target_sources(amrex
+      PRIVATE
+      OpenBC/AMReX_OpenBC.H
+      OpenBC/AMReX_OpenBC_K.H
+      OpenBC/AMReX_OpenBC.cpp
+      )
+endif ()

Src/LinearSolvers/MLMG/AMReX_MLPoisson.H

@@ -70,6 +70,10 @@ public:
 
     virtual void copyNSolveSolution (MultiFab& dst, MultiFab const& src) const final override;
 
+    //! Compute dphi/dn on domain faces after the solver has converged.
+    void get_dpdn_on_domain_faces (Array<MultiFab*,AMREX_SPACEDIM> const& dpdn,
+                                   MultiFab const& phi);
+
 private:
 
     Vector<int> m_is_singular;

Src/LinearSolvers/MLMG/AMReX_MLPoisson.cpp

@@ -702,4 +702,63 @@ MLPoisson::copyNSolveSolution (MultiFab& dst, MultiFab const& src) const
     dst.ParallelCopy(src);
 }
 
+void
+MLPoisson::get_dpdn_on_domain_faces (Array<MultiFab*,AMREX_SPACEDIM> const& dpdn,
+                                     MultiFab const& phi)
+{
+    BL_PROFILE("MLPoisson::dpdn_faces()");
+
+    // We do not need to call applyBC because this function is used by the
+    // OpenBC solver after solver has converged.  That means the BC has been
+    // filled to check the residual.
+
+    Box const& domain0 = m_geom[0][0].Domain();
+    AMREX_D_TERM(const Real dxi = m_geom[0][0].InvCellSize(0);,
+                 const Real dyi = m_geom[0][0].InvCellSize(1);,
+                 const Real dzi = m_geom[0][0].InvCellSize(2);)
+
+#ifdef AMREX_USE_OMP
+#pragma omp parallel if (Gpu::notInLaunchRegion())
+#endif
+    for (MFIter mfi(phi);  mfi.isValid(); ++mfi)
+    {
+        Box const& vbx = mfi.validbox();
+        for (OrientationIter oit; oit.isValid(); ++oit) {
+            Orientation face = oit();
+            if (vbx[face] == domain0[face]) {
+                int dir = face.coordDir();
+                Array4<Real const> const& p = phi.const_array(mfi);
+                Array4<Real> const& gp = dpdn[dir]->array(mfi);
+                Box const& b2d = amrex::bdryNode(vbx,face);
+                if (dir == 0) {
+                    // because it's dphi/dn, not dphi/dx.
+                    Real fac = dxi * (face.isLow() ? -1.0_rt : 1._rt);
+                    AMREX_HOST_DEVICE_PARALLEL_FOR_3D(b2d, i, j, k,
+                    {
+                        gp(i,j,k) = fac * (p(i,j,k) - p(i-1,j,k));
+                    });
+                }
+#if (AMREX_SPACEDIM > 1)
+                else if (dir == 1) {
+                    Real fac = dyi * (face.isLow() ? -1.0_rt : 1._rt);
+                    AMREX_HOST_DEVICE_PARALLEL_FOR_3D(b2d, i, j, k,
+                    {
+                        gp(i,j,k) = fac * (p(i,j,k) - p(i,j-1,k));
+                    });
+                }
+#if (AMREX_SPACEDIM > 2)
+                else {
+                    Real fac = dzi * (face.isLow() ? -1.0_rt : 1._rt);
+                    AMREX_HOST_DEVICE_PARALLEL_FOR_3D(b2d, i, j, k,
+                    {
+                        gp(i,j,k) = fac * (p(i,j,k) - p(i,j,k-1));
+                    });
+                }
+#endif
+#endif
+            }
+        }
+    }
+}
+
 }

Src/LinearSolvers/OpenBC/AMReX_OpenBC.H

@@ -0,0 +1,136 @@
+#ifndef AMREX_OPENBC_H_
+#define AMREX_OPENBC_H_
+#include <AMReX_Config.H>
+
+#include <AMReX_MLMG.H>
+#include <AMReX_MLPoisson.H>
+
+namespace amrex
+{
+
+namespace openbc {
+
+    static constexpr int M = 7; // highest order of moments
+    static constexpr int P = 3;
+
+    struct Moments
+    {
+        typedef GpuArray<Real,(M+2)*(M+1)/2> array_type;
+        array_type mom;
+        Real x, y, z;
+        Orientation face;
+    };
+
+    struct MomTag
+    {
+        Array4<Real const> gp;
+        Box b2d;
+        Orientation face;
+        int offset;
+    };
+
+    std::ostream& operator<< (std::ostream& os, Moments const& mom);
+}
+
+#if defined(AMREX_USE_CUDA) || defined(AMREX_USE_HIP)
+template<>
+struct Gpu::SharedMemory<openbc::Moments::array_type>
+{
+    AMREX_GPU_DEVICE openbc::Moments::array_type* dataPtr () noexcept {
+        AMREX_HIP_OR_CUDA(HIP_DYNAMIC_SHARED(openbc::Moments::array_type,amrex_openbc_momarray);,
+                          extern __shared__  openbc::Moments::array_type amrex_openbc_momarray[];)
+            return amrex_openbc_momarray;
+    }
+};
+#endif
+
+/**
+ * \brief Open Boundary Poisson Solver
+ *
+ * References:
+ *    (1) The Solution of Poisson's Equation for Isolated Source
+ *        Distributions, R. A. James, 1977, JCP 25, 71
+ *    (2) A Local Corrections Algorithm for Solving Poisson's Equation in Three
+ *        Dimensions, P. McCorquodale, P. Colella, G. T. Balls, & S. B. Baden,
+ *        2007, Communications in Applied Mathematics and Computational Science,
+ *        2, 1, 57-81
+ */
+class OpenBCSolver
+{
+public:
+    OpenBCSolver ();
+
+    OpenBCSolver (const Vector<Geometry>& a_geom,
+                  const Vector<BoxArray>& a_grids,
+                  const Vector<DistributionMapping>& a_dmap,
+                  const LPInfo& a_info = LPInfo());
+
+    ~OpenBCSolver ();
+
+    OpenBCSolver (const OpenBCSolver&) = delete;
+    OpenBCSolver (OpenBCSolver&&) = delete;
+    OpenBCSolver& operator= (const OpenBCSolver&) = delete;
+    OpenBCSolver& operator= (OpenBCSolver&&) = delete;
+
+    void define (const Vector<Geometry>& a_geom,
+                 const Vector<BoxArray>& a_grids,
+                 const Vector<DistributionMapping>& a_dmap,
+                 const LPInfo& a_info = LPInfo());
+
+    void setVerbose (int v) noexcept;
+
+    Real solve (const Vector<MultiFab*>& a_sol, const Vector<MultiFab const*>& a_rhs,
+                Real a_tol_rel, Real a_tol_abs);
+
+public: // public for cuda
+
+    void compute_moments (Gpu::DeviceVector<openbc::Moments>& moments);
+    void compute_potential (Gpu::DeviceVector<openbc::Moments> const& moments);
+    void interpolate_potential (MultiFab& solg);
+
+private:
+
+#ifdef AMREX_USE_MPI
+    void bcast_moments (Gpu::DeviceVector<openbc::Moments>& moments);
+#endif
+
+    int m_verbose = 0;
+    Vector<Geometry> m_geom;
+    Vector<BoxArray> m_grids;
+    Vector<DistributionMapping> m_dmap;
+    LPInfo m_info;
+    std::unique_ptr<MLPoisson> m_poisson_1;
+    std::unique_ptr<MLPoisson> m_poisson_2;
+    std::unique_ptr<MLMG> m_mlmg_1;
+    std::unique_ptr<MLMG> m_mlmg_2;
+
+    int m_coarsen_ratio = 0;
+    Array<MultiFab,AMREX_SPACEDIM> m_dpdn;
+    Gpu::PinnedVector<openbc::MomTag> m_momtags_h;
+#ifdef AMREX_USE_GPU
+    Gpu::DeviceVector<openbc::MomTag> m_momtags_d;
+    Gpu::PinnedVector<int> m_ngpublocks_h;
+    Gpu::DeviceVector<int> m_ngpublocks_d;
+    int m_nthreads_momtag;
+#endif
+
+    int m_nblocks_local = 0;
+    int m_nblocks = 0;
+#ifdef AMREX_USE_MPI
+    Vector<int> m_countvec;
+    Vector<int> m_offset;
+#endif
+
+    IntVect m_ngrowdomain;
+    MultiFab m_crse_grown_faces_phi;
+    MultiFab m_phind;
+    BoxArray m_bag;
+
+    BoxArray m_ba_all;
+    DistributionMapping m_dm_all;
+    Geometry m_geom_all;
+};
+
+}
+
+#endif