Performance: drop redundant work in greedy-additive, BFS, label accumulation

constantinpape · claude · constantinpape · commit 1946cb4066ef · 2026-06-12T00:06:16.000+02:00
I1: GreedyAdditiveWorkspace::reset called heap.reset_capacity, and so did
detail::initialize_dynamic_graph which every caller runs immediately after.
The second call re-wiped the heap locator vector (O(E)) on every solve — hot
under fusion-move (per fuse x slot x iteration). Drop it from reset() in both
the multicut and lifted-multicut workspaces; initialize_dynamic_graph owns it.

I3: BfsWorkspace::reset cleared the visited/distance buffers (O(N) over the
whole graph) on every call, making the "k-hop neighborhood from every node"
pattern O(N^2) of memset. Replace the boolean visited buffer with a uint32
generation stamp: reset just increments the generation (O(1)), clearing only on
the rare wraparound. distance_ is sized but written before read.

I4: accumulate_labels used an n_threads x n_nodes map-of-maps plus a per-node
merge map. Switch to one combined-key (node, other) histogram per thread,
merged once, with a single-pass per-node argmax. Far fewer container
allocations; the argmax/tie-break/empty-node semantics are unchanged (covered
by the existing accumulate_labels tests, incl. parallel-vs-single and nifty
cross-checks).

Co-Authored-By: Claude Fable 5 &lt;noreply@anthropic.com&gt;
diff --git a/include/bioimage_cpp/graph/breadth_first_search.hxx b/include/bioimage_cpp/graph/breadth_first_search.hxx
@@ -2,6 +2,7 @@
 
 #include "bioimage_cpp/graph/undirected_graph.hxx"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <limits>
@@ -19,23 +20,44 @@ struct BfsEntry {
 };
 
 // Reusable scratch state for `breadth_first_search`. Reset via `reset(graph)`
-// before each call so the visited-flag and distance buffers grow once and stay
+// before each call so the visited and distance buffers grow once and stay
 // allocated across many BFS invocations on the same graph.
+//
+// Visited tracking uses a per-call generation stamp rather than a boolean
+// buffer cleared each call: `reset` just increments the generation (O(1)),
+// doing a full clear only on the rare 32-bit wraparound. This matters for the
+// "k-hop neighborhood from every node" pattern, where an O(N)-per-call clear
+// would otherwise make the whole sweep O(N^2). `distance_` is sized but not
+// cleared; it is written before being read for every visited node.
 class BfsWorkspace {
 public:
     BfsWorkspace() = default;
 
     void reset(const UndirectedGraph &graph) {
         const auto n_nodes = static_cast<std::size_t>(graph.number_of_nodes());
-        visited_.assign(n_nodes, 0);
-        distance_.assign(n_nodes, 0);
+        if (visited_.size() != n_nodes) {
+            visited_.assign(n_nodes, 0);
+            generation_ = 0;
+        }
+        distance_.resize(n_nodes);
+        if (generation_ == std::numeric_limits<std::uint32_t>::max()) {
+            std::fill(visited_.begin(), visited_.end(), std::uint32_t{0});
+            generation_ = 0;
+        }
+        ++generation_;
     }
 
-    [[nodiscard]] std::vector<unsigned char> &visited() { return visited_; }
+    [[nodiscard]] bool is_visited(const std::uint64_t node) const {
+        return visited_[static_cast<std::size_t>(node)] == generation_;
+    }
+    void mark_visited(const std::uint64_t node) {
+        visited_[static_cast<std::size_t>(node)] = generation_;
+    }
     [[nodiscard]] std::vector<std::uint64_t> &distance() { return distance_; }
 
 private:
-    std::vector<unsigned char> visited_;
+    std::vector<std::uint32_t> visited_;
+    std::uint32_t generation_ = 0;
     std::vector<std::uint64_t> distance_;
 };
 
@@ -71,13 +93,12 @@ inline std::vector<BfsEntry> breadth_first_search(
         );
     }
     workspace.reset(graph);
