Implement optimized kernel dispatch for flow computation

Author: constantinpape

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/PERFORMANCE_NOTES.md

@@ -90,6 +90,15 @@ table and the iteration-major loop (with its per-iteration alive scan) with:
    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`.
+
 Chosen defaults (see `src/bioimage_cpp/flow/_flow.py`) are unchanged:
 
 ```python
@@ -200,9 +209,46 @@ threads   runtime   speedup vs 1T
 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.
+
 ## Correctness
 
-- All 1067 tests in the package pass; the 17 `tests/test_flow.py` cases cover
+- 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,
@@ -251,6 +297,10 @@ 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.
 - **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
@@ -288,11 +338,12 @@ Lower priority than the landed work; all need a fixed-clock host with `perf`
    several kernel calls) to confirm where the remaining cycles go (front-end vs
    back-end memory vs back-end core) and to provide a low-noise environment for
    evaluating (1).
-3. **Cross-architecture confirmation** of the truncation gain magnitude on arm64
-   (macOS AppleClang, Linux aarch64) and Windows MSVC. The particle-major and
-   direct-interpolation gains are portable C++20; the truncation magnitude is the
-   most microarchitecture-dependent (it depends on how `floor` is lowered without
-   SSE4.1 / on the arm equivalent).
+3. **Cross-architecture confirmation** on macOS x86-64 and Windows x86-64 for
+   the FMA dispatch, plus the truncation gain magnitude on arm64 (macOS
+   AppleClang and Linux aarch64). The particle-major and direct-interpolation
+   gains are portable C++20; the x86 dispatch is compiler-specific and the
+   truncation magnitude depends on how `floor` is lowered without SSE4.1 / on
+   the arm equivalent.
 
 ## Reproducing these measurements
 

include/bioimage_cpp/flow/flow_density.hxx

@@ -10,6 +10,11 @@
 #include <cstdint>
 #include <vector>
 
+#if defined(BIOIMAGE_FLOW_FMA_DISPATCH) && defined(_MSC_VER)
+#include <immintrin.h>
+#include <intrin.h>
+#endif
+
 namespace bioimage_cpp::flow {
 namespace detail {
 
@@ -233,7 +238,12 @@ inline void trace_particle(
     }
 }
 
-template <std::size_t D, bool UseRK2, bool CheckConvergence, bool RestrictToMask>
+// 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<std::array<float, D>> &positions,
     const std::array<const float *, D> &channels,
@@ -261,6 +271,90 @@ void trace_all(
     );
 }
 
+#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<std::array<float, 2>> &positions,
+    const std::array<const float *, 2> &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<std::array<float, 3>> &positions,
+    const std::array<const float *, 3> &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 <std::size_t D>
+bool try_trace_all_fma(
+    std::vector<std::array<float, D>> &positions,
+    const std::array<const float *, D> &channels,
+    const GridLayout<D> &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):
@@ -346,7 +440,25 @@ void compute_flow_density(
             BIOIMAGE_FLOW_TRACE(4, true,  false, false)
             BIOIMAGE_FLOW_TRACE(5, true,  false, true)
             BIOIMAGE_FLOW_TRACE(6, true,  true,  false)
-            BIOIMAGE_FLOW_TRACE(7, true,  true,  true)
+            case 7:
+#if defined(BIOIMAGE_FLOW_FMA_DISPATCH)
+                if (detail::try_trace_all_fma<D>(
+                        positions,
+                        channels,
+                        grid,
+                        fg_mask.data,
+                        n_threads,
+                        n_iter,
+                        dt,
+                        tol
+                    )) {
+                    break;
+                }
+#endif
+                detail::trace_all<D, true, true, true>(
+                    positions, channels, grid, fg_mask.data, n_threads, n_iter, dt, tol
+                );
+                break;
         }
 #undef BIOIMAGE_FLOW_TRACE
     }

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<std::array<float, 2>> &positions,
+    const std::array<const float *, 2> &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<std::array<float, 3>> &positions,
+    const std::array<const float *, 3> &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