Coalesced Buffer Communication

src/CMakeLists.txt

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

src/bvals/comms/bnd_info.cpp

@@ -24,6 +24,7 @@
 
 #include "basic_types.hpp"
 #include "bvals/comms/bnd_info.hpp"
+#include "bvals/comms/bvals_utils.hpp"
 #include "bvals/neighbor_block.hpp"
 #include "config.hpp"
 #include "globals.hpp"
@@ -251,7 +252,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
@@ -277,7 +278,7 @@ BndInfo::BndInfo(MeshBlock *pmb, const NeighborBlock &nb,
   allocated = v->IsAllocated();
   alloc_status = v->GetAllocationStatus();
 
-  buf = combuf->buffer();
+  if (combuf != nullptr) buf = combuf->buffer();
   same_to_same = pmb->gid == nb.gid && nb.offsets.IsCell();
   lcoord_trans = nb.lcoord_trans;
   if (!allocated) return;
@@ -305,6 +306,24 @@ BndInfo::BndInfo(MeshBlock *pmb, const NeighborBlock &nb,
   }
 }
 
+BndId BndId::GetSend(MeshBlock *pmb, const NeighborBlock &nb,
+                     std::shared_ptr<Variable<Real>> v, BoundaryType b_type,
+                     int partition, int start_idx) {
+  auto [send_gid, recv_gid, vlabel, loc, extra_id] = SendKey(pmb, nb, v, b_type);
+  BndId out;
+  out.send_gid() = send_gid;
+  out.recv_gid() = recv_gid;
+  out.loc_idx() = loc;
+  out.var_id() = v->GetUniqueID();
+  out.extra_id() = extra_id;
+  out.rank_send() = Globals::my_rank;
+  out.rank_recv() = nb.rank;
+  out.partition() = partition;
+  out.size() = BndInfo::GetSendBndInfo(pmb, nb, v, nullptr).size();
+  out.start_idx() = start_idx;
+  return out;
+}
+
 BndInfo BndInfo::GetSendBndInfo(MeshBlock *pmb, const NeighborBlock &nb,
                                 std::shared_ptr<Variable<Real>> v,
                                 CommBuffer<buf_pool_t<Real>::owner_t> *buf) {
@@ -326,7 +345,7 @@ BndInfo BndInfo::GetSetBndInfo(MeshBlock *pmb, const NeighborBlock &nb,
   if (nb.offsets.IsCell()) idx_range_type = IndexRangeType::InteriorRecv;
   BndInfo out(pmb, nb, v, buf, idx_range_type);
 
-  auto buf_state = buf->GetState();
+  auto buf_state = buf != nullptr ? buf->GetState() : BufferState::received;
   if (buf_state == BufferState::received) {
     out.buf_allocated = true;
   } else if (buf_state == BufferState::received_null) {

src/bvals/comms/bnd_info.hpp

@@ -48,13 +48,65 @@ enum class IndexRangeType {
   InteriorRecv
 };
 
+struct BndId {
+  constexpr static std::size_t NDAT = 10;
+  int data[NDAT];
+
+  // Information for identifying the buffer with a communication
+  // channel, variable, and the ranks it is communicated across
+  int &send_gid() { return data[0]; }
+  int &recv_gid() { return data[1]; }
+  int &loc_idx() { return data[2]; }
+  int &var_id() { return data[3]; }
+  int &extra_id() { return data[4]; }
+  int &rank_send() { return data[5]; }
+  int &rank_recv() { return data[6]; }
+  BoundaryType bound_type;
+
+  // MeshData partition id of the *sender*
+  // not set by constructors and only necessary for coalesced comms
+  int &partition() { return data[7]; }
+  int &size() { return data[8]; }
+  int &start_idx() { return data[9]; }
+
+  KOKKOS_DEFAULTED_FUNCTION
+  BndId() = default;
+  KOKKOS_DEFAULTED_FUNCTION
+  BndId(const BndId &) = default;
+
+  explicit BndId(const int *const data_in) {
+    for (int i = 0; i < NDAT; ++i) {
+      data[i] = data_in[i];
+    }
+  }
+
+  void Serialize(int *data_out) {
+    for (int i = 0; i < NDAT; ++i) {
+      data_out[i] = data[i];
+    }
+  }
+
+  static BndId GetSend(MeshBlock *pmb, const NeighborBlock &nb,
+                       std::shared_ptr<Variable<Real>> v, BoundaryType b_type,
+                       int partition, int start_idx);
+};
+
 struct BndInfo {
   int ntopological_elements = 1;
   using TE = TopologicalElement;
   TE topo_idx[3]{TE::CC, TE::CC, TE::CC};
   SpatiallyMaskedIndexer6D idxer[3];
   forest::LogicalCoordinateTransformation lcoord_trans;
 
+  KOKKOS_FORCEINLINE_FUNCTION
+  int size() const {
+    int s = 0;
+    for (int n = 0; n < ntopological_elements; ++n) {
+      s += idxer[n].size();
+    }
+    return s;
+  }
+
   CoordinateDirection dir{CoordinateDirection::X0DIR};
   bool allocated = true;
   bool buf_allocated = true;
@@ -124,7 +176,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,19 +111,25 @@ 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);
+
 #else
       // Setting to zero is fine here since this doesn't actually get used when everything
       // is on the same rank
       mpi_comm_t comm = 0;
 #endif
 
     bool use_sparse_buffers = v->IsSet(Metadata::Sparse);
-    auto get_resource_method = [pmesh, buf_size]() {
+    auto get_resource_method = [pmesh, buf_size](int size) {
+      PARTHENON_REQUIRE(size <= buf_size,
+                        "Asking for a buffer that is larger than size of pool.");
       return buf_pool_t<Real>::owner_t(pmesh->pool_map.at(buf_size).Get());
     };
 
     // Build send buffer (unless this is a receiving flux boundary)
     if constexpr (IsSender(BTYPE)) {
+      // Register this buffer with the combined buffers
+      if (receiver_rank != sender_rank)
+        pmesh->pcombined_buffers->AddSendBuffer(md->partition, pmb, nb, v, BTYPE);
       auto s_key = SendKey(pmb, nb, v, BTYPE);
       if (buf_map.count(s_key) == 0)
         buf_map[s_key] = CommBuffer<buf_pool_t<Real>::owner_t>(
@@ -133,6 +140,7 @@ void BuildBoundaryBufferSubset(std::shared_ptr<MeshData<Real>> &md,
     // Also build the non-local receive buffers here
     if constexpr (IsReceiver(BTYPE)) {
       if (sender_rank != receiver_rank) {
+        pmesh->pcombined_buffers->AddRecvBuffer(pmb, nb, v, BTYPE);
         auto r_key = ReceiveKey(pmb, nb, v, BTYPE);
         if (buf_map.count(r_key) == 0)
           buf_map[r_key] = CommBuffer<buf_pool_t<Real>::owner_t>(

src/bvals/comms/bvals_utils.hpp

@@ -66,6 +66,11 @@ inline Mesh::channel_key_t ReceiveKey(const MeshBlock *pmb, const NeighborBlock
   return {sender_id, receiver_id, pcv->label(), location_idx, other};
 }
 
+inline Mesh::channel_key_t GetChannelKey(BndId &in) {
+  return {in.send_gid(), in.recv_gid(), Variable<Real>::GetLabel(in.var_id()),
+          in.loc_idx(), in.extra_id()};
+}
+
 // Build a vector of pointers to all of the sending or receiving communication buffers on
 // MeshData md. This cache is important for performance, since this elides a map look up
 // for the buffer every time the bvals code iterates over boundaries.

src/bvals/comms/combined_buffers.cpp

@@ -0,0 +1,125 @@
+//========================================================================================
+// (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.
+//========================================================================================
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "basic_types.hpp"
+#include "bvals/comms/bvals_utils.hpp"
+#include "bvals/comms/combined_buffers.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 {
+
+CombinedBuffersRank::CombinedBuffersRank(int o_rank, BoundaryType b_type, bool send)
+    : other_rank(o_rank), sender(send), buffers_built(false) {
+  if (sender) {
+    message = com_buf_t(1234, Globals::my_rank, other_rank, MPI_COMM_WORLD, [](int size) {
+      PARTHENON_FAIL("Comms should not be allocating sender.");
+      return std::vector<int>(size);
+    });
+  } else {
+    message = com_buf_t(1234, other_rank, Globals::my_rank, MPI_COMM_WORLD,
+                        [](int size) { return std::vector<int>(size); });
+  }
+  PARTHENON_REQUIRE(other_rank != Globals::my_rank, "Should only build for other ranks.");
+}
+
+void CombinedBuffersRank::AddSendBuffer(int partition, MeshBlock *pmb,
+                                        const NeighborBlock &nb,
+                                        const std::shared_ptr<Variable<Real>> &var,
+                                        BoundaryType b_type) {
+  if (current_size.count(partition) == 0) current_size[partition] = 0;
+  auto &cur_size = current_size[partition];
+  combined_info[partition].push_back(
+      BndId::GetSend(pmb, nb, var, b_type, partition, cur_size));
+  cur_size += combined_info[partition].back().size();
+}
+
+bool CombinedBuffersRank::TryReceiveBufInfo() {
+  PARTHENON_REQUIRE(!sender, "Trying to receive on a combined sender.");
+  if (buffers_built) return buffers_built;
+
+  bool received = message.TryReceive();
+  if (received) {
+    auto &mess_buf = message.buffer();
+    int npartitions = mess_buf[0];
+    // Unpack into per combined buffer information
+    int idx{nglobal};
+    for (int p = 0; p < npartitions; ++p) {
+      const int partition = mess_buf[idx++];
+      const int nbuf = mess_buf[idx++];
+      const int total_size = mess_buf[idx++];
+      combined_buffers[partition] = buf_t("combined recv buffer", total_size);
+      auto &cr_info = combined_info[partition];
+      for (int b = 0; b < nbuf; ++b) {
+        cr_info.emplace_back(&(mess_buf[idx]));
+        idx += BndId::NDAT;
+      }
+    }
+    message.Stale();
+    buffers_built = true;
+    return true;
+  }
+  return false;
+}
+
+void CombinedBuffersRank::ResolveSendBuffersAndSendInfo() {
+  // First calculate the total size of the message
+  int total_buffers{0};
+  for (auto &[partition, buf_struct_vec] : combined_info)
+    total_buffers += buf_struct_vec.size();
+  int total_partitions = combined_info.size();
+
+  auto &mess_buf = message.buffer();
+  int mesg_size = nglobal + nper_part * total_partitions + BndId::NDAT * total_buffers;
+  mess_buf.resize(mesg_size);
+
+  mess_buf[0] = total_partitions;
+
+  // Pack the data
+  int idx{nglobal};
+  for (auto &[partition, buf_struct_vec] : combined_info) {
+    mess_buf[idx++] = partition;               // Used as the comm tag
+    mess_buf[idx++] = buf_struct_vec.size();   // Number of buffers
+    mess_buf[idx++] = current_size[partition]; // combined size of buffers
+    for (auto &buf_struct : buf_struct_vec) {
+      buf_struct.Serialize(&(mess_buf[idx]));
+      idx += BndId::NDAT;
+    }
+  }
+
+  message.Send();
+
+  // Allocate the combined buffers
+  int total_size{0};
+  for (auto &[partition, size] : current_size)
+    total_size += size;
+
+  buf_t alloc_at_once("shared combined buffer", total_size);
+  int current_position{0};
+  for (auto &[partition, size] : current_size) {
+    combined_buffers[partition] =
+        buf_t(alloc_at_once, std::make_pair(current_position, current_position + size));
+    current_position += size;
+  }
+  buffers_built = true;
+}
+} // namespace parthenon