Implement parallek dijkstra

Author: constantinpape

.gitignore

@@ -28,6 +28,7 @@ venv/
 # Data
 *.h5
 *.npz
+*.mrc
 
 CLAUDE.md
 pypi-release.txt

MIGRATION_GUIDE.md

@@ -2355,6 +2355,7 @@ field = bic.distance.dijkstra_distance_field(
     sources,                         # (N, mask.ndim) integer coordinates
     connectivity=None,              # full connectivity by default
     spacing=(2.0, 1.0, 1.0),
+    number_of_threads=4,
 )
 
 # Flat C-order predecessor indices can be requested with the same solve.
@@ -2379,7 +2380,13 @@ All costs must be finite and non-negative. Distances are `float64`; background
 and unreachable voxels are `+inf`. Predecessors are `int64` with the mask shape:
 sources point to their own flat C-order index and background/unreachable voxels
 contain `-1`. Paths are `int64` NumPy-axis-order coordinates from the source to
-the selected target. Ties are deterministic and use flat C-order indices.
+the selected target. Target ties are deterministic and use flat C-order indices.
+`number_of_threads=1` retains the specialized sequential heaps; `0` uses
+hardware concurrency. Threaded solves use exact FP64 delta stepping: distance
+fields are bitwise-identical to the sequential result, while an equal-cost
+predecessor or path may differ from the sequential heap but remains deterministic.
+Small workloads and underperforming cost modes automatically retain the
+sequential implementation.
 
 This Dijkstra field is deliberately distinct from
 `geodesic_distance_field` below. Dijkstra is exact for the chosen 4/8- or
