perf(graph): ⚡ make per-rank graph memory fall as the world grows

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;
 }
 

cpp/monoprop/detail/graph/MPGraphLayers.h

@@ -64,48 +64,67 @@ 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 <typename Func>
+    auto for_each_occupied_slot(Func &&func) const -> void {
+        detail::for_each_occupied_slot(core_->cross_rank, std::forward<Func>(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 <typename Func>
     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 <typename Func>
     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;
     }
 
     // Endpoints are counted in+out across ranks. Every endpoint is also in cos_data, so cosine-only
     // 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;
     }
 

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;
     }
 };