start on combined communication

src/CMakeLists.txt

@@ -101,6 +101,7 @@ add_library(parthenon
   bvals/comms/bnd_info.cpp
   bvals/comms/bnd_info.hpp
   bvals/comms/boundary_communication.cpp
+  bvals/comms/combined_buffers.hpp
   bvals/comms/tag_map.cpp
   bvals/comms/tag_map.hpp
 

src/bvals/comms/bnd_info.cpp

@@ -251,7 +251,7 @@ CalcIndices(const NeighborBlock &nb, MeshBlock *pmb,
                                   {s[2], e[2]}, {s[1], e[1]}, {s[0], e[0]});
 }
 
-int GetBufferSize(MeshBlock *pmb, const NeighborBlock &nb,
+int GetBufferSize(const MeshBlock *const pmb, const NeighborBlock &nb,
                   std::shared_ptr<Variable<Real>> v) {
   // This does not do a careful job of calculating the buffer size, in many
   // cases there will be some extra storage that is not required, but there

src/bvals/comms/bnd_info.hpp

@@ -124,7 +124,7 @@ struct ProResInfo {
                                         std::shared_ptr<Variable<Real>> v);
 };
 
-int GetBufferSize(MeshBlock *pmb, const NeighborBlock &nb,
+int GetBufferSize(const MeshBlock *const pmb, const NeighborBlock &nb,
                   std::shared_ptr<Variable<Real>> v);
 
 using BndInfoArr_t = ParArray1D<BndInfo>;

src/bvals/comms/boundary_communication.cpp

@@ -138,6 +138,15 @@ TaskStatus SendBoundBufs(std::shared_ptr<MeshData<Real>> &md) {
           sending_nonzero_flags(b) = non_zero[0] || non_zero[1] || non_zero[2];
         });
       });
+  // 1. Parallel scan per rank to get the starting indices of the buffers
+
+  // 2. Check the size of the buffer (how do you do this without extra DtoH call?) and
+  // possibly allocate more storage
+  //    [Alternatively could just allocate to maximal size initially]
+
+  // 3. Pack the combined buffers
+
+  // 4. Send the combined buffers
 
   // Send buffers
   if (Globals::sparse_config.enabled)

src/bvals/comms/build_boundary_buffers.cpp

@@ -27,6 +27,7 @@
 
 #include "bvals_in_one.hpp"
 #include "bvals_utils.hpp"
+#include "combined_buffers.hpp"
 #include "config.hpp"
 #include "globals.hpp"
 #include "interface/variable.hpp"
