Optimize fusion moves

Author: constantinpape

AGENTS.md

@@ -120,6 +120,25 @@ Tests run against the installed Python package. For each public function, cover:
 
 Correct first, fast second. Benchmark before adding complexity. Prefer algorithmic improvements over build-level tweaks. Release the GIL for expensive kernels. Add threading only after the single-threaded implementation is stable; it must be portable and user-controllable.
 
+### Profiling
+
+Measure before optimizing. The codebase carries a lightweight per-phase profiling utility for exactly this:
+
+- Header: `include/bioimage_cpp/detail/profile.hxx`. Macros: `BIOIMAGE_PROFILE_INIT(name)`, `BIOIMAGE_PROFILE_SCOPE(name, "label")`, `BIOIMAGE_PROFILE_REPORT(name)`.
+- Gated behind the `BIOIMAGE_PROFILE` compile-time flag. Outside of profile builds the macros expand to no-ops and a `NullProfiler` stub so the same code compiles unchanged.
+- Enable via CMake option: `pip install -e . --no-build-isolation -C cmake.define.BIOIMAGE_PROFILE=ON`. Rebuild without the flag for production work.
+- Reports per-phase wall-clock totals to stderr at the end of the instrumented scope (e.g., at the end of `optimize`). Same labels accumulate across multiple invocations of the scope (e.g., per-iteration phases).
+
+Workflow when adding or chasing a performance issue:
+
+1. **Compare standalone primitives first** (`development/.../check_*.py` scripts vs. nifty). If a primitive is already fast, the gap is elsewhere — don't optimize it speculatively.
+2. **Instrument the suspect function** by wrapping each logical phase in a `BIOIMAGE_PROFILE_SCOPE`. Pick labels that map to one operation each ("agreement_contract", "sub_solve", "energy_eval", ...), not full call paths.
+3. **Build with `BIOIMAGE_PROFILE=ON` and run a realistic problem** — typically the external multicut instance loaded by the comparison scripts. Run with `--repeats 1` so the report isn't drowned out.
+4. **Optimize the largest phase** (50% phase beats two 10% phases combined). Re-measure after each change; verify no other phase regressed.
+5. **Strip the instrumentation when done** only if it adds clutter; otherwise leave it in place — it's free when the flag is off.
+
+Don't add `std::chrono` snippets ad hoc; use the existing macros so future profiling sessions land in a consistent format.
+
 ## Documentation
 
 `README.md` covers: what `bioimage-cpp` is and isn't, install/build, minimal examples, design philosophy. Public functions need concise docstrings documenting input shapes, supported dtypes, output shapes/dtypes, copy behavior, background-label behavior (if relevant), and axis/coordinate conventions.

CMakeLists.txt

@@ -27,4 +27,9 @@ target_compile_options(_core PRIVATE
     $<$<NOT:$<CONFIG:Debug>>:-O3>
 )
 
+option(BIOIMAGE_PROFILE "Enable per-phase profiling instrumentation (development only)" OFF)
+if(BIOIMAGE_PROFILE)
+    target_compile_definitions(_core PRIVATE BIOIMAGE_PROFILE)
+endif()
+
 install(TARGETS _core LIBRARY DESTINATION bioimage_cpp)

include/bioimage_cpp/detail/indexed_heap.hxx

@@ -133,6 +133,28 @@ public:
         sift_up(pos);
     }
 
