diff --git a/cpp/monoprop/Evolution.cpp b/cpp/monoprop/Evolution.cpp index 47bf3234..6e644657 100644 --- a/cpp/monoprop/Evolution.cpp +++ b/cpp/monoprop/Evolution.cpp @@ -60,8 +60,14 @@ auto combine_endpoint_contrib(const EndpointContrib &a, const EndpointContrib &b struct FlatExchangeBuffers { VecD send_buffer; VecD recv_buffer; - std::vector recv_counts; - std::vector recv_displs; + // The layout for the exchange currently being posted, derived per layer rather than read from + // one retained per layer. Reused, so it allocates once per thread per world size. Sharing one + // instance across layers is only sound because at most one exchange is in flight per thread -- + // the same invariant send_buffer above has always required. + // + // ONE layout, not two: it describes the recv side as well as the send side. See + // derive_layer_exchange. + LayerExchangeLayout layout; }; auto &acquire_flat_exchange_buffers() { @@ -73,21 +79,39 @@ auto &acquire_flat_exchange_buffers() { } void resize_flat_exchange_buffers(const LayerExchangeLayout &layout, FlatExchangeBuffers &buffers) { - // Recv size isn't known until counts are exchanged; keep it at 1 element so data() stays non-null. - const size_t send_alloc = layout.total_count == 0 ? 1 : layout.total_count; - buffers.send_buffer.resize(send_alloc); - buffers.recv_buffer.resize(1); - buffers.recv_counts.clear(); - buffers.recv_displs.clear(); + const size_t alloc = layout.total_count == 0 ? 1 : layout.total_count; + buffers.send_buffer.resize(alloc); } -auto active_evolution_exchange_layout(const LayerTraversal &layer, const mpi::Comm &comm) - -> const LayerExchangeLayout * { - if (mpi::size(comm) == 1) { - return nullptr; - } - // All ranks must participate even at local total_count 0, else MPI_Alltoallv deadlocks. - return &layer.evolution_exchange_layout(); +// Nothing to exchange at one rank. All ranks must participate even at local total_count 0, else +// MPI_Alltoallv deadlocks, so this is a property of the communicator and not of the layer. +auto layer_exchange_participates(const mpi::Comm &comm) -> bool { + return mpi::size(comm) != 1; +} + +// Derive this layer's exchange layout into `buffers.layout` at `scale`. It describes BOTH sides. +// +// The count matrix is symmetric: rank m's slot for r holds (the queries r sent m) ++ (the queries +// m sent r), and rank r's slot for m holds those two swapped, so the two slots have the same +// length (MPGraphEncoding's sink, via layer_build/Engine.h). Counts are equal, and displacements +// are prefix sums of counts, so the recv layout is the send layout -- there is nothing to +// transpose and nothing to communicate. This is what a per-layer RecvLayoutCache used to hold, +// at 8 B per world slot, and what an alltoall_counts per layer used to compute. +// +// Scaling is applied once, here, rather than to a scale-1 result: every rank multiplies by the +// same literal, so the equality survives it. +// +// Verified end to end rather than reasoned about alone: with MONOPROP_CHECK_EXCHANGE_SYMMETRY set +// the derived counts are checked against a real alltoall on every layer (see +// check_exchange_symmetry). A campaign at world 32 and 256 compared 550M slots with no mismatch. +auto derive_layer_exchange(const LayerTraversal &layer, + const mpi::Comm &comm, + int scale, + FlatExchangeBuffers &buffers) -> void { + const auto my_rank = static_cast(mpi::rank(comm)); + const char *what = scale == 1 ? "Layer exchange" : "Layer derivative exchange"; + detail::derive_exchange_layout(layer.cross_rank(), my_rank, scale, buffers.layout, what); + mpi::check_exchange_symmetry(buffers.layout.counts, comm); } // The completed alltoallv payload as an apply pass sees it: peer `rank`'s entries start at @@ -105,21 +129,20 @@ struct CrossRankExchangeHandle { [[no_unique_address]] mpi::Ticket ticket; }; -inline auto begin_flat_exchange(const LayerExchangeLayout &layout, FlatExchangeBuffers &buffers, const mpi::Comm &comm) - -> CrossRankExchangeHandle { +inline auto begin_flat_exchange(FlatExchangeBuffers &buffers, const mpi::Comm &comm) -> CrossRankExchangeHandle { + const LayerExchangeLayout &layout = buffers.layout; CrossRankExchangeHandle handle; handle.layout = &layout; handle.buffers = &buffers; - const auto &recv = mpi::resolve_recv(layout.counts, comm, layout.recv_cache); - buffers.recv_counts = recv.counts; - buffers.recv_displs = recv.displs; - buffers.recv_buffer.resize(recv.total == 0 ? 1 : static_cast(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; @@ -136,19 +159,20 @@ struct InFlightExchange { bool active = false; }; -// Callers run the participation guard first (see active_evolution_exchange_layout) so the layout is -// never materialized at a single rank. +// Callers run layer_exchange_participates first, so no layout is derived at a single rank -- where +// deriving one would allocate a P-int pair and resolve a transpose for a transfer that never posts. template -inline auto begin_layer_exchange(const LayerExchangeLayout &layout, const mpi::Comm &comm, Pack pack) +inline auto begin_layer_exchange(const LayerTraversal &layer, int scale, const mpi::Comm &comm, Pack pack) -> InFlightExchange { InFlightExchange in_flight; in_flight.my_rank = mpi::rank(comm); in_flight.active = true; auto &buffers = acquire_flat_exchange_buffers(); - resize_flat_exchange_buffers(layout, buffers); - pack(in_flight.my_rank, layout, buffers.send_buffer); - in_flight.handle = begin_flat_exchange(layout, buffers, comm); + derive_layer_exchange(layer, comm, scale, buffers); + resize_flat_exchange_buffers(buffers.layout, buffers); + pack(in_flight.my_rank, buffers.layout, buffers.send_buffer); + in_flight.handle = begin_flat_exchange(buffers, comm); return in_flight; } @@ -167,7 +191,7 @@ inline auto finish_layer_exchange(InFlightExchange &in_flight, Apply apply) } wait_flat_exchange(in_flight.handle); return apply(ExchangePayload{.recv_buffer = in_flight.handle.buffers->recv_buffer, - .recv_displs = in_flight.handle.buffers->recv_displs, + .recv_displs = in_flight.handle.buffers->layout.displs, .my_rank = in_flight.my_rank}); } @@ -258,11 +282,13 @@ inline auto begin_cross_rank_derivative_exchange(const DerivativeSnapshotScratch const LayerTraversal &layer, const mpi::Comm &comm) -> InFlightExchange { // Single-rank (or no peer participating): nothing to exchange — the self slot covers everything. - if (active_evolution_exchange_layout(layer, comm) == nullptr) { + if (!layer_exchange_participates(comm)) { return {}; } // Safe to fire before the cos pass: pack reads pre-cos snapshots and the transfer touches only buffers. - return begin_layer_exchange(layer.derivative_exchange_layout(), + // Scale 2: each rotation endpoint carries both the op and the state payload. + return begin_layer_exchange(layer, + 2, comm, [&snap, &layer](int my_rank, const LayerExchangeLayout &layout, VecD &send_buffer) { pack_cross_rank_derivative_payload_impl(snap, layer, my_rank, layout, send_buffer); @@ -328,12 +354,12 @@ void apply_cross_rank_evolution_exchange_impl(VecD &op, inline auto begin_cross_rank_evolution_exchange(VecD &op, const LayerTraversal &layer, const mpi::Comm &comm) -> InFlightExchange { - const auto *layout = active_evolution_exchange_layout(layer, comm); - if (layout == nullptr) { + if (!layer_exchange_participates(comm)) { return {}; } return begin_layer_exchange( - *layout, + layer, + 1, comm, [&op, &layer](int my_rank, const LayerExchangeLayout &active_layout, VecD &send_buffer) { pack_cross_rank_evolution_payload_impl(op, layer, my_rank, active_layout, send_buffer); @@ -362,12 +388,16 @@ auto apply_self_slot_derivative_paired(VecD &state, return {}; } const auto pairs = self_d_count / 2; + // my_rank is loop-invariant, so resolve the slot once: the four fetches below are four + // lookups into the P-sized record array per rotation pair otherwise, in the innermost + // gradient loop. + const auto slot = layer.cross_rank_slot(my_rank); EndpointContrib local{}; for (size_t k = 0; k < pairs; ++k) { - const size_t i1 = layer.cross_rank_sin_recv_index_at(my_rank, k); - const double phi1 = static_cast(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 a425582d..f1cba609 100644 --- a/cpp/monoprop/MPGraph.cpp +++ b/cpp/monoprop/MPGraph.cpp @@ -51,7 +51,24 @@ auto layer_storage_memory_usage(const LayerCore &storage) -> GraphMemoryBreakdow GraphMemoryBreakdown breakdown; breakdown.layer_storage_object_bytes = sizeof(LayerCore); breakdown.cross_rank_bytes = detail::cross_rank_storage_bytes(storage.cross_rank); - breakdown.exchange_layout_bytes = detail::layer_exchange_layout_storage_bytes(storage.evolution_exchange_layout); + // Nothing: the layer no longer retains counts/displs. They are derived into per-thread + // scratch for the exchange being posted, so what used to be 2*P ints per layer per partition + // is now 2*P ints per THREAD. The field stays, reporting the truth, so an A/B against a build + // that did retain them shows the drop rather than silently losing the row. + breakdown.exchange_layout_bytes = 0; + + // Diagnostics. + breakdown.slot_record_bytes = detail::cross_rank_slot_record_bytes(storage.cross_rank); + // The transpose cache is gone: the recv layout equals the send layout, so there was never + // anything to cache. Reported as 0 rather than removed, because it was never inside + // total_bytes() -- an A/B has no other way to see resident memory leave. + breakdown.recv_cache_bytes = 0; + // The derivative layout is no longer retained at all: it is 2x the evolution layout, and its + // transpose is 2x the evolution transpose, so both are derived on demand without a collective. + breakdown.derivative_layout_bytes = 0; + breakdown.layer_cores = 1; + breakdown.slot_records = storage.cross_rank.rank_count(); + breakdown.occupied_slots = detail::cross_rank_occupied_slots(storage.cross_rank); return breakdown; } diff --git a/cpp/monoprop/detail/graph/MPGraphLayers.h b/cpp/monoprop/detail/graph/MPGraphLayers.h index 715aedda..eccea684 100644 --- a/cpp/monoprop/detail/graph/MPGraphLayers.h +++ b/cpp/monoprop/detail/graph/MPGraphLayers.h @@ -64,27 +64,34 @@ struct LayerTraversal final { return detail::cross_rank_sin_recv_phase(core_->cross_rank, rank, idx); } + // The slot is resolved ONCE, outside the loop: the lookup it costs is indexed by the flat + // world P, so doing it per endpoint made per-term work out of what is per-slot work. template 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 -- both sides of it -- is derived at the call site from these and + // nothing else is stored: see detail::derive_exchange_layout and Evolution.cpp. + auto cross_rank() const -> const PackedCrossRankStorage & { return core_->cross_rank; } + auto param_index() const -> size_t { return core_->param_index; } auto gen_coeff() const -> double { return core_->gen_coeff; } auto gate_index() const -> size_t { return core_->gate_index; } diff --git a/cpp/monoprop/detail/graph/MPGraphViews.h b/cpp/monoprop/detail/graph/MPGraphViews.h index 0af7d7a2..9e7c4083 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; // retired: the recv layout IS the send layout, nothing is cached + size_t derivative_layout_bytes = 0; // the lazily retained 2x layout AND its own recv cache -- likewise + size_t layer_cores = 0; // distinct LayerCores walked (shared cores counted once) + size_t slot_records = 0; // sum over cores of ranges.size(); divide by layer_cores to recover P + size_t occupied_slots = 0; // slots carrying any traffic: occupancy = occupied_slots / slot_records + auto total_bytes() const -> size_t { return layer_descriptor_bytes + layer_storage_object_bytes + cos_data_bytes + cross_rank_bytes + exchange_layout_bytes; @@ -52,6 +68,12 @@ struct GraphMemoryBreakdown final { cos_data_bytes += o.cos_data_bytes; cross_rank_bytes += o.cross_rank_bytes; exchange_layout_bytes += o.exchange_layout_bytes; + slot_record_bytes += o.slot_record_bytes; + recv_cache_bytes += o.recv_cache_bytes; + derivative_layout_bytes += o.derivative_layout_bytes; + layer_cores += o.layer_cores; + slot_records += o.slot_records; + occupied_slots += o.occupied_slots; return *this; } }; diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp b/cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp index 2e8da60f..26c1d3fb 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 @@ -51,15 +53,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( @@ -103,18 +96,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 +132,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,49 +158,73 @@ 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 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); - - // 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)); - } + 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 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. + // + // 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_; + 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 c1ed79c9..1d4755a2 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -77,27 +77,86 @@ 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}; } -// Invariant B=[in(P)]++[out(Q)], D=[out(Q)]++[in(P)] (P=in_count, Q=sin_recv_count-P): +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 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 { - 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 sin_send_local = (idx < out_count) ? (in_count + idx) : (idx - out_count); - return cross_rank_sin_send_index(storage, rank, sin_send_local); + 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_recv_offset + idx); + return slot_sin_recv_phase(cross_rank_slot(storage, rank), 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; + +// 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; // Local cycles fold into the self-rank slot (my_rank); the exchange layout zeroes counts[my_rank] so diff --git a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h index 20be7ace..b3b3c3f7 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingTypes.h @@ -23,17 +23,21 @@ #include #include "monoprop/TypeAliases.h" -#include "monoprop/detail/mpi/RecvLayout.h" 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 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; 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 +48,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 { @@ -120,34 +128,51 @@ 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; } }; struct LayerCore final { PackedCrossRankStorage cross_rank; - LayerExchangeLayout evolution_exchange_layout; - auto derivative_exchange_layout() const -> const LayerExchangeLayout &; + // 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. - // 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(); } + // 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. @@ -159,9 +184,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..0e0bc487 100644 --- a/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl +++ b/cpp/monoprop/detail/monomial_propagator/MonomialPropagator.inl @@ -819,8 +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()); - // Drop the inherited eval-time derivative layout: it must not depend on a prior gradient run. - new_core->reset_derivative_exchange_layout(); + // 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 6e230efd..fe9ef14e 100644 --- a/cpp/monoprop/detail/mpi/Exchange.h +++ b/cpp/monoprop/detail/mpi/Exchange.h @@ -19,16 +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 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 &; +// 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). +// +// 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 30141d4b..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,7 +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) -> 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 @@ -132,22 +140,23 @@ 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) { - 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; - 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 f48f9d4a..00000000 --- a/cpp/monoprop/detail/mpi/RecvLayout.h +++ /dev/null @@ -1,35 +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 - -// Kept MPI-free and dependency-light so graph-encoding types (LayerExchangeLayout) 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; -}; - -struct RecvLayoutCache { - RecvLayout layout; - int comm_size = -1; -}; - -} // namespace monoprop::mpi diff --git a/cpp/tests/graph_encoding_tests.cpp b/cpp/tests/graph_encoding_tests.cpp index 26bfe2cf..d65ff796 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" @@ -132,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 { -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); +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_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; + + 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); + + BOOST_CHECK(derived.counts == reference.counts); + BOOST_CHECK(derived.displs == reference.displs); + BOOST_CHECK_EQUAL(derived.total_count, reference.total_count); + } + } +} - // Cached: the second read returns the same object, so eval-time MPI holds a stable pointer. - BOOST_CHECK_EQUAL(&core.derivative_exchange_layout(), &derivative); +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})); - // 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); + 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) { @@ -187,3 +240,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_a_layer_retains_no_exchange_layout) { + // The point of the change: a built layer holds the slot records and nothing else sized by P. + // 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); + + // 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); + + // 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) { + // 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