Communications arena implementation (#3388)

## Summary
Implement a communications arena for comm buffers to replace
`the_fa_arena`. It creates a separate arena when GPU-aware MPI is used
and `the_arena` is not managed.

## Additional background
The motivation for this is a communication performance degradation that
is observed for GPU-aware MPI with `amrex.the_arena_is_managed=0`.
@WeiqunZhang has a hypothesis that this may be due to the need for
frequent re-registering of comm buffer pointers when using the same
device arena as the other compute data. Hence a separate arena in this
case would alleviate this issue.

`the_fa_arena` is eliminated in this PR and the communication buffer
directly uses `the_comms_arena` to simplify the code.

## Performance tests

The above stated performance degradation is particularly observed with
the `GPU/CNS/Exec/Sod` code under `Tests` and is alleviated by using a
separate comms arena as seen in the performance data below. `original`
refers to the state before we made the change in #3362 related to
`the_fa_arena` pointing to the device arena which allowed
`amrex.the_arena_is_managed=1` with GPU-aware MPI without a significant
performance hit. It is compared with the current development branch and
the proposed comms arena implementation. The data pointing to the
performance improvement from this PR is highlighted.

![Screenshot 2023-06-27 at 4 11 54
PM](https://github.com/AMReX-Codes/amrex/assets/18251677/ae16b822-0178-4679-a90f-255cad6c5451)

In other tests such as the `ABecLaplacian` linear solve or the ERF code,
using `amrex.the_arena_is_managed=0` did not show a significant
performance hit and using this comms arena implementation did not harm
the performance either. More comprehensive tests would be required to
determine the effect on other codes and platforms.

---------

Co-authored-by: Mukul Dave <[email protected]>
Co-authored-by: Weiqun Zhang <[email protected]>

Author: mukul1992

Src/Base/AMReX_Arena.H

@@ -31,6 +31,7 @@ Arena* The_Async_Arena ();
 Arena* The_Device_Arena ();
 Arena* The_Managed_Arena ();
 Arena* The_Pinned_Arena ();
+Arena* The_Comms_Arena ();
 Arena* The_Cpu_Arena ();
 
 struct ArenaInfo

Src/Base/AMReX_Arena.cpp

@@ -34,15 +34,18 @@ namespace {
     Arena* the_managed_arena = nullptr;
     Arena* the_pinned_arena = nullptr;
     Arena* the_cpu_arena = nullptr;
+    Arena* the_comms_arena = nullptr;
 
     Long the_arena_init_size = 0L;
     Long the_device_arena_init_size = 1024*1024*8;
     Long the_managed_arena_init_size = 1024*1024*8;
     Long the_pinned_arena_init_size = 1024*1024*8;
+    Long the_comms_arena_init_size = 1024*1024*8;
     Long the_arena_release_threshold = std::numeric_limits<Long>::max();
     Long the_device_arena_release_threshold = std::numeric_limits<Long>::max();
     Long the_managed_arena_release_threshold = std::numeric_limits<Long>::max();
     Long the_pinned_arena_release_threshold = std::numeric_limits<Long>::max();
+    Long the_comms_arena_release_threshold = std::numeric_limits<Long>::max();
     Long the_async_arena_release_threshold = std::numeric_limits<Long>::max();
 #ifdef AMREX_USE_HIP
     bool the_arena_is_managed = false; // xxxxx HIP FIX HERE
@@ -276,6 +279,7 @@ Arena::Initialize ()
     BL_ASSERT(the_managed_arena == nullptr || the_managed_arena == The_BArena());
     BL_ASSERT(the_pinned_arena == nullptr);
     BL_ASSERT(the_cpu_arena == nullptr || the_cpu_arena == The_BArena());
+    BL_ASSERT(the_comms_arena == nullptr || the_comms_arena == The_BArena());
 
 #ifdef AMREX_USE_GPU
 #ifdef AMREX_USE_SYCL
@@ -292,10 +296,12 @@ Arena::Initialize ()
     pp.queryAdd( "the_device_arena_init_size",  the_device_arena_init_size);
     pp.queryAdd("the_managed_arena_init_size", the_managed_arena_init_size);
     pp.queryAdd( "the_pinned_arena_init_size",  the_pinned_arena_init_size);
+    pp.queryAdd( "the_comms_arena_init_size",  the_comms_arena_init_size);
     pp.queryAdd(       "the_arena_release_threshold" ,         the_arena_release_threshold);
     pp.queryAdd( "the_device_arena_release_threshold",  the_device_arena_release_threshold);
     pp.queryAdd("the_managed_arena_release_threshold", the_managed_arena_release_threshold);
     pp.queryAdd( "the_pinned_arena_release_threshold",  the_pinned_arena_release_threshold);
+    pp.queryAdd("the_comms_arena_release_threshold", the_comms_arena_release_threshold);
     pp.queryAdd(  "the_async_arena_release_threshold",   the_async_arena_release_threshold);
     pp.queryAdd("the_arena_is_managed", the_arena_is_managed);
     pp.queryAdd("abort_on_out_of_gpu_memory", abort_on_out_of_gpu_memory);
@@ -361,6 +367,22 @@ Arena::Initialize ()
                                   (the_pinned_arena_release_threshold));
     the_pinned_arena->registerForProfiling("Pinned Memory");
 
+#ifdef AMREX_USE_GPU
+    if (ParallelDescriptor::UseGpuAwareMpi()) {
+        if (!(the_arena->isDevice())) {
+            the_comms_arena = the_device_arena;
+        } else {
+            the_comms_arena = new CArena(0, ArenaInfo{}.SetDeviceMemory().SetReleaseThreshold
+                                        (the_comms_arena_release_threshold));
+            the_comms_arena->registerForProfiling("Comms Memory");
+        }
+    } else {
+        the_comms_arena = the_pinned_arena;
+    }
+#else
+    the_comms_arena = The_BArena();
+#endif
+
     if (the_device_arena_init_size > 0 && the_device_arena != the_arena) {
         BL_PROFILE("The_Device_Arena::Initialize()");
         void *p = the_device_arena->alloc(the_device_arena_init_size);
@@ -379,6 +401,13 @@ Arena::Initialize ()
         the_pinned_arena->free(p);
     }
 
+    if (the_comms_arena_init_size > 0 && the_comms_arena != the_arena
+        && the_comms_arena != the_device_arena && the_comms_arena != the_pinned_arena) {
+        BL_PROFILE("The_Comms_Arena::Initialize()");
+        void *p = the_comms_arena->alloc(the_comms_arena_init_size);
+        the_comms_arena->free(p);
+    }
+
     the_cpu_arena = The_BArena();
 
     // Initialize the null arena
@@ -440,6 +469,13 @@ Arena::PrintUsage ()
             p->PrintUsage("The  Pinned Arena");
         }
     }