+    // Bulk-load the heap from a list of entries in O(n) via Floyd's heapify.
+    // Replaces any current heap contents. Preconditions:
+    //   - the heap is empty before this call (call `clear()`/`reset_capacity`
+    //     first if needed);
+    //   - every key in `entries` is unique.
+    // Compared to N successive `push` calls (O(n log n)) this matters when
+    // initializing a heap from a large edge set in one shot.
+    void build_heap(std::vector<Entry> entries) {
+        heap_ = std::move(entries);
+        for (std::size_t pos = 0; pos < heap_.size(); ++pos) {
+            locator_.set(heap_[pos].key, pos);
+        }
+        if (heap_.size() < 2) {
+            return;
+        }
+        // Sift down from the last internal node back to the root. After each
+        // sift_down, the subtree rooted at `pos` is heap-ordered.
+        for (std::size_t pos = heap_.size() / 2; pos-- > 0;) {
+            sift_down(pos);
+        }
+    }
+
     // Precondition: `key` is currently in the heap.
     void change(const KeyT &key, PriorityT priority) {
         const auto pos = locator_.at(key);

include/bioimage_cpp/detail/profile.hxx

@@ -0,0 +1,82 @@
+#pragma once
+
+// Lightweight per-phase profiling utility. Active only when the translation
+// unit is compiled with `-DBIOIMAGE_PROFILE`; otherwise every macro is a no-op
+// and adds no overhead. Use it in development/profile-mode builds to find
+// hotspots; do not enable it in production wheels.
+
+#ifdef BIOIMAGE_PROFILE
+#include <chrono>
+#include <cstdio>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace bioimage_cpp::detail {
+
+class Profiler {
+public:
+    using Clock = std::chrono::steady_clock;
+    using Duration = std::chrono::duration<double>;
+
+    void accumulate(const char *name, const Duration delta) {
+        auto it = totals_.find(name);
+        if (it == totals_.end()) {
+            order_.push_back(name);
+            totals_.emplace(name, delta.count());
+        } else {
+            it->second += delta.count();
+        }
+    }
+
+    void report() const {
+        std::fprintf(stderr, "[bioimage profile]\n");
+        double total = 0.0;
+        for (const auto *name : order_) {
+            total += totals_.at(name);
+        }
+        for (const auto *name : order_) {
+            const auto seconds = totals_.at(name);
+            const auto fraction = total > 0.0 ? (100.0 * seconds / total) : 0.0;
+            std::fprintf(stderr, "  %-22s %8.4f s  (%5.1f%%)\n", name, seconds, fraction);
+        }
+        std::fprintf(stderr, "  %-22s %8.4f s\n", "total", total);
+    }
+
+private:
+    std::vector<const char *> order_;
+    std::unordered_map<const char *, double> totals_;
+};
+
+class ProfileTimer {
+public:
+    ProfileTimer(Profiler &profiler, const char *name)
+        : profiler_(profiler), name_(name), start_(Profiler::Clock::now()) {}
+
+    ~ProfileTimer() {
+        profiler_.accumulate(name_, Profiler::Clock::now() - start_);
+    }
+
+private:
+    Profiler &profiler_;
+    const char *name_;
+    Profiler::Clock::time_point start_;
+};
+
+} // namespace bioimage_cpp::detail
+
+#define BIOIMAGE_PROFILE_INIT(var) ::bioimage_cpp::detail::Profiler var;
+#define BIOIMAGE_PROFILE_SCOPE(var, name) ::bioimage_cpp::detail::ProfileTimer _bp_##__LINE__(var, name);
+#define BIOIMAGE_PROFILE_REPORT(var) (var).report();
+
+#else
+
+namespace bioimage_cpp::detail {
+struct NullProfiler {};
+} // namespace bioimage_cpp::detail
+
+#define BIOIMAGE_PROFILE_INIT(var) ::bioimage_cpp::detail::NullProfiler var;
+#define BIOIMAGE_PROFILE_SCOPE(var, name) (void)var;
+#define BIOIMAGE_PROFILE_REPORT(var) (void)var;
+
+#endif

include/bioimage_cpp/detail/union_find.hxx

@@ -70,6 +70,16 @@ public:
         return parents_.size();
     }
 
+    // Re-initialise the union-find to `n` singletons. Reuses the existing
+    // vectors' capacity when possible, so the workspace pattern (reusing a
+    // single `UnionFind` across many solver invocations on graphs of varying
+    // size) avoids reallocations.
+    void reset(const std::size_t n) {
+        parents_.resize(n);
+        ranks_.assign(n, 0);
+        std::iota(parents_.begin(), parents_.end(), std::uint64_t{0});
+    }
+
 private:
     std::vector<std::uint64_t> parents_;
     std::vector<std::uint64_t> ranks_;

include/bioimage_cpp/graph/detail/fusion_contract.hxx

@@ -1,13 +1,15 @@
 #pragma once
 
-#include "bioimage_cpp/detail/relabel.hxx"
 #include "bioimage_cpp/detail/union_find.hxx"
 #include "bioimage_cpp/graph/undirected_graph.hxx"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include <limits>
 #include <stdexcept>
 #include <string>
