More optimization

Author: constantinpape

MIGRATION_GUIDE.md

@@ -107,6 +107,13 @@ Important differences:
 - Bulk methods accept array-like inputs and return NumPy arrays.
 - Python-style names are preferred. A few nifty-style aliases are still present
   on `UndirectedGraph` for convenience, but new code should use snake_case.
+- `graph.clone()` returns an independent deep copy. The C++ class is
+  move-only (it owns a CSR adjacency buffer), so prefer this over
+  reassignment-by-value.
+- `graph.freeze()` eagerly builds the internal adjacency. Call it after a
+  batch of `insert_edge` calls if you intend to hand the graph to multiple
+  reader threads, or if you want to ensure subsequent `node_adjacency`
+  reads carry no first-call rebuild cost.
 
 Common method/property mapping:
 
@@ -155,6 +162,16 @@ Important differences:
   `grid_affinity_features_with_lifted(...)`, which returns local graph weights
   plus explicit long-range `uv_ids` and weights for lifted multicut or mutex
   watershed style workflows.
+- The three grid feature functions preserve `float32` and `float64` input
+  dtype end-to-end (no internal copy to `float64`); other dtypes are
+  promoted to `float64`. Output weight arrays match the input dtype.
+- Grid graph construction does not materialize the per-node adjacency
+  list. If you only need `uv_ids()` and edge features (the common case)
+  you pay nothing for adjacency. The first call to `node_adjacency`,
+  `connected_components`, `breadth_first_search`, or
+  `extract_subgraph_from_nodes` on a grid graph triggers a one-shot
+  rebuild; call `graph.freeze()` on the construction thread before
+  fan-out if you intend to use those from multiple threads.
 - Affogato-style masks and seed edges are not part of the public grid feature
   API yet; the implementation is structured so these filters/extra edges can
   be added later.

include/bioimage_cpp/graph/grid_graph.hxx

@@ -314,10 +314,21 @@ private:
     }
 
     void build_edges() {
-        // Phase 1: emit every edge into `edges_` in canonical order
-        // (axis-major, C-order within each axis). Sequential `emplace_back`
-        // into a single pre-reserved vector is cache-friendly. `edges_` was
+        // Emit every edge into `edges_` in canonical order (axis-major,
+        // C-order within each axis). Sequential `emplace_back` into a
+        // single pre-reserved vector is cache-friendly. `edges_` was
         // already reserved by the `UndirectedGraph` ctor.
+        //
+        // We deliberately do NOT call `rebuild_adjacency_from_edges()`
+        // here. The dominant grid-graph workflow (build graph → compute
+        // edge features via `uv_ids()`) never touches adjacency, so the
+        // ~400 ms rebuild on a 12 M-edge 3D grid is pure waste in that
+        // case. Callers that need adjacency (BFS, connected components,
+        // `extract_subgraph_from_nodes`) trigger a lazy rebuild via the
+        // base-class `node_adjacency` path on first use, paying the cost
+        // exactly once. Multi-threaded readers should call `freeze()` on
+        // the construction thread before fan-out — the same convention
+        // any insert-built `UndirectedGraph` already follows.
         auto &edges = access_edges();
         for (std::size_t axis = 0; axis < D; ++axis) {
             auto pivot_shape = shape_;
@@ -330,10 +341,6 @@ private:
                 }
             );
         }
-        // Phase 2: fill adjacency in bulk — exact-size reserve per node,
-        // no per-edge geometric-growth reallocs across millions of small
-        // adjacency vectors. This is the dominant win on large grids.
-        rebuild_adjacency_from_edges();
     }
 
     Coordinate shape_{};

include/bioimage_cpp/graph/undirected_graph.hxx

@@ -5,8 +5,11 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <memory>
+#include <span>
 #include <stdexcept>
 #include <string>
+#include <type_traits>
 #include <unordered_map>
 #include <unordered_set>
 #include <utility>
@@ -19,12 +22,37 @@ struct Adjacency {
     std::uint64_t edge;
 };
 