+    if (The_Comms_Arena() && The_Comms_Arena() != The_Device_Arena()
+         && The_Comms_Arena() != The_Pinned_Arena()) {
+        auto* p = dynamic_cast<CArena*>(The_Comms_Arena());
+        if (p) {
+            p->PrintUsage("The   Comms Arena");
+        }
+    }
 }
 
 void
@@ -485,6 +521,13 @@ Arena::PrintUsageToFiles (const std::string& filename, const std::string& messag
             p->PrintUsage(ofs, "The  Pinned Arena", "    ");
         }
     }
+    if (The_Comms_Arena() && The_Comms_Arena() != The_Device_Arena()
+        && The_Comms_Arena() != The_Pinned_Arena()) {
+        auto* p = dynamic_cast<CArena*>(The_Comms_Arena());
+        if (p) {
+            p->PrintUsage(ofs, "The   Comms Arena", "    ");
+        }
+    }
 
     ofs << "\n";
 }
@@ -509,6 +552,13 @@ Arena::Finalize ()
     //   MultiFab mf(...);  // this should be scoped in { ... }
     //   amrex::Finalize();
     // mf cannot be used now, but it can at least be freed without a segfault
+    if (!dynamic_cast<BArena*>(the_comms_arena)) {
+        if (the_comms_arena != the_device_arena && the_comms_arena != the_pinned_arena) {
+            delete the_comms_arena;
+        }
+        the_comms_arena = nullptr;
+    }
+
     if (!dynamic_cast<BArena*>(the_device_arena)) {
         if (the_device_arena != the_arena) {
             delete the_device_arena;
@@ -600,4 +650,14 @@ The_Cpu_Arena ()
     }
 }
 
