Rag optimization (#17)

* Optimize lifted affinity features

* More optimization

Author: constantinpape

development/graph/_rag_compatibility.py

@@ -85,6 +85,10 @@ def make_watershed_labels(
 
 
 def time_call(function, repeats: int):
+    # One untimed warm-up call before the measured loop so the first sample
+    # doesn't carry nanobind tuple-shape caching, numpy ufunc init, code-page
+    # faults, etc. Mirrors the grid-affinity helper for consistency.
+    function()
     timings = []
     result = None
     for _ in range(repeats):
@@ -136,7 +140,9 @@ def compare_boundary_features(
     import bioimage_cpp as bic
     import nifty.graph.rag as nrag
 
-    block_shape = list(labels.shape)
+    # Don't force blockShape on the nifty side — its default block layout is
+    # what its parallelism is tuned for, and forcing a single block (== labels
+    # shape) starves nifty's worker pool.
     bic_timings, bic_features = time_call(
         lambda: bic.graph.edge_map_features(
             bic_rag, labels, boundary_map, number_of_threads=threads
@@ -147,7 +153,6 @@ def compare_boundary_features(
         lambda: nrag.accumulateEdgeMeanAndLength(
             nifty_rag,
             np.ascontiguousarray(boundary_map, dtype=np.float32),
-            blockShape=block_shape,
             numberOfThreads=threads,
         ),
         repeats,
@@ -193,16 +198,20 @@ def compare_affinity_features(
         ),
         repeats,
     )
-    # The installed Nifty wrapper can crash or fail when an explicit
-    # numberOfThreads is passed here. The default path is still the reference
-    # implementation and keeps this compatibility script robust.
+    # nifty's accumulateAffinityStandartFeatures crashes inside vigra's
+    # UserRangeHistogram when numberOfThreads=1 (the accumulators never get
+    # setMinMax). For threads >= 2 we can pass the value through; for the
+    # single-thread case we fall back to nifty's default (-1, all cores) and
+    # surface the unfairness in the printed report.
+    nifty_threads_effective = threads if threads != 1 else -1
     nifty_timings, nifty_features_full = time_call(
         lambda: nrag.accumulateAffinityStandartFeatures(
             nifty_rag,
             affs_float32,
             offsets_for_nifty,
             min_val,
             max_val,
+            numberOfThreads=nifty_threads_effective,
         ),
         repeats,
     )
@@ -213,6 +222,8 @@ def compare_affinity_features(
     np.testing.assert_allclose(aligned_bic, nifty_features, rtol=1.0e-5, atol=1.0e-6)
     return bic_timings, nifty_timings, {
         "max_abs_diff": float(np.max(np.abs(aligned_bic - nifty_features))),
+        "nifty_threads_effective": nifty_threads_effective,
+        "nifty_threads_requested": threads,
     }
 
 
@@ -290,7 +301,15 @@ def run_compatibility_check(
     print(f"boundary-map size convention: {boundary_summary['size_convention']}")
     _print_timing("affinity features", affinity_bic_timings, affinity_nifty_timings)
     print(f"affinity feature max abs diff: {affinity_summary['max_abs_diff']:.6g}")
-    print("nifty affinity timing uses the wrapper default thread handling")
+    requested = affinity_summary["nifty_threads_requested"]
+    effective = affinity_summary["nifty_threads_effective"]
+    if requested != effective:
+        print(
+            f"WARNING: affinity features — requested {requested} thread(s) but "
+            f"nifty was called with numberOfThreads={effective} "
+            "(its single-thread path crashes inside vigra's UserRangeHistogram). "
+            f"bioimage-cpp used {requested} thread(s); the timing comparison is NOT apples-to-apples."
+        )
 
 
 def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float]):
@@ -303,7 +322,9 @@ def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float
 
 
 def add_common_arguments(parser: argparse.ArgumentParser) -> None:
-    parser.add_argument("--repeats", type=int, default=3)
+    # 5 matches the grid-affinity helper; median of 3 is noisy because one
+    # GC stall in the middle sample becomes the median.
+    parser.add_argument("--repeats", type=int, default=5)
     parser.add_argument("--threads", type=int, default=1)
     parser.add_argument("--watershed-min-distance", type=int, default=5)
     parser.add_argument("--watershed-grid-spacing", type=int, default=12)

include/bioimage_cpp/graph/feature_accumulation.hxx

@@ -2,6 +2,7 @@
 
 #include "bioimage_cpp/array_view.hxx"
 #include "bioimage_cpp/detail/grid.hxx"
+#include "bioimage_cpp/detail/profile.hxx"
 #include "bioimage_cpp/detail/threading.hxx"
 #include "bioimage_cpp/graph/region_adjacency_graph.hxx"
 
@@ -204,31 +205,168 @@ void scan_edge_map_3d_chunk(
     }
 }
 
+// Given an offset along one axis and the axis length, return the half-open
+// range of axis coordinates `[lo, hi)` for which `coord + delta` stays in
+// `[0, length)`. Returns `lo >= hi` if the offset is larger than the axis.
+inline void valid_axis_range(
+    const std::ptrdiff_t delta,
+    const std::size_t length,
+    std::size_t &lo,
+    std::size_t &hi
+) {
+    if (delta >= 0) {
+        lo = 0;
+        const auto d = static_cast<std::size_t>(delta);
+        hi = (d >= length) ? 0 : (length - d);
+    } else {
+        const auto d = static_cast<std::size_t>(-delta);
+        lo = (d >= length) ? length : d;
+        hi = length;
+    }
+}
+
+// Sweep every (node, target) pair on a 2D grid for which `node + offset` stays
+// in bounds, restricted to the half-open y-slab [y_begin, y_end). The body
+// receives flat C-order indices for both endpoints and is expected to inline
+// at -O2 since this is a header-only template with a fully-known callable
+// type at instantiation.
+template <class Body>
+void sweep_offset_box_2d(
+    const std::ptrdiff_t dy,
+    const std::ptrdiff_t dx,
+    const std::size_t height,
+    const std::size_t width,
+    const std::size_t y_begin,
+    const std::size_t y_end,
+    const Body &body
+) {
+    std::size_t y_lo_full, y_hi_full, x_lo, x_hi;
+    valid_axis_range(dy, height, y_lo_full, y_hi_full);
+    valid_axis_range(dx, width, x_lo, x_hi);
+    const auto y_lo = std::max(y_lo_full, y_begin);
+    const auto y_hi = std::min(y_hi_full, y_end);
+    if (y_lo >= y_hi || x_lo >= x_hi) {
+        return;
+    }
+    const auto offset_stride = dy * static_cast<std::ptrdiff_t>(width) + dx;
+    for (std::size_t y = y_lo; y < y_hi; ++y) {
+        const auto row_offset = y * width;
+        for (std::size_t x = x_lo; x < x_hi; ++x) {
+            const auto node = row_offset + x;
+            const auto target = static_cast<std::uint64_t>(
+                static_cast<std::ptrdiff_t>(node) + offset_stride
+            );
+            body(static_cast<std::uint64_t>(node), target);
+        }
+    }
+}
+
+// 3D variant of `sweep_offset_box_2d`. Restricts the sweep to a z-slab.
+template <class Body>
+void sweep_offset_box_3d(
+    const std::ptrdiff_t dz,
+    const std::ptrdiff_t dy,
+    const std::ptrdiff_t dx,
+    const std::size_t depth,
+    const std::size_t height,
+    const std::size_t width,
+    const std::size_t z_begin,
+    const std::size_t z_end,
+    const Body &body
+) {
+    std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi;
+    valid_axis_range(dz, depth, z_lo_full, z_hi_full);
+    valid_axis_range(dy, height, y_lo, y_hi);
+    valid_axis_range(dx, width, x_lo, x_hi);
+    const auto z_lo = std::max(z_lo_full, z_begin);
+    const auto z_hi = std::min(z_hi_full, z_end);
+    if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) {
+        return;
+    }
+    const auto slice_size = height * width;
+    const auto offset_stride =
+        dz * static_cast<std::ptrdiff_t>(slice_size) +
+        dy * static_cast<std::ptrdiff_t>(width) + dx;
+    for (std::size_t z = z_lo; z < z_hi; ++z) {
+        const auto slice_offset = z * slice_size;
+        for (std::size_t y = y_lo; y < y_hi; ++y) {
+            const auto row_offset = slice_offset + y * width;
+            for (std::size_t x = x_lo; x < x_hi; ++x) {
+                const auto node = row_offset + x;
+                const auto target = static_cast<std::uint64_t>(
+                    static_cast<std::ptrdiff_t>(node) + offset_stride
+                );
+                body(static_cast<std::uint64_t>(node), target);
+            }
+        }
+    }
+}
+
 template <class LabelT, class ValueT, class Stats>
-void scan_affinity_chunk(
+void scan_affinity_2d_chunk(
     const RegionAdjacencyGraph &rag,
     const LabelT *labels,
     const ValueT *affinities,
     const std::vector<std::vector<std::ptrdiff_t>> &offsets,
-    const std::vector<std::ptrdiff_t> &shape,
-    const std::size_t node_begin,
-    const std::size_t node_end,
+    const std::size_t height,
+    const std::size_t width,
+    const std::size_t y_begin,
+    const std::size_t y_end,
     std::vector<Stats> &stats
 ) {
-    const auto spatial_strides = bioimage_cpp::detail::c_order_strides(shape);
-    const auto number_of_nodes = static_cast<std::uint64_t>(number_of_pixels(shape));
+    const auto number_of_nodes = static_cast<std::uint64_t>(height * width);
     for (std::size_t channel = 0; channel < offsets.size(); ++channel) {
-        const auto channel_offset = static_cast<std::uint64_t>(channel) * number_of_nodes;
-        for (std::uint64_t node = node_begin; node < node_end; ++node) {
-            std::uint64_t target = 0;
-            if (!bioimage_cpp::detail::valid_offset_target(node, offsets[channel], shape, spatial_strides, target)) {
-                continue;
+        const auto &off = offsets[channel];
+        const auto channel_offset =
+            static_cast<std::uint64_t>(channel) * number_of_nodes;
+        sweep_offset_box_2d(
+            off[0], off[1], height, width, y_begin, y_end,
+            [&](const std::uint64_t node, const std::uint64_t target) {
+                const auto u = label_at(labels, node);
+                const auto v = label_at(labels, target);
+                const auto edge = edge_for_labels(rag, u, v);
+                if (edge >= 0) {
+                    stats[static_cast<std::size_t>(edge)].add(
+                        affinities[channel_offset + node]
+                    );
+                }
             }
-            const auto edge = edge_for_labels(rag, label_at(labels, node), label_at(labels, target));
-            if (edge >= 0) {
-                stats[static_cast<std::size_t>(edge)].add(affinities[channel_offset + node]);
+        );
+    }
+}
+
+template <class LabelT, class ValueT, class Stats>
+void scan_affinity_3d_chunk(
+    const RegionAdjacencyGraph &rag,
+    const LabelT *labels,
+    const ValueT *affinities,
+    const std::vector<std::vector<std::ptrdiff_t>> &offsets,
+    const std::size_t depth,
+    const std::size_t height,
+    const std::size_t width,
+    const std::size_t z_begin,
+    const std::size_t z_end,
+    std::vector<Stats> &stats
+) {
+    const auto slice_size = height * width;
+    const auto number_of_nodes = static_cast<std::uint64_t>(depth * slice_size);
+    for (std::size_t channel = 0; channel < offsets.size(); ++channel) {
+        const auto &off = offsets[channel];
+        const auto channel_offset =
+            static_cast<std::uint64_t>(channel) * number_of_nodes;
+        sweep_offset_box_3d(
+            off[0], off[1], off[2], depth, height, width, z_begin, z_end,
+            [&](const std::uint64_t node, const std::uint64_t target) {
+                const auto u = label_at(labels, node);
+                const auto v = label_at(labels, target);
+                const auto edge = edge_for_labels(rag, u, v);
+                if (edge >= 0) {
+                    stats[static_cast<std::size_t>(edge)].add(
+                        affinities[channel_offset + node]
+                    );
+                }
             }
-        }
+        );
     }
 }
 
@@ -408,40 +546,79 @@ void accumulate_affinity_features(
         throw std::invalid_argument("out shape must be (number_of_edges, number_of_features)");
     }
 
-    const auto number_of_nodes = detail_features::number_of_pixels(labels.shape);
-    const auto n_threads = detail::normalize_thread_count(number_of_threads, number_of_nodes);
+    const auto work_items = static_cast<std::size_t>(labels.shape[0]);
+    const auto n_threads = detail::normalize_thread_count(number_of_threads, work_items);
     const auto number_of_edges = static_cast<std::size_t>(rag.number_of_edges());
 
+    BIOIMAGE_PROFILE_INIT(aff_profiler);
+
     const auto run_scan = [&](auto &per_thread_stats) {
+        BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:scan");
         bioimage_cpp::detail::parallel_for_chunks(
             n_threads,
-            number_of_nodes,
+            work_items,
             [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) {
-                detail_features::scan_affinity_chunk(
-                    rag, labels.data, affinities.data, offsets, labels.shape,
-                    begin, end, per_thread_stats[thread_id]
-                );
+                if (labels.ndim() == 2) {
+                    detail_features::scan_affinity_2d_chunk(
+                        rag, labels.data, affinities.data, offsets,
+                        static_cast<std::size_t>(labels.shape[0]),
+                        static_cast<std::size_t>(labels.shape[1]),
+                        begin, end, per_thread_stats[thread_id]
+                    );
+                } else {
+                    detail_features::scan_affinity_3d_chunk(
+                        rag, labels.data, affinities.data, offsets,
+                        static_cast<std::size_t>(labels.shape[0]),
+                        static_cast<std::size_t>(labels.shape[1]),
+                        static_cast<std::size_t>(labels.shape[2]),
+                        begin, end, per_thread_stats[thread_id]
+                    );
+                }
             }
         );
     };
 
     if (compute_complex_features) {
-        std::vector<std::vector<detail_features::ComplexStats>> per_thread_stats(
-            n_threads,
-            std::vector<detail_features::ComplexStats>(number_of_edges)
-        );
+        std::vector<std::vector<detail_features::ComplexStats>> per_thread_stats;
+        {
+            BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:alloc");
+            per_thread_stats.assign(
+                n_threads,
+                std::vector<detail_features::ComplexStats>(number_of_edges)
+            );
+        }
         run_scan(per_thread_stats);
-        auto stats = detail_features::merge_stats(per_thread_stats, number_of_edges);
-        detail_features::write_complex_features(stats, out);
+        std::vector<detail_features::ComplexStats> stats;
+        {
+            BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:merge");
+            stats = detail_features::merge_stats(per_thread_stats, number_of_edges);
+        }
+        {
+            BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:write");
+            detail_features::write_complex_features(stats, out);
+        }
     } else {
-        std::vector<std::vector<detail_features::SimpleStats>> per_thread_stats(
-            n_threads,
-            std::vector<detail_features::SimpleStats>(number_of_edges)
-        );
+        std::vector<std::vector<detail_features::SimpleStats>> per_thread_stats;
+        {
+            BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:alloc");
+            per_thread_stats.assign(
+                n_threads,
+                std::vector<detail_features::SimpleStats>(number_of_edges)
+            );
+        }
         run_scan(per_thread_stats);
-        auto stats = detail_features::merge_stats(per_thread_stats, number_of_edges);
-        detail_features::write_simple_features(stats, out);
+        std::vector<detail_features::SimpleStats> stats;
+        {
+            BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:merge");
+            stats = detail_features::merge_stats(per_thread_stats, number_of_edges);
+        }
+        {
+            BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:write");
+            detail_features::write_simple_features(stats, out);
+        }
     }
+
+    BIOIMAGE_PROFILE_REPORT(aff_profiler);
 }
 
 } // namespace bioimage_cpp::graph