perf(evolution): ⚡ stop retaining the P-sized exchange state per layer

cpp/monoprop/Evolution.cpp

@@ -60,8 +60,14 @@ auto combine_endpoint_contrib(const EndpointContrib &a, const EndpointContrib &b
 struct FlatExchangeBuffers {
     VecD send_buffer;
     VecD recv_buffer;
-    std::vector<int> recv_counts;
-    std::vector<int> recv_displs;
+    // The layout for the exchange currently being posted, derived per layer rather than read from
+    // one retained per layer. Reused, so it allocates once per thread per world size. Sharing one
+    // instance across layers is only sound because at most one exchange is in flight per thread --
+    // the same invariant send_buffer above has always required.
+    //
+    // ONE layout, not two: it describes the recv side as well as the send side. See
+    // derive_layer_exchange.
+    LayerExchangeLayout layout;
 };
 
 auto &acquire_flat_exchange_buffers() {
@@ -73,21 +79,39 @@ auto &acquire_flat_exchange_buffers() {
 }
 
 void resize_flat_exchange_buffers(const LayerExchangeLayout &layout, FlatExchangeBuffers &buffers) {
-    // Recv size isn't known until counts are exchanged; keep it at 1 element so data() stays non-null.
-    const size_t send_alloc = layout.total_count == 0 ? 1 : layout.total_count;
-    buffers.send_buffer.resize(send_alloc);
-    buffers.recv_buffer.resize(1);
-    buffers.recv_counts.clear();
-    buffers.recv_displs.clear();
+    const size_t alloc = layout.total_count == 0 ? 1 : layout.total_count;
+    buffers.send_buffer.resize(alloc);
 }
 
-auto active_evolution_exchange_layout(const LayerTraversal &layer, const mpi::Comm &comm)
-    -> const LayerExchangeLayout * {
-    if (mpi::size(comm) == 1) {
-        return nullptr;
-    }
-    // All ranks must participate even at local total_count 0, else MPI_Alltoallv deadlocks.
-    return &layer.evolution_exchange_layout();
+// Nothing to exchange at one rank. All ranks must participate even at local total_count 0, else
+// MPI_Alltoallv deadlocks, so this is a property of the communicator and not of the layer.
+auto layer_exchange_participates(const mpi::Comm &comm) -> bool {
+    return mpi::size(comm) != 1;
+}
+
+// Derive this layer's exchange layout into `buffers.layout` at `scale`. It describes BOTH sides.
+//
+// The count matrix is symmetric: rank m's slot for r holds (the queries r sent m) ++ (the queries
+// m sent r), and rank r's slot for m holds those two swapped, so the two slots have the same
+// length (MPGraphEncoding's sink, via layer_build/Engine.h). Counts are equal, and displacements
+// are prefix sums of counts, so the recv layout is the send layout -- there is nothing to
+// transpose and nothing to communicate. This is what a per-layer RecvLayoutCache used to hold,
+// at 8 B per world slot, and what an alltoall_counts per layer used to compute.
+//
+// Scaling is applied once, here, rather than to a scale-1 result: every rank multiplies by the
+// same literal, so the equality survives it.
+//
+// Verified end to end rather than reasoned about alone: with MONOPROP_CHECK_EXCHANGE_SYMMETRY set
+// the derived counts are checked against a real alltoall on every layer (see
+// check_exchange_symmetry). A campaign at world 32 and 256 compared 550M slots with no mismatch.
+auto derive_layer_exchange(const LayerTraversal &layer,
+                           const mpi::Comm &comm,
+                           int scale,
+                           FlatExchangeBuffers &buffers) -> void {
+    const auto my_rank = static_cast<size_t>(mpi::rank(comm));
+    const char *what = scale == 1 ? "Layer exchange" : "Layer derivative exchange";
+    detail::derive_exchange_layout(layer.cross_rank(), my_rank, scale, buffers.layout, what);
+    mpi::check_exchange_symmetry(buffers.layout.counts, comm);
 }
 
 // The completed alltoallv payload as an apply pass sees it: peer `rank`'s entries start at
@@ -105,21 +129,20 @@ struct CrossRankExchangeHandle {
     [[no_unique_address]] mpi::Ticket ticket;
 };
 
-inline auto begin_flat_exchange(const LayerExchangeLayout &layout, FlatExchangeBuffers &buffers, const mpi::Comm &comm)
-    -> CrossRankExchangeHandle {
+inline auto begin_flat_exchange(FlatExchangeBuffers &buffers, const mpi::Comm &comm) -> CrossRankExchangeHandle {
+    const LayerExchangeLayout &layout = buffers.layout;
     CrossRankExchangeHandle handle;
     handle.layout = &layout;
     handle.buffers = &buffers;
-    const auto &recv = mpi::resolve_recv(layout.counts, comm, layout.recv_cache);
-    buffers.recv_counts = recv.counts;
-    buffers.recv_displs = recv.displs;
-    buffers.recv_buffer.resize(recv.total == 0 ? 1 : static_cast<size_t>(recv.total));
+    buffers.recv_buffer.resize(layout.total_count == 0 ? 1 : layout.total_count);
+    // Same arrays on both sides. MPI reads recvcounts/recvdispls, it does not write them, so
+    // aliasing them onto the send layout is legal as well as correct here.
     handle.ticket = mpi::post_flat_alltoallv<double>({.send = buffers.send_buffer.data(),
                                                       .send_counts = layout.counts.data(),
                                                       .send_displs = layout.displs.data(),
                                                       .recv = buffers.recv_buffer.data(),
-                                                      .recv_counts = buffers.recv_counts.data(),
-                                                      .recv_displs = buffers.recv_displs.data()},
+                                                      .recv_counts = layout.counts.data(),
+                                                      .recv_displs = layout.displs.data()},
                                                      mpi::size(comm),
                                                      comm);
     return handle;
@@ -136,19 +159,20 @@ struct InFlightExchange {
     bool active = false;
 };
 
-// Callers run the participation guard first (see active_evolution_exchange_layout) so the layout is
-// never materialized at a single rank.
+// Callers run layer_exchange_participates first, so no layout is derived at a single rank -- where
+// deriving one would allocate a P-int pair and resolve a transpose for a transfer that never posts.
 template <typename Pack>
-inline auto begin_layer_exchange(const LayerExchangeLayout &layout, const mpi::Comm &comm, Pack pack)
+inline auto begin_layer_exchange(const LayerTraversal &layer, int scale, const mpi::Comm &comm, Pack pack)
     -> InFlightExchange {
     InFlightExchange in_flight;
     in_flight.my_rank = mpi::rank(comm);
     in_flight.active = true;
 
     auto &buffers = acquire_flat_exchange_buffers();
-    resize_flat_exchange_buffers(layout, buffers);
-    pack(in_flight.my_rank, layout, buffers.send_buffer);
-    in_flight.handle = begin_flat_exchange(layout, buffers, comm);
+    derive_layer_exchange(layer, comm, scale, buffers);
+    resize_flat_exchange_buffers(buffers.layout, buffers);
+    pack(in_flight.my_rank, buffers.layout, buffers.send_buffer);
+    in_flight.handle = begin_flat_exchange(buffers, comm);
     return in_flight;
 }
 
@@ -167,7 +191,7 @@ inline auto finish_layer_exchange(InFlightExchange &in_flight, Apply apply)
     }
     wait_flat_exchange(in_flight.handle);
     return apply(ExchangePayload{.recv_buffer = in_flight.handle.buffers->recv_buffer,
-                                 .recv_displs = in_flight.handle.buffers->recv_displs,
+                                 .recv_displs = in_flight.handle.buffers->layout.displs,
                                  .my_rank = in_flight.my_rank});
 }
 
@@ -258,11 +282,13 @@ inline auto begin_cross_rank_derivative_exchange(const DerivativeSnapshotScratch
                                                  const LayerTraversal &layer,
                                                  const mpi::Comm &comm) -> InFlightExchange {
     // Single-rank (or no peer participating): nothing to exchange — the self slot covers everything.
-    if (active_evolution_exchange_layout(layer, comm) == nullptr) {
+    if (!layer_exchange_participates(comm)) {
         return {};
     }
     // Safe to fire before the cos pass: pack reads pre-cos snapshots and the transfer touches only buffers.
-    return begin_layer_exchange(layer.derivative_exchange_layout(),
+    // Scale 2: each rotation endpoint carries both the op and the state payload.
+    return begin_layer_exchange(layer,
+                                2,
                                 comm,
                                 [&snap, &layer](int my_rank, const LayerExchangeLayout &layout, VecD &send_buffer) {
                                     pack_cross_rank_derivative_payload_impl(snap, layer, my_rank, layout, send_buffer);
@@ -328,12 +354,12 @@ void apply_cross_rank_evolution_exchange_impl(VecD &op,
 
 inline auto begin_cross_rank_evolution_exchange(VecD &op, const LayerTraversal &layer, const mpi::Comm &comm)
     -> InFlightExchange {
-    const auto *layout = active_evolution_exchange_layout(layer, comm);
-    if (layout == nullptr) {
+    if (!layer_exchange_participates(comm)) {
         return {};
     }
     return begin_layer_exchange(
-        *layout,
+        layer,
+        1,
         comm,
         [&op, &layer](int my_rank, const LayerExchangeLayout &active_layout, VecD &send_buffer) {
             pack_cross_rank_evolution_payload_impl(op, layer, my_rank, active_layout, send_buffer);
@@ -362,12 +388,16 @@ auto apply_self_slot_derivative_paired(VecD &state,
         return {};
     }
     const auto pairs = self_d_count / 2;
+    // my_rank is loop-invariant, so resolve the slot once: the four fetches below are four
+    // lookups into the P-sized record array per rotation pair otherwise, in the innermost
+    // gradient loop.
+    const auto slot = layer.cross_rank_slot(my_rank);
     EndpointContrib local{};
     for (size_t k = 0; k < pairs; ++k) {
-        const size_t i1 = layer.cross_rank_sin_recv_index_at(my_rank, k);
-        const double phi1 = static_cast<double>(layer.cross_rank_sin_recv_phase_at(my_rank, k));
-        const size_t i2 = layer.cross_rank_sin_recv_index_at(my_rank, k + pairs);
-        const auto phi2 = static_cast<double>(layer.cross_rank_sin_recv_phase_at(my_rank, k + pairs));
+        const size_t i1 = detail::slot_sin_recv_index(slot, k);
+        const double phi1 = static_cast<double>(detail::slot_sin_recv_phase(slot, k));
+        const size_t i2 = detail::slot_sin_recv_index(slot, k + pairs);
+        const auto phi2 = static_cast<double>(detail::slot_sin_recv_phase(slot, k + pairs));
         // Recover pre-cos values.
         const double s1 = state[i1] * trig.sec_val;
         const double h1 = op[i1] * trig.cos_val;

cpp/monoprop/MPGraph.cpp

@@ -51,7 +51,24 @@ auto layer_storage_memory_usage(const LayerCore &storage) -> GraphMemoryBreakdow
     GraphMemoryBreakdown breakdown;
     breakdown.layer_storage_object_bytes = sizeof(LayerCore);
     breakdown.cross_rank_bytes = detail::cross_rank_storage_bytes(storage.cross_rank);
-    breakdown.exchange_layout_bytes = detail::layer_exchange_layout_storage_bytes(storage.evolution_exchange_layout);
+    // Nothing: the layer no longer retains counts/displs. They are derived into per-thread
+    // scratch for the exchange being posted, so what used to be 2*P ints per layer per partition
+    // is now 2*P ints per THREAD. The field stays, reporting the truth, so an A/B against a build
+    // that did retain them shows the drop rather than silently losing the row.
+    breakdown.exchange_layout_bytes = 0;
+
+    // Diagnostics.
+    breakdown.slot_record_bytes = detail::cross_rank_slot_record_bytes(storage.cross_rank);
+    // The transpose cache is gone: the recv layout equals the send layout, so there was never
+    // anything to cache. Reported as 0 rather than removed, because it was never inside
+    // total_bytes() -- an A/B has no other way to see resident memory leave.
+    breakdown.recv_cache_bytes = 0;
+    // The derivative layout is no longer retained at all: it is 2x the evolution layout, and its
+    // transpose is 2x the evolution transpose, so both are derived on demand without a collective.
+    breakdown.derivative_layout_bytes = 0;
+    breakdown.layer_cores = 1;
+    breakdown.slot_records = storage.cross_rank.rank_count();
+    breakdown.occupied_slots = detail::cross_rank_occupied_slots(storage.cross_rank);
     return breakdown;
 }
 

cpp/monoprop/detail/graph/MPGraphLayers.h

@@ -64,27 +64,34 @@ struct LayerTraversal final {
         return detail::cross_rank_sin_recv_phase(core_->cross_rank, rank, idx);
     }
 
+    // The slot is resolved ONCE, outside the loop: the lookup it costs is indexed by the flat
+    // world P, so doing it per endpoint made per-term work out of what is per-slot work.
     template <typename Func>
     auto for_each_cross_rank_sin_send_range(size_t rank, size_t begin, size_t end, Func &&func) const -> void {
+        const auto slot = detail::cross_rank_slot(core_->cross_rank, rank);
         for (size_t idx = begin; idx < end; ++idx) {
-            func(idx, detail::cross_rank_sin_send_index(core_->cross_rank, rank, idx));
+            func(idx, detail::slot_sin_send_index(slot, idx));
         }
     }
 
     template <typename Func>
     auto for_each_cross_rank_sin_recv_range(size_t rank, size_t begin, size_t end, Func &&func) const -> void {
+        const auto slot = detail::cross_rank_slot(core_->cross_rank, rank);
         for (size_t idx = begin; idx < end; ++idx) {
-            func(idx,
-                 detail::cross_rank_sin_recv_index(core_->cross_rank, rank, idx),
-                 detail::cross_rank_sin_recv_phase(core_->cross_rank, rank, idx));
+            func(idx, detail::slot_sin_recv_index(slot, idx), detail::slot_sin_recv_phase(slot, idx));
         }
     }
 
-    auto evolution_exchange_layout() const -> const LayerExchangeLayout & { return core_->evolution_exchange_layout; }
-    auto derivative_exchange_layout() const -> const LayerExchangeLayout & {
-        return core_->derivative_exchange_layout();
+    // For the paired self-slot derivative fetches, which read d[k] and d[k+pairs] together:
+    // resolve the slot once and hand the caller the view rather than four lookups per pair.
+    auto cross_rank_slot(size_t rank) const -> detail::CrossRankSlotView {
+        return detail::cross_rank_slot(core_->cross_rank, rank);
     }
 
+    // The exchange layout -- both sides of it -- is derived at the call site from these and
+    // nothing else is stored: see detail::derive_exchange_layout and Evolution.cpp.
+    auto cross_rank() const -> const PackedCrossRankStorage & { return core_->cross_rank; }
+
     auto param_index() const -> size_t { return core_->param_index; }
     auto gen_coeff() const -> double { return core_->gen_coeff; }
     auto gate_index() const -> size_t { return core_->gate_index; }

cpp/monoprop/detail/graph/MPGraphViews.h

@@ -40,6 +40,22 @@ struct GraphMemoryBreakdown final {
     size_t cross_rank_bytes = 0;
     size_t exchange_layout_bytes = 0;
 
+    // Diagnostics, deliberately EXCLUDED from total_bytes(): the first three are either a
+    // subset of a field above or memory that total_bytes() has never counted, and folding
+    // them in would silently redefine graph_memory_bytes() mid-flight, so an A/B against an
+    // older build would compare two different quantities. The rest are counts, not bytes.
+    //
+    // The point of the split: a per-layer array indexed by rank is sized by the FLAT world
+    // (mpi::size on a Hybrid comm is ranks x partitions), so it costs O(P) per layer per
+    // partition and O(P^2) across the job. slot_bytes is that part; traffic_bytes is the
+    // part that scales with terms actually crossing, which is real work.
+    size_t slot_record_bytes = 0;      // cross_rank ranges[]: one record per world slot, occupied or not
+    size_t recv_cache_bytes = 0;       // retired: the recv layout IS the send layout, nothing is cached
+    size_t derivative_layout_bytes = 0; // the lazily retained 2x layout AND its own recv cache -- likewise
+    size_t layer_cores = 0;            // distinct LayerCores walked (shared cores counted once)
+    size_t slot_records = 0;           // sum over cores of ranges.size(); divide by layer_cores to recover P
+    size_t occupied_slots = 0;         // slots carrying any traffic: occupancy = occupied_slots / slot_records
+
     auto total_bytes() const -> size_t {
         return layer_descriptor_bytes + layer_storage_object_bytes + cos_data_bytes + cross_rank_bytes
                + exchange_layout_bytes;
@@ -52,6 +68,12 @@ struct GraphMemoryBreakdown final {
         cos_data_bytes += o.cos_data_bytes;
         cross_rank_bytes += o.cross_rank_bytes;
         exchange_layout_bytes += o.exchange_layout_bytes;
+        slot_record_bytes += o.slot_record_bytes;
+        recv_cache_bytes += o.recv_cache_bytes;
+        derivative_layout_bytes += o.derivative_layout_bytes;
+        layer_cores += o.layer_cores;
+        slot_records += o.slot_records;
+        occupied_slots += o.occupied_slots;
         return *this;
     }
 };