+Arena*
+The_Comms_Arena ()
+{
+    if        (the_comms_arena) {
+        return the_comms_arena;
+    } else {
+        return The_Null_Arena();
+    }
 }
+
+}

Src/Base/AMReX_FabArray.H

@@ -144,7 +144,7 @@ struct MFInfo {
 struct TheFaArenaDeleter {
     using pointer = char*;
     void operator()(pointer p) const noexcept {
-        The_FA_Arena()->free(p);
+        The_Comms_Arena()->free(p);
     }
 };
 using TheFaArenaPointer = std::unique_ptr<char, TheFaArenaDeleter>;

Src/Base/AMReX_FabArrayBase.H

@@ -721,8 +721,6 @@ bool CheckRcvStats (Vector<MPI_Status>& recv_stats, const Vector<std::size_t>& r
 
 std::ostream& operator<< (std::ostream& os, const FabArrayBase::BDKey& id);
 
-Arena* The_FA_Arena ();
-
 }
 
 #endif

Src/Base/AMReX_FabArrayBase.cpp

@@ -88,7 +88,6 @@ std::vector<std::string>                    FabArrayBase::m_region_tag;
 
 namespace
 {
-    Arena* the_fa_arena = nullptr;
     bool initialized = false;
 }
 
@@ -123,16 +122,6 @@ FabArrayBase::Initialize ()
         MaxComp = 1;
     }
 
-#ifdef AMREX_USE_GPU
-    if (ParallelDescriptor::UseGpuAwareMpi()) {
-        the_fa_arena = The_Device_Arena();
-    } else {
-        the_fa_arena = The_Pinned_Arena();
-    }
-#else
-    the_fa_arena = The_Cpu_Arena();
-#endif
-
     amrex::ExecOnFinalize(FabArrayBase::Finalize);
 
 #ifdef AMREX_MEM_PROFILING
@@ -159,12 +148,6 @@ FabArrayBase::Initialize ()
 #endif
 }
 
-Arena*
-The_FA_Arena ()
-{
-    return the_fa_arena;
-}
-
 FabArrayBase::FabArrayBase (const BoxArray&            bxs,
                             const DistributionMapping& dm,
                             int                        nvar,
@@ -2245,8 +2228,6 @@ FabArrayBase::Finalize ()
 
     m_FA_stats = FabArrayStats();
 
-    the_fa_arena = nullptr;
-
     initialized = false;
 }
 

Src/Base/AMReX_FabArrayCommI.H

@@ -228,7 +228,7 @@ FabArray<FAB>::FillBoundary_finish ()
 
         if (fbd->the_recv_data)
         {
-            amrex::The_FA_Arena()->free(fbd->the_recv_data);
+            amrex::The_Comms_Arena()->free(fbd->the_recv_data);
             fbd->the_recv_data = nullptr;
         }
     }
@@ -237,7 +237,7 @@ FabArray<FAB>::FillBoundary_finish ()
     if (N_snds > 0) {
         Vector<MPI_Status> stats(fbd->send_reqs.size());
         ParallelDescriptor::Waitall(fbd->send_reqs, stats);
-        amrex::The_FA_Arena()->free(fbd->the_send_data);
+        amrex::The_Comms_Arena()->free(fbd->the_send_data);
         fbd->the_send_data = nullptr;
     }
 
