OverlapSync nowait and finish. (AMReX-Codes#2346)

kaiserpeng · Sep 24, 2021 · 66c9080 · 66c9080
1 parent 777a292
commit 66c9080
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 8 deletions.
diff --git a/Src/Base/AMReX_FabArray.H b/Src/Base/AMReX_FabArray.H
@@ -151,7 +151,6 @@ using TheFaArenaPointer = std::unique_ptr<char, TheFaArenaDeleter>;
 template <class FAB>
 struct FBData {
 
-    //! Data used in non-blocking FillBoundary
     const FabArrayBase::FB*  fb = nullptr;
     int                 scomp;
     int                 ncomp;
@@ -1128,8 +1127,10 @@ public:
 #endif
 
     std::unique_ptr<FBData<FAB>> fbd;
-
     std::unique_ptr<PCData<FAB>> pcd;
+
+    // Pointer to temporary fab used in non-blocking OverrideSync
+    std::unique_ptr< FabArray<FAB> > os_temp;
 };
 
 

diff --git a/Src/Base/AMReX_FabArrayUtility.H b/Src/Base/AMReX_FabArrayUtility.H
@@ -1358,13 +1358,27 @@ prefetchToDevice (FabArray<FAB> const& fa, const bool synchronous = true)
 #endif
 }
 
+
 template <class FAB, class IFAB, class bar = std::enable_if_t<IsBaseFab<FAB>::value
                                                                && IsBaseFab<IFAB>::value> >
 void
 OverrideSync (FabArray<FAB> & fa, FabArray<IFAB> const& msk, const Periodicity& period)
 {
     BL_PROFILE("OverrideSync()");
 
+    OverrideSync_nowait(fa, msk, period);
+    OverrideSync_finish(fa);
+}
+
+
+template <class FAB, class IFAB, class bar = std::enable_if_t<IsBaseFab<FAB>::value
+                                                               && IsBaseFab<IFAB>::value> >
+void
+OverrideSync_nowait (FabArray<FAB> & fa, FabArray<IFAB> const& msk, const Periodicity& period)
+{
+    BL_PROFILE("OverrideSync_nowait()");
+    AMREX_ASSERT_WITH_MESSAGE(!fa.os_temp, "OverrideSync_nowait() called when already in progress.");
+
     if (fa.ixType().cellCentered()) return;
 
     const int ncomp = fa.nComp();
@@ -1397,12 +1411,24 @@ OverrideSync (FabArray<FAB> & fa, FabArray<IFAB> const& msk, const Periodicity&
         }
     }
 
-    FabArray<FAB> tmpmf(fa.boxArray(), fa.DistributionMap(), ncomp, 0,
-                        MFInfo(), fa.Factory());
-    tmpmf.setVal(0);
-    tmpmf.ParallelCopy(fa, period, FabArrayBase::ADD);
+    fa.os_temp = std::make_unique< FabArray<FAB> > ( fa.boxArray(), fa.DistributionMap(),
+                                                     ncomp, 0, MFInfo(), fa.Factory() );
+    fa.os_temp->setVal(0);
+    fa.os_temp->ParallelCopy_nowait(fa, period, FabArrayBase::ADD);
+}
+
+template <class FAB, class bar = std::enable_if_t<IsBaseFab<FAB>::value> >
+void
+OverrideSync_finish (FabArray<FAB> & fa)
+{
+    BL_PROFILE("OverrideSync_finish()");
+
+    if (fa.ixType().cellCentered()) return;
+
+    fa.os_temp->ParallelCopy_finish();
+    amrex::Copy(fa, *(fa.os_temp), 0, 0, fa.nComp(), 0);
 
-    amrex::Copy(fa, tmpmf, 0, 0, ncomp, 0);
+    fa.os_temp.reset();
 }
 
 template <class FAB, class foo = std::enable_if_t<IsBaseFab<FAB>::value> >

diff --git a/Src/Base/AMReX_MultiFab.H b/Src/Base/AMReX_MultiFab.H
@@ -665,6 +665,11 @@ public:
     void OverrideSync (const Periodicity& period = Periodicity::NonPeriodic());
     void OverrideSync (const iMultiFab& msk, const Periodicity& period = Periodicity::NonPeriodic());
 
+    void OverrideSync_nowait (const Periodicity& period = Periodicity::NonPeriodic());
+    void OverrideSync_nowait (const iMultiFab& msk, const Periodicity& period = Periodicity::NonPeriodic());
+    void OverrideSync_finish ();
+
+
     static void Initialize ();
     static void Finalize ();
 

diff --git a/Src/Base/AMReX_MultiFab.cpp b/Src/Base/AMReX_MultiFab.cpp
@@ -1778,7 +1778,7 @@ MultiFab::OverrideSync (const Periodicity& period)
 {
     if (ixType().cellCentered()) return;
     auto msk = this->OwnerMask(period);
-    this->OverrideSync(*msk, period);
+    amrex::OverrideSync(*this, *msk, period);
 }
 
 void
@@ -1787,4 +1787,24 @@ MultiFab::OverrideSync (const iMultiFab& msk, const Periodicity& period)
     amrex::OverrideSync(*this, msk, period);
 }
 
+void
+MultiFab::OverrideSync_nowait (const Periodicity& period)
+{
+    if (ixType().cellCentered()) return;
+    auto msk = this->OwnerMask(period);
+    amrex::OverrideSync_nowait(*this, *msk, period);
+}
+
+void
+MultiFab::OverrideSync_nowait (const iMultiFab& msk, const Periodicity& period)
+{
+    amrex::OverrideSync_nowait(*this, msk, period);
+}
+
+void
+MultiFab::OverrideSync_finish ()
+{
+    amrex::OverrideSync_finish(*this);
+}
+
 }