feat(graph): 🔍 separate the graph's flat-world slot cost from its traffic

cpp/monoprop/MPGraph.cpp

@@ -52,6 +52,14 @@ auto layer_storage_memory_usage(const LayerCore &storage) -> GraphMemoryBreakdow
     breakdown.layer_storage_object_bytes = sizeof(LayerCore);
     breakdown.cross_rank_bytes = detail::cross_rank_storage_bytes(storage.cross_rank);
     breakdown.exchange_layout_bytes = detail::layer_exchange_layout_storage_bytes(storage.evolution_exchange_layout);
+
+    // Diagnostics. 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);
+    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;
 }
 

cpp/monoprop/detail/graph/MPGraphViews.h

@@ -40,6 +40,14 @@ struct GraphMemoryBreakdown final {
     size_t cross_rank_bytes = 0;
     size_t exchange_layout_bytes = 0;
 
+    // Diagnostics, deliberately OUTSIDE total_bytes(): each is a count or a slice of cross_rank_bytes,
+    // separating the part of the graph that scales with the flat world size P from the traffic part.
+    size_t slot_record_bytes = 0;    // the slice of cross_rank_bytes that is one record per world slot
+    size_t layer_cores = 0;          // distinct LayerCores; slot_records / layer_cores recovers P
+    size_t slot_records = 0;         // summed over cores, so P per core
+    size_t occupied_slots = 0;       // slots carrying any traffic; / slot_records is the occupancy
+    size_t cross_rank_endpoints = 0; // traffic itself; P-independent ceiling on occupied_slots (>= 1 endpoint each)
+
     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 +60,11 @@ 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;
+        layer_cores += o.layer_cores;
+        slot_records += o.slot_records;
+        occupied_slots += o.occupied_slots;
+        cross_rank_endpoints += o.cross_rank_endpoints;
         return *this;
     }
 };

cpp/monoprop/detail/graph_encoding/MPGraphEncoding.cpp

@@ -14,6 +14,7 @@
 
 #include "monoprop/detail/graph_encoding/MPGraphEncodingStorage.h"
 
+#include <algorithm>
 #include <format>
 #include <limits>
 #include <memory>
@@ -148,14 +149,34 @@
 }
 
 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.ranges.capacity() * sizeof(CrossRankPartnerRange);
+}
+
+auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t {
+    // sin_send_count alone is the predicate. B and D hold the same endpoint set in two orders (see
+    // cross_rank_sin_recv_index), so a slot cannot carry D entries while carrying no B entries, and
+    // counting either gives the same answer.
+    return static_cast<size_t>(std::ranges::count_if(storage.ranges, [](const CrossRankPartnerRange &range) {
+        return range.sin_send_count != 0;
+    }));
+}
+
+auto cross_rank_endpoint_count(const PackedCrossRankStorage &storage) -> size_t {
+    size_t count = 0;
+    for (const auto &range : storage.ranges) {
+        count += range.sin_send_count;
+    }
+    return count;
+}
+
 auto layer_exchange_layout_storage_bytes(const LayerExchangeLayout &layout) -> size_t {
     return layout.counts.capacity() * sizeof(int) + layout.displs.capacity() * sizeof(int);
 }
 
 auto build_layer_storage_unified(std::vector<CrossRankPartnerData> all_partners, size_t my_rank)
@@ -176,7 +197,7 @@
         static_cast<void>(build_derivative_exchange_layout(storage->evolution_exchange_layout));
     }
 
     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()) {

cpp/monoprop/detail/graph_encoding/MPGraphEncodingStorage.h

@@ -98,6 +98,15 @@ inline auto cross_rank_sin_recv_phase(const PackedCrossRankStorage &storage, siz
 
 auto cross_rank_storage_bytes(const PackedCrossRankStorage &storage) -> size_t;
 
+// The slot-proportional slice of cross_rank_storage_bytes: one record per world slot, whether or not
+// that slot carries traffic; the remainder scales with terms crossing rather than with the world size.
+auto cross_rank_slot_record_bytes(const PackedCrossRankStorage &storage) -> size_t;
+
+auto cross_rank_occupied_slots(const PackedCrossRankStorage &storage) -> size_t;
+
+// The traffic itself, and the world-size-independent ceiling on cross_rank_occupied_slots.
+auto cross_rank_endpoint_count(const PackedCrossRankStorage &storage) -> size_t;
+
 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

cpp/tests/graph_encoding_tests.cpp

@@ -19,8 +19,10 @@
 
 #include <cstdint>
 #include <limits>
+#include <memory>
 #include <vector>
 
+#include "monoprop/MPGraph.h"
 #include "monoprop/detail/graph_encoding/MPGraphEncodingStorage.h"
 
 using namespace monoprop;
@@ -187,3 +189,99 @@ 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);
 }