@@ -548,7 +548,7 @@ FabArray<FAB>::ParallelCopy_finish ()
 
         if (pcd->the_recv_data)
         {
-            amrex::The_FA_Arena()->free(pcd->the_recv_data);
+            amrex::The_Comms_Arena()->free(pcd->the_recv_data);
             pcd->the_recv_data = nullptr;
         }
     }
@@ -558,7 +558,7 @@ FabArray<FAB>::ParallelCopy_finish ()
             Vector<MPI_Status> stats(pcd->send_reqs.size());
             ParallelDescriptor::Waitall(pcd->send_reqs, stats);
         }
-        amrex::The_FA_Arena()->free(pcd->the_send_data);
+        amrex::The_Comms_Arena()->free(pcd->the_send_data);
         pcd->the_send_data = nullptr;
     }
 
@@ -685,7 +685,7 @@ FabArray<FAB>::PrepareSendBuffers (const MapOfCopyComTagContainers&     SndTags,
 
     if (total_volume > 0)
     {
-        the_send_data = static_cast<char*>(amrex::The_FA_Arena()->alloc(total_volume));
+        the_send_data = static_cast<char*>(amrex::The_Comms_Arena()->alloc(total_volume));
         for (int i = 0, N = static_cast<int>(send_size.size()); i < N; ++i) {
             send_data[i] = the_send_data + offset[i];
         }
@@ -783,7 +783,7 @@ FabArray<FAB>::PostRcvs (const MapOfCopyComTagContainers&  RcvTags,
     }
     else
     {
-        the_recv_data = static_cast<char*>(amrex::The_FA_Arena()->alloc(TotalRcvsVolume));
+        the_recv_data = static_cast<char*>(amrex::The_Comms_Arena()->alloc(TotalRcvsVolume));
 
         for (int i = 0; i < nrecv; ++i)
         {
@@ -1004,7 +1004,7 @@ FillBoundary (Vector<MF*> const& mf, Vector<int> const& scomp,
             recv_size.push_back(nbytes);
         }
 
-        the_recv_data = static_cast<char*>(amrex::The_FA_Arena()->alloc(TotalRcvsVolume));
+        the_recv_data = static_cast<char*>(amrex::The_Comms_Arena()->alloc(TotalRcvsVolume));
 
         int k = 0;
         for (int i = 0; i < nrecv; ++i) {
@@ -1077,7 +1077,7 @@ FillBoundary (Vector<MF*> const& mf, Vector<int> const& scomp,
             send_size.push_back(nbytes);
         }
 
-        the_send_data = static_cast<char*>(amrex::The_FA_Arena()->alloc(TotalSndsVolume));
+        the_send_data = static_cast<char*>(amrex::The_Comms_Arena()->alloc(TotalSndsVolume));
         int k = 0;
         for (int i = 0; i < nsend; ++i) {
             send_data[i] = the_send_data + offset[i];
@@ -1113,13 +1113,13 @@ FillBoundary (Vector<MF*> const& mf, Vector<int> const& scomp,
 
         detail::fbv_copy(recv_tags);
 
-        amrex::The_FA_Arena()->free(the_recv_data);
+        amrex::The_Comms_Arena()->free(the_recv_data);
     }
 
     if (N_snds > 0) {
         Vector<MPI_Status> stats(send_reqs.size());
         ParallelDescriptor::Waitall(send_reqs, stats);
-        amrex::The_FA_Arena()->free(the_send_data);
+        amrex::The_Comms_Arena()->free(the_send_data);
     }
 
 #endif  // #ifdef AMREX_USE_MPI

Src/Base/AMReX_NonLocalBC.cpp

@@ -95,7 +95,7 @@ void PrepareCommBuffers(CommData& comm,
     }
     else
     {
-        comm.the_data.reset(static_cast<char*>(amrex::The_FA_Arena()->alloc(total_volume)));
+        comm.the_data.reset(static_cast<char*>(amrex::The_Comms_Arena()->alloc(total_volume)));
         for (int i = 0; i < N_comms; ++i) {
             comm.data[i] = comm.the_data.get() + comm.offset[i];
         }