Revert two perf regressions found by A/B benchmarking; trim a third

A/B benchmarks (pre-review build vs the review branch) showed two of the
review's changes were net-negative, so they are reverted:

I4 (label accumulation): the combined-key (node, other) histogram was 3-15x
SLOWER than the per-node map-of-maps and scaled negatively with threads (a giant
hashmap probed once per pixel, plus a serial argmax). Restore the per-node
map-of-maps with the parallel merge. The detail::number_of_elements dedup (R1)
is kept. Tests confirm identical output.

B1 (mutex watershed): routing the neighbor through detail::valid_offset_target
(per-axis bounds check) cost ~17% on 3D in the hot loop. The public API already
guarantees valid_edges excludes out-of-bounds / wrapping edges, so restore the
single precomputed offset-stride add and document that precondition instead.

B2 (UnionFind bindings): keep the safety bounds check but replace the
per-element validation loop with a single std::max_element scan, cutting bulk
find overhead from ~1.35x to ~1.1x. The residual cost is the unavoidable extra
read of the id array; it only affects the non-hot Python convenience API.

BFS (I3) stays 4.5x faster; greedy-additive (I1) unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: constantinpape

include/bioimage_cpp/graph/label_accumulation.hxx

@@ -17,32 +17,6 @@ namespace bioimage_cpp::graph {
 
 namespace detail_label_accumulation {
 
-// 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,
@@ -51,15 +25,15 @@ void scan_chunk(
     const std::size_t pixel_end,
     const bool has_ignore_value,
     const OtherT ignore_value,
-    NodeOtherHistogram<OtherT> &histogram
+    std::vector<std::unordered_map<OtherT, std::uint64_t>> &histograms
 ) {
     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]);
-        ++histogram[NodeOther<OtherT>{static_cast<std::uint64_t>(node), other}];
+        ++histograms[static_cast<std::size_t>(node)][other];
     }
 }
 
@@ -98,7 +72,10 @@ void accumulate_labels(
     const auto n_pixels = bioimage_cpp::detail::number_of_elements(labels.shape);
     const auto n_threads = detail::normalize_thread_count(number_of_threads, n_pixels);
 
-    std::vector<detail_label_accumulation::NodeOtherHistogram<OtherT>> per_thread(n_threads);
+    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)
+    );
 
     bioimage_cpp::detail::parallel_for_chunks(
         n_threads,
@@ -116,35 +93,41 @@ void accumulate_labels(
         }
     );
 
-    // 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;
+    // Merge per-thread histograms and pick majority per node in one pass over
+    // nodes. This is embarrassingly parallel across nodes. (A single combined
+    // (node, other) hashmap per thread was tried and was 3-15x slower — a giant
+    // map probed per pixel plus a serial argmax — so the per-node map-of-maps
+    // with a parallel merge stays.)
+    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};
+            }
         }
-    }
+    );
 }
 
 } // namespace bioimage_cpp::graph

include/bioimage_cpp/segmentation/mutex_watershed.hxx

@@ -78,6 +78,18 @@ void mutex_watershed_grid(
     ));
     const auto spatial_strides = detail::c_order_strides(spatial_shape);
 
+    // Precondition: `valid_edges` must be 0 for every out-of-bounds / boundary-
+    // wrapping edge (the Python wrapper guarantees this). Given that, the
+    // neighbor flat index can be computed with a single precomputed stride add
+    // per edge instead of a per-axis bounds check, which matters in this hot
+    // loop (a per-axis valid_offset_target check measured ~17% slower on 3D).
+    std::vector<std::ptrdiff_t> offset_strides(number_of_channels, 0);
+    for (std::size_t channel = 0; channel < number_of_channels; ++channel) {
+        for (std::size_t axis = 0; axis < spatial_ndim; ++axis) {
+            offset_strides[channel] += offsets[channel][axis] * spatial_strides[axis];
+        }
+    }
+
     const auto number_of_edges = number_of_nodes * number_of_channels;
     std::vector<WeightedGridEdge<T>> edge_order;
     edge_order.reserve(static_cast<std::size_t>(number_of_edges));