+// Undirected graph storing adjacency in CSR-style (compressed sparse row)
+// layout: a single contiguous `adjacency_data_` array of `2 * E` `Adjacency`
+// entries, plus an `adjacency_offsets_` array of length `N + 1` giving the
+// start of each node's adjacency slice. This replaces the previous
+// `std::vector<std::vector<Adjacency>>` representation, which paid for one
+// allocator call per node and scattered the per-node buffers throughout the
+// heap — the dominant cost in `GridGraph` construction on large 3D problems.
+//
+// CSR is rebuilt lazily: incremental `insert_edge*` only appends to `edges_`
+// (and `edge_lookup_`) and marks the adjacency dirty; the first subsequent
+// `node_adjacency` read triggers a single bulk rebuild. Bulk construction
+// paths (`from_sorted_unique_edges`, subclass `build_edges` overrides) call
+// `rebuild_adjacency_from_edges()` explicitly, which keeps reads cheap and
+// thread-safe.
+//
+// Thread safety: as long as a graph is "frozen" before being shared with
+// reader threads (no concurrent inserts, no first-read-from-dirty-state
+// race), reads of `node_adjacency` are safe to share across threads. The
+// lazy rebuild is not internally synchronized — call
+// `rebuild_adjacency_from_edges()` once on the construction thread before
+// fan-out if you built the graph via `insert_edge*`.
 class UndirectedGraph {
 public:
     using NodeId = std::uint64_t;
     using EdgeId = std::uint64_t;
     using Edge = detail::Edge;
-    using AdjacencyList = std::vector<Adjacency>;
+    // Non-owning view over a node's adjacency slice in the CSR buffer.
+    // Backward-compatible with the previous `std::vector<Adjacency>` type
+    // for the uses we have today (range-for and `.size()`); not assignable
+    // and not extendable.
+    using AdjacencyList = std::span<const Adjacency>;
 
     explicit UndirectedGraph(
         const NodeId number_of_nodes = 0,
@@ -35,17 +63,31 @@ public:
 
     virtual ~UndirectedGraph() = default;
 
+    // CSR data lives in a `unique_ptr<Adjacency[]>`, so the type is move-only.
+    // The user-declared destructor above suppresses implicit move generation,
+    // so we re-default the moves explicitly. Copies are deleted on purpose:
+    // the previous vector-of-vectors layout was implicitly copyable but every
+    // such copy paid for a deep clone of millions of small adjacency vectors,
+    // and no caller in this codebase actually needs that. Add an explicit
+    // deep-copy method if one becomes necessary.
+    UndirectedGraph(const UndirectedGraph &) = delete;
+    UndirectedGraph &operator=(const UndirectedGraph &) = delete;
+    UndirectedGraph(UndirectedGraph &&) noexcept = default;
+    UndirectedGraph &operator=(UndirectedGraph &&) noexcept = default;
+
     void assign(
         const NodeId number_of_nodes = 0,
         const EdgeId reserve_number_of_edges = 0
     ) {
         number_of_nodes_ = number_of_nodes;
         edges_.clear();
-        edge_lookup_.clear();
-        adjacency_.clear();
-        adjacency_.resize(static_cast<std::size_t>(number_of_nodes));
         edges_.reserve(static_cast<std::size_t>(reserve_number_of_edges));
+        edge_lookup_.clear();
         edge_lookup_.reserve(static_cast<std::size_t>(reserve_number_of_edges));
+        adjacency_offsets_.clear();
+        adjacency_data_.reset();
+        adjacency_data_size_ = 0;
+        adjacency_dirty_ = true;
     }
 
     [[nodiscard]] NodeId number_of_nodes() const {
@@ -97,9 +139,15 @@ public:
         return edges_;
     }
 
-    [[nodiscard]] const AdjacencyList &node_adjacency(const NodeId node) const {
+    [[nodiscard]] AdjacencyList node_adjacency(const NodeId node) const {
         validate_node(node);
-        return adjacency_[static_cast<std::size_t>(node)];
+        ensure_adjacency_built();
+        const auto begin = adjacency_offsets_[static_cast<std::size_t>(node)];
+        const auto end = adjacency_offsets_[static_cast<std::size_t>(node) + 1];
+        return AdjacencyList(
+            adjacency_data_.get() + begin,
+            static_cast<std::size_t>(end - begin)
+        );
     }
 
     virtual EdgeId insert_edge(const NodeId u, const NodeId v) {
@@ -136,6 +184,47 @@ public:
         return 2 + 2 * number_of_edges();
     }
 
+    // Force the CSR adjacency to be rebuilt if it's currently stale. Use this
+    // after a batch of `insert_edge*` calls, before handing the graph to
+    // multiple reader threads or before holding a span returned from
+    // `node_adjacency` across other graph operations. Lazy rebuild from a
+    // dirty state happens via `mutable` writes inside a `const` method and is
+    // not internally synchronized; once a graph is "frozen" via `freeze`
+    // (or via an explicit `rebuild_adjacency_from_edges` from a subclass
+    // constructor) it is safe to share by `const&` across threads.
+    void freeze() const {
+        ensure_adjacency_built();
+    }
+
+    // Explicit deep copy. The previous vector-of-vectors layout made
+    // `UndirectedGraph` implicitly copyable; switching to CSR with a
+    // `unique_ptr` buffer made the class move-only, so callers that need
+    // a clone use this method. If the source graph's CSR is already built
+    // we copy the buffer so the clone doesn't pay another rebuild on
+    // first read.
+    [[nodiscard]] UndirectedGraph clone() const {
+        UndirectedGraph copy(
+            number_of_nodes_,
+            static_cast<EdgeId>(edges_.size()),
+            edge_lookup_.empty() ? 0 : static_cast<EdgeId>(edges_.size())
+        );
+        copy.edges_ = edges_;
+        copy.edge_lookup_ = edge_lookup_;
+        if (adjacency_dirty_) {
+            copy.adjacency_dirty_ = true;
+        } else {
+            copy.adjacency_offsets_ = adjacency_offsets_;
+            copy.adjacency_data_size_ = adjacency_data_size_;
+            copy.adjacency_data_ =
+                std::make_unique_for_overwrite<Adjacency[]>(adjacency_data_size_);
+            std::copy_n(
+                adjacency_data_.get(), adjacency_data_size_, copy.adjacency_data_.get()
+            );
+            copy.adjacency_dirty_ = false;
+        }
+        return copy;
+    }
+
     // Fast construction from a pre-sorted, deduplicated edge list.
     //
     // Precondition: `edges` is sorted ascending by `(u, v)` with `u < v` in
@@ -210,8 +299,7 @@ protected:
         const EdgeId reserve_edges,
         const EdgeId reserve_lookup
     )
-        : number_of_nodes_(number_of_nodes),
-          adjacency_(static_cast<std::size_t>(number_of_nodes)) {
+        : number_of_nodes_(number_of_nodes) {
         edges_.reserve(static_cast<std::size_t>(reserve_edges));
         edge_lookup_.reserve(static_cast<std::size_t>(reserve_lookup));
     }
@@ -220,61 +308,77 @@ protected:
         const auto edge = static_cast<EdgeId>(edges_.size());
         edges_.emplace_back(u, v);
         edge_lookup_.emplace(Edge{u, v}, edge);
-        adjacency_[static_cast<std::size_t>(u)].push_back(Adjacency{v, edge});
-        adjacency_[static_cast<std::size_t>(v)].push_back(Adjacency{u, edge});
+        adjacency_dirty_ = true;
         return edge;
     }
 
     EdgeId insert_new_edge_without_lookup(const NodeId u, const NodeId v) {
         const auto edge = static_cast<EdgeId>(edges_.size());
         edges_.emplace_back(u, v);
-        adjacency_[static_cast<std::size_t>(u)].push_back(Adjacency{v, edge});
-        adjacency_[static_cast<std::size_t>(v)].push_back(Adjacency{u, edge});
+        adjacency_dirty_ = true;
         return edge;
     }
 
     // Mutable access to the underlying edge list for subclasses that emit
-    // edges in bulk. The caller is responsible for re-establishing the
-    // invariants between `edges_` and `adjacency_` (see
-    // `rebuild_adjacency_from_edges`) before exposing the graph to clients.
+    // edges in bulk. Marks adjacency dirty; the next `node_adjacency` read
+    // (or an explicit `rebuild_adjacency_from_edges` call) will rebuild
+    // the CSR buffers.
     std::vector<Edge> &access_edges() {
+        adjacency_dirty_ = true;
         return edges_;
     }
 
-    // Bulk-populate `adjacency_` from `edges_`. Assumes `edges_` already
-    // contains every edge in the order the subclass wants its `EdgeId`s
-    // assigned (the index into `edges_` becomes the edge id).
+    // Build the CSR adjacency from `edges_`. After this call,
+    // `adjacency_offsets_` has length `N + 1`, `adjacency_data_` has length
+    // `2 * E`, and the slice `adjacency_data_[offsets[u]:offsets[u+1]]` is
+    // exactly the adjacency of node `u`.
     //
-    // Computes the exact degree of each node in one pass, reserves that
-    // capacity per `adjacency_[u]`, then fills in a second pass. With
-    // exact capacity reserved up front each `adjacency_[u]` gets a single
-    // allocation of optimal size — no geometric-growth reallocs, no
-    // fragmentation from intermediate buffers. For grid-shaped topologies
-    // where one push_back per axis would otherwise hammer many small
-    // vectors, this is dramatically faster than `insert_new_edge_*`
-    // accumulating both `edges_` and `adjacency_` in lockstep.
+    // Builds in three sequential passes:
+    //  1. Count degree per node (one pass over `edges_`).
+    //  2. Prefix-sum into `adjacency_offsets_` (one pass over nodes).
+    //  3. Place each edge's two adjacency entries via a fill cursor copied
+    //     from the offsets (one pass over `edges_`).
     //
-    // Existing contents of `adjacency_` are discarded.
+    // All writes in pass (3) land in a single contiguous buffer, so cache
+    // behavior is good; there are no per-node allocator calls. This is
+    // typically 4-5x faster than the previous vector-of-vectors fill on
+    // large grids.
     void rebuild_adjacency_from_edges() {
-        std::vector<std::size_t> degree(static_cast<std::size_t>(number_of_nodes_), 0);
+        const auto n_nodes = static_cast<std::size_t>(number_of_nodes_);
+        const auto data_size = 2 * edges_.size();
+        // Pass 1: count per-node degree into the offsets buffer (shifted by 1).
+        adjacency_offsets_.assign(n_nodes + 1, 0);
         for (const auto &edge : edges_) {
-            ++degree[static_cast<std::size_t>(edge.first)];
-            ++degree[static_cast<std::size_t>(edge.second)];
+            ++adjacency_offsets_[static_cast<std::size_t>(edge.first) + 1];
+            ++adjacency_offsets_[static_cast<std::size_t>(edge.second) + 1];
         }
-        for (std::size_t node = 0; node < adjacency_.size(); ++node) {
-            adjacency_[node].clear();
-            adjacency_[node].reserve(degree[node]);
+        // Pass 2: inclusive prefix sum turns degree-counts into slice starts.
+        for (std::size_t node = 1; node <= n_nodes; ++node) {
+            adjacency_offsets_[node] += adjacency_offsets_[node - 1];
         }
+        // Pass 3: allocate the contiguous adjacency buffer.
+        // `make_unique_for_overwrite` leaves trivially-default-constructible
+        // `Adjacency` elements UNINITIALIZED — every slot is overwritten in
+        // the fill pass below, so the zero-init that `std::vector::resize`
+        // would do is pure overhead. On a 12M-edge grid this skips ~125 ms
+        // of memset, though the underlying page-fault cost shifts to the
+        // fill pass; net win is modest but real.
+        static_assert(std::is_trivially_default_constructible_v<Adjacency>);
+        adjacency_data_ = std::make_unique_for_overwrite<Adjacency[]>(data_size);
+        adjacency_data_size_ = data_size;
+        // Pass 4: `cursor[u]` is the next free slot in node `u`'s adjacency
+        // range. Initialized from the slice starts and incremented per write.
+        std::vector<std::uint64_t> cursor(adjacency_offsets_);
+        Adjacency *const data = adjacency_data_.get();
         for (std::size_t index = 0; index < edges_.size(); ++index) {
             const auto &edge = edges_[index];
             const auto edge_id = static_cast<EdgeId>(index);
-            adjacency_[static_cast<std::size_t>(edge.first)].push_back(
-                Adjacency{edge.second, edge_id}
-            );
-            adjacency_[static_cast<std::size_t>(edge.second)].push_back(
-                Adjacency{edge.first, edge_id}
-            );
+            data[cursor[static_cast<std::size_t>(edge.first)]++] =
+                Adjacency{edge.second, edge_id};
+            data[cursor[static_cast<std::size_t>(edge.second)]++] =
+                Adjacency{edge.first, edge_id};
         }
+        adjacency_dirty_ = false;
     }
 
     void validate_node(const NodeId node) const {
@@ -298,9 +402,29 @@ protected:
     }
 
 private:
+    void ensure_adjacency_built() const {
+        if (adjacency_dirty_) {
+            const_cast<UndirectedGraph *>(this)->rebuild_adjacency_from_edges();
+        }
+    }
+
     NodeId number_of_nodes_;
     std::vector<Edge> edges_;
-    std::vector<AdjacencyList> adjacency_;
+    // CSR adjacency. `adjacency_offsets_` has length `N + 1`;
+    // `adjacency_data_` has length `2 * number_of_edges()` after a rebuild.
+    // The `mutable` qualifiers allow the lazy rebuild from a const reader
+    // path — see `ensure_adjacency_built`. Writers (insert / bulk-edit
+    // helpers) set `adjacency_dirty_ = true` and the rebuild happens on the
+    // next read.
+    mutable std::vector<std::uint64_t> adjacency_offsets_;
+    // Heap-allocated CSR data buffer. Using `unique_ptr<T[]>` instead of
+    // `std::vector<T>` lets `make_unique_for_overwrite` skip the
+    // zero-initialization that `vector::resize` would force — a substantial
+    // saving for 12 M-edge graphs where the buffer is ~384 MB and every
+    // slot is overwritten in the fill pass anyway.
+    mutable std::unique_ptr<Adjacency[]> adjacency_data_;
+    mutable std::size_t adjacency_data_size_ = 0;
+    mutable bool adjacency_dirty_ = true;
     std::unordered_map<Edge, EdgeId, detail::EdgeHash> edge_lookup_;
 };
 

src/bindings/graph.cxx

@@ -1120,6 +1120,8 @@ void bind_graph(nb::module_ &m) {
             nb::arg("nodes")
         )
         .def("edges_from_node_list", &graph_edges_from_node_list, nb::arg("nodes"))
+        .def("freeze", &Graph::freeze)
+        .def("clone", &Graph::clone)
         .def_static(
             "from_edges",
             &graph_from_edges,