Curl Curl Solver: Option to use PCG instead of LU (AMReX-Codes#3812)

Add support for variable beta coefficient. --------- Co-authored-by: Revathi Jambunathan <[email protected]>
ajnonaka · Aug 27, 2024 · df0d594 · df0d594
1 parent 12002e7
commit df0d594
Show file tree

Hide file tree

Showing 8 changed files with 277 additions and 117 deletions.
diff --git a/Src/LinearSolvers/CMakeLists.txt b/Src/LinearSolvers/CMakeLists.txt
@@ -21,6 +21,7 @@ foreach(D IN LISTS AMReX_SPACEDIM)
        MLMG/AMReX_MLCellABecLap_K.H
        MLMG/AMReX_MLCellABecLap_${D}D_K.H
        MLMG/AMReX_MLCGSolver.H
+       MLMG/AMReX_PCGSolver.H
        MLMG/AMReX_MLABecLaplacian.H
        MLMG/AMReX_MLABecLap_K.H
        MLMG/AMReX_MLABecLap_${D}D_K.H

diff --git a/Src/LinearSolvers/MLMG/AMReX_MLCurlCurl.H b/Src/LinearSolvers/MLMG/AMReX_MLCurlCurl.H
@@ -58,6 +58,8 @@ public:
         return std::string("curl of curl");
     }
 
+    bool setUsePCG (bool flag) { return std::exchange(m_use_pcg, flag); }
+
     void setLevelBC (int amrlev, const MF* levelbcdata,
                      const MF* robinbc_a = nullptr,
                      const MF* robinbc_b = nullptr,
@@ -137,6 +139,7 @@ private:
     Vector<Vector<std::unique_ptr<Gpu::DeviceScalar
                                   <LUSolver<AMREX_SPACEDIM*2,RT>>>>> m_lusolver;
     Vector<Vector<Array<std::unique_ptr<MultiFab>,3>>> m_bcoefs;
+    bool m_use_pcg = false;
 };
 
 }

diff --git a/Src/LinearSolvers/MLMG/AMReX_MLCurlCurl.cpp b/Src/LinearSolvers/MLMG/AMReX_MLCurlCurl.cpp
@@ -353,22 +353,36 @@ void MLCurlCurl::smooth4 (int amrlev, int mglev, MF& sol, MF const& rhs,
         auto* plusolver = m_lusolver[amrlev][mglev]->dataPtr();
         ParallelFor(nmf, [=] AMREX_GPU_DEVICE (int bno, int i, int j, int k)
         {
-            mlcurlcurl_gs4(i,j,k,ex[bno],ey[bno],ez[bno],rhsx[bno],rhsy[bno],rhsz[bno],
+            mlcurlcurl_gs4_lu(i,j,k,ex[bno],ey[bno],ez[bno],
+                              rhsx[bno],rhsy[bno],rhsz[bno],
 #if (AMREX_SPACEDIM == 2)
-                           b,
+                              b,
 #endif
-                           adxinv,color,*plusolver,dinfo,sinfo);
+                              adxinv,color,*plusolver,dinfo,sinfo);
         });
     } else {
         auto const& bcx = m_bcoefs[amrlev][mglev][0]->const_arrays();
         auto const& bcy = m_bcoefs[amrlev][mglev][1]->const_arrays();
         auto const& bcz = m_bcoefs[amrlev][mglev][2]->const_arrays();
-        ParallelFor(nmf, [=] AMREX_GPU_DEVICE (int bno, int i, int j, int k)
-        {
+        if (m_use_pcg) {
+            ParallelFor(nmf, [=] AMREX_GPU_DEVICE (int bno, int i, int j, int k)
+            {
 
-            mlcurlcurl_gs4(i,j,k,ex[bno],ey[bno],ez[bno],rhsx[bno],rhsy[bno],rhsz[bno],
-                           adxinv,color,bcx[bno],bcy[bno],bcz[bno],dinfo,sinfo);
-        });
+                mlcurlcurl_gs4<true>(i,j,k,ex[bno],ey[bno],ez[bno],
+                                     rhsx[bno],rhsy[bno],rhsz[bno],
+                                     adxinv,color,bcx[bno],bcy[bno],bcz[bno],
+                                     dinfo,sinfo);
+            });
+        } else {
+            ParallelFor(nmf, [=] AMREX_GPU_DEVICE (int bno, int i, int j, int k)
+            {
+
+                mlcurlcurl_gs4<false>(i,j,k,ex[bno],ey[bno],ez[bno],
+                                      rhsx[bno],rhsy[bno],rhsz[bno],
+                                      adxinv,color,bcx[bno],bcy[bno],bcz[bno],
+                                      dinfo,sinfo);
+            });
+        }
     }
     Gpu::streamSynchronize();
 }