@@ -110,6 +111,11 @@ void BuildBoundaryBufferSubset(std::shared_ptr<MeshData<Real>> &md,
     tag = pmesh->tag_map.GetTag(pmb, nb);
     auto comm_label = v->label();
     mpi_comm_t comm = pmesh->GetMPIComm(comm_label);
+
+    // Register this buffer with the combined buffers
+    if (receiver_rank != sender_rank) {
+      pmesh->pcombined_buffers->AddSendBuffer(md->partition, pmb, nb, v, BTYPE);
+    }
 #else
       // Setting to zero is fine here since this doesn't actually get used when everything
       // is on the same rank

src/bvals/comms/combined_buffers.hpp

@@ -0,0 +1,108 @@
+//========================================================================================
+// (C) (or copyright) 2020-2024. Triad National Security, LLC. All rights reserved.
+//
+// This program was produced under U.S. Government contract 89233218CNA000001 for Los
+// Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC
+// for the U.S. Department of Energy/National Nuclear Security Administration. All rights
+// in the program are reserved by Triad National Security, LLC, and the U.S. Department
+// of Energy/National Nuclear Security Administration. The Government is granted for
+// itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide
+// license in this material to reproduce, prepare derivative works, distribute copies to
+// the public, perform publicly and display publicly, and to permit others to do so.
+//========================================================================================
+
+#ifndef BVALS_COMMS_COMBINED_BUFFERS_HPP_
+#define BVALS_COMMS_COMBINED_BUFFERS_HPP_
+
+#include <map>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "basic_types.hpp"
+#include "bvals/comms/bvals_utils.hpp"
+#include "bvals/neighbor_block.hpp"
+#include "coordinates/coordinates.hpp"
+#include "interface/variable.hpp"
+#include "mesh/mesh.hpp"
+#include "mesh/meshblock.hpp"
+#include "utils/communication_buffer.hpp"
+
+namespace parthenon {
+
+struct BufferStructure {
+  // These first five variables should be enough information to
+  // uniquely identify the buffer
+  int tag;    // Tag defining communication channel between blocks
+              // (which subsumes send_gid, recv_gid, location_on_block)
+              // within a given MPI rank pair
+  int var_id; // We use an int for the Uid_t since we will be sending via MPI
+  int extra_id;
+  int rank_send; // MPI rank of sender
+  int rank_recv; // MPI rank of receiver
+
+  // Other information that could be useful to sending messages
+  int size;                 // Size of the buffer
+  Mesh::channel_key_t key;  // Actual key
+  bool currently_allocated; // Current allocation status of the buffer
+  int partition;            // Partition of sender
+  BoundaryType btype;       // Type of boundary this was registered for
+
+  static BufferStructure Send(int partition, const MeshBlock *const pmb,
+                              const NeighborBlock &nb,
+                              const std::shared_ptr<Variable<Real>> &var,
+                              BoundaryType b_type) {
+    BufferStructure out;
+    out.tag = pmb->pmy_mesh->tag_map.GetTag(pmb, nb);
+    out.var_id = var->GetUniqueID();
+    out.extra_id = static_cast<int>(b_type);
+    out.rank_send = Globals::my_rank;
+    out.rank_recv = nb.rank;
+
+    out.key = SendKey(pmb, nb, var, b_type);
+    out.size = GetBufferSize(pmb, nb, var);
+    out.currently_allocated = true;
+    out.partition = partition;
+    out.btype = b_type;
+    return out;
+  }
+};
+
+// Structure containing the information required for sending coalesced
+// messages between ranks
+struct CombinedBuffersRank {
+  using coalesced_message_structure_t = std::vector<BufferStructure>;
+
+  // Rank that these buffers communicate with
+  int other_rank;
+
+  // map from partion id to coalesced message structure for communication
+  // from this rank to other_rank
+  std::map<int, coalesced_message_structure_t> combined_send_info;
+  std::map<int, ParArray1D<Real>> combined_send_buffers;
+
+  // map from neighbor partition id to coalesced message structures that
+  // this rank can receive from other_rank. We will use the partition id
+  // as the mpi tag
+  std::map<int, coalesced_message_structure_t> combined_recv_info;
+  std::map<int, ParArray1D<Real>> combined_recv_buffers;
+
+  void AddSendBuffer(int partition, const MeshBlock *const &pmb, const NeighborBlock &nb,
+                     const std::shared_ptr<Variable<Real>> &var, BoundaryType b_type) {
+    combined_send_info[partition].push_back(
+        BufferStructure::Send(partition, pmb, nb, var, b_type));
+  }
+};
+
+struct CombinedBuffers {
+  // Combined buffers for each rank
+  std::vector<CombinedBuffersRank> combined_buffers;
+  void AddSendBuffer(int partition, const MeshBlock *const pmb, const NeighborBlock &nb,
+                     const std::shared_ptr<Variable<Real>> &var, BoundaryType b_type) {
+    combined_buffers[nb.rank].AddSendBuffer(partition, pmb, nb, var, b_type);
+  }
+};
+
+} // namespace parthenon
+
+#endif // BVALS_COMMS_COMBINED_BUFFERS_HPP_

src/mesh/mesh.cpp

@@ -42,6 +42,7 @@
 #include "application_input.hpp"
 #include "bvals/boundary_conditions.hpp"
 #include "bvals/bvals.hpp"
+#include "bvals/comms/combined_buffers.hpp"
 #include "defs.hpp"
 #include "globals.hpp"
 #include "interface/packages.hpp"
@@ -85,7 +86,8 @@ Mesh::Mesh(ParameterInput *pin, ApplicationInput *app_in, Packages_t &packages,
       lb_manual_(), nslist(Globals::nranks), nblist(Globals::nranks),
       nref(Globals::nranks), nderef(Globals::nranks), rdisp(Globals::nranks),
       ddisp(Globals::nranks), bnref(Globals::nranks), bnderef(Globals::nranks),
-      brdisp(Globals::nranks), bddisp(Globals::nranks) {
+      brdisp(Globals::nranks), bddisp(Globals::nranks),
+      pcombined_buffers(std::make_shared<CombinedBuffers>()) {
   // Allow for user overrides to default Parthenon functions
   if (app_in->InitUserMeshData != nullptr) {
     InitUserMeshData = app_in->InitUserMeshData;

src/mesh/mesh.hpp

@@ -59,6 +59,7 @@ namespace parthenon {
 
 // Forward declarations
 class ApplicationInput;
+class CombinedBuffers;
 class MeshBlock;
 class MeshRefinement;
 class Packages_t;
@@ -218,6 +219,9 @@ class Mesh {
 
   // Ordering here is important to prevent deallocation of pools before boundary
   // communication buffers
+  // channel_key_t is tuple of (gid_sender, gid_receiver, variable_name,
+  // block_location_idx, extra_delineater) which uniquely define a communication channel
+  // between two blocks for a given variable
   using channel_key_t = std::tuple<int, int, std::string, int, int>;
   using comm_buf_t = CommBuffer<buf_pool_t<Real>::owner_t>;
   std::unordered_map<int, buf_pool_t<Real>> pool_map;
@@ -226,6 +230,8 @@ class Mesh {
   comm_buf_map_t boundary_comm_map;
   TagMap tag_map;
 
+  std::shared_ptr<CombinedBuffers> pcombined_buffers;
+
 #ifdef MPI_PARALLEL
   MPI_Comm GetMPIComm(const std::string &label) const { return mpi_comm_map_.at(label); }
 #endif