From 1d6277344c48c41809fac037d6dfa309de1fc484 Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sat, 15 Aug 2026 15:03:33 +0100 Subject: [PATCH 1/6] 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/6] 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/6] 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) { From 020e3a8b067d0e73d8fa3afd58a6f9349e5b8c76 Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sun, 16 Aug 2026 12:01:45 +0100 Subject: [PATCH 4/6] feat(graph): :chart_with_upwards_trend: separate world-slot metadata from traffic in the graph breakdown The graph does not partition. Its per-layer arrays are indexed by rank, and on a partitioned run that index space is the FLAT world P = ranks x partitions, so they grow with a P the MPI rank count never shows. `graph_memory_bytes` is a single scalar and cannot say how much of it is that. Split the two growth laws so a measurement can separate them: d_slot_record_bytes the slice of cross_rank_bytes that is one record per world slot, carried whether or not the slot has traffic d_slot_records P per layer core; / d_layer_cores recovers P d_occupied_slots slots carrying any traffic; / d_slot_records is occupancy d_cross_rank_endpoints the traffic itself, and the ceiling on d_occupied_slots The last one is the point of the exercise. An occupied slot holds at least one endpoint, so endpoints bound occupied slots from above -- and endpoints do not depend on P at all. Together the two say how much of the slot array is information and how much is reserved-and-empty. All of them sit OUTSIDE total_bytes(): each is a count or a slice of a field already summed there, so adding them would double-count. Behaviour is unchanged; this only reports. --- cpp/monoprop/MPGraph.cpp | 4 +- cpp/monoprop/detail/graph/MPGraphViews.h | 27 +++++----- .../detail/graph_encoding/MPGraphEncoding.cpp | 17 +++++-- .../graph_encoding/MPGraphEncodingStorage.h | 14 ++++-- cpp/tests/graph_encoding_tests.cpp | 49 +++++++------------ src/monoprop/bindings/binder.h | 13 ++--- 6 files changed, 67 insertions(+), 57 deletions(-) diff --git a/cpp/monoprop/MPGraph.cpp b/cpp/monoprop/MPGraph.cpp index f1cba609..655a6630 100644 --- a/cpp/monoprop/MPGraph.cpp +++ b/cpp/monoprop/MPGraph.cpp @@ -57,7 +57,8 @@ 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. + // Diagnostics. The graph does not partition -- its per-layer arrays are indexed by the FLAT world + // (ranks x partitions), so on a partitioned run these grow with a P the rank count never shows. breakdown.slot_record_bytes = detail::cross_rank_slot_record_bytes(storage.cross_rank); // The transpose cache is gone: the recv layout equals the send layout, so there was never // anything to cache. Reported as 0 rather than removed, because it was never inside @@ -69,6 +70,7 @@ auto layer_storage_memory_usage(const LayerCore &storage) -> GraphMemoryBreakdow breakdown.layer_cores = 1; breakdown.slot_records = storage.cross_rank.rank_count(); breakdown.occupied_slots = detail::cross_rank_occupied_slots(storage.cross_rank); + breakdown.cross_rank_endpoints = detail::cross_rank_endpoint_count(storage.cross_rank); return breakdown; } diff --git a/cpp/monoprop/detail/graph/MPGraphViews.h b/cpp/monoprop/detail/graph/MPGraphViews.h index 9e7c4083..bbb2bdc8 100644 --- a/cpp/monoprop/detail/graph/MPGraphViews.h +++ b/cpp/monoprop/detail/graph/MPGraphViews.h @@ -40,21 +40,25 @@ 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. + // Diagnostics, deliberately EXCLUDED from total_bytes(): each is either a count, a subset of a + // field above, or memory total_bytes() has never counted. Folding any of them in would silently + // redefine graph_memory_bytes() mid-flight, so an A/B against an older build would compare two + // different quantities. // // The point of the split: a per-layer array indexed by rank is sized by the FLAT world // (mpi::size on a Hybrid comm is ranks x partitions), so it costs O(P) per layer per - // partition and O(P^2) across the job. slot_bytes is that part; traffic_bytes is the - // part that scales with terms actually crossing, which is real work. - size_t slot_record_bytes = 0; // cross_rank ranges[]: one record per world slot, occupied or not - size_t recv_cache_bytes = 0; // retired: the recv layout IS the send layout, nothing is cached + // partition and O(P^2) across the job. slot_record_bytes is that part; the endpoint count below + // is the part that scales with terms actually crossing, which is real work. + size_t slot_record_bytes = 0; // one record per STORED world slot -- occupied only, once sparse + size_t recv_cache_bytes = 0; // retired: the recv layout IS the send layout, nothing is cached size_t derivative_layout_bytes = 0; // the lazily retained 2x layout AND its own recv cache -- likewise - size_t layer_cores = 0; // distinct LayerCores walked (shared cores counted once) - size_t slot_records = 0; // sum over cores of ranges.size(); divide by layer_cores to recover P - size_t occupied_slots = 0; // slots carrying any traffic: occupancy = occupied_slots / slot_records + size_t layer_cores = 0; // distinct LayerCores walked (shared cores counted once) + size_t slot_records = 0; // the flat world P per core, so slot_records / layer_cores == P + size_t occupied_slots = 0; // slots carrying any traffic: occupancy = occupied_slots / slot_records + // Cross-rank endpoints -- the traffic itself, and the ceiling on occupied_slots, since an + // occupied slot holds at least one endpoint. Unlike slot_records it does not depend on P, so the + // two together say how much of the slot array is information and how much is reserved-and-empty. + size_t cross_rank_endpoints = 0; auto total_bytes() const -> size_t { return layer_descriptor_bytes + layer_storage_object_bytes + cos_data_bytes + cross_rank_bytes @@ -74,6 +78,7 @@ struct GraphMemoryBreakdown final { layer_cores += o.layer_cores; slot_records += o.slot_records; occupied_slots += o.occupied_slots; + cross_rank_endpoints += o.cross_rank_endpoints; return *this; } }; diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index 26c1d3fb..9becff8a 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp @@ -152,8 +152,7 @@ auto build_packed_cross_rank_storage(const std::vector &da } auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t { - size_t bytes = - storage.ranges.capacity() * sizeof(CrossRankPartnerRange) + packed_phase_storage_bytes(storage.sin_recv_phases); + size_t bytes = cross_rank_slot_record_bytes(storage) + packed_phase_storage_bytes(storage.sin_recv_phases); bytes += storage.sin_send_indices.capacity() * sizeof(TermIndex); return bytes; } @@ -163,12 +162,22 @@ auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size } 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. + // sin_send_count alone is the predicate. B and D hold the same endpoint set in two orders (see + // cross_rank_sin_recv_index), so a slot cannot carry D entries while carrying no B entries, and + // counting either gives the same answer. return static_cast(std::ranges::count_if( storage.ranges, [](const CrossRankPartnerRange &range) { return range.sin_send_count != 0; })); } +auto cross_rank_endpoint_count(const PackedCrossRankStorage &storage) -> size_t { + size_t count = 0; + for (const auto &range : storage.ranges) { + count += range.sin_send_count; + } + return count; +} + + auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> size_t { return layout.counts.capacity() * sizeof(int) + layout.displs.capacity() * sizeof(int); } diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h index 1d4755a2..a89fcbf9 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -137,14 +137,20 @@ inline auto cross_rank_sin_recv_phase(const PackedCrossRankStorage &storage, siz 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. +// The slot-proportional slice of cross_rank_storage_bytes: one record per STORED world slot. Once the +// storage is sparse that is one record per OCCUPIED slot, so this stops tracking the world size and +// starts tracking traffic -- which is what makes the two growth laws separable in a measurement. 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. +// World slots carrying any traffic. Read against rank_count() to get occupancy: low occupancy means a +// sparse layout pays, high occupancy means only a narrower record would. auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t; +// Total cross-rank endpoints in this layer -- the length of the B array, so the traffic itself. It +// bounds cross_rank_occupied_slots from above (an occupied slot holds at least one endpoint) and, +// unlike the slot count, does not depend on the world size at all. +auto cross_rank_endpoint_count(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 -- diff --git a/cpp/tests/graph_encoding_tests.cpp b/cpp/tests/graph_encoding_tests.cpp index d65ff796..543c5a46 100644 --- a/cpp/tests/graph_encoding_tests.cpp +++ b/cpp/tests/graph_encoding_tests.cpp @@ -243,44 +243,31 @@ BOOST_AUTO_TEST_CASE(graph_encoding_d_from_b_derivation_both_arms) { // 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. +// when the traffic through it does. (slot_partners is defined above, with the layout tests.) 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; + // Zeros at the front, in the interior and at the back -- the three places a scan loses count. + const auto storage = detail::build_packed_cross_rank_storage(slot_partners({0, 3, 0, 0, 7, 0})); - const auto storage = detail::build_packed_cross_rank_storage(data); - - BOOST_CHECK_EQUAL(storage.rank_count(), 5U); + BOOST_CHECK_EQUAL(storage.rank_count(), 6U); BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(storage), 2U); + BOOST_CHECK_EQUAL(detail::cross_rank_endpoint_count(storage), 10U); + // The ceiling this instrument exists to expose: an occupied slot holds at least one endpoint. + BOOST_CHECK_LE(detail::cross_rank_occupied_slots(storage), detail::cross_rank_endpoint_count(storage)); } 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)); + // Same traffic, four times the world. The record array grows; the payload does not. + const auto narrow = detail::build_packed_cross_rank_storage(slot_partners({5, 0, 0, 0})); + const auto wide = detail::build_packed_cross_rank_storage(slot_partners({5, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0})); + + BOOST_CHECK_EQUAL(detail::cross_rank_endpoint_count(narrow), detail::cross_rank_endpoint_count(wide)); + BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(narrow), detail::cross_rank_occupied_slots(wide)); + BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(narrow), 4U * sizeof(CrossRankPartnerRange)); + BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(wide), 16U * sizeof(CrossRankPartnerRange)); + // And the slot records are a slice of cross_rank_bytes, not an addition to it. + BOOST_CHECK_LT(detail::cross_rank_slot_record_bytes(wide), detail::cross_rank_storage_bytes(wide)); } BOOST_AUTO_TEST_CASE(graph_encoding_a_layer_retains_no_exchange_layout) { diff --git a/src/monoprop/bindings/binder.h b/src/monoprop/bindings/binder.h index 0c9b606a..9118fadd 100644 --- a/src/monoprop/bindings/binder.h +++ b/src/monoprop/bindings/binder.h @@ -272,11 +272,11 @@ auto bind_monomial_propagator(nb::module_ &mod) -> void { {"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. + // 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). So these grow with a + // P that the MPI rank count never reveals. d_slot_records / d_layer_cores recovers P; + // d_occupied_slots / d_slot_records is the occupancy that says whether a sparse layout would pay; + // and d_cross_rank_endpoints is the P-independent ceiling on d_occupied_slots. 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}, @@ -291,7 +291,8 @@ auto bind_monomial_propagator(nb::module_ &mod) -> void { {"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}}; + {"d_occupied_slots", b.occupied_slots}, + {"d_cross_rank_endpoints", b.cross_rank_endpoints}}; }); } } // namespace monoprop::bindings::detail From bbfa9185fe445dbf5f63cbab04e7d26d1b368e44 Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sun, 16 Aug 2026 12:34:00 +0100 Subject: [PATCH 5/6] perf(graph): :zap: store only the world slots that carry traffic The graph's last array indexed by the flat world size P. Each layer held one record per POSSIBLE partner, so with P participants each holding a P-length array the job carried L x P-squared records whether or not anything was ever sent between them. At L=5,420 and P=512 that is 22.7 GB of slot records against 3.7 GB of actual traffic -- 6.1 bytes of addressing per byte of data. Store the occupied slots instead, ascending by slot id. That is bounded by something with no P in it: an occupied slot holds at least one endpoint, so occupied_slots <= total cross-rank endpoints and the endpoint count is a property of the operator and the circuit, measured flat in P to 0.096% across a 4x change in it. The quadratic is not merely smaller, it is capped by the traffic it describes. The record is 12 B, and two things are absent from it by design: * the D range, already dropped -- B and D are one endpoint set in two orders; * the B/D offset, which is the running prefix over stored entries in ascending order. Empty slots contributed zero to the dense prefix, so the derived value equals the stored one exactly. A size_t offset would have padded the record to 24 B, so deriving it is worth 2x on its own. Access changes shape rather than getting slower. Every partner sweep in production was already `for r in 0..P { if empty continue }` -- walking the whole world to find the part of it with anything in it -- and becomes for_each_occupied_slot, which carries the derived offset and never visits an empty slot. The self slot keeps O(1) through a position resolved once at build: it is read per rotation pair in the innermost gradient loop and cannot afford a search. Converted: the four packing loops and the snapshot pass in Evolution.cpp, both totals in MPGraphLayers.h, endpoint marking in PareGraph.cpp, and the layer export in MonomialPropagator.inl (still dense in its output, since callers index it by rank, but now scattered into rather than interrogated for). graph_encoding_slot_record_bytes_track_the_world_not_the_traffic asserted precisely the property being removed, so it is inverted rather than repaired: quadrupling the world must now leave the record array byte-identical. 213/213 serial. --- cpp/monoprop/Evolution.cpp | 185 +++++++----------- cpp/monoprop/detail/graph/MPGraphLayers.h | 30 ++- .../detail/graph_encoding/MPGraphEncoding.cpp | 71 +++++-- .../graph_encoding/MPGraphEncodingStorage.h | 59 +++++- .../graph_encoding/MPGraphEncodingTypes.h | 74 +++++-- .../MonomialPropagator.inl | 43 ++-- cpp/monoprop/detail/pare/PareGraph.cpp | 27 +-- cpp/tests/graph_encoding_tests.cpp | 75 ++++++- cpp/tests/large_cosine_storage_tests.cpp | 8 +- 9 files changed, 355 insertions(+), 217 deletions(-) diff --git a/cpp/monoprop/Evolution.cpp b/cpp/monoprop/Evolution.cpp index 6e644657..4dede1eb 100644 --- a/cpp/monoprop/Evolution.cpp +++ b/cpp/monoprop/Evolution.cpp @@ -216,23 +216,18 @@ void pack_cross_rank_derivative_payload_impl(const DerivativeSnapshotScratch &sn int my_rank, const LayerExchangeLayout &layout, VecD &send_buffer) { - const size_t num_ranks = layer.cross_rank_rank_count(); - for (size_t rank = 0; rank < num_ranks; ++rank) { + layer.for_each_occupied_slot([&](size_t rank, const detail::CrossRankSlotView &slot) { if (static_cast(rank) == my_rank) { - continue; - } - const size_t end = layer.cross_rank_sin_send_size(rank); - if (end == 0) { - continue; + return; } const auto base = static_cast(layout.displs[rank]); const auto &bs = snap.sin_send_state[rank]; const auto &bh = snap.sin_send_op[rank]; - layer.for_each_cross_rank_sin_send_range(rank, 0, end, [&send_buffer, &base, &bs, &bh](size_t k, size_t /*i*/) { + for (size_t k = 0; k < slot.sin_send_count; ++k) { send_buffer[base + 2 * k] = bs[k]; send_buffer[base + 2 * k + 1] = bh[k]; - }); - } + } + }); } // Remote endpoint pass: own pre-cos values come from the sin_recv snapshots, partner values from the @@ -243,37 +238,29 @@ auto apply_cross_rank_derivative_exchange_impl(VecD &state, const DerivativeSnapshotScratch &snap, const TrigValues &trig, const ExchangePayload &payload) -> EndpointContrib { - const size_t num_ranks = layer.cross_rank_rank_count(); EndpointContrib local{}; - for (size_t rank = 0; rank < num_ranks; ++rank) { + layer.for_each_occupied_slot([&](size_t rank, const detail::CrossRankSlotView &slot) { if (static_cast(rank) == payload.my_rank) { - continue; - } - const size_t end = layer.cross_rank_sin_recv_size(rank); - if (end == 0) { - continue; + return; } const auto *rv = payload.recv_buffer.data() + payload.recv_displs[rank]; const auto &ds = snap.sin_recv_state[rank]; const auto &dh = snap.sin_recv_op[rank]; - layer.for_each_cross_rank_sin_recv_range( - rank, - 0, - end, - [&trig, &ds, &dh, &rv, &local, &op, &state](size_t k, size_t i, int phi_signed) { - const auto phi = static_cast(phi_signed); - // Inverse-rotation write-back (−sin): un-evolves state/op for the next reverse layer. - const double ps = -trig.sin_val * phi; - const double s_old = ds[k]; - const double h_old = dh[k]; - const double s_p = rv[2 * k]; - const double h_p = rv[2 * k + 1]; - local.cos_terms += s_old * h_old; - local.sin_terms += phi * s_old * h_p; - op[i] = (h_old * trig.cos_val) + (ps * h_p); - state[i] = (s_old * trig.cos_val) + (ps * s_p); - }); - } + for (size_t k = 0; k < slot.sin_send_count; ++k) { + const size_t i = detail::slot_sin_recv_index(slot, k); + const auto phi = static_cast(detail::slot_sin_recv_phase(slot, k)); + // Inverse-rotation write-back (−sin): un-evolves state/op for the next reverse layer. + const double ps = -trig.sin_val * phi; + const double s_old = ds[k]; + const double h_old = dh[k]; + const double s_p = rv[2 * k]; + const double h_p = rv[2 * k + 1]; + local.cos_terms += s_old * h_old; + local.sin_terms += phi * s_old * h_p; + op[i] = (h_old * trig.cos_val) + (ps * h_p); + state[i] = (s_old * trig.cos_val) + (ps * s_p); + } + }); return local; } @@ -312,20 +299,15 @@ void pack_cross_rank_evolution_payload_impl(VecD &op, int my_rank, const LayerExchangeLayout &layout, VecD &send_buffer) { - const size_t num_ranks = layer.cross_rank_rank_count(); - for (size_t rank = 0; rank < num_ranks; ++rank) { + layer.for_each_occupied_slot([&](size_t rank, const detail::CrossRankSlotView &slot) { if (static_cast(rank) == my_rank) { - continue; - } - const size_t end = layer.cross_rank_sin_send_size(rank); - if (end == 0) { - continue; + return; } const auto base = static_cast(layout.displs[rank]); - layer.for_each_cross_rank_sin_send_range(rank, 0, end, [&send_buffer, &base, &op](size_t k, size_t i) { - send_buffer[base + k] = op[i]; - }); - } + for (size_t k = 0; k < slot.sin_send_count; ++k) { + send_buffer[base + k] = op[detail::slot_sin_send_index(slot, k)]; + } + }); } void apply_cross_rank_evolution_exchange_impl(VecD &op, @@ -333,23 +315,16 @@ void apply_cross_rank_evolution_exchange_impl(VecD &op, double sin_val, const ExchangePayload &payload) { // op[i] is already cos-scaled, so only the sine term is added; rv[k] is the partner's pre-cos value. - const size_t num_ranks = layer.cross_rank_rank_count(); - for (size_t rank = 0; rank < num_ranks; ++rank) { + layer.for_each_occupied_slot([&](size_t rank, const detail::CrossRankSlotView &slot) { if (static_cast(rank) == payload.my_rank) { - continue; - } - const size_t end = layer.cross_rank_sin_recv_size(rank); - if (end == 0) { - continue; + return; } const auto *rv = payload.recv_buffer.data() + payload.recv_displs[rank]; - layer.for_each_cross_rank_sin_recv_range(rank, - 0, - end, - [&op, &sin_val, &rv](size_t k, size_t i, int phi_signed) { - op[i] += sin_val * static_cast(phi_signed) * rv[k]; - }); - } + for (size_t k = 0; k < slot.sin_send_count; ++k) { + const size_t i = detail::slot_sin_recv_index(slot, k); + op[i] += sin_val * static_cast(detail::slot_sin_recv_phase(slot, k)) * rv[k]; + } + }); } inline auto begin_cross_rank_evolution_exchange(VecD &op, const LayerTraversal &layer, const mpi::Comm &comm) @@ -383,15 +358,13 @@ auto apply_self_slot_derivative_paired(VecD &state, const LayerTraversal &layer, size_t my_rank, const TrigValues &trig) -> EndpointContrib { - const size_t self_d_count = layer.cross_rank_sin_recv_size(my_rank); + // O(1): the self slot's position is recorded at build precisely so this loop does not search for it. + const auto slot = layer.cross_rank_self_slot(); + const size_t self_d_count = slot.sin_send_count; if (self_d_count == 0) { 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 = detail::slot_sin_recv_index(slot, k); @@ -428,40 +401,37 @@ void snapshot_remote_endpoints(const VecD &state, snap.sin_send_op.resize(R); snap.sin_recv_state.resize(R); snap.sin_recv_op.resize(R); + // The scratch is indexed by world slot, so it is still sized R -- but only the occupied slots have + // anything to put in it. Clear first, then fill those: an empty slot's snapshot is an empty vector, + // which is what the dense sweep produced for it anyway. for (size_t r = 0; r < R; ++r) { + snap.sin_send_state[r].clear(); + snap.sin_send_op[r].clear(); + snap.sin_recv_state[r].clear(); + snap.sin_recv_op[r].clear(); + } + layer.for_each_occupied_slot([&](size_t r, const detail::CrossRankSlotView &slot) { if (r == my_rank) { - snap.sin_send_state[r].clear(); - snap.sin_send_op[r].clear(); - snap.sin_recv_state[r].clear(); - snap.sin_recv_op[r].clear(); - continue; - } - const size_t bc = layer.cross_rank_sin_send_size(r); - snap.sin_send_state[r].resize(bc); - snap.sin_send_op[r].resize(bc); - if (bc > 0) { - auto &bs = snap.sin_send_state[r]; - auto &bh = snap.sin_send_op[r]; - layer.for_each_cross_rank_sin_send_range(r, 0, bc, [&bs, &bh, &state, &op](size_t k, size_t i) { - bs[k] = state[i]; - bh[k] = op[i]; - }); + return; // the self slot recovers live; it needs no snapshot } - const size_t dc = layer.cross_rank_sin_recv_size(r); - snap.sin_recv_state[r].resize(dc); - snap.sin_recv_op[r].resize(dc); - if (dc > 0) { - auto &ds = snap.sin_recv_state[r]; - auto &dh = snap.sin_recv_op[r]; - layer.for_each_cross_rank_sin_recv_range(r, - 0, - dc, - [&ds, &dh, &state, &op](size_t k, size_t i, int /*phi*/) { - ds[k] = state[i]; - dh[k] = op[i]; - }); + const size_t count = slot.sin_send_count; + auto &bs = snap.sin_send_state[r]; + auto &bh = snap.sin_send_op[r]; + bs.resize(count); + bh.resize(count); + auto &ds = snap.sin_recv_state[r]; + auto &dh = snap.sin_recv_op[r]; + ds.resize(count); + dh.resize(count); + for (size_t k = 0; k < count; ++k) { + const size_t bi = detail::slot_sin_send_index(slot, k); + bs[k] = state[bi]; + bh[k] = op[bi]; + const size_t di = detail::slot_sin_recv_index(slot, k); + ds[k] = state[di]; + dh[k] = op[di]; } - } + }); } } // namespace @@ -513,14 +483,12 @@ auto evolve_step_traversal_impl(VecD &op, // Snapshot my_rank's own sin_send values before the cos pass; runs unconditionally (the remote pack // skips my_rank) so single-rank works. - const size_t self_b_count = (my_rank < layer.cross_rank_rank_count()) ? layer.cross_rank_sin_send_size(my_rank) : 0; + const auto self_slot = layer.cross_rank_self_slot(); + const size_t self_b_count = self_slot.sin_send_count; VecD self_b_snapshot; self_b_snapshot.resize(self_b_count); - if (self_b_count > 0) { - auto &snap = self_b_snapshot; - layer.for_each_cross_rank_sin_send_range(my_rank, 0, self_b_count, [&snap, &op](size_t k, size_t i) { - snap[k] = op[i]; - }); + for (size_t k = 0; k < self_b_count; ++k) { + self_b_snapshot[k] = op[detail::slot_sin_send_index(self_slot, k)]; } // Pack + start the exchange before the cos scan so partner values are pre-cos and the transfer overlaps. @@ -528,16 +496,11 @@ auto evolve_step_traversal_impl(VecD &op, cos_scale(layer_idx, op_data, cos_val); finish_cross_rank_evolution_exchange(op, layer, sin_val, in_flight); - // Self-slot sin_recv entries: op[i] is already cos-scaled, so only the sine term is added. - if (self_b_count > 0) { - const size_t self_d_count = layer.cross_rank_sin_recv_size(my_rank); - layer.for_each_cross_rank_sin_recv_range(my_rank, - 0, - self_d_count, - [&op, &sin_val, &self_b_snapshot](size_t k, size_t i, int phi_signed) { - op[i] += - sin_val * static_cast(phi_signed) * self_b_snapshot[k]; - }); + // Self-slot sin_recv entries: op[i] is already cos-scaled, so only the sine term is added. B and D + // have the same count, so self_b_count serves both. + for (size_t k = 0; k < self_b_count; ++k) { + const size_t i = detail::slot_sin_recv_index(self_slot, k); + op[i] += sin_val * static_cast(detail::slot_sin_recv_phase(self_slot, k)) * self_b_snapshot[k]; } } diff --git a/cpp/monoprop/detail/graph/MPGraphLayers.h b/cpp/monoprop/detail/graph/MPGraphLayers.h index eccea684..901f5d2e 100644 --- a/cpp/monoprop/detail/graph/MPGraphLayers.h +++ b/cpp/monoprop/detail/graph/MPGraphLayers.h @@ -64,8 +64,22 @@ 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. + // O(1); the self slot is read per rotation pair in the innermost gradient loop. + auto cross_rank_self_slot() const -> detail::CrossRankSlotView { + return detail::cross_rank_self_slot(core_->cross_rank); + } + + // Every slot carrying traffic, ascending, each with its offset. func(slot_id, view). + // + // This is what a partner sweep should use. The old shape -- loop 0..P, ask each slot its size, + // `continue` on zero -- walked the whole world to find the part of it that had anything in it. + template + auto for_each_occupied_slot(Func &&func) const -> void { + detail::for_each_occupied_slot(core_->cross_rank, std::forward(func)); + } + + // 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); @@ -97,12 +111,11 @@ struct LayerTraversal final { auto gate_index() const -> size_t { return core_->gate_index; } // Rotations (Givens cycles) = sum of per-rank in-counts (one in-entry per rotation). sin_recv_size - // would double-count self-rank rotations (in+out). + // would double-count self-rank rotations (in+out). Empty slots contribute nothing, so summing over + // the occupied ones is the same total the full sweep gave. auto total_cycles() const -> size_t { size_t count = 0; - for (size_t rank = 0; rank < cross_rank_rank_count(); ++rank) { - count += cross_rank_in_count(rank); - } + for_each_occupied_slot([&count](size_t, const detail::CrossRankSlotView &slot) { count += slot.in_count; }); return count; } @@ -110,9 +123,8 @@ struct LayerTraversal final { // indices = num_cos_inds() - total_rotation_endpoints(). auto total_rotation_endpoints() const -> size_t { size_t count = 0; - for (size_t rank = 0; rank < cross_rank_rank_count(); ++rank) { - count += cross_rank_sin_recv_size(rank); - } + for_each_occupied_slot( + [&count](size_t, const detail::CrossRankSlotView &slot) { count += slot.sin_send_count; }); return count; } diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index 9becff8a..a1240f7c 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp @@ -34,6 +34,18 @@ auto checked_mpi_int(size_t value, const char *what) -> int { return static_cast(value); } +// The slot id is stored as uint32 to keep the occupied record at 12 B. That bounds the flat world at +// ~4.3e9 participants, which is not a limit anyone will meet, but it is a narrowing conversion and so +// it is checked rather than cast. +auto checked_world_slot(size_t rank) -> uint32_t { + if (rank > static_cast(std::numeric_limits::max())) { + throw std::overflow_error(std::format("World slot {} exceeds the {} the occupied-slot record can hold.", + rank, + std::numeric_limits::max())); + } + return static_cast(rank); +} + auto build_layer_exchange_layout(const std::vector &send_counts, int scale, const char *what) -> LayerExchangeLayout { const std::string count_label = std::format("{} count", what); @@ -93,7 +105,7 @@ auto packed_phase_storage_bytes(const PackedPhaseStorage &storage) -> size_t { auto build_packed_cross_rank_storage(const std::vector &data) -> PackedCrossRankStorage { PackedCrossRankStorage storage; const size_t num_ranks = data.size(); - storage.ranges.resize(num_ranks); + storage.world_size = num_ranks; size_t total_b = 0; for (size_t rank = 0; rank < num_ranks; ++rank) { @@ -110,12 +122,18 @@ auto build_packed_cross_rank_storage(const std::vector &da 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.in_count = static_cast(partner.in_count); + // The whole point: a slot with no traffic gets no record. Ascending rank order makes `occupied` + // sorted by construction, which is what lets readers binary-search it and lets the offset be a + // running prefix rather than a stored field. + if (partner.sin_send_indices.empty()) { + continue; + } + storage.occupied.push_back({.slot = checked_world_slot(rank), + .sin_send_count = static_cast(partner.sin_send_indices.size()), + .in_count = static_cast(partner.in_count)}); total_b += partner.sin_send_indices.size(); } + storage.occupied.shrink_to_fit(); // push_back overshoots, and this array is the thing being shrunk const size_t total_d = total_b; bool uses_binary_phases = true; @@ -130,12 +148,13 @@ auto build_packed_cross_rank_storage(const std::vector &da storage.sin_send_indices.resize(total_b); storage.sin_recv_phases = make_packed_phase_storage(total_d, uses_binary_phases); - 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 = b_off; + // Fill in the same ascending order the offsets were accumulated in, so the running prefix here is + // the one for_each_occupied_slot will reconstruct on every later read. + size_t offset = 0; + for (const auto &entry : storage.occupied) { + const auto &partner = data[entry.slot]; + const size_t b_off = offset; + const size_t d_off = b_off; // equal counts, so equal prefix sums 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"); @@ -146,11 +165,27 @@ auto build_packed_cross_rank_storage(const std::vector &da (void)i; store_packed_phase(storage.sin_recv_phases, d_off + k, phi, "Cross-rank D phase"); } + offset += entry.sin_send_count; } return storage; } +auto resolve_self_slot(PackedCrossRankStorage &storage, size_t my_rank) -> void { + storage.self_pos = kNoSelfSlot; + storage.self_offset = 0; + size_t offset = 0; + for (size_t pos = 0; pos < storage.occupied.size(); ++pos) { + const auto &entry = storage.occupied[pos]; + if (entry.slot == my_rank) { + storage.self_pos = pos; + storage.self_offset = offset; + return; + } + offset += entry.sin_send_count; + } +} + auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t { size_t bytes = cross_rank_slot_record_bytes(storage) + packed_phase_storage_bytes(storage.sin_recv_phases); bytes += storage.sin_send_indices.capacity() * sizeof(TermIndex); @@ -158,21 +193,18 @@ auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t { } auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size_t { - return storage.ranges.capacity() * sizeof(CrossRankPartnerRange); + return storage.occupied.capacity() * sizeof(CrossRankOccupiedSlot); } auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t { - // sin_send_count alone is the predicate. B and D hold the same endpoint set in two orders (see - // cross_rank_sin_recv_index), so a slot cannot carry D entries while carrying no B entries, and - // counting either gives the same answer. - return static_cast(std::ranges::count_if( - storage.ranges, [](const CrossRankPartnerRange &range) { return range.sin_send_count != 0; })); + // No predicate and no scan any more: an entry exists only if the slot carries traffic. + return storage.occupied.size(); } auto cross_rank_endpoint_count(const PackedCrossRankStorage &storage) -> size_t { size_t count = 0; - for (const auto &range : storage.ranges) { - count += range.sin_send_count; + for (const auto &entry : storage.occupied) { + count += entry.sin_send_count; } return count; } @@ -211,6 +243,7 @@ auto build_layer_storage_unified(std::vector all_partners, const size_t num_ranks = all_partners.size(); storage->cross_rank = build_packed_cross_rank_storage(std::move(all_partners)); + resolve_self_slot(storage->cross_rank, my_rank); // 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. diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h index a89fcbf9..6835a04d 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -77,6 +77,10 @@ inline auto store_packed_phase(PackedPhaseStorage &storage, size_t idx, int phas auto build_packed_cross_rank_storage(const std::vector &data) -> PackedCrossRankStorage; +// Record where my_rank's own slot sits, so the gradient's self-slot reads are O(1). Called once per +// layer at build; separate from the builder because only build_layer_storage_unified knows my_rank. +auto resolve_self_slot(PackedCrossRankStorage &storage, size_t my_rank) -> void; + // 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 @@ -94,13 +98,56 @@ struct CrossRankSlotView final { 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, +namespace slot_detail { +inline auto view_at(const PackedCrossRankStorage &storage, const CrossRankOccupiedSlot &entry, size_t offset) + -> CrossRankSlotView { + return CrossRankSlotView{.sin_send_indices = storage.sin_send_indices.data() + 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}; + .phase_offset = offset, + .sin_send_count = entry.sin_send_count, + .in_count = entry.in_count}; +} +} // namespace slot_detail + +// Every occupied slot in ascending order, each with its B/D offset -- the running prefix, which is +// exactly what the dense layout stored per slot. func(slot_id, view). +// +// This is the shape production code should use. It is O(occupied) for the whole sweep rather than +// O(P), and it never visits a slot with nothing in it: under a dense layout those were visited and +// skipped, so on a large world most of the loop was the skip. +template +auto for_each_occupied_slot(const PackedCrossRankStorage &storage, Func &&func) -> void { + size_t offset = 0; + for (const auto &entry : storage.occupied) { + func(static_cast(entry.slot), slot_detail::view_at(storage, entry, offset)); + offset += entry.sin_send_count; + } +} + +// O(1). The self slot is read per rotation pair in the innermost gradient loop, so it cannot pay the +// search or the prefix walk that an arbitrary slot does. An all-zero view when it carries no traffic. +inline auto cross_rank_self_slot(const PackedCrossRankStorage &storage) -> CrossRankSlotView { + if (storage.self_pos == kNoSelfSlot) { + return CrossRankSlotView{.sin_recv_phases = &storage.sin_recv_phases}; + } + return slot_detail::view_at(storage, storage.occupied[storage.self_pos], storage.self_offset); +} + +// Arbitrary slot, and O(occupied): the offset is a prefix over preceding entries, so resolving one +// slot in isolation walks them. Diagnostic and test use -- a production loop wants +// for_each_occupied_slot, and the self slot wants cross_rank_self_slot. +inline auto cross_rank_slot(const PackedCrossRankStorage &storage, size_t rank) -> CrossRankSlotView { + size_t offset = 0; + for (const auto &entry : storage.occupied) { + if (entry.slot == rank) { + return slot_detail::view_at(storage, entry, offset); + } + if (entry.slot > rank) { + break; + } + offset += entry.sin_send_count; + } + return CrossRankSlotView{.sin_recv_phases = &storage.sin_recv_phases}; } inline auto slot_sin_send_index(const CrossRankSlotView &slot, size_t idx) -> size_t { diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h index b3b3c3f7..8f966a9e 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include @@ -128,34 +129,69 @@ 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 AND sin_recv_phases; cumulative, so may exceed 2^32 - // == the D count; TermIndex-wide so one rank/layer can exceed 2^32. +// One world slot that carries traffic. +// +// Slots with none are not stored at all. The dense array this replaces reserved a record for every +// POSSIBLE partner, so its size was the flat world P and, summed over the P participants that each +// hold one, the graph carried P-squared records regardless of how much was ever sent. What is stored +// instead is bounded by the traffic, which does not depend on P. +// +// The D range is not stored: it IS the B range. B and D hold the same endpoint set in two orders, so +// their counts are equal by construction (Engine.h resizes both to P+Q) and the two prefix sums are +// therefore identical. The equality is enforced in build_packed_cross_rank_storage, not assumed. +// +// The B/D offset is not a field either. It is the running prefix over stored entries in ascending slot +// order, and empty slots contributed zero to the dense prefix, so the derived value equals the stored +// one exactly. Deriving it is what keeps this record at 12 B: a size_t offset would pad it to 24. +struct CrossRankOccupiedSlot final { + uint32_t slot = 0; // flat world slot id -- what the dense array encoded by position + // TermIndex-wide so one slot in one layer can exceed 2^32 endpoints. TermIndex sin_send_count = 0; TermIndex in_count = 0; }; +static_assert(sizeof(TermIndex) != sizeof(uint32_t) || sizeof(CrossRankOccupiedSlot) == 12, + "narrow build: the occupied-slot record is the graph's P coefficient; keep it padding-free."); -// 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."); +// Sentinel for self_pos: this rank's own slot carries no traffic in this layer. +inline constexpr size_t kNoSelfSlot = static_cast(-1); struct PackedCrossRankStorage final { - std::vector ranges; // size == the flat world P, not the MPI rank count + // Ascending by slot, unique, every entry carrying traffic. + std::vector occupied; 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; } - // 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; } + // P. No longer recoverable from the array length -- that is the point -- but the exchange still + // needs it, since MPI_Alltoallv wants dense counts and displacements. + size_t world_size = 0; + + // The self slot is read in the innermost gradient loop, so it gets O(1) access instead of the + // search the general case pays. Resolved once at build; see resolve_self_slot. + size_t self_pos = kNoSelfSlot; + size_t self_offset = 0; + + auto rank_count() const -> size_t { return world_size; } + + // O(log occupied). Callers walking every partner should use for_each_occupied_slot instead, which + // is O(occupied) for the whole sweep and carries the offset with it. + auto find(size_t rank) const -> const CrossRankOccupiedSlot * { + const auto it = std::lower_bound(occupied.begin(), + occupied.end(), + rank, + [](const CrossRankOccupiedSlot &e, size_t r) { return e.slot < r; }); + return (it == occupied.end() || it->slot != rank) ? nullptr : &*it; + } + + auto sin_send_size(size_t rank) const -> size_t { + const auto *e = find(rank); + return e == nullptr ? 0 : e->sin_send_count; + } + // Same count as the send side: B and D are the same endpoint set, permuted. + auto sin_recv_size(size_t rank) const -> size_t { return sin_send_size(rank); } + auto in_count(size_t rank) const -> size_t { + const auto *e = find(rank); + return e == nullptr ? 0 : e->in_count; + } }; struct LayerCore final { diff --git a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl index 0e0bc487..ac46b5aa 100644 --- a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl +++ b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl @@ -389,31 +389,24 @@ auto MonomialPropagator::graph_data() const -> std::vector // Always empty: local cycles are folded into cross_rank[my_rank]. std::vector local_cyc_data; - std::vector b_data, d_data; - b_data.reserve(rank_count); - d_data.reserve(rank_count); - for (size_t rank = 0; rank < rank_count; ++rank) { - VecZ sin_send_indices(traversal.cross_rank_sin_send_size(rank)); - VecI b_phases(traversal.cross_rank_sin_send_size(rank), 0); - VecZ d_indices(traversal.cross_rank_sin_recv_size(rank)); - VecI sin_recv_phases(traversal.cross_rank_sin_recv_size(rank)); - - traversal.for_each_cross_rank_sin_send_range( - rank, - 0, - traversal.cross_rank_sin_send_size(rank), - [&](size_t logical_idx, size_t value_idx) { sin_send_indices[logical_idx] = value_idx; }); - traversal.for_each_cross_rank_sin_recv_range(rank, - 0, - traversal.cross_rank_sin_recv_size(rank), - [&](size_t logical_idx, size_t value_idx, int phase) { - d_indices[logical_idx] = value_idx; - sin_recv_phases[logical_idx] = phase; - }); - - b_data.emplace_back(std::move(sin_send_indices), std::move(b_phases)); - d_data.emplace_back(std::move(d_indices), std::move(sin_recv_phases)); - } + // The exported shape stays dense in the world size -- callers index it by rank -- but only the + // occupied slots have anything to put in it, so it is filled by scattering into a sized-empty + // array rather than by asking every possible slot how much it holds. + std::vector b_data(rank_count), d_data(rank_count); + traversal.for_each_occupied_slot([&](size_t rank, const detail::CrossRankSlotView &slot) { + const size_t count = slot.sin_send_count; + VecZ sin_send_indices(count); + VecI b_phases(count, 0); + VecZ d_indices(count); + VecI sin_recv_phases(count); + for (size_t k = 0; k < count; ++k) { + sin_send_indices[k] = detail::slot_sin_send_index(slot, k); + d_indices[k] = detail::slot_sin_recv_index(slot, k); + sin_recv_phases[k] = detail::slot_sin_recv_phase(slot, k); + } + b_data[rank] = CrossRankData{std::move(sin_send_indices), std::move(b_phases)}; + d_data[rank] = CrossRankData{std::move(d_indices), std::move(sin_recv_phases)}; + }); // Same two-way read as cos_index_count_(): a pared layer's stored set is authoritative, and // recomputing the fold over it would report the indices the pare removed. VecZ cos_inds; diff --git a/cpp/monoprop/detail/pare/PareGraph.cpp b/cpp/monoprop/detail/pare/PareGraph.cpp index 4d6dccec..63a2be93 100644 --- a/cpp/monoprop/detail/pare/PareGraph.cpp +++ b/cpp/monoprop/detail/pare/PareGraph.cpp @@ -26,15 +26,6 @@ namespace monoprop { namespace { -template -auto for_each_remote_rank(const LayerTraversal &layer, size_t my_rank, Func &&func) -> void { - for (size_t rank = 0; rank < layer.cross_rank_rank_count(); ++rank) { - if (rank != my_rank) { - func(rank); - } - } -} - // Bits of `present` whose absolute index is in the keep set. inline auto keep_mask_for_block(const std::vector &keep, size_t base, uint64_t present) -> uint64_t { uint64_t mask = 0; @@ -84,17 +75,13 @@ auto mark_cross_rank_endpoints_kept(const LayerTraversal &layer, size_t my_rank, nodes_to_keep[idx] = 1; } }; - for (size_t rank = 0; rank < layer.cross_rank_rank_count(); ++rank) { - layer.for_each_cross_rank_sin_recv_range(rank, - 0, - layer.cross_rank_sin_recv_size(rank), - [&mark](size_t, size_t tgt_idx, int) { mark(tgt_idx); }); - } - for_each_remote_rank(layer, my_rank, [&layer, &mark](size_t rank) { - layer.for_each_cross_rank_sin_send_range(rank, - 0, - layer.cross_rank_sin_send_size(rank), - [&mark](size_t, size_t src_idx) { mark(src_idx); }); + layer.for_each_occupied_slot([&](size_t rank, const detail::CrossRankSlotView &slot) { + for (size_t k = 0; k < slot.sin_send_count; ++k) { + mark(detail::slot_sin_recv_index(slot, k)); + if (rank != my_rank) { + mark(detail::slot_sin_send_index(slot, k)); + } + } }); } diff --git a/cpp/tests/graph_encoding_tests.cpp b/cpp/tests/graph_encoding_tests.cpp index 543c5a46..db3aa967 100644 --- a/cpp/tests/graph_encoding_tests.cpp +++ b/cpp/tests/graph_encoding_tests.cpp @@ -256,16 +256,20 @@ BOOST_AUTO_TEST_CASE(graph_encoding_occupied_slots_counts_only_slots_carrying_tr BOOST_CHECK_LE(detail::cross_rank_occupied_slots(storage), detail::cross_rank_endpoint_count(storage)); } -BOOST_AUTO_TEST_CASE(graph_encoding_slot_record_bytes_track_the_world_not_the_traffic) { - // Same traffic, four times the world. The record array grows; the payload does not. +BOOST_AUTO_TEST_CASE(graph_encoding_slot_record_bytes_track_the_traffic_not_the_world) { + // Same traffic, four times the world. This is the whole claim of the sparse layout: the record + // array is a function of what is sent, not of how many participants could have been sent to. const auto narrow = detail::build_packed_cross_rank_storage(slot_partners({5, 0, 0, 0})); const auto wide = detail::build_packed_cross_rank_storage(slot_partners({5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0})); + BOOST_CHECK_EQUAL(narrow.rank_count(), 4U); + BOOST_CHECK_EQUAL(wide.rank_count(), 16U); BOOST_CHECK_EQUAL(detail::cross_rank_endpoint_count(narrow), detail::cross_rank_endpoint_count(wide)); BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(narrow), detail::cross_rank_occupied_slots(wide)); - BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(narrow), 4U * sizeof(CrossRankPartnerRange)); - BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(wide), 16U * sizeof(CrossRankPartnerRange)); + // The world quadrupled and the record array did not move at all. + BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(narrow), detail::cross_rank_slot_record_bytes(wide)); + BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(wide), 1U * sizeof(CrossRankOccupiedSlot)); // And the slot records are a slice of cross_rank_bytes, not an addition to it. BOOST_CHECK_LT(detail::cross_rank_slot_record_bytes(wide), detail::cross_rank_storage_bytes(wide)); } @@ -301,3 +305,66 @@ BOOST_AUTO_TEST_CASE(graph_encoding_skewed_endpoint_counts_are_refused) { BOOST_CHECK_THROW(detail::build_packed_cross_rank_storage(data), std::logic_error); } + +BOOST_AUTO_TEST_CASE(graph_encoding_occupied_sweep_matches_the_dense_sweep_it_replaces) { + // Zeros at the front, interior and back, and a slot whose in_count splits the B list, so the + // derived offset has somewhere to go wrong. + auto data = slot_partners({0, 3, 0, 0, 7, 0}); + data[1].in_count = 1; + data[4].in_count = 4; + const auto storage = detail::build_packed_cross_rank_storage(data); + + // What the dense layout would have produced for the occupied slots. The offsets are the prefix over + // ALL slots -- which the empty ones contributed zero to, which is why dropping them is lossless. + struct Expected { + size_t slot, offset, count, in_count; + }; + const std::vector expected{{1, 0, 3, 1}, {4, 3, 7, 4}}; + + std::vector seen; + detail::for_each_occupied_slot(storage, [&](size_t slot, const detail::CrossRankSlotView &view) { + seen.push_back({slot, view.phase_offset, view.sin_send_count, view.in_count}); + }); + + BOOST_REQUIRE_EQUAL(seen.size(), expected.size()); + for (size_t k = 0; k < expected.size(); ++k) { + BOOST_CHECK_EQUAL(seen[k].slot, expected[k].slot); + BOOST_CHECK_EQUAL(seen[k].offset, expected[k].offset); + BOOST_CHECK_EQUAL(seen[k].count, expected[k].count); + BOOST_CHECK_EQUAL(seen[k].in_count, expected[k].in_count); + } + + // The general single-slot resolver must agree with the sweep, including on an absent slot. + for (const auto &e : expected) { + const auto view = detail::cross_rank_slot(storage, e.slot); + BOOST_CHECK_EQUAL(view.phase_offset, e.offset); + BOOST_CHECK_EQUAL(view.sin_send_count, e.count); + } + BOOST_CHECK_EQUAL(detail::cross_rank_slot(storage, 0).sin_send_count, 0U); + BOOST_CHECK_EQUAL(detail::cross_rank_slot(storage, 5).sin_send_count, 0U); + BOOST_CHECK_EQUAL(storage.sin_send_size(3), 0U); + BOOST_CHECK_EQUAL(storage.sin_send_size(4), 7U); +} + +BOOST_AUTO_TEST_CASE(graph_encoding_self_slot_is_resolved_without_a_search) { + auto storage = detail::build_packed_cross_rank_storage(slot_partners({0, 3, 0, 7, 0})); + + detail::resolve_self_slot(storage, 3); + BOOST_CHECK_EQUAL(storage.self_offset, 3U); // slot 1's three endpoints precede it + const auto self = detail::cross_rank_self_slot(storage); + BOOST_CHECK_EQUAL(self.sin_send_count, 7U); + BOOST_CHECK_EQUAL(self.phase_offset, 3U); + + // A rank whose own slot carries nothing must resolve to an empty view, not to a neighbour's. + detail::resolve_self_slot(storage, 2); + BOOST_CHECK_EQUAL(storage.self_pos, kNoSelfSlot); + BOOST_CHECK_EQUAL(detail::cross_rank_self_slot(storage).sin_send_count, 0U); +} + +BOOST_AUTO_TEST_CASE(graph_encoding_a_layer_with_no_cross_rank_traffic_stores_no_slots) { + const auto storage = detail::build_packed_cross_rank_storage(slot_partners({0, 0, 0, 0, 0, 0, 0, 0})); + + BOOST_CHECK_EQUAL(storage.rank_count(), 8U); // the world is still eight wide + BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(storage), 0U); + BOOST_CHECK_EQUAL(detail::cross_rank_slot_record_bytes(storage), 0U); +} diff --git a/cpp/tests/large_cosine_storage_tests.cpp b/cpp/tests/large_cosine_storage_tests.cpp index d9dd20b8..6ebe11aa 100644 --- a/cpp/tests/large_cosine_storage_tests.cpp +++ b/cpp/tests/large_cosine_storage_tests.cpp @@ -79,12 +79,12 @@ BOOST_AUTO_TEST_CASE(pruned_layer_supports_cos_counts_above_u32) { // The per-rank cross-rank counts index into one layer's term set, so under the wide build they must // be TermIndex-wide; uint32_t would silently cap a single partition/layer at ~2^32 terms. BOOST_AUTO_TEST_CASE(cross_rank_partner_range_counts_track_term_index_width) { - CrossRankPartnerRange r{}; + CrossRankOccupiedSlot r{}; BOOST_CHECK_EQUAL(sizeof(r.sin_send_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)); + // One record per OCCUPIED slot, so its width scales the traffic-bounded term, not a P-squared one. + // The B/D offset is derived rather than stored precisely to keep it this narrow. + BOOST_CHECK_EQUAL(sizeof(CrossRankOccupiedSlot), sizeof(uint32_t) + 2 * sizeof(TermIndex)); } #if defined(monoprop_WIDE_TERM_INDEX) From 9e22925e8f6b3c4560bb574d14743fa9b3a90ade Mon Sep 17 00:00:00 2001 From: Aaron Miller Date: Sun, 16 Aug 2026 14:28:40 +0100 Subject: [PATCH 6/6] perf(graph): :zap: scatter the derived exchange layout over the occupied slots Only needed once the two halves coexist, which is why neither branch carries it. #237 derives counts[r] by asking cross_rank.sin_send_size(r) for every r < P. That was O(1) against the dense range array it was written for. Under the sparse storage sin_send_size is a binary search over the occupied slots, so the same loop became O(P log occupied) -- per layer, per exchange -- to fill an array that is ~82% zeros at P=512 by construction, and whose zero fraction only grows with P. So fill it the other way round: zero the counts, walk the slots that actually carry traffic via for_each_occupied_slot, and scatter. O(P) + O(occupied) with no search at all. The displacement prefix stays dense because MPI_Alltoallv wants an entry per rank and an empty slot still needs a valid, repeated displacement. assign() rather than resize() for the counts: `out` is scratch reused across layers, and a slot carrying nothing this layer must read zero rather than inherit the last layer's count. graph_encoding_derived_layout_reuses_its_scratch pins exactly that. The self slot is skipped by slot id, not by the old r == my_rank test on the loop variable: under sparse storage this rank's own slot is simply one of the stored entries, and it may or may not be present at all. Equivalence is asserted elementwise against build_layer_exchange_layout, for every my_rank and both scales, by graph_encoding_derived_layout_matches_the_layout_it_replaces. --- .../detail/graph_encoding/MPGraphEncoding.cpp | 25 ++++++++++++++----- 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index a1240f7c..28be0cc4 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp @@ -223,16 +223,29 @@ auto derive_exchange_layout(const PackedCrossRankStorage &cross_rank, const std::string displacement_label = std::format("{} displacement", what); const size_t num_ranks = cross_rank.rank_count(); - out.counts.resize(num_ranks); + // assign() over resize(): every slot that carries nothing must read zero, and `out` is reused + // across layers, so last layer's counts would otherwise survive into this one's. + out.counts.assign(num_ranks, 0); out.displs.resize(num_ranks); + + // Scatter over the slots that carry traffic instead of interrogating every possible partner. + // sin_send_size(r) is a binary search once the storage is sparse, so the dense probe this + // replaces would cost O(P log occupied) per layer per exchange -- to fill an array that is + // ~82% zeros at P=512 by construction. Occupancy only falls as P grows, so the gap widens. + for_each_occupied_slot(cross_rank, [&](size_t slot, const CrossRankSlotView &view) { + if (slot == my_rank) { + return; // excluded from the transfer and handled locally, as the stored layout did + } + out.counts[slot] = + checked_mpi_int(static_cast(scale) * view.sin_send_count, count_label.c_str()); + }); + + // The prefix sum stays dense: MPI_Alltoallv wants a displacement for every rank, and an empty + // slot still needs a valid (repeated) one. 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; + total += static_cast(out.counts[r]); } out.total_count = total; }