diff --git a/cpp/monoprop/Evolution.cpp b/cpp/monoprop/Evolution.cpp index 47bf3234..4dede1eb 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}); } @@ -192,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 @@ -219,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; } @@ -258,11 +269,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); @@ -286,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, @@ -307,33 +315,26 @@ 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) -> 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); @@ -357,17 +358,19 @@ 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; 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; @@ -398,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; + return; // the self slot recovers live; it needs no snapshot } - 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]; - }); + 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]; } - 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]; - }); - } - } + }); } } // namespace @@ -483,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. @@ -498,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/MPGraph.cpp b/cpp/monoprop/MPGraph.cpp index a425582d..655a6630 100644 --- a/cpp/monoprop/MPGraph.cpp +++ b/cpp/monoprop/MPGraph.cpp @@ -51,7 +51,26 @@ 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. 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 + // 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); + breakdown.cross_rank_endpoints = detail::cross_rank_endpoint_count(storage.cross_rank); return breakdown; } diff --git a/cpp/monoprop/detail/graph/MPGraphLayers.h b/cpp/monoprop/detail/graph/MPGraphLayers.h index 715aedda..901f5d2e 100644 --- a/cpp/monoprop/detail/graph/MPGraphLayers.h +++ b/cpp/monoprop/detail/graph/MPGraphLayers.h @@ -64,38 +64,58 @@ struct LayerTraversal final { return detail::cross_rank_sin_recv_phase(core_->cross_rank, rank, idx); } + // 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); 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; } // 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; } @@ -103,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/MPGraphViews.h b/cpp/monoprop/detail/graph/MPGraphViews.h index 0af7d7a2..bbb2bdc8 100644 --- a/cpp/monoprop/detail/graph/MPGraphViews.h +++ b/cpp/monoprop/detail/graph/MPGraphViews.h @@ -40,6 +40,26 @@ struct GraphMemoryBreakdown final { size_t cross_rank_bytes = 0; size_t exchange_layout_bytes = 0; + // 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_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; // 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 + exchange_layout_bytes; @@ -52,6 +72,13 @@ 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; + 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 2e8da60f..28be0cc4 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 @@ -32,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); @@ -51,15 +65,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( @@ -100,21 +105,36 @@ 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; - size_t total_d = 0; for (size_t rank = 0; rank < num_ranks; ++rank) { const auto &partner = data[rank]; - 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); + // 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())); + } + // 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(); - total_d += partner.sin_recv_entries.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; for (const auto &partner : data) { @@ -128,10 +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]; - const size_t b_off = storage.ranges[rank].sin_send_offset; - const size_t d_off = storage.ranges[rank].sin_recv_offset; + // 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"); @@ -142,61 +165,121 @@ 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 = - 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; } +auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size_t { + return storage.occupied.capacity() * sizeof(CrossRankOccupiedSlot); +} + +auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t { + // 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 &entry : storage.occupied) { + count += entry.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); } -auto build_layer_storage_unified(std::vector all_partners, size_t my_rank) - -> std::shared_ptr { - auto storage = std::make_shared(); +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(); + // 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); - { - 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()); + // 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 } - storage->evolution_exchange_layout = build_layer_exchange_layout(send_counts, 1); + out.counts[slot] = + checked_mpi_int(static_cast(scale) * view.sin_send_count, count_label.c_str()); + }); - // 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)); + // 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) { + out.displs[r] = checked_mpi_int(total, displacement_label.c_str()); + total += static_cast(out.counts[r]); } + 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(); + 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. - 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..6835a04d 100644 --- a/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h +++ b/cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h @@ -77,27 +77,139 @@ 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]); +// 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 +// 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; +}; + +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 = 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}; } -// 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 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. 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 -- +// 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..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 @@ -23,17 +24,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 +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 { @@ -120,34 +129,86 @@ struct CrossRankPartnerData { size_t in_count = 0; }; -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; +// 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."); + +// 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 == R + // 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; } - auto sin_recv_size(size_t rank) const -> size_t { return ranges[rank].sin_recv_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 { 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 +220,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..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; @@ -819,8 +812,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/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 26bfe2cf..db3aa967 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 - 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); +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)); - // 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 (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; - // 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); + 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); + } + } +} + +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) { @@ -187,3 +240,131 @@ 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. (slot_partners is defined above, with the layout tests.) + +BOOST_AUTO_TEST_CASE(graph_encoding_occupied_slots_counts_only_slots_carrying_traffic) { + // 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})); + + 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_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)); + // 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)); +} + +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); +} + +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 2110b355..6ebe11aa 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; @@ -79,10 +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.sin_recv_count), sizeof(TermIndex)); BOOST_CHECK_EQUAL(sizeof(r.in_count), 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) @@ -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..9118fadd 100644 --- a/src/monoprop/bindings/binder.h +++ b/src/monoprop/bindings/binder.h @@ -271,5 +271,28 @@ 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). 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}, + {"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}, + {"d_cross_rank_endpoints", b.cross_rank_endpoints}}; + }); } } // namespace monoprop::bindings::detail