diff --git a/CMakeLists.txt b/CMakeLists.txt index 7cc5323..4e9324f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -36,6 +36,36 @@ target_compile_options(_core PRIVATE $<$>:-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) diff --git a/development/flow/PERFORMANCE_NOTES.md b/development/flow/PERFORMANCE_NOTES.md index aa37054..9781f27 100644 --- a/development/flow/PERFORMANCE_NOTES.md +++ b/development/flow/PERFORMANCE_NOTES.md @@ -10,297 +10,448 @@ Test data, here and below, refers to the registered `flow_data_2d.h5` / Hardware on which these numbers were taken: **11th Gen Intel Core i7-1185G7 @ 3.00 GHz** (Tiger Lake, 4 physical cores / 8 SMT threads, 12 MB L3, AVX2 + -AVX-512 capable). +AVX-512 capable). The 2026-07-12 pass below was cross-checked against +GPT 5.6-Sol's independent investigation, which measured the same changes on an +AMD EPYC 7513 (Zen 3); see `FLOW_OPTIM.md` in the repo root. Absolute timings +differ between the hosts; the relative gains agree closely across both. ## Headline numbers +Current state after the 2026-07-12 scalar-kernel rewrite (particle-major + +direct interpolation + truncation + compile-time specialization). Timings are +`check_flow_density.py` medians of 5, measured in one session on the Tiger Lake +host. The "pre-rewrite" column is the previous kernel (generic `2^D` corner +table, iteration-major with a per-iteration alive scan): + ``` - baseline current gain -3D, 1 thread : 5.55 s → 5.40 s 1.03× -3D, 8 threads : 5.55 s → 1.51 s 3.67× -2D, 1 thread : 0.43 s → 0.45 s 0.96× -2D, 8 threads : 0.43 s → 0.18 s 2.4× - -Accuracy (3D vs reference): - rel_diff = mean(|ours - ref|) / mean(ref) = 0.021 (gate ≤ 0.15) - pearson = 0.998 + pre-rewrite current gain +3D, 1 thread : 5.66 s → 2.45 s 2.31× (−57 %) +3D, 8 threads : 1.51 s → 0.64 s 2.36× (−58 %) +2D, 1 thread : 0.63 s → 0.29 s 2.18× (−54 %) +2D, 8 threads : 0.18 s → 0.057 s 3.07× (−67 %) + +Accuracy vs reference (unchanged — the rewrite is bitwise identical): + 3D rel_diff = 0.0470 pearson = 0.9690 (gate ≤ 0.15) + 2D rel_diff = 0.0201 pearson = 0.9975 ``` -The big win is threading. Algorithmic changes only marginally help the -single-threaded path; they pay off because they cap the work and give the -threaded path something to scale. +The rewrite produces a **byte-for-byte identical** density on both fixtures: +`max_abs_diff`, `mean_abs_diff`, `rel_diff`, and `pearson` are unchanged from the +pre-rewrite build. It was also verified bitwise identical across 1056 randomized +small/edge-case inputs (see "Correctness" below). + +Two eras of work are recorded here: + +1. **Threading + early-exit era** (earlier). Threading was the big lever; the + algorithmic changes capped work so the threaded path had something to scale. +2. **Scalar-kernel rewrite** (2026-07-12). Roughly halved single-threaded *and* + multi-threaded runtime by cutting per-step instruction count and removing the + per-iteration synchronization, on the existing channel-first layout with no + SIMD, no relayout, and no precision change. ## Changes that landed -In rough order of impact: +In rough order of when they landed: 1. **Threading via `detail::parallel_for_chunks`** — particle tracing is embarrassingly parallel; the final density scatter stays single-threaded (`<1 %` of time per the profile). Determinism preserved. -2. **Convergence-based early exit** (`tol` parameter) — per-particle alive - flag; freeze a particle once `max_axis(|dt·step|) < tol`. Outer iteration - loop breaks once all particles are frozen. -3. **Mask-restricted tracing** (`restrict_to_mask` parameter) — a particle - that rounds to a background voxel freezes in place. Combined with (2) it - sheds slow trajectories that would otherwise run the full `n_iter`. -4. **RK2 (midpoint) integrator** (`method="rk2"`) — two flow samples per - step but allows larger `dt`; combined with early-exit it converges in - fewer iterations. -5. **Corner-table sharing inside `sample_channel`** — `compute_corners` - builds the `2^D` (offset, weight) table once per particle per iteration; - all `D` channel samples reuse it. Cuts the redundant lower/frac work that - the original `sample_linear_nearest` did per axis. -6. **Hoisted invariants** out of the per-particle loop (`upper[axis]`, - channel base pointers, channel stride). -7. **Profiling instrumentation** (`BIOIMAGE_PROFILE_SCOPE`) around `init`, - `iter_loop`, `scatter`, `mask_zero`. Free when the CMake flag is off; live - when `-C cmake.define.BIOIMAGE_PROFILE=ON`. - -New chosen defaults (see `src/bioimage_cpp/flow/_flow.py`): +2. **Convergence-based early exit** (`tol` parameter) — freeze a particle once + `max_axis(|dt·step|) < tol`. +3. **Mask-restricted tracing** (`restrict_to_mask` parameter) — a particle whose + proposed endpoint rounds to a background voxel freezes in place. Combined with + (2) it sheds slow trajectories that would otherwise run the full `n_iter`. +4. **RK2 (midpoint) integrator** (`method="rk2"`) — two flow samples per step but + allows larger `dt`; combined with early-exit it converges in fewer iterations. +5. **Profiling instrumentation** (`BIOIMAGE_PROFILE_SCOPE`) around `init`, + `iter_loop`, `scatter`, `mask_zero`. Free when the CMake flag is off. + +The 2026-07-12 scalar-kernel rewrite then replaced the generic `2^D` corner +table and the iteration-major loop (with its per-iteration alive scan) with: + +6. **Direct bilinear / trilinear interpolation** (`detail::sample_flow`) — + replaces the generic `2^D` corner-offset/weight table with explicit nested + lerps. The lower/upper index and fractional weight are computed once per axis + and shared across channels; no `2^D` product weights or offsets are + materialized. `if constexpr` selects the 2D vs 3D path. +7. **Truncation instead of `std::floor`** for the (clipped, nonnegative) sampling + and rounding coordinates. See "Why truncation is worth ~10-35 %" below. +8. **Particle-major traversal** (`detail::trace_particle` / `trace_all`) — the + integration loop moved *inside* one `parallel_for_chunks` fan-out. Each worker + traces its contiguous range of particles across all `n_iter` steps. Removes + the global `alive` vector, the up-to-`n_iter` thread create/join cycles, and + the per-iteration sequential alive scan. Trajectories are independent until the + sequential scatter, so results stay deterministic. +9. **Compile-time specialization** of the trace loop on the three loop-invariant + flags `(use_rk2, check_convergence, restrict_to_mask)` via an 8-way switch + dispatching to `trace_particle`. Hoists the per-step + branches on these flags out of the innermost loop. When `restrict_to_mask` is + set, the specialization also drops the start-of-step clip, which is provably + redundant there (committed endpoints already passed the in-bounds mask test, + and the seed is an in-bounds integer voxel). + +10. **Runtime-dispatched FMA specialization** for the common default mode + `(RK2, convergence, restrict_to_mask)`. On supported x86 CPUs the extension + dispatches once, before thread fan-out, to a separately compiled tracer that + uses scalar FMA contraction in the nested lerps. All other modes and CPUs use + the portable kernel. Only `src/cpp/flow/flow_density_fma.cxx` receives the + ISA flags (`-mavx -mfma` for GCC/Clang, `/arch:AVX2` for MSVC); the extension + and binding layer remain baseline-safe. The build can disable this path with + `-C cmake.define.BIOIMAGE_FLOW_FMA_DISPATCH=OFF`. + +11. **3-way lockstep trajectory interleave for RK2** + (`detail::trace_particle_block`, 2026-07-12). Each worker traces three + particles per group in lockstep instead of one trajectory to completion. + The trajectories are independent, so the out-of-order core fills one + lane's serial sample→lerp→update latency chain with the other lanes' + work; a converged/frozen lane costs one predictable branch per remaining + group iteration. Euler keeps the plain per-particle loop — its shorter + single-sample chain measured ~10 % *slower* interleaved. Results are + bitwise identical (per-trajectory arithmetic is untouched). See + "Trajectory interleave (2026-07-12)" below for measurements and K choice. + +Chosen defaults (see `src/bioimage_cpp/flow/_flow.py`) are unchanged: ```python n_iter=50, dt=0.2, tol=0.005, method="rk2", restrict_to_mask=True, number_of_threads=1 ``` -A caller who wants the threaded speedup passes `number_of_threads=8` (or -whatever value fits their box). +## Attribution of the 2026-07-12 rewrite -## Profile breakdown (1-thread, 3D, profile build) +Standalone C++ kernel harness, `-O3 -DNDEBUG` gcc-14, min-of-N, Tiger Lake. +Every row was verified bitwise identical to the pre-rewrite kernel on both +fixtures. `PM` = particle-major, `DIRECT` = direct interpolation, `TRUNC` = +truncation (the three toggles were measured in every combination): ``` -init ~0.4 % density-zero + positions collection -iter_loop ~99 % the particle tracing loop -scatter ~0.1 % final density write -mask_zero ~0.1 % final mask zeroing +PM DIRECT TRUNC | 3D 1T 3D 8T | 2D 1T 2D 8T + 0 0 0 | 5.45 1.34 | 0.598 0.118 (pre-rewrite baseline) + 0 0 1 | 4.89 1.23 | 0.391 0.103 + 0 1 0 | 4.35 1.01 | 0.501 0.125 + 0 1 1 | 3.79 0.90 | 0.374 0.088 + 1 0 0 | 4.48 1.05 | 0.469 0.085 + 1 0 1 | 3.95 0.87 | 0.319 0.059 + 1 1 0 | 3.18 0.68 | 0.424 0.072 + 1 1 1 | 2.51 0.55 | 0.303 0.054 (all three) ``` -The iteration loop dominates so completely that none of the other phases -were worth touching. A "single-pass init" cleanup and per-thread scatter -buffers were planned and skipped after the profile. - -## Thread scaling (3D, current defaults) +Single-thread marginal contributions from the pre-rewrite baseline: ``` -threads runtime speedup vs 1T - 1 5.40 s 1.00× - 2 3.76 s 1.44× - 4 2.07 s 2.61× - 6 1.71 s 3.17× - 8 1.52 s 3.57× + 3D 2D +direct interpolation ~20 % ~16 % +particle-major ~18 % ~22 % +truncation ~10 % ~35 % +all three combined ~54 % ~49 % ``` -Sub-linear from `1→4` (2.61× on 4 physical cores) and a marginal -improvement from `4→8` (SMT pair on each core). Likely contributors: - -- The per-iter `alive` scan is sequential. -- Single-thread TurboBoost is higher than all-core sustained clock. -- Memory subsystem contention even before bandwidth saturation - (L3 sharing, prefetcher load). - -It is **not** memory-bandwidth saturation — see the next section. - -## Bandwidth diagnosis (no `perf` available, used a streaming probe) - -Sustained streaming read bandwidth, measured via `numpy.sum` on float32 -arrays larger than L3: +- All three compose well and are independent wins. +- **In 2D, truncation is the single biggest lever** (~35 %): fewer corners (4 vs + 8) and channels (2 vs 3) make the floor operations a larger share of the work. +- Particle-major helps 8T more than 1T (barrier + alive-scan removal), so it + improves both raw throughput and scaling. + +On top of the three changes, **compile-time specialization (item 9) adds a +further ~6-9 % on 3D** (3D 1T 2.51 → ~2.29 s; total ~59 %). Isolating it: ~5.6 % +is from hoisting the three per-step flag branches, ~1 % from dropping the +redundant start-of-step clip. This is the same start-clip removal Sol tried as a +*runtime* `if (!restrict_to_mask)` and found to regress (~3.87 → 4.97 s): the +regression was the added invariant branch, not the removed clip. Done at compile +time there is no added branch, so it is a small positive instead. + +The rounding used by the per-step mask test (`round_to_flat_index`) also had a +`std::floor(x + 0.5f)`; it is now a truncating cast (0.5-2.6 %, bitwise safe by +the same nonnegativity argument). + +## Why truncation is worth ~10-35 % + +The gain is real *because of*, not despite, the portable wheel build. On the +baseline x86-64 / SSE2 target the wheels compile for (the project forbids +`-march=native`), gcc-14 lowers `std::floor` to an **inlined ~15-instruction +software routine with a branch** — `roundss` requires SSE4.1, which the baseline +target does not assume. A truncating `static_cast` is a single +`cvttss2si`. Every sampling coordinate is clipped to `[0, shape-1]` before use, +so it is nonnegative and truncation equals `std::floor` exactly. + +Consequences: + +- The win **holds on the shipped wheels** (same SSE2 target). It would only + shrink if the build enabled SSE4.1. +- Keep truncation local to helpers whose contract requires clipped, nonnegative + coordinates (`sample_flow`, `round_to_flat_index`). Do not push it into a + general sampler that might later be called with negative positions. + +## Particle-major load balance (adversarial check) + +Particle-major assigns each thread a static contiguous range for the whole trace. +The concern is that spatially-clustered slow trajectories could leave threads +idle. Measured on a worst-case synthetic fixture — `(48,512,512)`, all +foreground, zero flow for `z<24` (converges at iteration 0) and a small constant +in-mask drift for `z>=24` (runs all 50 steps). In C-order every slow particle is +in the upper half of the index range, so with 8 static chunks the lower 4 threads +are idle: ``` -array size best wall GB/s - 16 MB 1.44 ms 11.67 - 64 MB 5.83 ms 11.52 - 144 MB 13.08 ms 11.54 (matches the 3D flow array layout) - 256 MB 23.63 ms 11.36 + 1T 8T 8T scaling +baseline (iter-major) 44.4 s 12.2 s 3.65× +opt (particle-major) 22.2 s 6.0 s 3.70× ``` -Sustained read bandwidth on this machine: **≈ 11.5 GB/s**. - -Upper-bound bytes read from the flow array per `compute_flow_density` call, -3D RK2 path: - -``` -729 236 particles × 50 iter × 2 samples (RK2) × 3 channels × 8 corners × 4 B - ≈ 7.0 GB (upper bound — convergence early-exit reduces this in practice) -``` +Particle-major scales **as well** as iteration-major even in this worst case, and +is ~2× faster overall. This matches the theory: iteration-major pays a +create/join barrier on every one of the 50 steps, so it is gated by +`Σ_iter max_thread(work)`, whereas particle-major has one barrier and is gated by +`max_thread(Σ_iter work)`, and `max-of-sums ≤ sum-of-maxes`. The ~3.7× ceiling is +imposed by the shared static partition (4 idle threads), identical for both, not +by particle-major. **A dynamic scheduler is not warranted** on this evidence. -Effective kernel bandwidth, against that upper bound: +Realistic-fixture scaling (optimized kernel, 3D, `check_flow_density.py` median): ``` -1 thread @ 5.40 s → 1.3 GB/s = 11 % of sustained streaming -8 threads @ 1.52 s → 4.6 GB/s = 40 % of sustained streaming +threads runtime speedup vs 1T + 1 2.46 s 1.00× + 2 1.33 s 1.85× + 4 0.73 s 3.37× + 6 0.66 s 3.73× + 8 0.54 s 4.56× ``` -So even at 8 threads the kernel only utilises ~40 % of the memory subsystem. -The conclusion is that **the hot loop is compute-bound**, not -bandwidth-bound. Almost certainly bound by gather latency / per-particle -serial instruction count, not RAM throughput. - -## SIMD experiment (rolled back) - -GCC/Clang `__attribute__((target_clones("default,arch=haswell,arch=skylake-avx512")))` -was wired up end-to-end so the IFUNC resolver would pick the AVX2 / -AVX-512F clone at first call on capable CPUs. Wheel-portability was -preserved: the attribute is only enabled on Linux x86_64 GCC/Clang; macOS -(Mach-O lacks IFUNC) and MSVC fall back to a single default clone. - -The dispatcher was confirmed to work — the IFUNC resolver was emitted, and -the three clones diverged in their object code once the templated kernel was -forced inline via `[[gnu::always_inline]]` + `[[gnu::flatten]]`. Result on -the AVX2 clone: - -- ✓ `vfmadd231ss` + ymm registers in scalar FMA encoding. -- ✗ No packed-vector ops (`*ps` instructions, `vfmadd*ps`). -- ✗ No `vpgatherdd` or `vgatherqps`. - -GCC's autovectorizer evaluated the `sample_channel` corner-sum (8 gather -loads + multiply-add) and decided the gather pattern wasn't profitable. -Tried in sequence: - -1. `target_clones` alone — all clones folded to the same code (the template - was called via a normal function boundary, target attributes did not - propagate). Net: no change. -2. `[[gnu::always_inline]] inline` on the template + `[[gnu::flatten]]` on - the dispatch wrappers — clones diverged but only emitted scalar FMA. - 8-thread runtime regressed from 1.7 s to 2.3 s (icache pressure from - three full inlined copies in the same TU). -3. `#pragma omp simd reduction(+:value)` on the corner sum + `-fopenmp-simd` - in CMake — packed SIMD still did not materialise. -4. Narrowing corner offsets from `ptrdiff_t` to `int32_t` to enable 8-wide - `vpgatherdd` — no codegen change observed. - -Final decision: **roll back**. The complexity was non-zero (a new TU, IFUNC -plumbing, force-inline annotations on the templated kernel, an `omp-simd` -flag) and the 8-thread regression was real. - -What would actually work, but was out of scope for the session: hand-written -AVX2 intrinsics for `sample_channel<3>` (and `<2>`), specifically: - -```cpp -__attribute__((target("avx2,fma"))) -inline float sample_channel_avx2_3d(const float *channel, - const SamplingCorners<3> &c) { - const __m256 w = _mm256_loadu_ps(c.weights.data()); - const __m256i o = _mm256_loadu_si256( - reinterpret_cast(c.offsets.data())); - const __m256 v = _mm256_i32gather_ps(channel, o, sizeof(float)); - const __m256 p = _mm256_mul_ps(w, v); - // horizontal reduce - __m128 lo = _mm256_castps256_ps128(p); - __m128 hi = _mm256_extractf128_ps(p, 1); - __m128 s4 = _mm_add_ps(lo, hi); - __m128 s2 = _mm_add_ps(s4, _mm_movehl_ps(s4, s4)); - __m128 s1 = _mm_add_ss(s2, _mm_shuffle_ps(s2, s2, 1)); - return _mm_cvtss_f32(s1); -} -``` +Better than the pre-rewrite kernel's 3.57× at 8T; barrier and alive-scan removal +let it scale further before hitting the 4-physical-core / SMT ceiling. + +## Runtime FMA dispatch (2026-07-12) + +The direct bilinear/trilinear sampler is a chain of scalar lerps. On baseline +x86-64 those lerps compile as separate multiply/add operations; an FMA-enabled +build contracts them and provides a further portable-at-runtime speedup without +changing the flow layout or API. + +Paired `check_flow_density.py` medians on the AMD EPYC 7513, from the same source +revision and build environment: + +| Fixture | Threads | Dispatch OFF | Dispatch ON | Improvement | +|---|---:|---:|---:|---:| +| 3D | 1 | 1.4900 s | 1.3084 s | 12.2% | +| 3D | 8 | 0.3035 s | 0.2553 s | 15.9% | +| 2D | 1 | 0.2021 s | 0.1801 s | 10.9% | + +The full registered 3D density remained byte-for-byte identical, both stored +reference checks passed with unchanged metrics, and the complete test suite +reported 1057 passed / 8 skipped. A further 400 randomized 2D/3D default-mode +cases, including axis-of-length-one shapes and 1/4 threads, produced identical +aggregate output digests with dispatch enabled and disabled. + +Implementation details that matter: + +- Runtime feature detection happens in the baseline-compiled caller before the + specialized function is entered. +- GCC/Clang require AVX+FMA; MSVC uses `/arch:AVX2` and therefore also checks AVX2 + before dispatch. +- Only the common selector 7 specialization is duplicated. Less common Euler, + no-convergence, and unrestricted modes keep using the portable implementation. +- The FMA translation unit uses a distinct `CodegenVariant` template argument. + Without a distinct mangled name, linker COMDAT selection can silently retain + the portable `trace_all` instantiation and discard the FMA implementation. +- Automatic `target_clones` remains rejected: putting a target boundary around + the driver or chunk inhibited the hot-loop inlining and regressed runtime. + +## Trajectory interleave (2026-07-12) + +The rejected midpoint-reuse experiment showed the per-step loads are +latency-hidden *within* one trajectory — so the remaining 1T lever is ILP +*across* trajectories. `trace_particle_block` traces K consecutive +particles in lockstep (local position/alive lane state, per-step body identical +to `trace_particle`); the RK2 selectors use K=3, with a plain `trace_particle` +remainder loop. + +Measured with tightly paired `.so`-swap ABA runs (K1 = per-particle loop, +`paired_bench.py`, min-of-3 per invocation, Tiger Lake). Two flag regimes, +because the conda toolchain exports default `CXXFLAGS` +(`-march=nocona -mtune=haswell -ftree-vectorize ...`) while wheel builds get +plain flags — see the pitfall below: + +| Workload (RK2 defaults) | conda flags | plain flags | +|---|---:|---:| +| registered 3D, 1T | **−8.6 %** (1.93 → 1.76 s) | **−13 %** | +| registered 2D, 1T | **−27 %** (0.33 → 0.24 s) | **−28 %** | +| registered 3D, 8T | −3.5 % | ~parity | +| euler 3D, 1T (gated to K=1) | parity | — | +| stripes 1/4 long-lived, 1T (lane-divergence worst case) | +1.2 % | parity | +| stripes 3/4 long-lived, 1T | −1.6 % | — | +| random flow s=1 / s=10, 1T | −7 % / −10 % | — | + +- **K choice**: K=3 and K=4 both beat K=1 by 6–16 % on the 3D fixture in the + triage; K=2 was inside noise. Head-to-head, K=3 beat K=4 by ~4 % and carries + less lane state, so K=3 landed. GCC 14 does *not* unroll the K-lane loop — + the ILP comes from the out-of-order window overlapping loop iterations, and + that was measurably sufficient; manual unrolling was not needed. +- **2D vs 3D**: the 2D win (~27 %) dwarfs 3D — 2D has less per-lane state + (position + 4-corner sampling), so the interleave adds ILP without register + pressure. +- **Lane divergence** (the lockstep concern: a group lives as long as its + slowest lane) is benign: the stripe worst case — one never-converging + orbiter per group of 4, neighbours converging on step 1 — costs ~1 %, + because dead lanes skip their step body and only pay an alive-bit branch. +- **Adversarial cases** live in `development/flow/benchmark_interleave.py`. +- Verified bitwise identical: 320-case randomized digest harness, registered + 2D/3D fixtures at 1/4/8 threads, full suite (1070 tests). + +### Build-flag pitfall (affects all local benchmarking) + +The micromamba GCC 14 environment exports `CXXFLAGS`. CMake appends it to +every TU, so *normal* local builds are `-march=nocona -mtune=haswell +-ftree-vectorize ...` on top of the project's `-O3`. Setting `CXXFLAGS=` +for an experiment build **replaces** those defaults rather than adding to +them, silently changing codegen of every TU — an early triage here compared +mixed flag regimes before this was caught. Wheel builds (cibuildwheel, no +conda) get the plain flags, so plain-flag numbers are the shipping-relevant +ones. When A/B-benchmarking local builds: `echo $CXXFLAGS` first, and keep +the regime identical on both sides. + +## Correctness + +- The current suite reports 1057 passed / 8 skipped. The 17 + `tests/test_flow.py` cases cover + 2D/3D, Euler/RK2, mask on/off, convergence on/off, single-vs-multithread + equality, non-contiguous inputs, degenerate `n_iter`, and invalid inputs. +- A differential harness ran 1056 randomized cases — axis-of-length-1 shapes, + tiny grids, zero/integer-aligned flows, Euler and RK2, `tol=0` and `>0`, + `restrict_to_mask` on/off, 1 vs 4 threads — all bitwise identical to the + pre-rewrite kernel. + +Bitwise identity is robust here because the density scatter rounds each endpoint +to an integer voxel, which absorbs the sub-half-pixel differences from the +changed floating-point association order of the nested lerps. It is not a theorem +for every possible flow field, which is why the differential edge-case coverage +above matters. -Wired up behind a runtime dispatch (`__builtin_cpu_supports("avx2")`) so the -wheel still loads on SSE2-only CPUs. Expected speedup at 1T: ~20–30 %, less -at 8T due to closer-to-bandwidth saturation. Estimated effort: half a day -plus a careful microbenchmark of `vpgatherdd` cost on the target CPUs. - -## Interleaved (channel-last) layout + hand-written AVX2 FMA (rolled back, 2026-06-12) - -This implemented the two top "remaining" ideas below and measured them. **Net -result: no win; small regression. Rolled back.** Documented here so the avenue -is not re-attempted without new information. - -The idea, in two composable stages: - -1. **Channel-last (interleaved) flow buffer.** The input is channel-first - (`channels[axis] = flow.data + axis*n_pixels`), so the `D` components at one - position are gathers from regions `n_pixels` floats apart — in theory `D` - distinct cache lines per corner. A one-time `O(N)` transpose into a - `std::vector` where a voxel's components sit at `[v*VS + axis]` makes - each corner touch one line and yield all `D` components. `compute_corners` - is unchanged (its offset is already the voxel flat index); only the sampler - multiplies by the voxel stride `VS`. -2. **Hand-written AVX2+FMA sampler for 3D** on that interleaved buffer, with - `VS = 4` padding so each corner loads as one 128-bit vector. Lane-wise FMA - accumulates the weighted `D`-vector across the 8 corners — **no gather** - (the rejected approach above). Runtime-dispatched via - `__builtin_cpu_supports`, per-function `__attribute__((target("avx2,fma")))`, - scalar fallback for MSVC/arm64/non-AVX2. No CMake/global-arch-flag change. - -The codegen was confirmed ideal via `objdump`: 8× `vfmadd132ps` each fusing a -128-bit load (`(%rdi,%rax,1)`) with a `vbroadcastss` weight, zero gather -instructions. So the SIMD path was real and optimal — and still no faster than -scalar. - -Measurements (3D fixture, warm, min of 8 runs, same session, same machine as -the headline numbers): +## Profile breakdown (1-thread, 3D, profile build) ``` -variant 1T min 8T min -channel-first (baseline) 5.33 s 1.357 s -interleaved VS=3 (no pad, scalar) 5.44 s 1.46 s -interleaved VS=4 (scalar) 5.56 s 1.42 s -interleaved VS=4 + AVX2 FMA 5.58 s 1.46 s +init ~0.4 % density-zero + positions collection +iter_loop ~99 % the particle tracing loop +scatter ~0.1 % final density write +mask_zero ~0.1 % final mask zeroing ``` -Conclusions: - -- **The loop is memory-latency bound, as the roofline diagnosis said.** - Cutting the corner sum from 24 scalar mul-adds to 8 packed FMAs changes the - ALU count, not the load latency, so it does nothing. The 8 corner loads are - independent, so the out-of-order engine already overlaps them (high MLP); - channel-first's three independent channel loads overlap too. -- **Channel-first already has good locality.** A 64-byte line holds 16 - consecutive-x voxels of one channel; a trilinear cell's x-pair lives on one - line. Interleaving with `VS=4` cuts that to 4 voxels/line — *worse* density — - which is why VS=4 lost ~0.12 s/1T to VS=3. But even VS=3 (no waste, same - 144 MB footprint) did not beat channel-first. -- **Padding `VS=4` also costs +33 % flow memory** (144 → 192 MB for the 3D - fixture) for the AVX2 path's aligned 4-wide load. -- The differences are partly inside this laptop's **±~8 % thermal-throttle - noise** (baseline 1T alone ranged 5.33–5.85 s across the session, 8T - 1.357–1.59 s). That noise floor exceeds the expected gain of any remaining - micro-optimization, so small wins cannot be validated on this host — they - need a fixed-clock machine with `perf` counters. - -## What was tried and rejected - -- **Per-thread density scatter buffers** — scatter is <1 % of runtime, - not worth it. -- **Single-pass init** (fuse density-zero with positions collection) — init - is 0.4 % of runtime, not worth it. -- **Active-list compaction** (rebuild a `vector active_idx` to skip - dead-particle scans) — projected impact <1 %; the current alive-byte - check is one load + one branch. -- **`target_clones` SIMD** — see above, no net win, rolled back. -- **Half-precision (fp16) flow storage** — would halve memory traffic, but - the kernel isn't bandwidth-bound (see diagnosis), and it's a breaking API - change. -- **Interleaved channel-last layout + hand-written AVX2 FMA** — implemented and - measured 2026-06-12 (see the dedicated section above). No win, small - regression, rolled back. The SIMD path was confirmed optimal in codegen yet - did not beat scalar — the loop is load-latency bound, not ALU bound. - -## What remains worth trying - -In rough order of expected return on effort. **Note (2026-06-12):** items 1 and -2 below were implemented and measured — they gave no speedup (see "Interleaved -layout + AVX2 FMA" above). They are kept here only with that caveat; the live -candidate is now (3), and (4) is the prerequisite for evaluating it. - -1. ~~Hand-written AVX2 intrinsics for the corner sum~~ — **tried, no win.** The - corner sum on an interleaved buffer compiled to 8 packed FMAs with no - gathers, and was no faster than scalar: the loop waits on the 8 corner - loads, not the arithmetic. -2. ~~SoA position layout~~ — its only rationale was enabling the SIMD in (1); - with (1) shown not to help, SoA's batched clip/step/convergence has nothing - to pay for the extra coordinate gather/scatter. Not pursued separately. -3. **Reduce load latency** via software prefetching: prefetch the next-K - particles' corner cache lines before computing the current one. This is the - one lever that targets the actual (latency) bottleneck, and it works on the - existing channel-first layout — no relayout needed. Expected gain is modest - and **cannot be validated on the current laptop** (its ±~8 % thermal noise - floor swamps it); needs a fixed-clock host. -4. **A real `perf stat`** run on a host where `linux-perf-tools` is - installed, to confirm where cycles actually go (front-end, back-end - memory, back-end core, retired-FMA-ratio) and to provide a low-noise - measurement environment in which (3) could actually be evaluated. The - streaming-probe upper bound is necessary but not sufficient. +The iteration loop still dominates after the rewrite, so init/scatter/mask-zero +remain not worth touching (per-thread scatter buffers and single-pass init were +scoped and skipped for this reason). + +## Bandwidth diagnosis (streaming probe, earlier era) + +Sustained streaming read bandwidth on this host is **≈ 11.5 GB/s** (measured via +`numpy.sum` on float32 arrays larger than L3). Against the upper-bound flow bytes +read per 3D RK2 call (~7.0 GB), the pre-rewrite kernel used ~11 % of sustained +bandwidth at 1T and ~40 % at 8T — i.e. the hot loop was **compute-bound, not +bandwidth-bound**, dominated by per-particle serial instruction count and gather +latency rather than RAM throughput. + +The 2026-07-12 rewrite is consistent with and sharpens that diagnosis: it won by +**removing instructions and synchronization**, not by touching the memory layout. +The earlier conclusion was too narrow only if read as "only prefetching can +help" — cutting the coordinate/branch/weight/alive-state work around the loads +was the larger lever. + +## Rejected / did-not-help experiments + +Kept so these avenues are not blindly re-attempted. + +- **`target_clones` autovectorization SIMD** (earlier, rolled back) — the AVX2 + clone emitted only scalar FMA (gcc judged the 8-corner gather unprofitable) and + regressed 8T from icache pressure of three inlined clones. +- **Current-sample channel prefetching** after interpolation offsets were known + regressed the EPYC 3D fixture by about 5% at 1T and 2% at 8T. The warm-kernel + counter run reported only about 1.45% L1-data load misses, so the extra + prefetch instructions cost more than the avoided demand-load latency. +- **Next-sample software prefetching** (closed after the `perf` gate, + 2026-07-13; not implemented). A production, plain-flag build of `e426920` was + measured on an AMD EPYC 7513 (Zen 3) with the registered 3D fixture, pinned to + one core. Each `perf stat -r 3` process loaded the fixture once, made one warm + kernel call, then made eight more calls; process setup and loading were + included but amortized by the repeated kernel work. The performance governor + was active (boost remained enabled), so cycles and counter repeatability were + used as the gate rather than wall time alone. + - Measurement noise was low enough for the decision: cycles varied by + 0.03--0.12% and wall time by 0.3--0.7% across the three runs. + - The warm process sustained 2.77 instructions/cycle. L1 data-load misses + were about 1.54% (roughly 500 million misses from 32.3 billion loads). + - Of the observed demand data reads reaching L2, about 453.5 million hit and + 19.5 million missed (approximately a 4.1% L2 miss rate). An L2 fill was + pending for roughly 7% of counted cycles; this is an upper bound on cycles + prefetching might overlap, not a direct stall measurement. + These counters do not show enough residual cache-miss cost to justify the + extra instructions and code complexity. The existing 3-way RK2 interleave + already overlaps independent load chains, current-sample prefetching regressed + on this CPU, and a next-iteration address is only known late in the dependent + trajectory update. The prefetch target is therefore closed unless a future + workload or architecture supplies contrary profile evidence. +- **Interleaved (channel-last) layout + hand-written AVX2 FMA** (earlier, rolled + back, 2026-06-12) — codegen was confirmed optimal via `objdump` (8× packed + `vfmadd132ps`, zero gathers) yet was no faster than scalar, and `VS=4` padding + cost +33 % flow memory and worse line density than channel-first. Confirms the + loop is load-latency / instruction-count bound, not ALU bound. The scalar + rewrite targets exactly that (fewer instructions), which is why it succeeded + where SIMD did not. +- **Half-precision (fp16) flow storage** — would halve traffic, but the kernel is + not bandwidth-bound, and it is a breaking API change. +- **Per-thread density scatter buffers** / **single-pass init** — scatter and + init are each `<0.5 %` of runtime. +- **Active-list compaction** (superseded) — was projected `<1 %` against the old + alive-byte scan; particle-major removes the alive state entirely, so this is + moot. +- **Runtime `if (!restrict_to_mask)` start-clip removal** (Sol, reverted) — + regressed ~3.87 → 4.97 s from the added invariant branch. Now achieved for free + via compile-time specialization (item 9). +- **Explicit `dt·step` displacement reuse** (Sol) — neutral (~0.02 s); the + compiler already handles it. +- **16-particle iteration-major sub-blocking inside a persistent worker** (Sol) — + regressed ~34 %. Keeping one trajectory in registers beats interleaving 16 to + expose memory-level parallelism. This is also why a per-particle prefetch of the + *next* particle is unlikely to pay: the dependent chain is within one trajectory. +- **RK2 midpoint same-cell corner reuse** (2026-07-12, reverted). After the first + sample of a step, keep the 2^D·D corner values and skip the midpoint's offset + computation and reloads when `trunc(mid) == trunc(position)` (bitwise-exact by + construction: equal lower coords imply equal clamped offsets and corner + values). The mechanism worked as designed — measured with temporary counters, + the reuse hit rate is **96.7 %** on the registered 3D fixture (default + `dt=0.2`; the midpoint moves `0.1·|flow|` voxels) — and a 320-case randomized + digest harness plus the full suite confirmed bitwise-identical output. It was + still rejected on tightly A/B-paired benchmarks (alternating the two built + `.so` files per run; stash-and-rebuild rounds proved useless against this + laptop's 15–20 % thermal drift): + - Registered fixtures: 3D 1T ~2–3 % faster, 2D 1T ~5 %, 3D 8T ~2–5 % — the + midpoint loads are L1 hits whose latency the out-of-order core already hides + behind the trajectory's dependent chain, so removing them buys little. + - Adversarial random flows (`development/flow/benchmark_midpoint_reuse.py`, + scale sweep): at scale 10 (~30 % hit rate, maximally unpredictable branch) + the kernel regressed **13–15 %** at 1T, consistently across pairs; ~1–4 % + regression even at scale 1 (92 % hit). + - Two codegen traps worth remembering: (a) expressing `sample_flow` as + load-cell + interpolate-cell made GCC 14 materialize the corner array on the + stack in the single-sample Euler tracer (~20 % Euler regression; fixed by + keeping `sample_flow` monolithic); (b) the fatter RK2 `trace_particle` made + GCC stop inlining it, emitting it out-of-line as an *untagged* weak symbol + with AVX code in the FMA TU — the COMDAT/ODR hazard the dispatch design had + only avoided through full inlining. The `CodegenVariant` tag is therefore + now propagated through `trace_particle` (landed; instruction streams verified + identical to the pre-change build in both TUs). + +## Optimization status and future validation + +The 2026-07-13 Zen 3 counter run completed the final `perf` gate and did not +support another kernel implementation experiment. Flow-density optimization is +concluded for now at `e426920`; reopen it only for a demonstrated regression, a +new representative workload, or hardware-counter evidence of a substantial +bottleneck not covered above. + +Cross-architecture validation remains useful but is not an active optimization +target: confirm the FMA dispatch on macOS x86-64 and Windows x86-64, and confirm +the truncation gain magnitude on arm64 (macOS AppleClang and a Linux arm64 wheel +environment). The particle-major and direct-interpolation changes are portable +C++20; the x86 dispatch is compiler-specific and the truncation magnitude can +vary with how `floor` is lowered on each target. ## Reproducing these measurements @@ -316,7 +467,7 @@ pip install -e . --no-build-isolation # Thread scaling for nt in 1 2 4 6 8; do - python development/flow/check_flow_density.py --dim 3 --repeats 3 --threads $nt + python development/flow/check_flow_density.py --dim 3 --repeats 5 --threads $nt done # Accuracy gate at the chosen defaults @@ -325,6 +476,5 @@ python development/flow/check_flow_density.py --dim both --repeats 3 `check_flow_density.py` accepts `--method`, `--dt`, `--tol`, `--n-iter`, `--restrict-to-mask` / `--no-restrict-to-mask`, and `--threads` to override -defaults; useful for sweep work. The PASS gate is -`rel_diff = mean(|ours-ref|)/mean(ref) ≤ 0.15` by default +defaults. The PASS gate is `rel_diff = mean(|ours-ref|)/mean(ref) ≤ 0.15` (`--rel-tol` to override). diff --git a/development/flow/benchmark_interleave.py b/development/flow/benchmark_interleave.py new file mode 100644 index 0000000..a1d26c8 --- /dev/null +++ b/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()) diff --git a/development/flow/benchmark_midpoint_reuse.py b/development/flow/benchmark_midpoint_reuse.py new file mode 100644 index 0000000..9ab6981 --- /dev/null +++ b/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()) diff --git a/include/bioimage_cpp/flow/flow_density.hxx b/include/bioimage_cpp/flow/flow_density.hxx index 6eced35..3820aed 100644 --- a/include/bioimage_cpp/flow/flow_density.hxx +++ b/include/bioimage_cpp/flow/flow_density.hxx @@ -10,6 +10,11 @@ #include #include +#if defined(BIOIMAGE_FLOW_FMA_DISPATCH) && defined(_MSC_VER) +#include +#include +#endif + namespace bioimage_cpp::flow { namespace detail { @@ -34,66 +39,79 @@ GridLayout make_grid_layout( return layout; } -// Per-position cache of the 2^D corners used for linear interpolation. Offsets -// are in elements relative to the channel base pointer; weights sum to 1. +// Linearly interpolate all D flow channels at `position` and write the result +// to `out`. This is explicit bilinear (D==2) / trilinear (D==3) sampling rather +// than a generic 2^D corner table: the lower/upper index and fractional weight +// are computed once per axis and shared across channels, and the nested lerps +// avoid materializing 2^D product weights and offsets. +// +// Precondition: 0 <= position[axis] <= shape[axis]-1 for every axis (callers +// clip before sampling). Because the coordinate is nonnegative, truncation +// equals std::floor, so the integer cast is exact; this matters because on the +// portable (no -march, SSE2) wheel build std::floor is an inlined multi-branch +// software routine while the cast is a single instruction. At an upper boundary +// the lower and upper index coincide, matching nearest-boundary behavior. template -struct SamplingCorners { - std::array offsets{}; - std::array weights{}; -}; - -template -SamplingCorners compute_corners( +inline void sample_flow( + const std::array &channels, const std::array &position, - const GridLayout &grid + const GridLayout &grid, + std::array &out ) { - std::array lower{}; - std::array frac{}; - for (std::size_t axis = 0; axis < D; ++axis) { - const auto lo = static_cast(std::floor(position[axis])); - lower[axis] = lo; - frac[axis] = position[axis] - static_cast(lo); - } - - SamplingCorners corners{}; - constexpr std::size_t n_corners = std::size_t{1} << D; - for (std::size_t corner = 0; corner < n_corners; ++corner) { - std::ptrdiff_t offset = 0; - float weight = 1.0f; + if constexpr (D == 2) { + const std::ptrdiff_t sy = grid.strides[0]; + const std::ptrdiff_t sx = grid.strides[1]; + const std::ptrdiff_t y0 = static_cast(position[0]); + const std::ptrdiff_t x0 = static_cast(position[1]); + const float fy = position[0] - static_cast(y0); + const float fx = position[1] - static_cast(x0); + const std::ptrdiff_t y1 = (y0 + 1 < grid.shape[0]) ? y0 + 1 : y0; + const std::ptrdiff_t x1 = (x0 + 1 < grid.shape[1]) ? x0 + 1 : x0; + const std::ptrdiff_t o00 = y0 * sy + x0 * sx; + const std::ptrdiff_t o01 = y0 * sy + x1 * sx; + const std::ptrdiff_t o10 = y1 * sy + x0 * sx; + const std::ptrdiff_t o11 = y1 * sy + x1 * sx; for (std::size_t axis = 0; axis < D; ++axis) { - const bool upper_side = ((corner >> axis) & std::size_t{1}) != 0; - std::ptrdiff_t coord; - if (upper_side) { - weight *= frac[axis]; - coord = lower[axis] + 1; - } else { - weight *= 1.0f - frac[axis]; - coord = lower[axis]; - } - if (coord < 0) { - coord = 0; - } else if (coord >= grid.shape[axis]) { - coord = grid.shape[axis] - 1; - } - offset += coord * grid.strides[axis]; + const float *c = channels[axis]; + const float top = c[o00] + fx * (c[o01] - c[o00]); + const float bot = c[o10] + fx * (c[o11] - c[o10]); + out[axis] = top + fy * (bot - top); + } + } else { + const std::ptrdiff_t sz = grid.strides[0]; + const std::ptrdiff_t sy = grid.strides[1]; + const std::ptrdiff_t sx = grid.strides[2]; + const std::ptrdiff_t z0 = static_cast(position[0]); + const std::ptrdiff_t y0 = static_cast(position[1]); + const std::ptrdiff_t x0 = static_cast(position[2]); + const float fz = position[0] - static_cast(z0); + const float fy = position[1] - static_cast(y0); + const float fx = position[2] - static_cast(x0); + const std::ptrdiff_t z1 = (z0 + 1 < grid.shape[0]) ? z0 + 1 : z0; + const std::ptrdiff_t y1 = (y0 + 1 < grid.shape[1]) ? y0 + 1 : y0; + const std::ptrdiff_t x1 = (x0 + 1 < grid.shape[2]) ? x0 + 1 : x0; + const std::ptrdiff_t z0s = z0 * sz, z1s = z1 * sz; + const std::ptrdiff_t y0s = y0 * sy, y1s = y1 * sy; + const std::ptrdiff_t x0s = x0 * sx, x1s = x1 * sx; + const std::ptrdiff_t o000 = z0s + y0s + x0s; + const std::ptrdiff_t o001 = z0s + y0s + x1s; + const std::ptrdiff_t o010 = z0s + y1s + x0s; + const std::ptrdiff_t o011 = z0s + y1s + x1s; + const std::ptrdiff_t o100 = z1s + y0s + x0s; + const std::ptrdiff_t o101 = z1s + y0s + x1s; + const std::ptrdiff_t o110 = z1s + y1s + x0s; + const std::ptrdiff_t o111 = z1s + y1s + x1s; + for (std::size_t axis = 0; axis < D; ++axis) { + const float *c = channels[axis]; + const float c00 = c[o000] + fx * (c[o001] - c[o000]); + const float c01 = c[o010] + fx * (c[o011] - c[o010]); + const float c10 = c[o100] + fx * (c[o101] - c[o100]); + const float c11 = c[o110] + fx * (c[o111] - c[o110]); + const float c0 = c00 + fy * (c01 - c00); + const float c1 = c10 + fy * (c11 - c10); + out[axis] = c0 + fz * (c1 - c0); } - corners.offsets[corner] = offset; - corners.weights[corner] = weight; - } - return corners; -} - -template -inline float sample_channel( - const float *channel, - const SamplingCorners &corners -) { - constexpr std::size_t n_corners = std::size_t{1} << D; - float value = 0.0f; - for (std::size_t corner = 0; corner < n_corners; ++corner) { - value += corners.weights[corner] * channel[corners.offsets[corner]]; } - return value; } template @@ -110,9 +128,11 @@ inline std::ptrdiff_t round_to_flat_index( clipped = grid.upper[axis]; } // Round half up, matching the nearest-neighbor convention in - // transformation/affine.hxx and segmentation/watershed.hxx. - // std::nearbyint would honor the FP rounding mode (round-half-to-even). - const auto coord = static_cast(std::floor(clipped + 0.5f)); + // transformation/affine.hxx and segmentation/watershed.hxx. clipped is + // nonnegative, so (clipped + 0.5f) >= 0 and the truncating cast equals + // std::floor(clipped + 0.5f) exactly, while avoiding the comparatively + // expensive inlined floor on the portable (non-SSE4.1) wheel build. + const auto coord = static_cast(clipped + 0.5f); flat += coord * grid.strides[axis]; } return flat; @@ -139,6 +159,336 @@ enum class IntegrationMethod { RK2, }; +namespace detail { + +// Trace a single particle through its whole trajectory. The integration flags +// are compile-time parameters so the per-step branches on them fold away and +// the sampler/RK2/convergence/mask code inlines into one specialized loop. +// +// CodegenVariant must match the instantiating trace_all (see there): the +// compiler is free to emit this function out-of-line (observed with GCC 14 +// under mild size pressure), and without the tag the FMA translation unit +// would emit AVX code under the same weak symbol name as the portable +// instantiation, letting COMDAT selection ship AVX code to the portable +// fallback path (SIGILL on pre-AVX CPUs) or silently discard the FMA kernel. +template < + std::size_t D, bool UseRK2, bool CheckConvergence, bool RestrictToMask, + bool CodegenVariant = false> +inline void trace_particle( + std::array &position, + const std::array &channels, + const GridLayout &grid, + const std::uint8_t *mask, + const std::size_t n_iter, + const float dt, + const float tol +) { + const auto clip = [&grid](std::array &p) { + for (std::size_t axis = 0; axis < D; ++axis) { + if (p[axis] < 0.0f) { + p[axis] = 0.0f; + } else if (p[axis] > grid.upper[axis]) { + p[axis] = grid.upper[axis]; + } + } + }; + + // When restricting to the mask, only in-mask (hence in-bounds) endpoints are + // ever committed and the seed is an in-bounds integer voxel, so `position` + // is already inside the domain at the start of every step and the + // start-of-step clip is redundant. Without mask restriction an endpoint may + // leave the domain and must be clipped back before the next sample. + if constexpr (!RestrictToMask) { + clip(position); + } + + for (std::size_t iter = 0; iter < n_iter; ++iter) { + std::array step{}; + sample_flow(channels, position, grid, step); + + if constexpr (UseRK2) { + std::array mid{}; + for (std::size_t axis = 0; axis < D; ++axis) { + mid[axis] = position[axis] + 0.5f * dt * step[axis]; + } + clip(mid); + sample_flow(channels, mid, grid, step); + } + + if constexpr (CheckConvergence) { + float max_step = 0.0f; + for (std::size_t axis = 0; axis < D; ++axis) { + const float abs_step = std::fabs(dt * step[axis]); + if (abs_step > max_step) { + max_step = abs_step; + } + } + if (max_step < tol) { + break; + } + } + + std::array proposed = position; + for (std::size_t axis = 0; axis < D; ++axis) { + proposed[axis] += dt * step[axis]; + } + + if constexpr (RestrictToMask) { + // A particle whose proposed endpoint leaves the foreground is frozen + // at its last in-mask position (only the endpoint is mask-tested, + // not the RK2 midpoint). + if (!position_is_in_mask(proposed, grid, mask)) { + break; + } + } else { + clip(proposed); + } + position = proposed; + } +} + +// Trace K consecutive particles in lockstep. The trajectories are independent, +// so the out-of-order core can overlap one lane's serial sample->update chain +// with the other lanes' chains; a converged or mask-frozen lane costs one +// predictable branch per remaining group iteration, and the group ends when +// every lane is done. The per-step body is identical to trace_particle (with +// `break` expressed as clearing the lane's alive flag), so the traced +// positions are bitwise equal to K independent trace_particle calls. +template < + std::size_t D, std::size_t K, bool UseRK2, bool CheckConvergence, + bool RestrictToMask, bool CodegenVariant = false> +inline void trace_particle_block( + std::array *positions, + const std::array &channels, + const GridLayout &grid, + const std::uint8_t *mask, + const std::size_t n_iter, + const float dt, + const float tol +) { + static_assert(K >= 2, "use trace_particle for single trajectories"); + + const auto clip = [&grid](std::array &p) { + for (std::size_t axis = 0; axis < D; ++axis) { + if (p[axis] < 0.0f) { + p[axis] = 0.0f; + } else if (p[axis] > grid.upper[axis]) { + p[axis] = grid.upper[axis]; + } + } + }; + + // Local copies keep the lane state in registers across the group loop. + std::array, K> pos; + std::array alive; + for (std::size_t k = 0; k < K; ++k) { + pos[k] = positions[k]; + alive[k] = true; + } + if constexpr (!RestrictToMask) { + for (std::size_t k = 0; k < K; ++k) { + clip(pos[k]); + } + } + + for (std::size_t iter = 0; iter < n_iter; ++iter) { + for (std::size_t k = 0; k < K; ++k) { + if (!alive[k]) { + continue; + } + std::array step{}; + sample_flow(channels, pos[k], grid, step); + + if constexpr (UseRK2) { + std::array mid{}; + for (std::size_t axis = 0; axis < D; ++axis) { + mid[axis] = pos[k][axis] + 0.5f * dt * step[axis]; + } + clip(mid); + sample_flow(channels, mid, grid, step); + } + + if constexpr (CheckConvergence) { + float max_step = 0.0f; + for (std::size_t axis = 0; axis < D; ++axis) { + const float abs_step = std::fabs(dt * step[axis]); + if (abs_step > max_step) { + max_step = abs_step; + } + } + if (max_step < tol) { + alive[k] = false; + continue; + } + } + + std::array proposed = pos[k]; + for (std::size_t axis = 0; axis < D; ++axis) { + proposed[axis] += dt * step[axis]; + } + + if constexpr (RestrictToMask) { + if (!position_is_in_mask(proposed, grid, mask)) { + alive[k] = false; + continue; + } + } else { + clip(proposed); + } + pos[k] = proposed; + } + + bool any_alive = false; + for (std::size_t k = 0; k < K; ++k) { + any_alive = any_alive || alive[k]; + } + if (!any_alive) { + break; + } + } + + for (std::size_t k = 0; k < K; ++k) { + positions[k] = pos[k]; + } +} + +// CodegenVariant gives separately compiled ISA variants a distinct linker +// identity. Without it, COMDAT selection may replace an FMA instantiation with +// the portable instantiation that has the otherwise-identical template name. +template < + std::size_t D, bool UseRK2, bool CheckConvergence, bool RestrictToMask, + bool CodegenVariant = false> +void trace_all( + std::vector> &positions, + const std::array &channels, + const GridLayout &grid, + const std::uint8_t *mask, + const std::size_t n_threads, + const std::size_t n_iter, + const float dt, + const float tol +) { + // Particle-major: each worker traces its whole contiguous range of particles + // across all integration steps in one fan-out. Trajectories are independent + // until the sequential scatter, so no global alive state, per-step barrier, + // or per-step alive scan is needed. + ::bioimage_cpp::detail::parallel_for_chunks( + n_threads, + positions.size(), + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + // Lockstep interleaving only pays for RK2: its two dependent + // samples per step leave latency bubbles that other lanes fill. + // The shorter Euler chain measured ~10% slower when interleaved + // (extra lane state without enough latency to hide). K=3 beat + // K=2/4 and the per-particle loop on paired benchmarks (see + // development/flow/PERFORMANCE_NOTES.md, trajectory interleave). + constexpr std::size_t K = UseRK2 ? 3 : 1; + std::size_t i = begin; + if constexpr (K > 1) { + for (; i + K <= end; i += K) { + trace_particle_block< + D, K, UseRK2, CheckConvergence, RestrictToMask, + CodegenVariant>( + &positions[i], channels, grid, mask, n_iter, dt, tol + ); + } + } + for (; i < end; ++i) { + trace_particle< + D, UseRK2, CheckConvergence, RestrictToMask, CodegenVariant>( + positions[i], channels, grid, mask, n_iter, dt, tol + ); + } + } + ); +} + +#if defined(BIOIMAGE_FLOW_FMA_DISPATCH) + +inline bool runtime_fma_supported() noexcept { + static const bool supported = []() noexcept { +#if defined(_MSC_VER) + int registers[4]{}; + __cpuid(registers, 1); + constexpr int fma_bit = 1 << 12; + constexpr int osxsave_bit = 1 << 27; + constexpr int avx_bit = 1 << 28; + if ((registers[2] & (fma_bit | osxsave_bit | avx_bit)) != + (fma_bit | osxsave_bit | avx_bit)) { + return false; + } + if ((_xgetbv(0) & 0x6) != 0x6) { + return false; + } +#if defined(BIOIMAGE_FLOW_FMA_REQUIRES_AVX2) + __cpuidex(registers, 7, 0); + constexpr int avx2_bit = 1 << 5; + if ((registers[1] & avx2_bit) == 0) { + return false; + } +#endif + return true; +#elif defined(__GNUC__) || defined(__clang__) + return __builtin_cpu_supports("avx") && __builtin_cpu_supports("fma"); +#else + return false; +#endif + }(); + return supported; +} + +void trace_all_fma_2d( + std::vector> &positions, + const std::array &channels, + const GridLayout<2> &grid, + const std::uint8_t *mask, + std::size_t n_threads, + std::size_t n_iter, + float dt, + float tol +); + +void trace_all_fma_3d( + std::vector> &positions, + const std::array &channels, + const GridLayout<3> &grid, + const std::uint8_t *mask, + std::size_t n_threads, + std::size_t n_iter, + float dt, + float tol +); + +template +bool try_trace_all_fma( + std::vector> &positions, + const std::array &channels, + const GridLayout &grid, + const std::uint8_t *mask, + const std::size_t n_threads, + const std::size_t n_iter, + const float dt, + const float tol +) { + if (!runtime_fma_supported()) { + return false; + } + if constexpr (D == 2) { + trace_all_fma_2d( + positions, channels, grid, mask, n_threads, n_iter, dt, tol + ); + } else { + trace_all_fma_3d( + positions, channels, grid, mask, n_threads, n_iter, dt, tol + ); + } + return true; +} + +#endif + +} // namespace detail + // Preconditions (validated in the binding layer): // * flow.ndim() == D + 1, flow.shape[0] == D, flow.shape[1..] == fg_mask.shape // * fg_mask.ndim() == D and density.shape == fg_mask.shape @@ -199,95 +549,50 @@ void compute_flow_density( number_of_threads, positions.size() ); - std::vector alive(positions.size(), 1); const bool use_rk2 = (method == IntegrationMethod::RK2); const bool check_convergence = (tol > 0.0f); - auto clip_position = [&grid](std::array &p) { - for (std::size_t axis = 0; axis < D; ++axis) { - if (p[axis] < 0.0f) { - p[axis] = 0.0f; - } else if (p[axis] > grid.upper[axis]) { - p[axis] = grid.upper[axis]; - } - } - }; - { BIOIMAGE_PROFILE_SCOPE(profiler, "iter_loop"); - for (std::size_t iter = 0; iter < n_iter; ++iter) { - ::bioimage_cpp::detail::parallel_for_chunks( - n_threads, - positions.size(), - [&](const std::size_t, const std::size_t begin, const std::size_t end) { - for (std::size_t i = begin; i < end; ++i) { - if (alive[i] == 0) { - continue; - } - auto &position = positions[i]; - clip_position(position); - - const auto corners = detail::compute_corners(position, grid); - std::array step{}; - for (std::size_t axis = 0; axis < D; ++axis) { - step[axis] = detail::sample_channel(channels[axis], corners); - } - - if (use_rk2) { - std::array mid{}; - for (std::size_t axis = 0; axis < D; ++axis) { - mid[axis] = position[axis] + 0.5f * dt * step[axis]; - } - clip_position(mid); - const auto mid_corners = detail::compute_corners(mid, grid); - for (std::size_t axis = 0; axis < D; ++axis) { - step[axis] = detail::sample_channel(channels[axis], mid_corners); - } - } - - float max_step = 0.0f; - for (std::size_t axis = 0; axis < D; ++axis) { - const float abs_step = std::fabs(dt * step[axis]); - if (abs_step > max_step) { - max_step = abs_step; - } - } - if (check_convergence && max_step < tol) { - alive[i] = 0; - continue; - } - auto proposed = position; - for (std::size_t axis = 0; axis < D; ++axis) { - proposed[axis] += dt * step[axis]; - } - // A particle whose proposed endpoint leaves the - // foreground is frozen at its last in-mask position - // (only the endpoint is mask-tested, not the RK2 - // midpoint). Every alive particle is seeded in the mask - // and only commits in-mask endpoints, so its current - // position is always a valid last position and no - // separate start-of-step mask test is needed. - if (restrict_to_mask && !detail::position_is_in_mask( - proposed, grid, fg_mask.data - )) { - alive[i] = 0; - continue; - } - position = proposed; - } - } - ); - - if (check_convergence || restrict_to_mask) { - std::size_t still_alive = 0; - for (const auto a : alive) { - still_alive += a; - } - if (still_alive == 0) { + // Dispatch the three loop-invariant flags to a compile-time specialized + // tracer (see trace_particle). The flags are constant for the whole + // call, so this hoists their branches out of the innermost step loop. + const int selector = + (use_rk2 ? 4 : 0) | (check_convergence ? 2 : 0) | (restrict_to_mask ? 1 : 0); +#define BIOIMAGE_FLOW_TRACE(SELECTOR, RK2, CONV, RESTRICT) \ + case SELECTOR: \ + detail::trace_all( \ + positions, channels, grid, fg_mask.data, n_threads, n_iter, dt, tol); \ + break; + switch (selector) { + BIOIMAGE_FLOW_TRACE(0, false, false, false) + BIOIMAGE_FLOW_TRACE(1, false, false, true) + BIOIMAGE_FLOW_TRACE(2, false, true, false) + BIOIMAGE_FLOW_TRACE(3, false, true, true) + BIOIMAGE_FLOW_TRACE(4, true, false, false) + BIOIMAGE_FLOW_TRACE(5, true, false, true) + BIOIMAGE_FLOW_TRACE(6, true, true, false) + case 7: +#if defined(BIOIMAGE_FLOW_FMA_DISPATCH) + if (detail::try_trace_all_fma( + positions, + channels, + grid, + fg_mask.data, + n_threads, + n_iter, + dt, + tol + )) { break; } - } +#endif + detail::trace_all( + positions, channels, grid, fg_mask.data, n_threads, n_iter, dt, tol + ); + break; } +#undef BIOIMAGE_FLOW_TRACE } { diff --git a/src/cpp/flow/flow_density_fma.cxx b/src/cpp/flow/flow_density_fma.cxx new file mode 100644 index 0000000..f321d35 --- /dev/null +++ b/src/cpp/flow/flow_density_fma.cxx @@ -0,0 +1,39 @@ +#include "bioimage_cpp/flow/flow_density.hxx" + +#if !defined(BIOIMAGE_FLOW_FMA_DISPATCH) +#error "flow_density_fma.cxx must only be built with BIOIMAGE_FLOW_FMA_DISPATCH" +#endif + +namespace bioimage_cpp::flow::detail { + +void trace_all_fma_2d( + std::vector> &positions, + const std::array &channels, + const GridLayout<2> &grid, + const std::uint8_t *mask, + const std::size_t n_threads, + const std::size_t n_iter, + const float dt, + const float tol +) { + trace_all<2, true, true, true, true>( + positions, channels, grid, mask, n_threads, n_iter, dt, tol + ); +} + +void trace_all_fma_3d( + std::vector> &positions, + const std::array &channels, + const GridLayout<3> &grid, + const std::uint8_t *mask, + const std::size_t n_threads, + const std::size_t n_iter, + const float dt, + const float tol +) { + trace_all<3, true, true, true, true>( + positions, channels, grid, mask, n_threads, n_iter, dt, tol + ); +} + +} // namespace bioimage_cpp::flow::detail diff --git a/tests/test_flow.py b/tests/test_flow.py index 7b73cc1..afc68d5 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -168,6 +168,66 @@ def test_restrict_to_mask_freezes_at_last_valid_position(method): np.testing.assert_array_equal(density, expected) +def test_rk2_midpoint_cell_crossing_is_exact(): + # Constant flow of 3 px along x: the RK2 midpoint always lands outside the + # interpolation cell of the current position, while the constant flow keeps + # the exact trajectory independent of which cell is sampled. + shape = (3, 16) + flow = np.zeros((2,) + shape, dtype=np.float32) + flow[1] = 3.0 + mask = np.ones(shape, dtype=bool) + + density = bic.flow.compute_flow_density( + flow, mask, n_iter=2, dt=1.0, tol=0.0, method="rk2", restrict_to_mask=False + ) + + # Each particle advances 3 px per iteration and is clipped at x=15, so + # x_final = min(x + 6, 15): starts 0..8 land on 6..14, starts 9..15 pile + # up on the boundary column. + expected = np.zeros(shape, dtype=np.float32) + expected[:, 6:15] = 1.0 + expected[:, 15] = 7.0 + np.testing.assert_array_equal(density, expected) + + +def test_rk2_large_flow_multithreaded_matches_single_threaded(): + # Large flow steps make the RK2 midpoint cross interpolation cells; the + # result must stay independent of the thread count. + rng = np.random.default_rng(99) + flow = rng.normal(scale=8.0, size=(3, 6, 20, 20)).astype(np.float32) + mask = rng.random((6, 20, 20)) > 0.3 + + kwargs = dict(n_iter=20, dt=0.2, tol=0.005, method="rk2", restrict_to_mask=True) + single = bic.flow.compute_flow_density(flow, mask, number_of_threads=1, **kwargs) + multi = bic.flow.compute_flow_density(flow, mask, number_of_threads=4, **kwargs) + + np.testing.assert_array_equal(single, multi) + + +def test_rk2_mixed_lifetimes_thread_equality(): + # 65 foreground particles (not divisible by 2/3/4) with strongly divergent + # lifetimes: zero-flow columns converge on the first step while drift + # columns run until frozen at the mask border. Different thread counts + # split the particle range differently, so equality across them exercises + # the lockstep block/remainder boundaries of the interleaved RK2 tracer. + shape = (5, 13) + flow = np.zeros((2,) + shape, dtype=np.float32) + flow[1, :, 1::2] = 2.0 + mask = np.ones(shape, dtype=bool) + kwargs = dict(n_iter=50, dt=0.2, tol=0.005, method="rk2", restrict_to_mask=True) + + results = [ + bic.flow.compute_flow_density(flow, mask, number_of_threads=t, **kwargs) + for t in (1, 2, 3, 4, 5) + ] + for other in results[1:]: + np.testing.assert_array_equal(results[0], other) + # restrict_to_mask freezes particles instead of dropping them + assert results[0].sum() == mask.sum() + # zero-flow columns converge in place and keep their particle + assert (results[0][:, 0::2] >= 1.0).all() + + def test_flow_rejects_non_finite_values(): flow = np.zeros((2, 3, 3), dtype=np.float32) flow[0, 1, 1] = np.nan