+#include <utility>
 #include <vector>
 
 namespace bioimage_cpp::graph::detail {
@@ -30,6 +32,17 @@ struct AgreementContraction {
 // `proposals` is a row-major buffer of shape (n_proposals, number_of_nodes).
 // Passing 0 proposals collapses every edge (all proposals trivially agree)
 // and yields a single-node contracted graph; we reject that as a usage error.
+//
+// Implementation notes:
+//   - Dense relabeling uses a flat sentinel array indexed by root id rather
+//     than a hash map; roots live in [0, number_of_nodes) so this is O(N).
+//   - Surviving (min_root, max_root) pairs are collected and sorted by a
+//     packed 64-bit key, then deduped sequentially. The contracted graph is
+//     built in one pass via `UndirectedGraph::from_sorted_unique_edges`,
+//     bypassing per-edge hash insertion in `insert_edge`.
+//   - The contracted graph's `edge_lookup_` is left empty because the
+//     fusion-move sub-solver only iterates edges and adjacency, never calls
+//     `find_edge`.
 inline AgreementContraction contract_by_agreement(
     const UndirectedGraph &graph,
     const std::uint64_t *proposals,
@@ -67,36 +80,87 @@ inline AgreementContraction contract_by_agreement(
         }
     }
 
-    std::vector<std::uint64_t> raw_root(static_cast<std::size_t>(number_of_nodes));
+    // Dense-relabel UFD roots in one O(N) pass with a sentinel array.
+    constexpr std::uint64_t unset = std::numeric_limits<std::uint64_t>::max();
+    std::vector<std::uint64_t> dense_of_raw(
+        static_cast<std::size_t>(number_of_nodes), unset
+    );
+    std::vector<std::uint64_t> root_of_node(static_cast<std::size_t>(number_of_nodes));
+    std::uint64_t number_of_components = 0;
     for (std::uint64_t node = 0; node < number_of_nodes; ++node) {
-        raw_root[static_cast<std::size_t>(node)] = sets.find(node);
+        const auto raw = sets.find(node);
+        auto dense = dense_of_raw[static_cast<std::size_t>(raw)];
+        if (dense == unset) {
+            dense = number_of_components++;
+            dense_of_raw[static_cast<std::size_t>(raw)] = dense;
+        }
+        root_of_node[static_cast<std::size_t>(node)] = dense;
     }
-    auto root_of_node = bioimage_cpp::detail::dense_relabel(raw_root);
 
-    std::uint64_t number_of_components = 0;
-    for (const auto root : root_of_node) {
-        if (root + 1 > number_of_components) {
-            number_of_components = root + 1;
-        }
+    // Sort key for surviving edges. The lower 32 bits hold `max_root` and
+    // the upper 32 bits hold `min_root` so a single uint64 comparison
+    // suffices. This requires number_of_components to fit in 32 bits, which
+    // is always true for graphs we can fit in memory.
+    if (number_of_components > (std::uint64_t{1} << 32)) {
+        throw std::runtime_error(
+            "number_of_components exceeds 2^32 — contraction packing assumption violated"
+        );
     }
 
-    UndirectedGraph contracted_graph(number_of_components);
+    struct Survivor {
+        std::uint64_t key;             // (min_root << 32) | max_root
+        std::uint64_t original_edge;
+    };
+    std::vector<Survivor> survivors;
+    survivors.reserve(static_cast<std::size_t>(number_of_edges));
+
     std::vector<std::int64_t> contracted_edge_of_original(
         static_cast<std::size_t>(number_of_edges), -1
     );
 
     for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) {
         const auto uv = graph.uv(edge);
-        const auto ru = root_of_node[static_cast<std::size_t>(uv.first)];
-        const auto rv = root_of_node[static_cast<std::size_t>(uv.second)];
+        auto ru = root_of_node[static_cast<std::size_t>(uv.first)];
+        auto rv = root_of_node[static_cast<std::size_t>(uv.second)];
         if (ru == rv) {
             continue;
         }
-        const auto inserted = contracted_graph.insert_edge(ru, rv);
-        contracted_edge_of_original[static_cast<std::size_t>(edge)] =
-            static_cast<std::int64_t>(inserted);
+        if (ru > rv) {
+            std::swap(ru, rv);
+        }
+        survivors.push_back(Survivor{(ru << 32) | rv, edge});
     }
 
+    std::sort(
+        survivors.begin(),
+        survivors.end(),
+        [](const Survivor &a, const Survivor &b) {
+            return a.key < b.key;
+        }
+    );
+
+    std::vector<UndirectedGraph::Edge> contracted_edges;
+    contracted_edges.reserve(survivors.size());
+
+    constexpr std::uint64_t no_key = std::numeric_limits<std::uint64_t>::max();
+    std::uint64_t last_key = no_key;
+    std::int64_t current_contracted = -1;
+    for (const auto &survivor : survivors) {
+        if (survivor.key != last_key) {
+            const auto ru = survivor.key >> 32;
+            const auto rv = survivor.key & std::uint64_t{0xFFFFFFFF};
+            current_contracted = static_cast<std::int64_t>(contracted_edges.size());
+            contracted_edges.push_back(UndirectedGraph::Edge{ru, rv});
+            last_key = survivor.key;
+        }
+        contracted_edge_of_original[static_cast<std::size_t>(survivor.original_edge)] =
+            current_contracted;
+    }
+
+    auto contracted_graph = UndirectedGraph::from_sorted_unique_edges(
+        number_of_components, std::move(contracted_edges), /*populate_lookup=*/false
+    );
+
     return AgreementContraction{
         std::move(contracted_graph),
         std::move(contracted_edge_of_original),