@@ -2526,6 +2533,7 @@ vertices, edges, radii = bic.skeleton.teasar(
     constant=0,
     pdrf_scale=100000,
     pdrf_exponent=4,
+    number_of_threads=4,
 )
 ```
 
@@ -2548,10 +2556,11 @@ Important differences and current scope:
   cross-section metadata, or postprocessing heuristics. These differences can
   change branch positions and vertex counts, so output is not expected to be
   vertex-for-vertex identical to kimimaro.
-- The C++ core is dependency-free and single-threaded. TEASAR uses an internal
-  FP64 compact-foreground Dijkstra backend for root fields and penalized rail
-  paths. Compact IDs retain C-order tie-breaking, so it is bitwise identical to
-  the dense public grid solver on the tested TEASAR workloads while avoiding
+- The C++ core remains dependency-free. `number_of_threads=1` is the default;
+  `0` uses hardware concurrency. One thread budget is shared by the exact
+  distance transform and the internal FP64 compact-foreground Dijkstra backend.
+  Dependent root sweeps and rails remain ordered, while sufficiently large
+  shortest-path wavefronts use deterministic delta stepping. Compact IDs avoid
   full-volume shortest-path fields. The public Dijkstra functions remain dense
   and exact FP64.
 

development/distance/DIJKSTRA_PERFORMANCE_NOTES.md

@@ -542,3 +542,47 @@ predecessor parity is intentionally relaxed, retain the previously used TEASAR
 gate: matching contracted branch/leaf topology, bidirectional 95th-percentile
 distance at most one normalized voxel, Hausdorff distance at most two voxels,
 and total physical length within 2% of sequential compact FP64.
+
+## Implemented parallel follow-up
+
+Implemented after reviewing the draft above. The implementation scope was
+expanded to the dense public Dijkstra fields as well as the compact TEASAR
+backend. Both use one deterministic FP64 delta-stepping primitive in
+`distance/detail/delta_stepping.hxx`; the optimized heaps remain unchanged.
+
+The concrete design closes several gaps in the draft:
+
+- frontier nodes and proposals are processed in fixed global batches independent
+  of thread count, with at most 1,048,576 worst-case proposals per batch;
+- workers write only per-thread buffers, then the calling thread sorts, reduces,
+  and applies strict improvements in a deterministic order;
+- equal-distance updates do not rewrite existing predecessors, preventing cycles
+  in zero-weight components;
+- the heavy phase uses the unique set removed during the complete light closure;
+- early stopping happens only after both phases and after checking that FP64
+  rounding did not reinsert work into the current bucket;
+- generation-marked compact state avoids clearing every rail-search array;
+- invalid bucket widths or bucket-index overflow restart on the sequential heap.
+
+The bucket width is the minimum physical neighbor length for physical costs, the
+deterministically sampled positive-cost median for node costs, and their product
+for node-times-physical costs. All-zero costs use a positive sentinel width and
+remain entirely in the light closure.
+
+Initial measurements on the four-core development machine also rejected the
+assumption that every large one-source solve benefits from delta stepping.
+Representative one-shot default-tier dense physical fields took about 0.38 s on
+the heap versus 0.64 s with four-thread staged relaxation, while a full source
+plane on the same `64 x 160 x 160` domain improved from about 0.72 s to 0.30 s.
+The directed-node heap remained decisively faster. TEASAR's compact wavefronts
+were likewise too narrow through the `256^3` tier, although threading its EDT
+still improved end-to-end time.
+
+The production dispatch therefore follows the benchmark rather than forcing the
+new backend: one-thread calls, paths, directed-node costs, small fields, and
+narrow multi-source fields retain the heap. Broad multi-source physical and
+node-times-physical fields use delta stepping. Compact TEASAR keeps heap Dijkstra
+through the measured tiers while sharing the requested thread budget with EDT;
+the delta backend remains available for much larger compact domains. The
+development benchmarks now accept thread matrices so these thresholds can be
+revisited with evidence on other machines.

development/distance/benchmark_dijkstra.py

@@ -64,7 +64,7 @@ def time_call(function, repeats: int, warmup: int) -> list[float]:
     return samples
 
 
-def build_cases(shape: tuple[int, ...], include_geodesic: bool):
+def build_cases(shape: tuple[int, ...], include_geodesic: bool, number_of_threads: int):
     mask = make_mask(shape)
     source = np.full(len(shape), 2, dtype=np.int64)
     target = np.asarray(shape, dtype=np.int64) - 3
@@ -79,25 +79,48 @@ def build_cases(shape: tuple[int, ...], include_geodesic: bool):
         (
             "dijkstra/physical-field",
             lambda: bic.distance.dijkstra_distance_field(
-                mask, source, spacing=spacing
+                mask, source, spacing=spacing, number_of_threads=number_of_threads
             ),
         ),
         (
             "dijkstra/field+parents",
             lambda: bic.distance.dijkstra_distance_field(
-                mask, source, spacing=spacing, return_predecessors=True
+                mask,
+                source,
+                spacing=spacing,
+                return_predecessors=True,
+                number_of_threads=number_of_threads,
             ),
         ),
         (
             "dijkstra/node-field",
             lambda: bic.distance.dijkstra_distance_field(
-                mask, source, costs=costs, cost_mode="node"
+                mask,
+                source,
+                costs=costs,
+                cost_mode="node",
+                number_of_threads=number_of_threads,
+            ),
+        ),
+        (
+            "dijkstra/node-times-physical-field",
+            lambda: bic.distance.dijkstra_distance_field(
+                mask,
+                source,
+                costs=costs,
+                cost_mode="node_times_physical",
+                spacing=spacing,
+                number_of_threads=number_of_threads,
             ),
         ),
         (
             "dijkstra/early-path",
             lambda: bic.distance.dijkstra_path(
-                mask, source, target, spacing=spacing
+                mask,
+                source,
+                target,
+                spacing=spacing,
+                number_of_threads=number_of_threads,
             ),
         ),
     ]
@@ -106,7 +129,10 @@ def build_cases(shape: tuple[int, ...], include_geodesic: bool):
             (
                 "geodesic/FMM-field",
                 lambda: bic.distance.geodesic_distance_field(
-                    mask, source, sampling=spacing, number_of_threads=1
+                    mask,
+                    source,
+                    sampling=spacing,
+                    number_of_threads=number_of_threads,
                 ),
             )
         )
@@ -121,6 +147,10 @@ def main() -> int:
     parser.add_argument("--repeats", type=int, default=7)
     parser.add_argument("--warmup", type=int, default=1)
     parser.add_argument("--include-geodesic", action="store_true")
+    parser.add_argument(
+        "--threads", type=int, nargs="+", default=[1],
+        help="thread counts to benchmark (0 uses hardware concurrency)",
+    )
     parser.add_argument("--json", default="", help="optional JSON result path")
     args = parser.parse_args()
     if args.repeats < 1 or args.warmup < 0:
@@ -134,29 +164,34 @@ def main() -> int:
         shapes = ((1024, 1024), (64, 160, 160))
 
     rows = []
-    header = f"{'case':>29} {'shape':>16} {'median ms':>11} {'min ms':>10} {'ns/fg voxel':>13}"
+    header = f"{'case':>36} {'threads':>7} {'shape':>16} {'median ms':>11} {'min ms':>10} {'ns/fg voxel':>13}"
     print(header)
     print("-" * len(header))
     for shape in shapes:
-        mask, cases = build_cases(shape, args.include_geodesic)
-        foreground = int(np.count_nonzero(mask))
-        for name, function in cases:
-            samples = time_call(function, args.repeats, args.warmup)
-            median_s = median(samples)
-            row = {
-                "case": name,
-                "shape": list(shape),
-                "foreground_voxels": foreground,
-                "samples_s": samples,
-                "median_s": median_s,
-                "min_s": min(samples),
-                "ns_per_foreground_voxel": median_s * 1e9 / foreground,
-            }
-            rows.append(row)
-            print(
-                f"{name:>29} {str(shape):>16} {median_s * 1e3:11.2f} "
-                f"{min(samples) * 1e3:10.2f} {row['ns_per_foreground_voxel']:13.1f}"
+        for number_of_threads in args.threads:
+            mask, cases = build_cases(
+                shape, args.include_geodesic, number_of_threads
             )
+            foreground = int(np.count_nonzero(mask))
+            for name, function in cases:
+                samples = time_call(function, args.repeats, args.warmup)
+                median_s = median(samples)
+                row = {
+                    "case": name,
+                    "number_of_threads": number_of_threads,
+                    "shape": list(shape),
+                    "foreground_voxels": foreground,
+                    "samples_s": samples,
+                    "median_s": median_s,
+                    "min_s": min(samples),
+                    "ns_per_foreground_voxel": median_s * 1e9 / foreground,
+                }
+                rows.append(row)
+                print(
+                    f"{name:>36} {number_of_threads:7d} {str(shape):>16} "
+                    f"{median_s * 1e3:11.2f} {min(samples) * 1e3:10.2f} "
+                    f"{row['ns_per_foreground_voxel']:13.1f}"
+                )
 
     if args.json:
         with open(args.json, "w", encoding="utf-8") as file:

development/skeleton/benchmark_teasar.py

@@ -151,7 +151,7 @@ def kimimaro_call(mask, spacing, scale, constant, pdrf_scale, pdrf_exponent):
     return skeletons[1]
 
 
-def bic_backend_call(mask, spacing, parameters, backend):
+def bic_backend_call(mask, spacing, parameters, backend, number_of_threads=1):
     """Call a development-only C++ backend without changing the public API."""
     return _core._teasar_uint8_backend(
         mask,
@@ -161,6 +161,7 @@ def bic_backend_call(mask, spacing, parameters, backend):
         parameters["pdrf_scale"],
         parameters["pdrf_exponent"],
         backend,
+        number_of_threads,
     )
 
 
@@ -178,6 +179,10 @@ def main() -> int:
         help="compare dense FP64, compact on-the-fly/CSR FP64, and CSR FP32",
     )
     parser.add_argument("--json", default="", help="optional JSON result path")
+    parser.add_argument(
+        "--threads", type=int, nargs="+", default=[1],
+        help="public TEASAR thread counts (0 uses hardware concurrency)",
+    )
     args = parser.parse_args()
     if args.repeats < 1 or args.warmup < 0:
         parser.error("--repeats must be >= 1 and --warmup must be >= 0")
@@ -218,10 +223,16 @@ def main() -> int:
     else:
         backends = [
             (
-                "bioimage-cpp",
-                lambda mask: bic.skeleton.teasar(mask, spacing=spacing, **parameters),
+                f"bioimage-cpp/t{number_of_threads}",
+                lambda mask, number_of_threads=number_of_threads: bic.skeleton.teasar(
+                    mask,
+                    spacing=spacing,
+                    number_of_threads=number_of_threads,
+                    **parameters,
+                ),
                 count_bic,
             )
+            for number_of_threads in args.threads
         ]
     if args.kimimaro:
         backends.append(
@@ -293,12 +304,12 @@ def main() -> int:
                 f"hausdorff={quality['hausdorff_voxels']:.3f} vox "
                 f"length_error={quality['relative_length_error']:.3%}"
             )
-        if len(case_rows) == 2 and {row["backend"] for row in case_rows} == {
-            "bioimage-cpp", "kimimaro"
-        }:
+        if len(args.threads) == 1 and len(case_rows) == 2 and {
+            row["backend"] for row in case_rows
+        } == {f"bioimage-cpp/t{args.threads[0]}", "kimimaro"}:
             by_name = {row["backend"]: row for row in case_rows}
             ratio = (
-                by_name["bioimage-cpp"]["median_s"]
+                by_name[f"bioimage-cpp/t{args.threads[0]}"]["median_s"]
                 / by_name["kimimaro"]["median_s"]
             )
             print(f"{'bioimage-cpp / kimimaro':>54}: {ratio:.2f}x")