Flow optim

CMakeLists.txt

@@ -36,6 +36,36 @@ target_compile_options(_core PRIVATE
     $<$<NOT:$<CONFIG:Debug>>:-O3>
 )
 
+option(
+    BIOIMAGE_FLOW_FMA_DISPATCH
+    "Build an x86 FMA-specialized flow tracer with runtime CPU dispatch"
+    ON
+)
+if(BIOIMAGE_FLOW_FMA_DISPATCH)
+    string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _bioimage_system_processor)
+    if(_bioimage_system_processor MATCHES "^(x86_64|amd64|i[3-6]86)$")
+        if(MSVC)
+            target_sources(_core PRIVATE src/cpp/flow/flow_density_fma.cxx)
+            set_property(
+                SOURCE src/cpp/flow/flow_density_fma.cxx
+                APPEND PROPERTY COMPILE_OPTIONS /arch:AVX2
+            )
+            target_compile_definitions(
+                _core PRIVATE
+                BIOIMAGE_FLOW_FMA_DISPATCH=1
+                BIOIMAGE_FLOW_FMA_REQUIRES_AVX2=1
+            )
+        elseif(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
+            target_sources(_core PRIVATE src/cpp/flow/flow_density_fma.cxx)
+            set_property(
+                SOURCE src/cpp/flow/flow_density_fma.cxx
+                APPEND PROPERTY COMPILE_OPTIONS -mavx -mfma
+            )
+            target_compile_definitions(_core PRIVATE BIOIMAGE_FLOW_FMA_DISPATCH=1)
+        endif()
+    endif()
+endif()
+
 option(BIOIMAGE_PROFILE "Enable per-phase profiling instrumentation (development only)" OFF)
 if(BIOIMAGE_PROFILE)
     target_compile_definitions(_core PRIVATE BIOIMAGE_PROFILE)

development/flow/benchmark_interleave.py

@@ -0,0 +1,126 @@
+"""Benchmark the lockstep trajectory interleave against adversarial inputs.
+
+The RK2 tracer processes K=3 particles per worker in lockstep
+(``trace_particle_block``), so the out-of-order core can overlap the
+independent trajectories' dependent chains. Two input families stress it:
+
+- ``stripes``: lane-divergence worst case. Zero flow everywhere except every
+  fourth x-column, which carries an in-plane swirl of magnitude
+  ``2*tol/dt`` — those particles never converge and never leave the mask,
+  while their group neighbours converge on the first step, so a lockstep
+  group idles most of its lanes for all 50 iterations.
+- ``random``: irregular trajectories with unpredictable per-lane branching
+  (same generator as ``benchmark_midpoint_reuse.py``).
+
+Compare kernels by running this script under two installed builds (swap the
+prebuilt ``_core*.so`` between runs — rebuild rounds drift thermally).
+
+Run::
+
+    python development/flow/benchmark_interleave.py --repeats 3
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from statistics import median
+from time import perf_counter
+
+import numpy as np
+
+import bioimage_cpp as bic
+from bioimage_cpp._data import load_flow_data
+
+
+SHAPE_3D = (32, 192, 192)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Benchmark the lockstep trajectory interleave."
+    )
+    parser.add_argument("--repeats", type=int, default=3)
+    parser.add_argument("--threads", type=int, default=1)
+    parser.add_argument("--timeout", type=float, default=60.0)
+    parser.add_argument(
+        "--skip-fixture",
+        action="store_true",
+        help="Only run the synthetic cases (no registered-fixture download).",
+    )
+    return parser.parse_args()
+
+
+def _time(flow: np.ndarray, mask: np.ndarray, repeats: int, threads: int) -> tuple[np.ndarray, list[float]]:
+    timings = []
+    result = None
+    for _ in range(repeats):
+        start = perf_counter()
+        result = bic.flow.compute_flow_density(
+            flow, mask, sigma=None, number_of_threads=threads
+        )
+        timings.append(perf_counter() - start)
+    assert result is not None
+    return result, timings
+
+
+def _report(label: str, result: np.ndarray, timings: list[float]) -> None:
+    print(
+        f"{label}: median={median(timings):.4f}s min={min(timings):.4f}s "
+        f"particles={int(result.sum())}"
+    )
+
+
+def _stripe_flow(n_orbiter_columns_of_4: int) -> np.ndarray:
+    """Swirling flow on every n-of-4 x-columns, zero elsewhere."""
+    zz, yy, xx = np.indices(SHAPE_3D, dtype=np.float32)
+    cy = (SHAPE_3D[1] - 1) / 2.0
+    cx = (SHAPE_3D[2] - 1) / 2.0
+    dy, dx = yy - cy, xx - cx
+    radius = np.hypot(dy, dx)
+    radius[radius == 0] = 1.0
+    # |dt * step| = 0.01 >= tol = 0.005 at the defaults: orbiters never
+    # converge; the rotation keeps them inside the mask for all iterations.
+    speed = 0.05
+    flow = np.stack(
+        [np.zeros(SHAPE_3D, np.float32), -dx / radius * speed, dy / radius * speed]
+    ).astype(np.float32)
+    orbiter = (xx.astype(np.int64) % 4) < n_orbiter_columns_of_4
+    flow *= orbiter[None]
+    return flow
+
+
+def main() -> int:
+    args = parse_args()
+    if args.repeats < 1:
+        print("--repeats must be >= 1", file=sys.stderr)
+        return 2
+
+    if not args.skip_fixture:
+        for ndim in (2, 3):
+            dist, fg, _ = load_flow_data(ndim, timeout=args.timeout)
+            flow = np.ascontiguousarray(-dist, dtype=np.float32)
+            mask = np.ascontiguousarray(fg > 0.5)
+            result, timings = _time(flow, mask, args.repeats, args.threads)
+            _report(f"fixture {ndim}D threads={args.threads}", result, timings)
+
+    mask = np.ones(SHAPE_3D, dtype=bool)
+    for n_orbiters in (1, 3):
+        flow = _stripe_flow(n_orbiters)
+        result, timings = _time(flow, mask, args.repeats, args.threads)
+        _report(
+            f"stripes {n_orbiters}/4 long-lived threads={args.threads}",
+            result,
+            timings,
+        )
+
+    rng = np.random.default_rng(0)
+    for scale in (1.0, 10.0):
+        flow = rng.normal(scale=scale, size=(3,) + SHAPE_3D).astype(np.float32)
+        result, timings = _time(flow, mask, args.repeats, args.threads)
+        _report(f"random scale={scale:g} threads={args.threads}", result, timings)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

development/flow/benchmark_midpoint_reuse.py

@@ -0,0 +1,98 @@
+"""Benchmark the RK2 midpoint cell-reuse against adversarial flow magnitudes.
+
+The RK2 midpoint moves ``0.5 * dt * |flow|`` voxels from the current position.
+With the default ``dt=0.2`` the tracer's same-cell reuse branch is nearly
+always taken for unit-magnitude flows (scale 1), becomes unpredictable around
+scale 10 (midpoint displacement ~1 voxel), and is nearly always a reload for
+larger scales. This script sweeps that range with synthetic random flows and
+also times the registered fixtures, so the reuse can be compared against the
+previous kernel by running the script under both builds.
+
+Run::
+
+    python development/flow/benchmark_midpoint_reuse.py --repeats 3
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from statistics import median
+from time import perf_counter
+
+import numpy as np
+
+import bioimage_cpp as bic
+from bioimage_cpp._data import load_flow_data
+
+
+SCALES = (1.0, 5.0, 10.0, 20.0, 40.0)
+SHAPE_3D = (32, 192, 192)
+SHAPE_2D = (512, 512)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Benchmark RK2 midpoint cell-reuse on adversarial flow scales."
+    )
+    parser.add_argument("--repeats", type=int, default=3)
+    parser.add_argument("--threads", type=int, default=1)
+    parser.add_argument("--timeout", type=float, default=60.0)
+    parser.add_argument(
+        "--skip-fixture",
+        action="store_true",
+        help="Only run the synthetic sweep (no registered-fixture download).",
+    )
+    return parser.parse_args()
+
+
+def _time(flow: np.ndarray, mask: np.ndarray, repeats: int, threads: int) -> tuple[np.ndarray, list[float]]:
+    timings = []
+    result = None
+    for _ in range(repeats):
+        start = perf_counter()
+        result = bic.flow.compute_flow_density(
+            flow, mask, sigma=None, number_of_threads=threads
+        )
+        timings.append(perf_counter() - start)
+    assert result is not None
+    return result, timings
+
+
+def _report(label: str, result: np.ndarray, timings: list[float]) -> None:
+    print(
+        f"{label}: median={median(timings):.4f}s min={min(timings):.4f}s "
+        f"particles={int(result.sum())}"
+    )
+
+
+def main() -> int:
+    args = parse_args()
+    if args.repeats < 1:
+        print("--repeats must be >= 1", file=sys.stderr)
+        return 2
+
+    if not args.skip_fixture:
+        for ndim in (2, 3):
+            dist, fg, _ = load_flow_data(ndim, timeout=args.timeout)
+            flow = np.ascontiguousarray(-dist, dtype=np.float32)
+            mask = np.ascontiguousarray(fg > 0.5)
+            result, timings = _time(flow, mask, args.repeats, args.threads)
+            _report(f"fixture {ndim}D threads={args.threads}", result, timings)
+
+    rng = np.random.default_rng(0)
+    mask_3d = np.ones(SHAPE_3D, dtype=bool)
+    mask_2d = np.ones(SHAPE_2D, dtype=bool)
+    for scale in SCALES:
+        flow = rng.normal(scale=scale, size=(3,) + SHAPE_3D).astype(np.float32)
+        result, timings = _time(flow, mask_3d, args.repeats, args.threads)
+        _report(f"sweep 3D scale={scale:g} threads={args.threads}", result, timings)
+
+    flow = rng.normal(scale=10.0, size=(2,) + SHAPE_2D).astype(np.float32)
+    result, timings = _time(flow, mask_2d, args.repeats, args.threads)
+    _report(f"sweep 2D scale=10 threads={args.threads}", result, timings)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())