From 1d6277344c48c41809fac037d6dfa309de1fc484 Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sat, 15 Aug 2026 15:03:33 +0100 Subject: [PATCH 1/3] perf(evolution): :zap: stop storing the cross-rank D range twice Inside the engine `rank_count` is `mpi::size(comm)`, and on a partitioned run the comm is Hybrid, whose size() is the FLAT world P = ranks x partitions. Every per-rank array in a layer is therefore P long, each MPI rank holds one per partition, and the graph retains one per layer -- so a per-slot record costs O(P^2) across the job. Measured on pauli c14 at 91,273,861 terms, the graph goes 3.81 GB at P=16 to 61.94 GB at P=512 while the operator stays flat near 6.5 GB. Fitting graph = a + b*P^2 on each adjacent pair gives b = 235,709 / 223,891 / 220,345 B/P^2 -- three independent pairs agreeing to 7%, the upper two to 1.6%. CrossRankPartnerRange carried an offset and a count for each of B and D. They were always equal: GraphSink::finalize resizes both vectors from the same P + Q, so the counts match per slot and their prefix sums match with them. B and D are the two endpoints of the same rotation set. Keeping one pair drops the record from 32 to 16 bytes with no padding either way, pinned by a static_assert. The equality is now a checked precondition rather than a comment. Unchecked, a skew would not throw: cross_rank_sin_recv_index would mis-derive Q and read a wrong-but-valid endpoint, and Evolution's self-slot snapshot would run off the end of a B-sized buffer. Three consumers already bet on it silently. Also adds graph_memory_breakdown(). The operator partitions and the graph does not, and one total could not say which. It splits the fields, reports the slot occupancy that decides whether a sparse layout would pay, and counts two things total_bytes() never has: the resolve_recv transpose cache and the lazily retained derivative layout. Those stay as diagnostics rather than joining total_bytes, so graph_memory_bytes() means the same thing before and after and an A/B against an older build still compares one quantity. Assisted-by: ClaudeCode:claude-opus-5 --- cpp/monoprop/MPGraph.cpp | 10 +++ cpp/monoprop/detail/graph/MPGraphViews.h | 22 ++++++ .../detail/graph_encoding/MPGraphEncoding.cpp | 47 ++++++++++-- .../graph_encoding/MPGraphEncodingStorage.h | 21 +++++- .../graph_encoding/MPGraphEncodingTypes.h | 30 ++++++-- cpp/tests/graph_encoding_tests.cpp | 75 +++++++++++++++++++ cpp/tests/large_cosine_storage_tests.cpp | 8 +- src/monoprop/bindings/binder.h | 22 ++++++ 8 files changed, 217 insertions(+), 18 deletions(-) diff --git a/cpp/monoprop/MPGraph.cpp b/cpp/monoprop/MPGraph.cpp index a425582d..942c61d6 100644 --- a/cpp/monoprop/MPGraph.cpp +++ b/cpp/monoprop/MPGraph.cpp @@ -52,6 +52,16 @@ auto layer_storage_memory_usage(const LayerCore &storage) -> GraphMemoryBreakdow 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); + + // Diagnostics. recv_cache and the derivative layout are real resident memory that + // total_bytes() has never counted, which is why the reported graph size sits below the + // process RSS by a margin that itself widens with the world size. + breakdown.slot_record_bytes = detail::cross_rank_slot_record_bytes(storage.cross_rank); + breakdown.recv_cache_bytes = detail::layer_exchange_layout_cache_bytes(storage.evolution_exchange_layout); + breakdown.derivative_layout_bytes = storage.derivative_exchange_layout_bytes(); + 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; } diff --git a/cpp/monoprop/detail/graph/MPGraphViews.h b/cpp/monoprop/detail/graph/MPGraphViews.h index 0af7d7a2..525d6a80 100644 --- a/cpp/monoprop/detail/graph/MPGraphViews.h +++ b/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; // evolution layout's resolve_recv transpose cache -- never in total_bytes() + 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; } }; diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index 2e8da60f..65899577 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp @@ -14,9 +14,11 @@ #include "monoprop/detail/graph_encoding/MPGraphEncodingStorage.h" +#include #include #include #include +#include #include #include #include @@ -103,18 +105,27 @@ auto build_packed_cross_rank_storage(const std::vector &da storage.ranges.resize(num_ranks); size_t total_b = 0; - size_t total_d = 0; for (size_t rank = 0; rank < num_ranks; ++rank) { const auto &partner = data[rank]; + // B and D are the two endpoints of the same rotation set, so they must be the same + // length. The record stores one count and one offset for both; checking here is what + // makes that a precondition instead of a convention. Unchecked, a skew would not throw + // -- cross_rank_sin_recv_index would mis-derive Q and silently read the wrong endpoint, + // and Evolution's self-slot snapshot would run off the end of its B-sized buffer. + if (partner.sin_send_indices.size() != partner.sin_recv_entries.size()) { + throw std::logic_error(std::format( + "Cross-rank slot {} has {} send endpoints against {} recv endpoints; B and D are the same set.", + rank, + partner.sin_send_indices.size(), + partner.sin_recv_entries.size())); + } auto &range = storage.ranges[rank]; range.sin_send_offset = total_b; range.sin_send_count = static_cast(partner.sin_send_indices.size()); - range.sin_recv_offset = total_d; - range.sin_recv_count = static_cast(partner.sin_recv_entries.size()); range.in_count = static_cast(partner.in_count); total_b += partner.sin_send_indices.size(); - total_d += partner.sin_recv_entries.size(); } + const size_t total_d = total_b; bool uses_binary_phases = true; for (const auto &partner : data) { @@ -130,8 +141,10 @@ auto build_packed_cross_rank_storage(const std::vector &da for (size_t rank = 0; rank < num_ranks; ++rank) { const auto &partner = data[rank]; + // One offset addresses both arrays: the counts are equal per slot (checked above), so + // their prefix sums are too. const size_t b_off = storage.ranges[rank].sin_send_offset; - const size_t d_off = storage.ranges[rank].sin_recv_offset; + const size_t d_off = b_off; for (size_t k = 0; k < partner.sin_send_indices.size(); ++k) { storage.sin_send_indices[b_off + k] = checked_term_index(partner.sin_send_indices[k], "Cross-rank B index"); @@ -154,10 +167,26 @@ auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t { return bytes; } +auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size_t { + return storage.ranges.capacity() * sizeof(CrossRankPartnerRange); +} + +auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t { + // sin_send_count alone is the predicate: it is the whole endpoint set for the slot, + // in-block and out-block together, so in_count cannot be non-zero while it is zero. + return static_cast(std::ranges::count_if( + storage.ranges, [](const CrossRankPartnerRange &range) { return range.sin_send_count != 0; })); +} + auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> size_t { return layout.counts.capacity() * sizeof(int) + layout.displs.capacity() * sizeof(int); } +auto layer_exchange_layout_cache_bytes(const LayerExchangeLayout &layout) -> size_t { + const auto &cached = layout.recv_cache.layout; + return cached.counts.capacity() * sizeof(int) + cached.displs.capacity() * sizeof(int); +} + auto build_layer_storage_unified(std::vector all_partners, size_t my_rank) -> std::shared_ptr { auto storage = std::make_shared(); @@ -199,4 +228,12 @@ auto LayerCore::derivative_exchange_layout() const -> const LayerExchangeLayout return *derivative_exchange_layout_cache_; } +auto LayerCore::derivative_exchange_layout_bytes() const -> size_t { + if (!derivative_exchange_layout_cache_) { + return 0; + } + const auto &layout = *derivative_exchange_layout_cache_; + return detail::layer_exchange_layout_storage_bytes(layout) + detail::layer_exchange_layout_cache_bytes(layout); +} + } // namespace monoprop diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h index c1ed79c9..09c18ee3 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -82,24 +82,39 @@ inline auto cross_rank_sin_send_index(const PackedCrossRankStorage &storage, siz return static_cast(storage.sin_send_indices[offset]); } -// Invariant B=[in(P)]++[out(Q)], D=[out(Q)]++[in(P)] (P=in_count, Q=sin_recv_count-P): +// Invariant B=[in(P)]++[out(Q)], D=[out(Q)]++[in(P)] (P=in_count, Q=sin_send_count-P): // D[idx] = (idx size_t { const auto &range = storage.ranges[rank]; const size_t in_count = range.in_count; // P - const size_t out_count = range.sin_recv_count - in_count; // Q + const size_t out_count = range.sin_send_count - in_count; // Q const size_t sin_send_local = (idx < out_count) ? (in_count + idx) : (idx - out_count); return cross_rank_sin_send_index(storage, rank, sin_send_local); } +// The D phases run parallel to the B indices -- same count per slot, so the same prefix sum +// addresses both. They are still separate arrays; only the offset into them is shared. inline auto cross_rank_sin_recv_phase(const PackedCrossRankStorage &storage, size_t rank, size_t idx) -> int { - return packed_phase_at(storage.sin_recv_phases, storage.ranges[rank].sin_recv_offset + idx); + return packed_phase_at(storage.sin_recv_phases, storage.ranges[rank].sin_send_offset + idx); } auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t; +// The slot-proportional part of cross_rank_storage_bytes: one record per world slot whether or +// not that slot carries traffic. The remainder (indices and phases) scales with terms crossing. +auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size_t; + +// World slots carrying any traffic for this layer. Read against rank_count() to get occupancy: +// low occupancy would make a sparse layout pay, high occupancy means only a narrower record does. +auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t; + auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> size_t; +// The resolve_recv transpose cache hanging off a layout. Separate from +// layer_exchange_layout_storage_bytes because that function's result is already carried in a +// shipped metric; folding this in would redefine it. +auto layer_exchange_layout_cache_bytes(const LayerExchangeLayout &layout) -> size_t; + // Local cycles fold into the self-rank slot (my_rank); the exchange layout zeroes counts[my_rank] so // MPI_Alltoallv skips it (replay does a local copy). auto build_layer_storage_unified(std::vector all_partners, size_t my_rank) diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h index 20be7ace..ac60854b 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h @@ -120,23 +120,33 @@ struct CrossRankPartnerData { size_t in_count = 0; }; +// One record per world slot, occupied or not, retained for every layer -- so on a partitioned run +// this is sizeof(record) x P x layers x partitions per rank, i.e. O(P^2) across the job. That is +// why it holds only what cannot be recovered: B and D are the two endpoints of the same rotation +// set, so their counts are equal and their prefix sums therefore identical, and storing the D pair +// separately cost 16 bytes a slot to say twice what the B pair already said. +// build_packed_cross_rank_storage enforces the equality rather than trusting it. struct CrossRankPartnerRange final { - size_t sin_send_offset = 0; // into sin_send_indices; cumulative across ranks, so size_t (may exceed 2^32) - TermIndex sin_send_count = - 0; // == sin_recv_count (both endpoints); TermIndex-wide so one rank/layer can exceed 2^32 - size_t sin_recv_offset = 0; // into sin_recv_phases; cumulative across ranks, so size_t (see sin_send_offset) - TermIndex sin_recv_count = 0; + size_t sin_send_offset = 0; // into sin_send_indices AND sin_recv_phases; cumulative, so may exceed 2^32 + // == the D count; TermIndex-wide so one rank/layer can exceed 2^32. + TermIndex sin_send_count = 0; TermIndex in_count = 0; }; +// Pins the saving: 8 + 4 + 4 narrow, 8 + 8 + 8 wide, with no tail padding either way. A new field +// here is paid for once per world slot per layer per partition, so it should be a deliberate act. +static_assert(sizeof(CrossRankPartnerRange) == sizeof(size_t) + 2 * sizeof(TermIndex), + "CrossRankPartnerRange is the per-world-slot record; keep it free of padding."); + struct PackedCrossRankStorage final { - std::vector ranges; // size == R + std::vector ranges; // size == the flat world P, not the MPI rank count std::vector sin_send_indices; PackedPhaseStorage sin_recv_phases; // one phased entry per D index, sign baked in auto rank_count() const -> size_t { return ranges.size(); } auto sin_send_size(size_t rank) const -> size_t { return ranges[rank].sin_send_count; } - auto sin_recv_size(size_t rank) const -> size_t { return ranges[rank].sin_recv_count; } + // D holds the same endpoints as B in the other order, so it has the same length. + auto sin_recv_size(size_t rank) const -> size_t { return ranges[rank].sin_send_count; } auto in_count(size_t rank) const -> size_t { return ranges[rank].in_count; } }; @@ -146,6 +156,12 @@ struct LayerCore final { auto derivative_exchange_layout() const -> const LayerExchangeLayout &; + // Bytes held by the lazily built derivative layout, 0 while it has never been asked for. + // Deliberately NOT implemented as a call to derivative_exchange_layout(): that would + // allocate the very thing being measured, turning an accounting read into a 2*P-int + // allocation on every layer and making the instrument report its own footprint. + auto derivative_exchange_layout_bytes() const -> size_t; + // A copied core must not inherit the source's cache: it is eval-time state, not data. auto reset_derivative_exchange_layout() -> void { derivative_exchange_layout_cache_.reset(); } diff --git a/cpp/tests/graph_encoding_tests.cpp b/cpp/tests/graph_encoding_tests.cpp index 26bfe2cf..b2f521ae 100644 --- a/cpp/tests/graph_encoding_tests.cpp +++ b/cpp/tests/graph_encoding_tests.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include "monoprop/detail/graph_encoding/MPGraphEncodingStorage.h" @@ -187,3 +188,77 @@ BOOST_AUTO_TEST_CASE(graph_encoding_d_from_b_derivation_both_arms) { BOOST_CHECK_EQUAL(detail::cross_rank_sin_send_index(storage, 0, 0), 10U); BOOST_CHECK_EQUAL(detail::cross_rank_sin_send_index(storage, 0, 4), 22U); } + +// The accounting split behind graph_memory_breakdown(). These lock the property the split +// exists to expose: the slot-record cost is set by the size of the world, and does not move +// when the traffic through it does. + +BOOST_AUTO_TEST_CASE(graph_encoding_occupied_slots_counts_only_slots_carrying_traffic) { + std::vector data(5); // five world slots, two of them used + data[1].sin_send_indices.push_back(7); + data[1].sin_recv_entries.push_back({0, 1}); + data[1].in_count = 1; + data[3].sin_send_indices.push_back(9); + data[3].sin_recv_entries.push_back({0, 1}); + data[3].in_count = 1; + + const auto storage = detail::build_packed_cross_rank_storage(data); + + BOOST_CHECK_EQUAL(storage.rank_count(), 5U); + BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(storage), 2U); +} + +BOOST_AUTO_TEST_CASE(graph_encoding_slot_record_bytes_track_the_world_not_the_traffic) { + // Same single sender, two different world sizes: the traffic is identical, so anything + // that grows here is paid for the world rather than for the work. + std::vector narrow(2); + std::vector wide(8); + for (auto *data : {&narrow, &wide}) { + (*data)[0].sin_send_indices.push_back(1); + (*data)[0].sin_recv_entries.push_back({0, 1}); + (*data)[0].in_count = 1; + } + + const auto narrow_storage = detail::build_packed_cross_rank_storage(narrow); + const auto wide_storage = detail::build_packed_cross_rank_storage(wide); + + BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(narrow_storage), 1U); + BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(wide_storage), 1U); + // Four times the slots for the same one term crossing. + BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(wide_storage), + 4 * detail::cross_rank_slot_record_bytes(narrow_storage)); + BOOST_CHECK_LT(detail::cross_rank_slot_record_bytes(narrow_storage), + detail::cross_rank_storage_bytes(narrow_storage)); +} + +BOOST_AUTO_TEST_CASE(graph_encoding_lazy_layout_bytes_do_not_force_the_allocation) { + LayerCore core; + core.evolution_exchange_layout = detail::build_layer_exchange_layout({3, 0, 5}, /*scale=*/1); + + // Reading the size must not build the thing being sized, or the instrument reports its + // own footprint and every layer pays 2*P ints for having been measured. + BOOST_CHECK_EQUAL(core.derivative_exchange_layout_bytes(), 0U); + static_cast(core.derivative_exchange_layout()); + BOOST_CHECK_GT(core.derivative_exchange_layout_bytes(), 0U); + + core.reset_derivative_exchange_layout(); + BOOST_CHECK_EQUAL(core.derivative_exchange_layout_bytes(), 0U); + + // The transpose cache is eval-time state: nothing has resolved it, so it is not resident. + BOOST_CHECK_EQUAL(detail::layer_exchange_layout_cache_bytes(core.evolution_exchange_layout), 0U); +} + +BOOST_AUTO_TEST_CASE(graph_encoding_skewed_endpoint_counts_are_refused) { + // B and D are the two endpoints of the same rotation set, so the packed record keeps one + // count and one offset for both. GraphSink::finalize resizes the two vectors from the same + // expression, so the engine cannot produce a skew -- but nothing in the TYPE prevents one, + // and unchecked it would not throw: cross_rank_sin_recv_index would mis-derive Q and read a + // wrong-but-valid endpoint. Refusing at the choke point makes the assumption a precondition. + std::vector data(1); + data[0].sin_send_indices.push_back(1); + data[0].sin_send_indices.push_back(2); + data[0].sin_recv_entries.push_back({1, 1}); // one D against two B + data[0].in_count = 1; + + BOOST_CHECK_THROW(detail::build_packed_cross_rank_storage(data), std::logic_error); +} diff --git a/cpp/tests/large_cosine_storage_tests.cpp b/cpp/tests/large_cosine_storage_tests.cpp index 2110b355..d9dd20b8 100644 --- a/cpp/tests/large_cosine_storage_tests.cpp +++ b/cpp/tests/large_cosine_storage_tests.cpp @@ -64,7 +64,7 @@ BOOST_AUTO_TEST_CASE(pruned_layer_supports_cos_counts_above_u32) { lt.for_each_cross_rank_sin_send_range(1, 0, 1, [&](size_t, size_t i) { b_idx = i; }); BOOST_CHECK_EQUAL(b_idx, 200UL); - // D[0] is derived from B: Q = sin_recv_count - in_count = 20 - 12 = 8, so D[0] = out-block[0] = 100, + // D[0] is derived from B: Q = sin_send_count - in_count = 20 - 12 = 8, so D[0] = out-block[0] = 100, // stored phase = -(out_phases[0]) = -(+1) = -1. size_t d_idx = static_cast(-1); int d_phi = 0; @@ -81,8 +81,10 @@ BOOST_AUTO_TEST_CASE(pruned_layer_supports_cos_counts_above_u32) { BOOST_AUTO_TEST_CASE(cross_rank_partner_range_counts_track_term_index_width) { CrossRankPartnerRange r{}; BOOST_CHECK_EQUAL(sizeof(r.sin_send_count), sizeof(TermIndex)); - BOOST_CHECK_EQUAL(sizeof(r.sin_recv_count), sizeof(TermIndex)); BOOST_CHECK_EQUAL(sizeof(r.in_count), sizeof(TermIndex)); + // The record is paid once per world slot per layer per partition, so its width is a result, + // not an implementation detail. Padding here would be invisible and quadratically expensive. + BOOST_CHECK_EQUAL(sizeof(r), sizeof(size_t) + 2 * sizeof(TermIndex)); } #if defined(monoprop_WIDE_TERM_INDEX) @@ -105,7 +107,7 @@ BOOST_AUTO_TEST_CASE(cross_rank_sin_send_index_round_trips_above_u32) { BOOST_CHECK_EQUAL(detail::cross_rank_sin_send_index(storage, 1, 0), big_in); BOOST_CHECK_EQUAL(detail::cross_rank_sin_send_index(storage, 1, 1), big_out); - // D[0] is derived from B: Q = sin_recv_count - in_count = 1, so D[0] = out-block[0] = big_out. + // D[0] is derived from B: Q = sin_send_count - in_count = 1, so D[0] = out-block[0] = big_out. BOOST_CHECK_EQUAL(detail::cross_rank_sin_recv_index(storage, 1, 0), big_out); } #endif diff --git a/src/monoprop/bindings/binder.h b/src/monoprop/bindings/binder.h index 5728aca5..0c9b606a 100644 --- a/src/monoprop/bindings/binder.h +++ b/src/monoprop/bindings/binder.h @@ -271,5 +271,27 @@ auto bind_monomial_propagator(nb::module_ &mod) -> void { {"d_terms_slack_bytes", b.operator_terms_slack_bytes}, {"d_state_coeffs_nonzero", b.state_coeffs_nonzero}}; }); + + // The operator partitions but the graph does not: its per-layer arrays are indexed by rank, + // and on a partitioned run that index space is the FLAT world (ranks x partitions). Splitting + // the total is what separates memory that grows with the problem from memory that grows with + // the machine. d_occupied_slots / d_slot_records is the occupancy that says whether a sparse + // layout would pay; d_slot_records / d_layer_cores recovers the world size P. + cls.def("graph_memory_breakdown", [](const MonomialPropagator &self) { + const auto b = self.graph_memory_usage(); + return std::map{{"layer_descriptor_bytes", b.layer_descriptor_bytes}, + {"layer_storage_object_bytes", b.layer_storage_object_bytes}, + {"cos_data_bytes", b.cos_data_bytes}, + {"cross_rank_bytes", b.cross_rank_bytes}, + {"exchange_layout_bytes", b.exchange_layout_bytes}, + {"total_bytes", b.total_bytes()}, + // Diagnostics (not part of total_bytes; see the struct). + {"d_slot_record_bytes", b.slot_record_bytes}, + {"d_recv_cache_bytes", b.recv_cache_bytes}, + {"d_derivative_layout_bytes", b.derivative_layout_bytes}, + {"d_layer_cores", b.layer_cores}, + {"d_slot_records", b.slot_records}, + {"d_occupied_slots", b.occupied_slots}}; + }); } } // namespace monoprop::bindings::detail From 5f33a7176b7d491fab4db5fe8925844ce2c0e4fe Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sat, 15 Aug 2026 17:54:25 +0100 Subject: [PATCH 2/3] perf(evolution): :zap: derive the exchange layouts instead of retaining them A graph layer retained two `int[P]` arrays for the evolution exchange and, after the first gradient, two more for the derivative round. P is the FLAT WORLD SIZE (ranks x partitions), each MPI process holds one set per partition, and the graph holds one per layer -- so those arrays cost O(P^2) across a job for content that is a prefix sum of what the slot records already say. `counts[r]` is `(r == my_rank ? 0 : cross_rank.sin_send_size(r))` and `displs` is its running prefix. Both are now derived into the per-thread scratch that already owns the send and recv buffers, for the transfer being posted. That takes the retained slot-proportional footprint from 32 B/slot to 16 B/slot, and from 48 to 16 once a gradient has run: only the slot records survive. The derivative round needs no collective of its own. Its counts are the evolution counts at a hardcoded scale of 2, applied identically on every rank, and displacements are prefix sums of counts -- so scaling commutes with the transpose and the derivative recv layout is 2x the evolution recv layout. One `resolve_recv` per layer per evaluation now serves both rounds. The transpose cache stays retained (8 B/slot). It is the one piece that cannot be derived locally, and dropping it would cost an MPI_Alltoall per layer per evaluation. ## The predicate this needed first Sharing scratch across layers is only sound once `resolve_recv` can tell one send pattern from another. Its predicate was `comm_size == comm_size && counts.size() == n` -- effectively "have we ever resolved anything for a communicator this size", which is true for every layer after the first. Correct only while each cache belonged to the one layout that produced it; silently wrong the moment two patterns share a cache. The fix is NOT a checksum of the counts. A miss runs `alltoall_counts`, a collective, so two ranks disagreeing about validity is a distributed HANG rather than a wrong answer, and any rank-local key can collide on one rank and not on another. The cache now carries a `generation` assigned per LayerCore at build time. Build order is identical on every rank, so every rank misses on a layer's first resolve and hits afterwards -- the DECISION is uniform even though the id values are not. A `LayerCore` copy made by `set_parameter_mapping` now inherits that cache rather than dropping it. Relabelling changes which parameter drives the rotation, never which endpoints cross to which slot, so the cached transpose is still correct; clearing it would have cost one collective per layer (5,420 at the anchor) to rebuild an identical answer. ## Hoisting the slot resolution Resolving a world slot is an index into the P-sized `ranges` array, and the per-element accessors were doing it per ENDPOINT -- three times per term on the recv side, four times per rotation pair in the self-slot gradient loop. `cross_rank_slot()` resolves it once and the element accessors take that view, so walking a slot's endpoints pays for the P-sized lookup once. No behaviour change, and it is the precondition for ever storing slots sparsely, where resolving one stops being an array index. ## Notes for review `build_layer_exchange_layout` now has no production caller and is kept deliberately, as the reference the derivation is tested against: the new equivalence case asserts derived == built elementwise for every my_rank and both scales. Checking a derivation against an independent construction beats checking it against literals. `exchange_layout_bytes` and `derivative_layout_bytes` now report 0 rather than being removed from the breakdown, so an A/B against a build that did retain them shows the drop instead of losing the row. Build-time derivation is retained purely as eager validation: an int overflow has to throw from build_graph, not from inside the exchange where peers are already blocked in the count round. --- cpp/monoprop/Evolution.cpp | 107 +++++++++++----- cpp/monoprop/MPGraph.cpp | 19 +-- cpp/monoprop/detail/graph/MPGraphLayers.h | 24 ++-- .../detail/graph_encoding/MPGraphEncoding.cpp | 96 +++++++------- .../graph_encoding/MPGraphEncodingStorage.h | 73 +++++++++-- .../graph_encoding/MPGraphEncodingTypes.h | 59 +++++---- .../MonomialPropagator.inl | 8 +- cpp/monoprop/detail/mpi/Exchange.h | 12 +- cpp/monoprop/detail/mpi/MPICompat.cpp | 13 +- cpp/monoprop/detail/mpi/RecvLayout.h | 16 ++- cpp/tests/graph_encoding_tests.cpp | 117 +++++++++++++----- 11 files changed, 379 insertions(+), 165 deletions(-) diff --git a/cpp/monoprop/Evolution.cpp b/cpp/monoprop/Evolution.cpp index 47bf3234..cfd679b6 100644 --- a/cpp/monoprop/Evolution.cpp +++ b/cpp/monoprop/Evolution.cpp @@ -62,6 +62,12 @@ struct FlatExchangeBuffers { VecD recv_buffer; std::vector recv_counts; std::vector recv_displs; + // The send 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. + LayerExchangeLayout layout; + int recv_total = 0; }; auto &acquire_flat_exchange_buffers() { @@ -73,21 +79,56 @@ 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. + // Only the send side: the recv counts/displs are filled by derive_layer_exchange before this + // runs, so clearing them here (as this did when the recv side was resolved afterwards) would + // throw away the layout the transfer is about to be posted with. 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(); } -auto active_evolution_exchange_layout(const LayerTraversal &layer, const mpi::Comm &comm) - -> const LayerExchangeLayout * { - if (mpi::size(comm) == 1) { - return nullptr; +// 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 send layout into `buffers.layout` at `scale`, and its recv side into +// buffers.recv_counts/recv_displs at the same scale. +// +// The recv side is resolved ONCE per layer, at scale 1, into the layer's own cache. Scaling +// commutes with the transpose -- every rank multiplies its counts by the same literal, so peer q +// sends me exactly `scale` times what it sent at scale 1, and displacements are prefix sums of +// counts so they scale with them. That is what lets the derivative round reuse the evolution +// round's collective instead of paying an alltoall of its own. +auto derive_layer_exchange(const LayerTraversal &layer, + const mpi::Comm &comm, + int scale, + FlatExchangeBuffers &buffers) -> void { + const auto my_rank = static_cast(mpi::rank(comm)); + + // Resolve the transpose at scale 1, so evolution and derivative rounds share one cache entry + // and therefore one collective. + detail::derive_exchange_layout(layer.cross_rank(), my_rank, 1, buffers.layout); + const auto &recv = + mpi::resolve_recv(buffers.layout.counts, comm, layer.evolution_recv_cache(), layer.exchange_generation()); + + const size_t n = recv.counts.size(); + buffers.recv_counts.resize(n); + buffers.recv_displs.resize(n); + for (size_t i = 0; i < n; ++i) { + buffers.recv_counts[i] = detail::checked_mpi_int(static_cast(recv.counts[i]) * static_cast(scale), + "Layer exchange recv count"); + buffers.recv_displs[i] = detail::checked_mpi_int(static_cast(recv.displs[i]) * static_cast(scale), + "Layer exchange recv displacement"); + } + buffers.recv_total = detail::checked_mpi_int(static_cast(recv.total) * static_cast(scale), + "Layer exchange recv total"); + + // Now the send side at the requested scale. Re-derived rather than scaled in place so the + // overflow check runs against the value MPI actually receives. + if (scale != 1) { + detail::derive_exchange_layout(layer.cross_rank(), my_rank, scale, buffers.layout, "Layer derivative exchange"); } - // All ranks must participate even at local total_count 0, else MPI_Alltoallv deadlocks. - return &layer.evolution_exchange_layout(); } // The completed alltoallv payload as an apply pass sees it: peer `rank`'s entries start at @@ -105,15 +146,12 @@ 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(recv.total)); + buffers.recv_buffer.resize(buffers.recv_total == 0 ? 1 : static_cast(buffers.recv_total)); handle.ticket = mpi::post_flat_alltoallv({.send = buffers.send_buffer.data(), .send_counts = layout.counts.data(), .send_displs = layout.displs.data(), @@ -136,19 +174,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 -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; } @@ -258,11 +297,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 +369,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 +403,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(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(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(detail::slot_sin_recv_phase(slot, k)); + const size_t i2 = detail::slot_sin_recv_index(slot, k + pairs); + const auto phi2 = static_cast(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; diff --git a/cpp/monoprop/MPGraph.cpp b/cpp/monoprop/MPGraph.cpp index 942c61d6..6df4fca9 100644 --- a/cpp/monoprop/MPGraph.cpp +++ b/cpp/monoprop/MPGraph.cpp @@ -51,14 +51,19 @@ 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); - - // Diagnostics. recv_cache and the derivative layout are real resident memory that - // total_bytes() has never counted, which is why the reported graph size sits below the - // process RSS by a margin that itself widens with the world size. + // 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. recv_cache is real resident memory that total_bytes() has never counted, which + // is why the reported graph size sits below the process RSS. breakdown.slot_record_bytes = detail::cross_rank_slot_record_bytes(storage.cross_rank); - breakdown.recv_cache_bytes = detail::layer_exchange_layout_cache_bytes(storage.evolution_exchange_layout); - breakdown.derivative_layout_bytes = storage.derivative_exchange_layout_bytes(); + breakdown.recv_cache_bytes = detail::layer_exchange_layout_cache_bytes(storage.evolution_recv_cache); + // 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); diff --git a/cpp/monoprop/detail/graph/MPGraphLayers.h b/cpp/monoprop/detail/graph/MPGraphLayers.h index 715aedda..b2156385 100644 --- a/cpp/monoprop/detail/graph/MPGraphLayers.h +++ b/cpp/monoprop/detail/graph/MPGraphLayers.h @@ -64,27 +64,37 @@ 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 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 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 is derived at the call site from these, not stored: see + // detail::derive_exchange_layout. Only the transpose cache survives, because only it + // needs a collective. + auto cross_rank() const -> const PackedCrossRankStorage & { return core_->cross_rank; } + auto evolution_recv_cache() const -> mpi::RecvLayoutCache & { return core_->evolution_recv_cache; } + auto exchange_generation() const -> uint64_t { return core_->exchange_generation; } + 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; } diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index 65899577..662a9c56 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp @@ -15,6 +15,7 @@ #include "monoprop/detail/graph_encoding/MPGraphEncodingStorage.h" #include +#include #include #include #include @@ -53,15 +54,6 @@ auto build_layer_exchange_layout(const std::vector &send_counts, int sca return layout; } -auto build_derivative_exchange_layout(const LayerExchangeLayout &evolution) -> LayerExchangeLayout { - std::vector send_counts; - send_counts.reserve(evolution.counts.size()); - for (const int count : evolution.counts) { - send_counts.push_back(static_cast(count)); - } - return build_layer_exchange_layout(send_counts, 2, "Layer derivative exchange"); -} - auto checked_term_index(size_t value, const char *what) -> TermIndex { if (value > static_cast(std::numeric_limits::max())) { throw std::overflow_error( @@ -182,58 +174,68 @@ auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> s return layout.counts.capacity() * sizeof(int) + layout.displs.capacity() * sizeof(int); } -auto layer_exchange_layout_cache_bytes(const LayerExchangeLayout &layout) -> size_t { - const auto &cached = layout.recv_cache.layout; +auto layer_exchange_layout_cache_bytes(const mpi::RecvLayoutCache &cache) -> size_t { + const auto &cached = cache.layout; return cached.counts.capacity() * sizeof(int) + cached.displs.capacity() * sizeof(int); } +auto derive_exchange_layout(const PackedCrossRankStorage &cross_rank, + size_t my_rank, + int scale, + LayerExchangeLayout &out, + const char *what) -> void { + const std::string count_label = std::format("{} count", what); + const std::string displacement_label = std::format("{} displacement", what); + + const size_t num_ranks = cross_rank.rank_count(); + out.counts.resize(num_ranks); + out.displs.resize(num_ranks); + size_t total = 0; + for (size_t r = 0; r < num_ranks; ++r) { + // The self slot is excluded from the transfer and handled locally, exactly as the + // stored layout did; an empty slot still gets a valid (repeated) displacement. + const size_t count = (r == my_rank) ? size_t{0} : static_cast(scale) * cross_rank.sin_send_size(r); + out.counts[r] = checked_mpi_int(count, count_label.c_str()); + out.displs[r] = checked_mpi_int(total, displacement_label.c_str()); + total += count; + } + out.total_count = total; +} + auto build_layer_storage_unified(std::vector all_partners, size_t my_rank) -> std::shared_ptr { - auto storage = std::make_shared(); - - { - std::vector send_counts; - send_counts.reserve(all_partners.size()); - for (size_t r = 0; r < all_partners.size(); ++r) { - send_counts.push_back((r == my_rank) ? size_t{0} : all_partners[r].sin_send_indices.size()); - } - storage->evolution_exchange_layout = build_layer_exchange_layout(send_counts, 1); + // Identifies the send pattern this core holds, so its recv cache cannot be served to a + // different one. Starts at 1 because a default-constructed cache carries 0 = never populated. + static std::atomic next_generation{1}; - // The derivative layout (2x) is allocated lazily on first gradient read, but validated here: an - // overflow must throw during build_graph, not from inside the gradient collective window, where - // peers are already blocked in mpi::resolve_recv's count round -> a distributed hang, not an error. - static_cast(build_derivative_exchange_layout(storage->evolution_exchange_layout)); - } + auto storage = std::make_shared(); + storage->exchange_generation = next_generation.fetch_add(1, std::memory_order_relaxed); + const size_t num_ranks = all_partners.size(); storage->cross_rank = build_packed_cross_rank_storage(std::move(all_partners)); - // Both are indexed by the same rank space. - if (storage->evolution_exchange_layout.counts.size() != storage->cross_rank.rank_count()) { + // Both are indexed by the same rank space. Checked here because everything downstream now + // derives the layout from cross_rank, so this is the one place the two can still disagree. + if (num_ranks != storage->cross_rank.rank_count()) { throw ExchangeLayoutRankMismatch( std::format("Layer exchange layout covers {} ranks but cross-rank storage has {}.", - storage->evolution_exchange_layout.counts.size(), + num_ranks, storage->cross_rank.rank_count())); } - return storage; -} - -} // namespace monoprop::detail -namespace monoprop { + // Derive both scales once at build time and throw the result away. This is purely eager + // validation: an overflow of MPI's int has to throw from build_graph, not from inside the + // exchange, where peers are already blocked in resolve_recv's count round -- there it is a + // distributed hang rather than an error. Scale 2 is checked as well as 1 because the + // derivative round overflows first and a gradient may run long after the graph was built. + // + // The vectors are not kept. They are a prefix sum of what cross_rank already holds, and + // retaining them per layer per partition is the O(P^2) term this change removes. + LayerExchangeLayout scratch; + derive_exchange_layout(storage->cross_rank, my_rank, 1, scratch); + derive_exchange_layout(storage->cross_rank, my_rank, 2, scratch, "Layer derivative exchange"); -auto LayerCore::derivative_exchange_layout() const -> const LayerExchangeLayout & { - if (!derivative_exchange_layout_cache_) { - derivative_exchange_layout_cache_ = detail::build_derivative_exchange_layout(evolution_exchange_layout); - } - return *derivative_exchange_layout_cache_; -} - -auto LayerCore::derivative_exchange_layout_bytes() const -> size_t { - if (!derivative_exchange_layout_cache_) { - return 0; - } - const auto &layout = *derivative_exchange_layout_cache_; - return detail::layer_exchange_layout_storage_bytes(layout) + detail::layer_exchange_layout_cache_bytes(layout); + return storage; } -} // namespace monoprop +} // namespace monoprop::detail diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h index 09c18ee3..18518633 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -77,25 +77,62 @@ inline auto store_packed_phase(PackedPhaseStorage &storage, size_t idx, int phas auto build_packed_cross_rank_storage(const std::vector &data) -> PackedCrossRankStorage; -inline auto cross_rank_sin_send_index(const PackedCrossRankStorage &storage, size_t rank, size_t idx) -> size_t { - const size_t offset = storage.ranges[rank].sin_send_offset + idx; - return static_cast(storage.sin_send_indices[offset]); +// One world slot's position in the flat B/D arrays, resolved once. +// +// Resolving is per-SLOT work -- an index into `ranges`, which is the array sized by the flat +// world P. The per-element accessors below take this instead of a slot id so that walking a +// slot's endpoints pays that cost once rather than on every endpoint. It matters more than it +// looks: a recv endpoint reads three fields of the record, so the unhoisted form touched the +// P-sized array three times per term. It is also the precondition for ever storing the slots +// sparsely -- under a sparse layout resolving a slot stops being an array index, and anything +// that resolves per term rather than per slot would become unaffordable. +struct CrossRankSlotView final { + const TermIndex *sin_send_indices = nullptr; // B, already advanced to this slot's offset + const PackedPhaseStorage *sin_recv_phases = nullptr; + size_t phase_offset = 0; + size_t sin_send_count = 0; + size_t in_count = 0; +}; + +inline auto cross_rank_slot(const PackedCrossRankStorage &storage, size_t rank) -> CrossRankSlotView { + const auto &range = storage.ranges[rank]; + return CrossRankSlotView{.sin_send_indices = storage.sin_send_indices.data() + range.sin_send_offset, + .sin_recv_phases = &storage.sin_recv_phases, + .phase_offset = range.sin_send_offset, + .sin_send_count = range.sin_send_count, + .in_count = range.in_count}; +} + +inline auto slot_sin_send_index(const CrossRankSlotView &slot, size_t idx) -> size_t { + return static_cast(slot.sin_send_indices[idx]); } // Invariant B=[in(P)]++[out(Q)], D=[out(Q)]++[in(P)] (P=in_count, Q=sin_send_count-P): // D[idx] = (idx size_t { - const auto &range = storage.ranges[rank]; - const size_t in_count = range.in_count; // P - const size_t out_count = range.sin_send_count - in_count; // Q - const size_t sin_send_local = (idx < out_count) ? (in_count + idx) : (idx - out_count); - return cross_rank_sin_send_index(storage, rank, sin_send_local); +inline auto slot_sin_recv_index(const CrossRankSlotView &slot, size_t idx) -> size_t { + const size_t out_count = slot.sin_send_count - slot.in_count; // Q + const size_t sin_send_local = (idx < out_count) ? (slot.in_count + idx) : (idx - out_count); + return slot_sin_send_index(slot, sin_send_local); } // The D phases run parallel to the B indices -- same count per slot, so the same prefix sum // addresses both. They are still separate arrays; only the offset into them is shared. +inline auto slot_sin_recv_phase(const CrossRankSlotView &slot, size_t idx) -> int { + return packed_phase_at(*slot.sin_recv_phases, slot.phase_offset + idx); +} + +// Single-endpoint forms, for callers that genuinely touch one endpoint of one slot. A loop +// should resolve the slot once with cross_rank_slot() instead of calling these repeatedly. +inline auto cross_rank_sin_send_index(const PackedCrossRankStorage &storage, size_t rank, size_t idx) -> size_t { + return slot_sin_send_index(cross_rank_slot(storage, rank), idx); +} + +inline auto cross_rank_sin_recv_index(const PackedCrossRankStorage &storage, size_t rank, size_t idx) -> size_t { + return slot_sin_recv_index(cross_rank_slot(storage, rank), idx); +} + inline auto cross_rank_sin_recv_phase(const PackedCrossRankStorage &storage, size_t rank, size_t idx) -> int { - return packed_phase_at(storage.sin_recv_phases, storage.ranges[rank].sin_send_offset + idx); + return slot_sin_recv_phase(cross_rank_slot(storage, rank), idx); } auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t; @@ -108,12 +145,24 @@ auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size // low occupancy would make a sparse layout pay, high occupancy means only a narrower record does. auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t; +// Derive a layer's send layout into caller-owned scratch instead of reading a stored one. +// +// counts[r] = scale * (r == my_rank ? 0 : cross_rank.sin_send_size(r)), displs the prefix sum -- +// the same rule build_layer_storage_unified used to build the stored copy, so this reproduces it +// exactly rather than approximating it. `out` is resized, not reallocated, when reused across +// layers at a fixed world size. +auto derive_exchange_layout(const PackedCrossRankStorage &cross_rank, + size_t my_rank, + int scale, + LayerExchangeLayout &out, + const char *what = "Layer exchange") -> void; + auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> size_t; -// The resolve_recv transpose cache hanging off a layout. Separate from +// The resolve_recv transpose cache a layer retains. Separate from // layer_exchange_layout_storage_bytes because that function's result is already carried in a // shipped metric; folding this in would redefine it. -auto layer_exchange_layout_cache_bytes(const LayerExchangeLayout &layout) -> size_t; +auto layer_exchange_layout_cache_bytes(const mpi::RecvLayoutCache &cache) -> size_t; // Local cycles fold into the self-rank slot (my_rank); the exchange layout zeroes counts[my_rank] so // MPI_Alltoallv skips it (replay does a local copy). diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h index ac60854b..5a4e35a4 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h @@ -27,13 +27,18 @@ namespace monoprop { +// counts/displs for one alltoallv. Both are dense int[P] because MPI requires that at the call +// site, but this is a TRANSIENT: it is materialized into per-thread scratch for the exchange +// being posted, never retained per layer. A retained one costs P ints x2 x layers x partitions, +// which is O(P^2) across a job for something derivable in a prefix sum. +// +// It deliberately does NOT own a RecvLayoutCache any more. The cache is the transpose of one +// specific send pattern, so hanging it off a reused scratch object is precisely the way to serve +// layer A's transpose to layer B; it now lives on LayerCore, beside the pattern that produced it. struct LayerExchangeLayout final { std::vector counts; std::vector displs; size_t total_count = 0; - - // Cached recv counts/displs (see mpi::resolve_recv); mutable — filled through const handles at eval time. - mutable mpi::RecvLayoutCache recv_cache; }; } // namespace monoprop @@ -44,13 +49,17 @@ auto checked_mpi_int(size_t value, const char *what) -> int; // Per-rank MPI counts = send_counts[r] * scale, with prefix-sum displacements. send_counts is full-width // (size_t) so checked_mpi_int catches the narrowing to MPI's int. +// +// The engine no longer calls this: it derives the layout from the slot records instead (see +// derive_exchange_layout, declared in MPGraphEncodingStorage.h because it needs +// PackedCrossRankStorage). It is kept deliberately, as the REFERENCE the derivation is tested +// against -- graph_encoding_derived_layout_matches_the_layout_it_replaces asserts the two agree +// elementwise. Checking a derivation against an independent construction is worth more than +// checking it against literals, so this is a test oracle, not dead code. Do not delete it +// without replacing what it proves. auto build_layer_exchange_layout(const std::vector &send_counts, int scale, const char *what = "Layer exchange") -> LayerExchangeLayout; -// The derivative layout is the evolution layout at 2x (each rotation endpoint carries both the op and -// state payload). -auto build_derivative_exchange_layout(const LayerExchangeLayout &evolution) -> LayerExchangeLayout; - } // namespace monoprop::detail namespace monoprop { @@ -152,18 +161,29 @@ struct PackedCrossRankStorage final { struct LayerCore final { PackedCrossRankStorage cross_rank; - LayerExchangeLayout evolution_exchange_layout; - auto derivative_exchange_layout() const -> const LayerExchangeLayout &; - - // Bytes held by the lazily built derivative layout, 0 while it has never been asked for. - // Deliberately NOT implemented as a call to derivative_exchange_layout(): that would - // allocate the very thing being measured, turning an accounting read into a 2*P-int - // allocation on every layer and making the instrument report its own footprint. - auto derivative_exchange_layout_bytes() const -> size_t; - - // A copied core must not inherit the source's cache: it is eval-time state, not data. - auto reset_derivative_exchange_layout() -> void { derivative_exchange_layout_cache_.reset(); } + // The evolution layout is NOT stored at all: counts[r] is + // (r == my_rank ? 0 : cross_rank.sin_send_size(r)), displs is its prefix sum, and the total + // is the last displacement -- so the whole 2*P-int array was a second copy of what `ranges` + // already says, retained per layer per partition. detail::derive_exchange_layout rebuilds it + // into per-thread scratch for the transfer being posted. + + // The transpose of this layer's send pattern, which is the one part that cannot be derived + // locally -- it takes a collective. Cached per layer because the alternative is an + // MPI_Alltoall per layer per evaluation. mutable: filled through const handles at eval time. + mutable mpi::RecvLayoutCache evolution_recv_cache; + + // Rank-uniform identity for the send pattern this core holds, assigned in build order so + // every rank agrees on it. It is what makes reusing evolution_recv_cache safe: see + // mpi::resolve_recv, where a rank-LOCAL key would let one rank reuse while another rebuilds + // and rebuilding is a collective -- a hang, not a wrong answer. + uint64_t exchange_generation = 0; + + // There is deliberately no reset here. The retained derivative layout that used to need + // dropping on a copy no longer exists -- it is derived per exchange. The transpose cache + // that remains is a function of the send pattern, and a copy carries that pattern and its + // generation with it, so the copy inherits a cache that is still correct. Clearing it would + // cost an MPI_Alltoall per layer on the next evaluation to rebuild something already right. // Per-layer recompute metadata: generator_words = this layer's generator G as backing words; // scaled_count = fold truncation bound = operator size after this layer's partner inserts. @@ -175,9 +195,6 @@ struct LayerCore final { double gen_coeff = 0.0; // Shared by all layers of one multi-term gate; absolute across build_graph calls (parameter_mapping). size_t gate_index = 0; - -private: - mutable std::optional derivative_exchange_layout_cache_; }; } // namespace monoprop diff --git a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl index 3b529777..8f9a860f 100644 --- a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl +++ b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl @@ -819,8 +819,12 @@ auto MonomialPropagator::set_parameter_mapping(const VecZ ¶meter_m auto relabel = [this](size_t layer, size_t new_param_index) { auto &target = graph_.get_layer(layer); auto new_core = std::make_shared(target.core()); - // Drop the inherited eval-time derivative layout: it must not depend on a prior gradient run. - new_core->reset_derivative_exchange_layout(); + // The copy keeps the source's transpose cache on purpose. Relabelling changes only which + // parameter drives the rotation, never which endpoints cross to which slot, so the cached + // recv layout is still the transpose of this core's send pattern -- it carries the pattern + // and its generation along with it. Dropping it would cost one collective per layer to + // rebuild an identical answer. (The retained derivative layout that used to be dropped + // here no longer exists; it is derived per exchange from the evolution one.) new_core->param_index = new_param_index; if (const CosMask *pruned = target.pruned_cos()) { target = Layer(std::move(new_core), *pruned); diff --git a/cpp/monoprop/detail/mpi/Exchange.h b/cpp/monoprop/detail/mpi/Exchange.h index 6e230efd..72367f2f 100644 --- a/cpp/monoprop/detail/mpi/Exchange.h +++ b/cpp/monoprop/detail/mpi/Exchange.h @@ -25,10 +25,14 @@ namespace monoprop::mpi { -// Resolve the recv side of a send-count vector, reusing `cache` when comm size is unchanged: a -// replayed graph's send pattern is fixed, so a hit removes one blocking count round-trip per layer -// per evaluation. -auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayoutCache &cache) -> const RecvLayout &; +// Resolve the recv side of a send-count vector, reusing `cache` when the comm size and the send +// pattern are both unchanged: a replayed graph's send pattern is fixed, so a hit removes one +// blocking count round-trip per layer per evaluation. +// +// `generation` identifies the send pattern and MUST be rank-uniform -- see RecvLayoutCache, where +// the reason (a miss is a collective, so a split decision hangs) is spelled out. +auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayoutCache &cache, uint64_t generation) + -> const RecvLayout &; // Idempotent completion handle for a posted payload transfer; move-only, so a request is waited on // exactly once. wait() is a no-op on the blocking path and in non-MPI builds. Owns its request: the diff --git a/cpp/monoprop/detail/mpi/MPICompat.cpp b/cpp/monoprop/detail/mpi/MPICompat.cpp index 30141d4b..c744ff69 100644 --- a/cpp/monoprop/detail/mpi/MPICompat.cpp +++ b/cpp/monoprop/detail/mpi/MPICompat.cpp @@ -119,7 +119,8 @@ auto alltoall_counts(const int *send_counts, int *recv_counts, int n, Comm comm) #endif } -auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayoutCache &cache) -> const RecvLayout & { +auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayoutCache &cache, uint64_t generation) + -> const RecvLayout & { const auto n = static_cast(send_counts.size()); const int comm_size = mpi::size(comm); // alltoall_counts moves comm_size ints each way regardless of `n`, so a send vector that is not @@ -132,7 +133,14 @@ auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayout n, comm_size)); } - if (cache.comm_size == comm_size && static_cast(cache.layout.counts.size()) == n) { + // `generation` identifies the send pattern. Without it the predicate was "have we ever + // resolved anything for a communicator this size", which is true for every layer after the + // first -- correct only while each cache belonged to the one layout that produced it, and + // silently wrong the moment a caller resolves a second pattern through the same cache. + // Every rank holds the same generation for the same layer, so all ranks agree on the miss + // and enter alltoall_counts together. + if (cache.comm_size == comm_size && cache.generation == generation + && static_cast(cache.layout.counts.size()) == n) { return cache.layout; } @@ -147,6 +155,7 @@ auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayout } out.total = checked_mpi_count(total); cache.comm_size = comm_size; + cache.generation = generation; return out; } diff --git a/cpp/monoprop/detail/mpi/RecvLayout.h b/cpp/monoprop/detail/mpi/RecvLayout.h index f48f9d4a..ac6aadff 100644 --- a/cpp/monoprop/detail/mpi/RecvLayout.h +++ b/cpp/monoprop/detail/mpi/RecvLayout.h @@ -14,9 +14,10 @@ #pragma once +#include #include -// Kept MPI-free and dependency-light so graph-encoding types (LayerExchangeLayout) can embed the cache +// Kept MPI-free and dependency-light so graph-encoding types (LayerCore) can embed the cache // without pulling in or the exchange machinery (see Exchange.h). namespace monoprop::mpi { @@ -27,9 +28,22 @@ struct RecvLayout { int total = 0; }; +// The transpose of ONE send pattern. Reusing it for a different pattern returns wrong +// displacements silently, so the cache carries the identity of the pattern it was built from. +// +// What must be rank-uniform is the hit/miss DECISION, not the id's value. A miss runs +// alltoall_counts, which is a collective, so two ranks disagreeing about validity is a +// distributed HANG rather than a wrong answer. Binding the id to the layer that owns the cache +// gives that for free: every rank walks the same layers in the same order, so every rank misses +// on a layer's first resolve and hits afterwards, whatever the local id values happen to be. +// +// It must NOT be derived from the send counts themselves (a total, a checksum). Those are +// rank-local, so two patterns can collide on one rank and not on another -- and that difference +// is exactly the split decision that hangs. struct RecvLayoutCache { RecvLayout layout; int comm_size = -1; + uint64_t generation = 0; // 0 = never populated }; } // namespace monoprop::mpi diff --git a/cpp/tests/graph_encoding_tests.cpp b/cpp/tests/graph_encoding_tests.cpp index b2f521ae..472f3ea3 100644 --- a/cpp/tests/graph_encoding_tests.cpp +++ b/cpp/tests/graph_encoding_tests.cpp @@ -133,34 +133,86 @@ BOOST_AUTO_TEST_CASE(graph_encoding_exchange_layout_scale_and_displacements) { BOOST_CHECK_GT(detail::layer_exchange_layout_storage_bytes(s1), 0U); } -// Production only builds scale=1; the 2x layout reaches MPI through this accessor, which is -// unreachable at comm size 1, so the default non-MPI suite would otherwise never touch it. +// The layout is no longer stored, so the claim under test is EQUIVALENCE: what the exchange +// derives at the call site must equal, elementwise, what the retained copy used to hold. Asserted +// against build_layer_exchange_layout rather than against hand-written literals, so the two cannot +// drift apart in the same direction. +namespace { + +auto slot_partners(const std::vector &sin_send_counts) -> std::vector { + std::vector data(sin_send_counts.size()); + for (size_t r = 0; r < sin_send_counts.size(); ++r) { + for (size_t k = 0; k < sin_send_counts[r]; ++k) { + data[r].sin_send_indices.push_back(k); + data[r].sin_recv_entries.push_back({k, 1}); + } + } + return data; +} + +} // namespace -BOOST_AUTO_TEST_CASE(graph_encoding_derivative_exchange_layout_is_twice_the_evolution_layout) { - LayerCore core; - core.evolution_exchange_layout = detail::build_layer_exchange_layout({3, 0, 5}, /*scale=*/1); +BOOST_AUTO_TEST_CASE(graph_encoding_derived_layout_matches_the_layout_it_replaces) { + const std::vector counts{3, 0, 5, 2}; + const auto storage = detail::build_packed_cross_rank_storage(slot_partners(counts)); - const auto &derivative = core.derivative_exchange_layout(); - BOOST_CHECK((derivative.counts == std::vector{6, 0, 10})); - BOOST_CHECK((derivative.displs == std::vector{0, 6, 6})); - BOOST_CHECK_EQUAL(derivative.total_count, 16U); + for (size_t my_rank = 0; my_rank < counts.size(); ++my_rank) { + // The self slot is excluded from the transfer and handled locally. + std::vector expected_counts = counts; + expected_counts[my_rank] = 0; - // Cached: the second read returns the same object, so eval-time MPI holds a stable pointer. - BOOST_CHECK_EQUAL(&core.derivative_exchange_layout(), &derivative); + for (const int scale : {1, 2}) { + const auto reference = detail::build_layer_exchange_layout(expected_counts, scale); + LayerExchangeLayout derived; + detail::derive_exchange_layout(storage, my_rank, scale, derived); - // Reset drops the cache (relabel copies cores and must not inherit eval-time state). - core.reset_derivative_exchange_layout(); - BOOST_CHECK_EQUAL(core.derivative_exchange_layout().total_count, 16U); + BOOST_CHECK(derived.counts == reference.counts); + BOOST_CHECK(derived.displs == reference.displs); + BOOST_CHECK_EQUAL(derived.total_count, reference.total_count); + } + } +} + +BOOST_AUTO_TEST_CASE(graph_encoding_derived_layout_reuses_its_scratch) { + // Reused across layers, so it must overwrite rather than append -- a stale tail would be read + // by MPI as a real count for a slot this layer does not send to. + const auto wide = detail::build_packed_cross_rank_storage(slot_partners({1, 2, 3, 4})); + const auto narrow = detail::build_packed_cross_rank_storage(slot_partners({7, 7})); + + LayerExchangeLayout scratch; + detail::derive_exchange_layout(wide, /*my_rank=*/0, 1, scratch); + BOOST_CHECK_EQUAL(scratch.counts.size(), 4U); + detail::derive_exchange_layout(narrow, /*my_rank=*/0, 1, scratch); + BOOST_CHECK_EQUAL(scratch.counts.size(), 2U); + BOOST_CHECK((scratch.counts == std::vector{0, 7})); + BOOST_CHECK_EQUAL(scratch.total_count, 7U); +} + +BOOST_AUTO_TEST_CASE(graph_encoding_a_zero_traffic_slot_still_gets_a_valid_displacement) { + // Empty slots are where an off-by-one in a prefix sum hides: the count is 0 but the + // displacement must still be non-decreasing, or MPI reads a peer's payload at the wrong base. + const auto storage = detail::build_packed_cross_rank_storage(slot_partners({0, 4, 0, 0, 6})); + LayerExchangeLayout derived; + detail::derive_exchange_layout(storage, /*my_rank=*/3, 1, derived); + + BOOST_CHECK((derived.counts == std::vector{0, 4, 0, 0, 6})); + BOOST_CHECK((derived.displs == std::vector{0, 0, 4, 4, 4})); + BOOST_CHECK_EQUAL(derived.total_count, 10U); + for (size_t r = 1; r < derived.displs.size(); ++r) { + BOOST_CHECK_GE(derived.displs[r], derived.displs[r - 1]); + } } BOOST_AUTO_TEST_CASE(graph_encoding_derivative_exchange_layout_overflow_throws) { - // A count that fits int at 1x but not at 2x. build_layer_storage_unified runs this derivation - // eagerly, so the throw lands in build_graph and not inside the gradient collective window. + // A count that fits int at 1x but not at 2x. build_layer_storage_unified derives the 2x layout + // eagerly, so the throw lands in build_graph and not inside the gradient collective window, + // where peers are already blocked in resolve_recv's count round -- a hang, not an error. const size_t just_over_half = static_cast(std::numeric_limits::max()) / 2 + 1; + const auto storage = detail::build_packed_cross_rank_storage(slot_partners({just_over_half})); - LayerCore core; - core.evolution_exchange_layout = detail::build_layer_exchange_layout({just_over_half}, 1); - BOOST_CHECK_THROW(detail::build_derivative_exchange_layout(core.evolution_exchange_layout), std::overflow_error); + LayerExchangeLayout derived; + BOOST_CHECK_NO_THROW(detail::derive_exchange_layout(storage, /*my_rank=*/1, 1, derived)); + BOOST_CHECK_THROW(detail::derive_exchange_layout(storage, /*my_rank=*/1, 2, derived), std::overflow_error); } BOOST_AUTO_TEST_CASE(graph_encoding_d_from_b_derivation_both_arms) { @@ -231,21 +283,24 @@ BOOST_AUTO_TEST_CASE(graph_encoding_slot_record_bytes_track_the_world_not_the_tr detail::cross_rank_storage_bytes(narrow_storage)); } -BOOST_AUTO_TEST_CASE(graph_encoding_lazy_layout_bytes_do_not_force_the_allocation) { - LayerCore core; - core.evolution_exchange_layout = detail::build_layer_exchange_layout({3, 0, 5}, /*scale=*/1); +BOOST_AUTO_TEST_CASE(graph_encoding_a_layer_retains_no_exchange_layout) { + // The point of the change: a built layer holds the slot records and nothing else sized by P. + // The transpose cache is eval-time state, so on a freshly built layer nothing has resolved it + // and it is not resident either. + const auto core = detail::build_layer_storage_unified(slot_partners({3, 0, 5}), /*my_rank=*/1); - // Reading the size must not build the thing being sized, or the instrument reports its - // own footprint and every layer pays 2*P ints for having been measured. - BOOST_CHECK_EQUAL(core.derivative_exchange_layout_bytes(), 0U); - static_cast(core.derivative_exchange_layout()); - BOOST_CHECK_GT(core.derivative_exchange_layout_bytes(), 0U); + BOOST_CHECK_EQUAL(detail::layer_exchange_layout_cache_bytes(core->evolution_recv_cache), 0U); - core.reset_derivative_exchange_layout(); - BOOST_CHECK_EQUAL(core.derivative_exchange_layout_bytes(), 0U); + // Not stored, but not lost: the send total is still recoverable from the slot records alone, + // which is the whole claim. 3 + 5, with my_rank's own slot contributing nothing. + LayerExchangeLayout derived; + detail::derive_exchange_layout(core->cross_rank, /*my_rank=*/1, /*scale=*/1, derived); + BOOST_CHECK_EQUAL(derived.total_count, 8U); - // The transpose cache is eval-time state: nothing has resolved it, so it is not resident. - BOOST_CHECK_EQUAL(detail::layer_exchange_layout_cache_bytes(core.evolution_exchange_layout), 0U); + // Rank-uniform identity for the send pattern, so the cache above cannot be served to another. + BOOST_CHECK_GT(core->exchange_generation, 0U); + const auto other = detail::build_layer_storage_unified(slot_partners({3, 0, 5}), /*my_rank=*/1); + BOOST_CHECK_NE(core->exchange_generation, other->exchange_generation); } BOOST_AUTO_TEST_CASE(graph_encoding_skewed_endpoint_counts_are_refused) { From cff597e36329ffdc6418d0fba2ac54fd5ad45457 Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sun, 16 Aug 2026 10:35:33 +0100 Subject: [PATCH 3/3] perf(evolution): :zap: derive the recv layout too, instead of caching the transpose The previous commit stopped retaining the send layout but kept a RecvLayoutCache per layer -- 8 B per world slot, 10.59 GiB at P=512 -- on the grounds that a transpose is the one thing a rank cannot work out alone. That was wrong: this transpose carries data both sides already have. Layer build gives slot r on rank m the queries r sent m, followed by the queries m sent r; rank r's slot for m holds those two swapped. The counts are therefore equal, and displacements are prefix sums of counts, so the recv layout IS the send layout. MPI reads recvcounts/recvdispls rather than writing them, so the same two arrays now serve both sides of the alltoallv. What goes with the cache: the alltoall_counts on its miss path, and the rank-uniform `exchange_generation` that existed only to keep that miss rank uniform. The hazard the previous commit documented so carefully -- a split reuse decision hanging the job -- is removed rather than managed, because there is no longer a collective on any cache-miss path. sizeof(LayerCore) 248 -> 168 B. Symmetry is an invariant of the routing, not of this file, so it is checked where it can actually break: MONOPROP_CHECK_EXCHANGE_SYMMETRY=1 re-adds the alltoall and throws naming the slot and both counts. Unguarded, a future routing change that broke it would surface as a peer blocked in MPI_Alltoallv against a size nobody sends -- a hang with no line number. Evidence: a probe comparing derived counts against a real alltoall on every resolve saw 0 mismatches in 550M slot comparisons at world 32 and 256, over the full MPI suite and a pauli c12 energy+gradient run. Gate 1826413: 214 ctest serial, and 625 Python tests on each of four geometries TWICE -- once on the production path, once with the assertion live. Co-Authored-By: Claude Opus 5 --- cpp/monoprop/Evolution.cpp | 79 ++++++++----------- cpp/monoprop/MPGraph.cpp | 8 +- cpp/monoprop/detail/graph/MPGraphLayers.h | 7 +- cpp/monoprop/detail/graph/MPGraphViews.h | 2 +- .../detail/graph_encoding/MPGraphEncoding.cpp | 13 +-- .../graph_encoding/MPGraphEncodingStorage.h | 5 -- .../graph_encoding/MPGraphEncodingTypes.h | 29 +++---- .../MonomialPropagator.inl | 10 +-- cpp/monoprop/detail/mpi/CMakeLists.txt | 1 - cpp/monoprop/detail/mpi/Exchange.h | 18 +++-- cpp/monoprop/detail/mpi/MPICompat.cpp | 48 +++++------ cpp/monoprop/detail/mpi/RecvLayout.h | 49 ------------ cpp/tests/graph_encoding_tests.cpp | 13 ++- 13 files changed, 93 insertions(+), 189 deletions(-) delete mode 100644 cpp/monoprop/detail/mpi/RecvLayout.h diff --git a/cpp/monoprop/Evolution.cpp b/cpp/monoprop/Evolution.cpp index cfd679b6..6e644657 100644 --- a/cpp/monoprop/Evolution.cpp +++ b/cpp/monoprop/Evolution.cpp @@ -60,14 +60,14 @@ auto combine_endpoint_contrib(const EndpointContrib &a, const EndpointContrib &b struct FlatExchangeBuffers { VecD send_buffer; VecD recv_buffer; - std::vector recv_counts; - std::vector recv_displs; - // The send 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. + // 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; - int recv_total = 0; }; auto &acquire_flat_exchange_buffers() { @@ -79,11 +79,8 @@ auto &acquire_flat_exchange_buffers() { } void resize_flat_exchange_buffers(const LayerExchangeLayout &layout, FlatExchangeBuffers &buffers) { - // Only the send side: the recv counts/displs are filled by derive_layer_exchange before this - // runs, so clearing them here (as this did when the recv side was resolved afterwards) would - // throw away the layout the transfer is about to be posted with. - const size_t send_alloc = layout.total_count == 0 ? 1 : layout.total_count; - buffers.send_buffer.resize(send_alloc); + const size_t alloc = layout.total_count == 0 ? 1 : layout.total_count; + buffers.send_buffer.resize(alloc); } // Nothing to exchange at one rank. All ranks must participate even at local total_count 0, else @@ -92,43 +89,29 @@ auto layer_exchange_participates(const mpi::Comm &comm) -> bool { return mpi::size(comm) != 1; } -// Derive this layer's send layout into `buffers.layout` at `scale`, and its recv side into -// buffers.recv_counts/recv_displs at the same scale. +// Derive this layer's exchange layout into `buffers.layout` at `scale`. It describes BOTH sides. // -// The recv side is resolved ONCE per layer, at scale 1, into the layer's own cache. Scaling -// commutes with the transpose -- every rank multiplies its counts by the same literal, so peer q -// sends me exactly `scale` times what it sent at scale 1, and displacements are prefix sums of -// counts so they scale with them. That is what lets the derivative round reuse the evolution -// round's collective instead of paying an alltoall of its own. +// 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(mpi::rank(comm)); - - // Resolve the transpose at scale 1, so evolution and derivative rounds share one cache entry - // and therefore one collective. - detail::derive_exchange_layout(layer.cross_rank(), my_rank, 1, buffers.layout); - const auto &recv = - mpi::resolve_recv(buffers.layout.counts, comm, layer.evolution_recv_cache(), layer.exchange_generation()); - - const size_t n = recv.counts.size(); - buffers.recv_counts.resize(n); - buffers.recv_displs.resize(n); - for (size_t i = 0; i < n; ++i) { - buffers.recv_counts[i] = detail::checked_mpi_int(static_cast(recv.counts[i]) * static_cast(scale), - "Layer exchange recv count"); - buffers.recv_displs[i] = detail::checked_mpi_int(static_cast(recv.displs[i]) * static_cast(scale), - "Layer exchange recv displacement"); - } - buffers.recv_total = detail::checked_mpi_int(static_cast(recv.total) * static_cast(scale), - "Layer exchange recv total"); - - // Now the send side at the requested scale. Re-derived rather than scaled in place so the - // overflow check runs against the value MPI actually receives. - if (scale != 1) { - detail::derive_exchange_layout(layer.cross_rank(), my_rank, scale, buffers.layout, "Layer derivative exchange"); - } + 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 @@ -151,13 +134,15 @@ inline auto begin_flat_exchange(FlatExchangeBuffers &buffers, const mpi::Comm &c CrossRankExchangeHandle handle; handle.layout = &layout; handle.buffers = &buffers; - buffers.recv_buffer.resize(buffers.recv_total == 0 ? 1 : static_cast(buffers.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({.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; @@ -206,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}); } diff --git a/cpp/monoprop/MPGraph.cpp b/cpp/monoprop/MPGraph.cpp index 6df4fca9..f1cba609 100644 --- a/cpp/monoprop/MPGraph.cpp +++ b/cpp/monoprop/MPGraph.cpp @@ -57,10 +57,12 @@ auto layer_storage_memory_usage(const LayerCore &storage) -> GraphMemoryBreakdow // that did retain them shows the drop rather than silently losing the row. breakdown.exchange_layout_bytes = 0; - // Diagnostics. recv_cache is real resident memory that total_bytes() has never counted, which - // is why the reported graph size sits below the process RSS. + // Diagnostics. breakdown.slot_record_bytes = detail::cross_rank_slot_record_bytes(storage.cross_rank); - breakdown.recv_cache_bytes = detail::layer_exchange_layout_cache_bytes(storage.evolution_recv_cache); + // 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; diff --git a/cpp/monoprop/detail/graph/MPGraphLayers.h b/cpp/monoprop/detail/graph/MPGraphLayers.h index b2156385..eccea684 100644 --- a/cpp/monoprop/detail/graph/MPGraphLayers.h +++ b/cpp/monoprop/detail/graph/MPGraphLayers.h @@ -88,12 +88,9 @@ struct LayerTraversal final { return detail::cross_rank_slot(core_->cross_rank, rank); } - // The exchange layout is derived at the call site from these, not stored: see - // detail::derive_exchange_layout. Only the transpose cache survives, because only it - // needs a collective. + // 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 evolution_recv_cache() const -> mpi::RecvLayoutCache & { return core_->evolution_recv_cache; } - auto exchange_generation() const -> uint64_t { return core_->exchange_generation; } auto param_index() const -> size_t { return core_->param_index; } auto gen_coeff() const -> double { return core_->gen_coeff; } diff --git a/cpp/monoprop/detail/graph/MPGraphViews.h b/cpp/monoprop/detail/graph/MPGraphViews.h index 525d6a80..9e7c4083 100644 --- a/cpp/monoprop/detail/graph/MPGraphViews.h +++ b/cpp/monoprop/detail/graph/MPGraphViews.h @@ -50,7 +50,7 @@ struct GraphMemoryBreakdown final { // 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; // evolution layout's resolve_recv transpose cache -- never in total_bytes() + 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 diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index 662a9c56..26c1d3fb 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp @@ -15,7 +15,6 @@ #include "monoprop/detail/graph_encoding/MPGraphEncodingStorage.h" #include -#include #include #include #include @@ -174,11 +173,6 @@ auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> s return layout.counts.capacity() * sizeof(int) + layout.displs.capacity() * sizeof(int); } -auto layer_exchange_layout_cache_bytes(const mpi::RecvLayoutCache &cache) -> size_t { - const auto &cached = cache.layout; - return cached.counts.capacity() * sizeof(int) + cached.displs.capacity() * sizeof(int); -} - auto derive_exchange_layout(const PackedCrossRankStorage &cross_rank, size_t my_rank, int scale, @@ -204,12 +198,7 @@ auto derive_exchange_layout(const PackedCrossRankStorage &cross_rank, auto build_layer_storage_unified(std::vector all_partners, size_t my_rank) -> std::shared_ptr { - // Identifies the send pattern this core holds, so its recv cache cannot be served to a - // different one. Starts at 1 because a default-constructed cache carries 0 = never populated. - static std::atomic next_generation{1}; - auto storage = std::make_shared(); - storage->exchange_generation = next_generation.fetch_add(1, std::memory_order_relaxed); const size_t num_ranks = all_partners.size(); storage->cross_rank = build_packed_cross_rank_storage(std::move(all_partners)); @@ -225,7 +214,7 @@ auto build_layer_storage_unified(std::vector all_partners, // Derive both scales once at build time and throw the result away. This is purely eager // validation: an overflow of MPI's int has to throw from build_graph, not from inside the - // exchange, where peers are already blocked in resolve_recv's count round -- there it is a + // exchange, where peers are already committed to a transfer of that size -- there it is a // distributed hang rather than an error. Scale 2 is checked as well as 1 because the // derivative round overflows first and a gradient may run long after the graph was built. // diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h index 18518633..1d4755a2 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -159,11 +159,6 @@ auto derive_exchange_layout(const PackedCrossRankStorage &cross_rank, auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> size_t; -// The resolve_recv transpose cache a layer retains. Separate from -// layer_exchange_layout_storage_bytes because that function's result is already carried in a -// shipped metric; folding this in would redefine it. -auto layer_exchange_layout_cache_bytes(const mpi::RecvLayoutCache &cache) -> size_t; - // Local cycles fold into the self-rank slot (my_rank); the exchange layout zeroes counts[my_rank] so // MPI_Alltoallv skips it (replay does a local copy). auto build_layer_storage_unified(std::vector all_partners, size_t my_rank) diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h index 5a4e35a4..b3b3c3f7 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h @@ -23,7 +23,6 @@ #include #include "monoprop/TypeAliases.h" -#include "monoprop/detail/mpi/RecvLayout.h" namespace monoprop { @@ -32,9 +31,9 @@ namespace monoprop { // being posted, never retained per layer. A retained one costs P ints x2 x layers x partitions, // which is O(P^2) across a job for something derivable in a prefix sum. // -// It deliberately does NOT own a RecvLayoutCache any more. The cache is the transpose of one -// specific send pattern, so hanging it off a reused scratch object is precisely the way to serve -// layer A's transpose to layer B; it now lives on LayerCore, beside the pattern that produced it. +// It describes the recv side too: the count matrix is symmetric, so the transpose of a send +// pattern is that send pattern. There is no RecvLayoutCache anywhere any more -- not here, and +// not on LayerCore, which is where one briefly lived. struct LayerExchangeLayout final { std::vector counts; std::vector displs; @@ -168,22 +167,12 @@ struct LayerCore final { // already says, retained per layer per partition. detail::derive_exchange_layout rebuilds it // into per-thread scratch for the transfer being posted. - // The transpose of this layer's send pattern, which is the one part that cannot be derived - // locally -- it takes a collective. Cached per layer because the alternative is an - // MPI_Alltoall per layer per evaluation. mutable: filled through const handles at eval time. - mutable mpi::RecvLayoutCache evolution_recv_cache; - - // Rank-uniform identity for the send pattern this core holds, assigned in build order so - // every rank agrees on it. It is what makes reusing evolution_recv_cache safe: see - // mpi::resolve_recv, where a rank-LOCAL key would let one rank reuse while another rebuilds - // and rebuilding is a collective -- a hang, not a wrong answer. - uint64_t exchange_generation = 0; - - // There is deliberately no reset here. The retained derivative layout that used to need - // dropping on a copy no longer exists -- it is derived per exchange. The transpose cache - // that remains is a function of the send pattern, and a copy carries that pattern and its - // generation with it, so the copy inherits a cache that is still correct. Clearing it would - // cost an MPI_Alltoall per layer on the next evaluation to rebuild something already right. + // NOTHING about the exchange is retained here -- no send layout, no transpose, no identity + // for one. The recv layout equals the send layout (the count matrix is symmetric; see + // Evolution.cpp's derive_layer_exchange), so the transpose that used to be cached per layer + // at 8 B per world slot is not merely derivable, it is the same array. With it goes the + // rank-uniform generation id that existed only to make reusing that cache safe, and the + // hazard it managed: there is no longer a collective on any cache-miss path to split ranks on. // Per-layer recompute metadata: generator_words = this layer's generator G as backing words; // scaled_count = fold truncation bound = operator size after this layer's partner inserts. diff --git a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl index 8f9a860f..0e0bc487 100644 --- a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl +++ b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl @@ -819,12 +819,10 @@ auto MonomialPropagator::set_parameter_mapping(const VecZ ¶meter_m auto relabel = [this](size_t layer, size_t new_param_index) { auto &target = graph_.get_layer(layer); auto new_core = std::make_shared(target.core()); - // The copy keeps the source's transpose cache on purpose. Relabelling changes only which - // parameter drives the rotation, never which endpoints cross to which slot, so the cached - // recv layout is still the transpose of this core's send pattern -- it carries the pattern - // and its generation along with it. Dropping it would cost one collective per layer to - // rebuild an identical answer. (The retained derivative layout that used to be dropped - // here no longer exists; it is derived per exchange from the evolution one.) + // A plain copy, with nothing to invalidate: the core no longer retains any exchange + // state for a stale copy to serve. Relabelling changes only which parameter drives the + // rotation, never which endpoints cross to which slot, but that argument is no longer + // load-bearing -- both the send layout and its transpose are derived per exchange. new_core->param_index = new_param_index; if (const CosMask *pruned = target.pruned_cos()) { target = Layer(std::move(new_core), *pruned); diff --git a/cpp/monoprop/detail/mpi/CMakeLists.txt b/cpp/monoprop/detail/mpi/CMakeLists.txt index c8226a41..5b9969b0 100644 --- a/cpp/monoprop/detail/mpi/CMakeLists.txt +++ b/cpp/monoprop/detail/mpi/CMakeLists.txt @@ -12,7 +12,6 @@ target_sources( "MPICompat.h" "MPIUtils.h" "PartitionBarrier.h" - "RecvLayout.h" "ShmComm.h" ) diff --git a/cpp/monoprop/detail/mpi/Exchange.h b/cpp/monoprop/detail/mpi/Exchange.h index 72367f2f..fe9ef14e 100644 --- a/cpp/monoprop/detail/mpi/Exchange.h +++ b/cpp/monoprop/detail/mpi/Exchange.h @@ -19,20 +19,22 @@ #include #include "monoprop/detail/mpi/MPICompat.h" -#include "monoprop/detail/mpi/RecvLayout.h" // Keeps #ifdef monoprop_ENABLE_MPI out of the consumers; non-MPI builds get self-copy stubs. namespace monoprop::mpi { -// Resolve the recv side of a send-count vector, reusing `cache` when the comm size and the send -// pattern are both unchanged: a replayed graph's send pattern is fixed, so a hit removes one -// blocking count round-trip per layer per evaluation. +// Opt-in audit of the invariant the exchange now rests on: a layer's recv counts equal its send +// counts, so there is no transpose to compute or store (see Evolution.cpp's derive_layer_exchange). // -// `generation` identifies the send pattern and MUST be rank-uniform -- see RecvLayoutCache, where -// the reason (a miss is a collective, so a split decision hangs) is spelled out. -auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayoutCache &cache, uint64_t generation) - -> const RecvLayout &; +// Off unless MONOPROP_CHECK_EXCHANGE_SYMMETRY is set, because ON it costs exactly the collective +// the change exists to remove. It IS a collective, so the variable must be set identically on +// every rank -- setting it on one rank alone hangs. Read once, at first use. +// +// Worth having at all because of how the invariant fails if a future routing change breaks it: a +// peer blocks in MPI_Alltoallv against a size nobody sends, which is a hang with no line number. +// This turns that into an exception naming the slot. +auto check_exchange_symmetry(std::span send_counts, const Comm &comm) -> void; // Idempotent completion handle for a posted payload transfer; move-only, so a request is waited on // exactly once. wait() is a no-op on the blocking path and in non-MPI builds. Owns its request: the diff --git a/cpp/monoprop/detail/mpi/MPICompat.cpp b/cpp/monoprop/detail/mpi/MPICompat.cpp index c744ff69..cb0e7c11 100644 --- a/cpp/monoprop/detail/mpi/MPICompat.cpp +++ b/cpp/monoprop/detail/mpi/MPICompat.cpp @@ -14,9 +14,11 @@ #include "monoprop/detail/mpi/Exchange.h" +#include #include #include #include +#include namespace monoprop::mpi { @@ -119,8 +121,13 @@ auto alltoall_counts(const int *send_counts, int *recv_counts, int n, Comm comm) #endif } -auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayoutCache &cache, uint64_t generation) - -> const RecvLayout & { +auto check_exchange_symmetry(std::span send_counts, const Comm &comm) -> void { + // One read of the environment, not one per layer. Rank-uniform by assumption: this is a + // collective, so a variable set on some ranks and not others hangs rather than misreports. + static const bool enabled = std::getenv("MONOPROP_CHECK_EXCHANGE_SYMMETRY") != nullptr; + if (!enabled) { + return; + } const auto n = static_cast(send_counts.size()); const int comm_size = mpi::size(comm); // alltoall_counts moves comm_size ints each way regardless of `n`, so a send vector that is not @@ -133,30 +140,23 @@ auto resolve_recv(std::span send_counts, const Comm &comm, RecvLayout n, comm_size)); } - // `generation` identifies the send pattern. Without it the predicate was "have we ever - // resolved anything for a communicator this size", which is true for every layer after the - // first -- correct only while each cache belonged to the one layout that produced it, and - // silently wrong the moment a caller resolves a second pattern through the same cache. - // Every rank holds the same generation for the same layer, so all ranks agree on the miss - // and enter alltoall_counts together. - if (cache.comm_size == comm_size && cache.generation == generation - && static_cast(cache.layout.counts.size()) == n) { - return cache.layout; - } - - RecvLayout &out = cache.layout; - out.counts.resize(static_cast(n)); - alltoall_counts(send_counts.data(), out.counts.data(), n, comm); - out.displs.resize(static_cast(n)); - long long total = 0; + std::vector recv_counts(static_cast(n)); + alltoall_counts(send_counts.data(), recv_counts.data(), n, comm); for (int i = 0; i < n; ++i) { - out.displs[static_cast(i)] = checked_mpi_count(total); - total += out.counts[static_cast(i)]; + const int sent = send_counts[static_cast(i)]; + const int received = recv_counts[static_cast(i)]; + if (sent != received) { + // Naming the slot and both counts, because the whole point of the check is that the + // unguarded failure carries neither. + throw CollectiveArgumentError(std::format( + "Exchange count matrix is not symmetric at slot {}: this rank sends {} there but receives {} back. " + "The exchange derives its recv layout from its send layout on the strength of that equality, so a " + "routing change that breaks it must be caught here rather than as a hang in MPI_Alltoallv.", + i, + sent, + received)); + } } - out.total = checked_mpi_count(total); - cache.comm_size = comm_size; - cache.generation = generation; - return out; } } // namespace monoprop::mpi diff --git a/cpp/monoprop/detail/mpi/RecvLayout.h b/cpp/monoprop/detail/mpi/RecvLayout.h deleted file mode 100644 index ac6aadff..00000000 --- a/cpp/monoprop/detail/mpi/RecvLayout.h +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright 2026 Algorithmiq -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#pragma once - -#include -#include - -// Kept MPI-free and dependency-light so graph-encoding types (LayerCore) can embed the cache -// without pulling in or the exchange machinery (see Exchange.h). - -namespace monoprop::mpi { - -struct RecvLayout { - std::vector counts; - std::vector displs; - int total = 0; -}; - -// The transpose of ONE send pattern. Reusing it for a different pattern returns wrong -// displacements silently, so the cache carries the identity of the pattern it was built from. -// -// What must be rank-uniform is the hit/miss DECISION, not the id's value. A miss runs -// alltoall_counts, which is a collective, so two ranks disagreeing about validity is a -// distributed HANG rather than a wrong answer. Binding the id to the layer that owns the cache -// gives that for free: every rank walks the same layers in the same order, so every rank misses -// on a layer's first resolve and hits afterwards, whatever the local id values happen to be. -// -// It must NOT be derived from the send counts themselves (a total, a checksum). Those are -// rank-local, so two patterns can collide on one rank and not on another -- and that difference -// is exactly the split decision that hangs. -struct RecvLayoutCache { - RecvLayout layout; - int comm_size = -1; - uint64_t generation = 0; // 0 = never populated -}; - -} // namespace monoprop::mpi diff --git a/cpp/tests/graph_encoding_tests.cpp b/cpp/tests/graph_encoding_tests.cpp index 472f3ea3..d65ff796 100644 --- a/cpp/tests/graph_encoding_tests.cpp +++ b/cpp/tests/graph_encoding_tests.cpp @@ -285,22 +285,19 @@ BOOST_AUTO_TEST_CASE(graph_encoding_slot_record_bytes_track_the_world_not_the_tr BOOST_AUTO_TEST_CASE(graph_encoding_a_layer_retains_no_exchange_layout) { // The point of the change: a built layer holds the slot records and nothing else sized by P. - // The transpose cache is eval-time state, so on a freshly built layer nothing has resolved it - // and it is not resident either. + // Neither side of the exchange is retained -- not the send layout, and not a transpose of it. const auto core = detail::build_layer_storage_unified(slot_partners({3, 0, 5}), /*my_rank=*/1); - BOOST_CHECK_EQUAL(detail::layer_exchange_layout_cache_bytes(core->evolution_recv_cache), 0U); - // Not stored, but not lost: the send total is still recoverable from the slot records alone, // which is the whole claim. 3 + 5, with my_rank's own slot contributing nothing. LayerExchangeLayout derived; detail::derive_exchange_layout(core->cross_rank, /*my_rank=*/1, /*scale=*/1, derived); BOOST_CHECK_EQUAL(derived.total_count, 8U); - // Rank-uniform identity for the send pattern, so the cache above cannot be served to another. - BOOST_CHECK_GT(core->exchange_generation, 0U); - const auto other = detail::build_layer_storage_unified(slot_partners({3, 0, 5}), /*my_rank=*/1); - BOOST_CHECK_NE(core->exchange_generation, other->exchange_generation); + // A LayerCore is what gets held L x P times across a job, so its size is the thing the change + // is about. Pinned against the members it should have: three vectors' worth of storage plus + // the scalars. A new P-sized member here would be paid for once per layer per partition. + BOOST_CHECK_LE(sizeof(LayerCore), 256U); } BOOST_AUTO_TEST_CASE(graph_encoding_skewed_endpoint_counts_are_refused) {