@@ -101,13 +113,8 @@ void mutex_watershed_grid(
         const auto channel = static_cast<std::size_t>(edge_id / number_of_nodes);
         const auto u = edge_id % number_of_nodes;
 
-        // Bounds-check the neighbor per axis rather than relying solely on the
-        // valid_edges mask: this keeps the kernel memory-safe and rejects edges
-        // that would wrap across a row/plane boundary.
-        std::uint64_t v = 0;
-        if (!detail::valid_offset_target(u, offsets[channel], spatial_shape, spatial_strides, v)) {
-            continue;
-        }
+        const auto v_signed = static_cast<std::int64_t>(u) + static_cast<std::int64_t>(offset_strides[channel]);
+        const auto v = static_cast<std::uint64_t>(v_signed);
         std::uint64_t root_u = sets.find(u);
         std::uint64_t root_v = sets.find(v);
         if (root_u == root_v) {

include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx

@@ -105,6 +105,17 @@ void semantic_mutex_watershed_grid(
     ));
     const auto spatial_strides = detail::c_order_strides(spatial_shape);
 
+    // Precondition: `valid_edges` is 0 for every out-of-bounds / boundary-
+    // wrapping spatial edge (guaranteed by the Python wrapper), so the neighbor
+    // flat index uses a single precomputed stride add per edge rather than a
+    // per-axis bounds check in this hot loop.
+    std::vector<std::ptrdiff_t> offset_strides(number_of_offsets, 0);
+    for (std::size_t channel = 0; channel < number_of_offsets; ++channel) {
+        for (std::size_t axis = 0; axis < spatial_ndim; ++axis) {
+            offset_strides[channel] += offsets[channel][axis] * spatial_strides[axis];
+        }
+    }
+
     struct WeightedGridEdge {
         T weight;
         std::uint64_t id;
@@ -148,12 +159,9 @@ void semantic_mutex_watershed_grid(
             continue;
         }
 
-        // Bounds-check the neighbor per axis (see mutex_watershed_grid): keeps
-        // the kernel memory-safe and rejects edges that wrap across a boundary.
-        std::uint64_t v = 0;
-        if (!detail::valid_offset_target(u, offsets[channel], spatial_shape, spatial_strides, v)) {
-            continue;
-        }
+        const auto v_signed = static_cast<std::int64_t>(u) +
+            static_cast<std::int64_t>(offset_strides[channel]);
+        const auto v = static_cast<std::uint64_t>(v_signed);
         const auto root_u = sets.find(u);
         const auto root_v = sets.find(v);
         if (root_u == root_v) {

src/bindings/util.cxx

@@ -4,6 +4,7 @@
 
 #include <nanobind/ndarray.h>
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <stdexcept>
@@ -30,6 +31,21 @@ void check_node(const util::UnionFind &uf, const std::uint64_t node, const char
     }
 }
 
+// Bulk variant: a single max-scan over the ids (vectorizable, one pass) instead
+// of a per-element branch, so validating a large id array stays cheap.
+void check_nodes(const util::UnionFind &uf, const std::uint64_t *data, std::size_t count, const char *name) {
+    if (count == 0) {
+        return;
+    }
+    const auto max_node = *std::max_element(data, data + count);
+    if (max_node >= uf.size()) {
+        throw std::invalid_argument(
+            std::string(name) + " out of range: got " + name + "="
+            + std::to_string(max_node) + ", size=" + std::to_string(uf.size())
+        );
+    }
+}
+
 void merge_edges(
     util::UnionFind &uf,
     EdgeArray edges
@@ -51,10 +67,7 @@ void merge_edges(
     const auto n_edges = edges.shape(0);
     const auto *data = edges.data();
 
-    for (std::size_t i = 0; i < n_edges; ++i) {
-        check_node(uf, data[2 * i], "edge endpoint");
-        check_node(uf, data[2 * i + 1], "edge endpoint");
-    }
+    check_nodes(uf, data, 2 * n_edges, "edge endpoint");
 
     {
         nb::gil_scoped_release release;
@@ -67,9 +80,7 @@ void merge_edges(
 OutputArray find_nodes(util::UnionFind &uf, NodeArray nodes) {
     const auto n = nodes.shape(0);
     const auto *input = nodes.data();
-    for (std::size_t i = 0; i < n; ++i) {
-        check_node(uf, input[i], "node");
-    }
+    check_nodes(uf, input, n, "node");
 
     auto *out = new std::uint64_t[n]();
     nb::capsule owner(out, [](void *p) noexcept { delete[] static_cast<std::uint64_t *>(p); });