-    auto &visited = workspace.visited();
     auto &distance = workspace.distance();
 
     std::vector<BfsEntry> result;
     std::queue<std::uint64_t> queue;
     queue.push(source);
-    visited[static_cast<std::size_t>(source)] = 1;
+    workspace.mark_visited(source);
     distance[static_cast<std::size_t>(source)] = 0;
     if (include_source) {
         result.push_back({source, 0});
@@ -93,10 +114,10 @@ inline std::vector<BfsEntry> breadth_first_search(
         const auto next_distance = node_distance + 1;
         for (const auto adjacency : graph.node_adjacency(node)) {
             const auto neighbor = adjacency.node;
-            if (visited[static_cast<std::size_t>(neighbor)] != 0) {
+            if (workspace.is_visited(neighbor)) {
                 continue;
             }
-            visited[static_cast<std::size_t>(neighbor)] = 1;
+            workspace.mark_visited(neighbor);
             distance[static_cast<std::size_t>(neighbor)] = next_distance;
             result.push_back({neighbor, next_distance});
             queue.push(neighbor);
diff --git a/include/bioimage_cpp/graph/label_accumulation.hxx b/include/bioimage_cpp/graph/label_accumulation.hxx
@@ -26,6 +26,32 @@ inline std::size_t number_of_pixels(const std::vector<std::ptrdiff_t> &shape) {
     ));
 }
 
+// Combined (node, other) histogram key. A single map per thread keyed by this
+// pair avoids the n_threads * n_nodes map-of-maps (one std::unordered_map per
+// node per thread), which dominates allocation for RAGs with many nodes.
+template <class OtherT>
+struct NodeOther {
+    std::uint64_t node;
+    OtherT other;
+    bool operator==(const NodeOther &rhs) const noexcept {
+        return node == rhs.node && other == rhs.other;
+    }
+};
+
+template <class OtherT>
+struct NodeOtherHash {
+    std::size_t operator()(const NodeOther<OtherT> &key) const noexcept {
+        std::uint64_t h = key.node * 0x9E3779B97F4A7C15ULL;
+        h ^= static_cast<std::uint64_t>(key.other) + 0x9E3779B97F4A7C15ULL +
+             (h << 6) + (h >> 2);
+        return static_cast<std::size_t>(h);
+    }
+};
+
+template <class OtherT>
+using NodeOtherHistogram =
+    std::unordered_map<NodeOther<OtherT>, std::uint64_t, NodeOtherHash<OtherT>>;
+
 template <class LabelT, class OtherT>
 void scan_chunk(
     const LabelT *labels,
@@ -34,15 +60,15 @@ void scan_chunk(
     const std::size_t pixel_end,
     const bool has_ignore_value,
     const OtherT ignore_value,
-    std::vector<std::unordered_map<OtherT, std::uint64_t>> &histograms
+    NodeOtherHistogram<OtherT> &histogram
 ) {
     for (std::size_t index = pixel_begin; index < pixel_end; ++index) {
         const auto other = other_labels[index];
         if (has_ignore_value && other == ignore_value) {
             continue;
         }
         const auto node = detail::checked_label_to_node(labels[index]);
-        ++histograms[static_cast<std::size_t>(node)][other];
+        ++histogram[NodeOther<OtherT>{static_cast<std::uint64_t>(node), other}];
     }
 }
 
@@ -81,10 +107,7 @@ void accumulate_labels(
     const auto n_pixels = detail_label_accumulation::number_of_pixels(labels.shape);
     const auto n_threads = detail::normalize_thread_count(number_of_threads, n_pixels);
 
-    std::vector<std::vector<std::unordered_map<OtherT, std::uint64_t>>> per_thread(
-        n_threads,
-        std::vector<std::unordered_map<OtherT, std::uint64_t>>(number_of_nodes)
-    );
+    std::vector<detail_label_accumulation::NodeOtherHistogram<OtherT>> per_thread(n_threads);
 
     bioimage_cpp::detail::parallel_for_chunks(
         n_threads,
@@ -102,38 +125,35 @@ void accumulate_labels(
         }
     );
 
-    // Merge per-thread histograms and pick majority per node in one pass over
-    // nodes. This is embarrassingly parallel across nodes.
-    const auto node_threads = detail::normalize_thread_count(number_of_threads, number_of_nodes);
-    bioimage_cpp::detail::parallel_for_chunks(
-        node_threads,
-        number_of_nodes,
-        [&](const std::size_t, const std::size_t node_begin, const std::size_t node_end) {
-            for (std::size_t node = node_begin; node < node_end; ++node) {
-                std::unordered_map<OtherT, std::uint64_t> merged;
-                for (auto &thread_histograms : per_thread) {
-                    auto &node_hist = thread_histograms[node];
-                    for (const auto &entry : node_hist) {
-                        merged[entry.first] += entry.second;
-                    }
-                    node_hist.clear();
-                }
-                OtherT best_label = OtherT{0};
-                std::uint64_t best_count = 0;
-                bool has_best = false;
-                for (const auto &entry : merged) {
-                    if (!has_best ||
-                        entry.second > best_count ||
-                        (entry.second == best_count && entry.first < best_label)) {
-                        best_label = entry.first;
-                        best_count = entry.second;
-                        has_best = true;
-                    }
-                }
-                out.data[node] = has_best ? best_label : OtherT{0};
-            }
+    // Merge the per-thread histograms into one, then pick the majority `other`
+    // label per node (smaller label wins ties; nodes with no pixels stay 0).
+    // The argmax pass is single-threaded over distinct (node, other) pairs,
+    // which is cheap relative to the parallel pixel scan above.
+    auto &merged = per_thread[0];
+    for (std::size_t thread_id = 1; thread_id < per_thread.size(); ++thread_id) {
+        for (const auto &entry : per_thread[thread_id]) {
+            merged[entry.first] += entry.second;
         }
-    );
+        detail_label_accumulation::NodeOtherHistogram<OtherT>().swap(per_thread[thread_id]);
+    }
+
+    for (std::size_t node = 0; node < number_of_nodes; ++node) {
+        out.data[node] = OtherT{0};
+    }
+    std::vector<std::uint64_t> best_count(number_of_nodes, 0);
+    std::vector<unsigned char> has_best(number_of_nodes, 0);
+    for (const auto &entry : merged) {
+        const auto node = static_cast<std::size_t>(entry.first.node);
+        const auto other = entry.first.other;
+        const auto count = entry.second;
+        if (has_best[node] == 0 ||
+            count > best_count[node] ||
+            (count == best_count[node] && other < out.data[node])) {
+            out.data[node] = other;
+            best_count[node] = count;
+            has_best[node] = 1;
+        }
+    }
 }
 
 } // namespace bioimage_cpp::graph
diff --git a/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx b/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx
@@ -19,7 +19,9 @@ struct GreedyAdditiveWorkspace {
     void reset(const UndirectedGraph &lifted_graph) {
         dynamic_graph.reset(lifted_graph);
         union_find.reset(static_cast<std::size_t>(lifted_graph.number_of_nodes()));
-        heap.reset_capacity(static_cast<std::size_t>(lifted_graph.number_of_edges()));
+        // The heap is sized by detail::initialize_dynamic_graph, which every
+        // caller runs immediately after reset(); resizing it here too would
+        // wipe the locator vector twice.
     }
 };
 
diff --git a/include/bioimage_cpp/graph/multicut/greedy_additive.hxx b/include/bioimage_cpp/graph/multicut/greedy_additive.hxx
@@ -21,7 +21,9 @@ struct GreedyAdditiveWorkspace {
     void reset(const UndirectedGraph &graph) {
         dynamic_graph.reset(graph);
         union_find.reset(static_cast<std::size_t>(graph.number_of_nodes()));
-        heap.reset_capacity(static_cast<std::size_t>(graph.number_of_edges()));
+        // The heap is sized by detail::initialize_dynamic_graph, which every
+        // caller runs immediately after reset(); resizing it here too would
+        // wipe the locator vector twice.
     }
 };
 
