From ba95d4cd87580699e1b19b60a89f9ba006fa6665 Mon Sep 17 00:00:00 2001
From: Andrew Myers <atmyers2@gmail.com>
Date: Fri, 8 Mar 2024 09:43:28 -0800
Subject: [PATCH] Port WriteEBSurface by copying EB information to pinned tmps
 before writing. (#3793)

---------

Co-authored-by: Weiqun Zhang <WeiqunZhang@lbl.gov>
---
 Src/EB/AMReX_WriteEBSurface.cpp | 72 ++++++++++++++++++++++++++++-----
 1 file changed, 61 insertions(+), 11 deletions(-)

diff --git a/Src/EB/AMReX_WriteEBSurface.cpp b/Src/EB/AMReX_WriteEBSurface.cpp
index 75a0421c94d..f26420cdcc7 100644
--- a/Src/EB/AMReX_WriteEBSurface.cpp
+++ b/Src/EB/AMReX_WriteEBSurface.cpp
@@ -24,25 +24,62 @@ void WriteEBSurface (const BoxArray & ba, const DistributionMapping & dmap, cons
 
         const auto & sfab    = static_cast<EBFArrayBox const &>(mf_ba[mfi]);
         const auto & my_flag = sfab.getEBCellFlagFab();
+        const auto * my_flag_ptr = &my_flag;
 
         const Box & bx = mfi.validbox();
 
         if (my_flag.getType(bx) == FabType::covered ||
             my_flag.getType(bx) == FabType::regular) { continue; }
 
-        std::array<const MultiCutFab *, AMREX_SPACEDIM> areafrac;
-        const MultiCutFab * bndrycent;
-
-        areafrac  =   ebf->getAreaFrac();
-        bndrycent = &(ebf->getBndryCent());
+        std::array<const CutFab *, AMREX_SPACEDIM> areafrac;
+        const CutFab * bndrycent;
+
+        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
+            areafrac[d]  =   &(*ebf->getAreaFrac()[d])[mfi];
+        }
+        bndrycent = &(ebf->getBndryCent()[mfi]);
+
+#ifdef AMREX_USE_GPU
+        std::unique_ptr<EBCellFlagFab> host_flag;
+        if (my_flag.arena()->isManaged() || my_flag.arena()->isDevice()) {
+            host_flag = std::make_unique<EBCellFlagFab>(my_flag.box(), my_flag.nComp(),
+                                                  The_Pinned_Arena());
+            Gpu::dtoh_memcpy_async(host_flag->dataPtr(), my_flag.dataPtr(),
+                                   host_flag->nBytes());
+            Gpu::streamSynchronize();
+            my_flag_ptr = host_flag.get();
+        }
+
+        std::array<std::unique_ptr<CutFab>, AMREX_SPACEDIM> areafrac_h;
+        for (int d = 0; d < AMREX_SPACEDIM; ++d) {
+            if (areafrac[d]->arena()->isManaged() || areafrac[d]->arena()->isDevice()) {
+                areafrac_h[d] = std::make_unique<CutFab>(areafrac[d]->box(), areafrac[d]->nComp(),
+                                                           The_Pinned_Arena());
+                Gpu::dtoh_memcpy_async(areafrac_h[d]->dataPtr(), areafrac[d]->dataPtr(),
+                                       areafrac[d]->size()*sizeof(Real));
+                Gpu::streamSynchronize();
+                areafrac[d] = areafrac_h[d].get();
+            }
+        }
+
+        std::unique_ptr<CutFab> bndrycent_h;
+        if (bndrycent->arena()->isManaged() || bndrycent->arena()->isDevice()) {
+            bndrycent_h = std::make_unique<CutFab>(bndrycent->box(), bndrycent->nComp(),
+                                                        The_Pinned_Arena());
+            Gpu::dtoh_memcpy_async(bndrycent_h->dataPtr(), bndrycent->dataPtr(),
+                                   bndrycent->size()*sizeof(Real));
+            Gpu::streamSynchronize();
+            bndrycent = bndrycent_h.get();
+        }
+#endif
 
         eb_to_pvd.EBToPolygon(
                 problo, dx,
-                bx, my_flag.const_array(),
-                bndrycent->const_array(mfi),
-                areafrac[0]->const_array(mfi),
-                areafrac[1]->const_array(mfi),
-                areafrac[2]->const_array(mfi));
+                bx, my_flag_ptr->const_array(),
+                bndrycent->const_array(),
+                areafrac[0]->const_array(),
+                areafrac[1]->const_array(),
+                areafrac[2]->const_array());
     }
 
     int cpu = ParallelDescriptor::MyProc();
@@ -58,13 +95,26 @@ void WriteEBSurface (const BoxArray & ba, const DistributionMapping & dmap, cons
 
         const auto & sfab    = static_cast<EBFArrayBox const &>(mf_ba[mfi]);
         const auto & my_flag = sfab.getEBCellFlagFab();
+        const auto * my_flag_ptr = &my_flag;
 
         const Box & bx = mfi.validbox();
 
         if (my_flag.getType(bx) == FabType::covered ||
             my_flag.getType(bx) == FabType::regular) { continue; }
 
-        eb_to_pvd.EBGridCoverage(cpu, problo, dx, bx, my_flag.const_array());
+#ifdef AMREX_USE_GPU
+        std::unique_ptr<EBCellFlagFab> host_flag;
+        if (my_flag.arena()->isManaged() || my_flag.arena()->isDevice()) {
+            host_flag = std::make_unique<EBCellFlagFab>(my_flag.box(), my_flag.nComp(),
+                                                  The_Pinned_Arena());
+            Gpu::dtoh_memcpy_async(host_flag->dataPtr(), my_flag.dataPtr(),
+                                   host_flag->nBytes());
+            Gpu::streamSynchronize();
+            my_flag_ptr = host_flag.get();
+        }
+#endif
+
+        eb_to_pvd.EBGridCoverage(cpu, problo, dx, bx, my_flag_ptr->const_array());
     }
 }