+
+namespace {
+// `counts.size()` world slots, slot r carrying counts[r] endpoints; zero means reserved and empty.
+auto slot_partners(const std::vector<size_t> &counts) -> std::vector<CrossRankPartnerData> {
+    std::vector<CrossRankPartnerData> data(counts.size());
+    for (size_t r = 0; r < counts.size(); ++r) {
+        for (size_t k = 0; k < 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_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_world_not_the_traffic) {
+    const auto narrow = detail::build_packed_cross_rank_storage(slot_partners({5, 0, 0, 0}));
+    const auto wide =
+        detail::build_packed_cross_rank_storage(slot_partners({5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}));
+
+    BOOST_CHECK_EQUAL(detail::cross_rank_endpoint_count(narrow), detail::cross_rank_endpoint_count(wide));
+    BOOST_CHECK_EQUAL(detail::cross_rank_occupied_slots(narrow), detail::cross_rank_occupied_slots(wide));
+    BOOST_CHECK_EQUAL(narrow.rank_count(), 4U);
+    BOOST_CHECK_EQUAL(wide.rank_count(), 16U);
+    // One record per world slot is the FLOOR: the figure is capacity-derived, so slack is not a defect.
+    BOOST_CHECK_GE(detail::cross_rank_slot_record_bytes(narrow), narrow.rank_count() * sizeof(CrossRankPartnerRange));
+    BOOST_CHECK_GE(detail::cross_rank_slot_record_bytes(wide), wide.rank_count() * sizeof(CrossRankPartnerRange));
+    BOOST_CHECK_LT(detail::cross_rank_slot_record_bytes(narrow), detail::cross_rank_slot_record_bytes(wide));
+    // 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));
+}
+
+// The five d_-prefixed graph_memory_breakdown() keys read these fields; each is a count or a slice of
+// cross_rank_bytes, so a total including them double-counts.
+BOOST_AUTO_TEST_CASE(graph_memory_breakdown_diagnostics_sit_outside_total_bytes) {
+    GraphMemoryBreakdown b;
+    b.layer_descriptor_bytes = 1;
+    b.layer_storage_object_bytes = 2;
+    b.cos_data_bytes = 4;
+    b.cross_rank_bytes = 8;
+    b.exchange_layout_bytes = 16;
+    constexpr size_t kOwnedBytes = 1U + 2U + 4U + 8U + 16U;
+    BOOST_CHECK_EQUAL(b.total_bytes(), kOwnedBytes);
+
+    b.slot_record_bytes = 32;
+    b.layer_cores = 64;
+    b.slot_records = 128;
+    b.occupied_slots = 256;
+    b.cross_rank_endpoints = 512;
+    BOOST_CHECK_EQUAL(b.total_bytes(), kOwnedBytes);
+
+    // A partitioned propagator sums per-partition breakdowns, so the diagnostics have to add too.
+    GraphMemoryBreakdown acc;
+    acc += b;
+    acc += b;
+    BOOST_CHECK_EQUAL(acc.total_bytes(), 2U * kOwnedBytes);
+    BOOST_CHECK_EQUAL(acc.slot_record_bytes, 64U);
+    BOOST_CHECK_EQUAL(acc.layer_cores, 128U);
+    BOOST_CHECK_EQUAL(acc.slot_records, 256U);
+    BOOST_CHECK_EQUAL(acc.occupied_slots, 512U);
+    BOOST_CHECK_EQUAL(acc.cross_rank_endpoints, 1024U);
+}
+
+// Each d_ key against the helper it derives from; two layers share one core, so d_layer_cores is 1.
+BOOST_AUTO_TEST_CASE(graph_memory_breakdown_diagnostics_report_the_flat_world) {
+    auto core = std::make_shared<LayerCore>();
+    core->cross_rank = detail::build_packed_cross_rank_storage(slot_partners({0, 3, 0, 0, 7, 0}));
+
+    MPGraph graph(/*schrodinger=*/false);
+    graph.append(core);
+    graph.append(core);
+    BOOST_REQUIRE_EQUAL(graph.layers(), 2U);
+
+    const auto b = graph.storage_memory_usage();
+    BOOST_CHECK_EQUAL(b.layer_cores, 1U);
+    BOOST_CHECK_EQUAL(b.slot_records, core->cross_rank.rank_count());
+    BOOST_CHECK_EQUAL(b.slot_records, 6U);
+    BOOST_CHECK_EQUAL(b.occupied_slots, detail::cross_rank_occupied_slots(core->cross_rank));
+    BOOST_CHECK_EQUAL(b.cross_rank_endpoints, detail::cross_rank_endpoint_count(core->cross_rank));
+    BOOST_CHECK_EQUAL(b.slot_record_bytes, detail::cross_rank_slot_record_bytes(core->cross_rank));
+    BOOST_CHECK_EQUAL(b.cross_rank_bytes, detail::cross_rank_storage_bytes(core->cross_rank));
+    // slot_record_bytes is a slice of cross_rank_bytes, which total_bytes already counts in full.
+    BOOST_CHECK_LE(b.slot_record_bytes, b.cross_rank_bytes);
+    BOOST_CHECK_LE(b.cross_rank_bytes, b.total_bytes());
+}

src/monoprop/bindings/binder.h

@@ -271,5 +271,21 @@ 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}};
     });
+
+    cls.def("graph_memory_breakdown", [](const MonomialPropagator<NumModes> &self) {
+        const auto b = self.graph_memory_usage();
+        return std::map<std::string, size_t>{{"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_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