Kernel fusing in Geometry (AMReX-Codes#2606)

wjge · Jan 28, 2022 · fd6d649 · fd6d649
1 parent c62334c
commit fd6d649
Showing 1 changed file with 85 additions and 10 deletions.
diff --git a/Src/Base/AMReX_Geometry.cpp b/Src/Base/AMReX_Geometry.cpp
@@ -6,6 +6,7 @@
 #include <AMReX_MultiFab.H>
 #include <AMReX_Utility.H>
 #include <AMReX_SPACE.H>
+#include <AMReX_COORDSYS_C.H>
 
 #include <AMReX_OpenMP.H>
 
@@ -214,14 +215,38 @@ Geometry::GetVolume (MultiFab&       vol,
 }
 
 void
-Geometry::GetVolume (MultiFab&       vol) const
+Geometry::GetVolume (MultiFab& vol) const
 {
+    const auto a_dx = CellSizeArray();
+    if (IsCartesian()) {
+        vol.setVal(AMREX_D_TERM(a_dx[0],*a_dx[1],*a_dx[2]), 0, 1, vol.nGrowVect());
+    } else {
+#if (AMREX_SPACEDIM == 3)
+        amrex::Abort("Geometry::GetVolume: for 3d, only Cartesian is supported");
+#else
+#ifdef AMREX_USE_GPU
+        if (Gpu::inLaunchRegion()) {
+            GpuArray<Real,AMREX_SPACEDIM> a_offset{{AMREX_D_DECL(offset[0],offset[1],offset[2])}};
+            int coord = (int) c_sys;
+            auto const& ma = vol.arrays();
+            ParallelFor(vol, vol.nGrowVect(),
+            [=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k) noexcept
+            {
+                amrex_setvol(makeSingleCellBox(i,j,k), ma[box_no], a_offset, a_dx, coord);
+            });
+            Gpu::streamSynchronize();
+        } else
+#endif
+        {
 #ifdef AMREX_USE_OMP
 #pragma omp parallel if (Gpu::notInLaunchRegion())
 #endif
-    for (MFIter mfi(vol,TilingIfNotGPU()); mfi.isValid(); ++mfi)
-    {
-        CoordSys::SetVolume(vol[mfi], mfi.growntilebox());
+            for (MFIter mfi(vol,TilingIfNotGPU()); mfi.isValid(); ++mfi)
+            {
+                CoordSys::SetVolume(vol[mfi], mfi.growntilebox());
+            }
+        }
+#endif
     }
 }
 
@@ -243,12 +268,29 @@ Geometry::GetDLogA (MultiFab&       dloga,
                     int             ngrow) const
 {
     dloga.define(grds,dm,1,ngrow,MFInfo(),FArrayBoxFactory());
+
+#ifdef AMREX_USE_GPU
+    if (Gpu::inLaunchRegion()) {
+        auto const& ma = dloga.arrays();
+        GpuArray<Real,AMREX_SPACEDIM> a_offset{AMREX_D_DECL(offset[0],offset[1],offset[2])};
+        GpuArray<Real,AMREX_SPACEDIM> a_dx    {AMREX_D_DECL(    dx[0],    dx[1],    dx[2])};
+        int coord = (int) c_sys;
+        ParallelFor(dloga, IntVect(ngrow),
+        [=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k) noexcept
+        {
+            amrex_setdloga(makeSingleCellBox(i,j,k), ma[box_no], a_offset, a_dx, dir, coord);
+        });
+        Gpu::streamSynchronize();
+    } else
+#endif
+    {
 #ifdef AMREX_USE_OMP
 #pragma omp parallel if (Gpu::notInLaunchRegion())
 #endif
-    for (MFIter mfi(dloga,TilingIfNotGPU()); mfi.isValid(); ++mfi)
-    {
-        CoordSys::SetDLogA(dloga[mfi], mfi.growntilebox(), dir);
+        for (MFIter mfi(dloga,TilingIfNotGPU()); mfi.isValid(); ++mfi)
+        {
+            CoordSys::SetDLogA(dloga[mfi], mfi.growntilebox(), dir);
+        }
     }
 }
 #endif
@@ -271,12 +313,45 @@ void
 Geometry::GetFaceArea (MultiFab&       area,
                        int             dir) const
 {
+    const auto a_dx = CellSizeArray();
+
+    if (IsCartesian()) {
+#if (AMREX_SPACEDIM == 1)
+        const Real a0 = 1._rt;
+#elif (AMREX_SPACEDIM == 2)
+        const Real a0 = (dir == 0) ? a_dx[1] : a_dx[0];
+#else
+        const Real a0 = (dir == 0) ? a_dx[1]*a_dx[2]
+            : ((dir == 1) ? a_dx[0]*a_dx[2] : a_dx[0]*a_dx[1]);
+#endif
+        area.setVal(a0, 0, 1, area.nGrowVect());
+    } else {
+#if (AMREX_SPACEDIM == 3)
+        amrex::Abort("Geometry::GetFaceArea:: for 3d, only Cartesian is supported");
+#else
+#ifdef AMREX_USE_GPU
+        if (Gpu::inLaunchRegion()) {
+            GpuArray<Real,AMREX_SPACEDIM> a_offset{AMREX_D_DECL(offset[0],offset[1],offset[2])};
+            int coord = (int) c_sys;
+            auto const& ma = area.arrays();
+            ParallelFor(area, area.nGrowVect(),
+            [=] AMREX_GPU_DEVICE (int box_no, int i, int j, int k) noexcept
+            {
+                amrex_setarea(makeSingleCellBox(i,j,k), ma[box_no], a_offset, a_dx, dir, coord);
+            });
+            Gpu::streamSynchronize();
+        } else
+#endif
+        {
 #ifdef AMREX_USE_OMP
 #pragma omp parallel if (Gpu::notInLaunchRegion())
 #endif
-    for (MFIter mfi(area,TilingIfNotGPU()); mfi.isValid(); ++mfi)
-    {
-        CoordSys::SetFaceArea(area[mfi],mfi.growntilebox(),dir);
+            for (MFIter mfi(area,TilingIfNotGPU()); mfi.isValid(); ++mfi)
+            {
+                CoordSys::SetFaceArea(area[mfi],mfi.growntilebox(),dir);
+            }
+        }
+#endif
     }
 }