From 7757046aef5752697ed3a307b69596e411c68d9b Mon Sep 17 00:00:00 2001 From: anwai98 Date: Tue, 3 Nov 2026 15:08:59 +0100 Subject: [PATCH 1/2] Add a NEON packed flow tracer for arm64 --- CMakeLists.txt | 19 + development/flow/PERFORMANCE_NOTES.md | 152 +++++++ development/flow/_flow_cases.py | 5 + include/bioimage_cpp/flow/flow_density.hxx | 66 ++- src/bindings/flow.cxx | 2 +- src/cpp/flow/flow_density_neon.cxx | 494 +++++++++++++++++++++ tests/test_flow.py | 9 +- 7 files changed, 740 insertions(+), 7 deletions(-) create mode 100644 src/cpp/flow/flow_density_neon.cxx diff --git a/CMakeLists.txt b/CMakeLists.txt index 0f1047d..e529bc7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -69,6 +69,25 @@ if(BIOIMAGE_FLOW_FMA_DISPATCH) endif() endif() +option( + BIOIMAGE_FLOW_NEON + "Build the arm64 NEON packed flow tracer (NEON+FMA are baseline on arm64)" + ON +) +if(BIOIMAGE_FLOW_NEON AND CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" _bioimage_neon_processor) + # CMAKE_OSX_ARCHITECTURES, if set, overrides the host processor. + if(CMAKE_OSX_ARCHITECTURES) + set(_bioimage_neon_target "${CMAKE_OSX_ARCHITECTURES}") + else() + set(_bioimage_neon_target "${_bioimage_neon_processor}") + endif() + if(_bioimage_neon_target MATCHES "^(arm64|aarch64)$") + target_sources(_core PRIVATE src/cpp/flow/flow_density_neon.cxx) + target_compile_definitions(_core PRIVATE BIOIMAGE_FLOW_NEON=1) + endif() +endif() + option( BIOIMAGE_FILTERS_AVX2_DISPATCH "Build x86 AVX2/FMA filter kernels with runtime CPU dispatch" diff --git a/development/flow/PERFORMANCE_NOTES.md b/development/flow/PERFORMANCE_NOTES.md index 234edc7..e0cee1b 100644 --- a/development/flow/PERFORMANCE_NOTES.md +++ b/development/flow/PERFORMANCE_NOTES.md @@ -545,6 +545,158 @@ translation units. - **Padding away the boundary clamps**: only ~1 % (integer clamps were not on the critical path), kept because it simplifies the sampler. +### Reproduction on Apple Silicon (2026-11-03) + +Setup: + +- **Apple M2** (4 P-cores, 4 E-cores, 16 GB), macOS 26.5, Apple clang 21.0.0, + Python 3.11. +- Plain flags (`CXXFLAGS` and the other flag variables unset), Release `-O3`. +- The laptop ran on battery with Low Power Mode on. The absolute times are + therefore slow. Both sides of each pair ran in the same state. +- macOS has no CPU affinity API. The harness runs without `--cpu`, and + `pin_cpus` does not pin on macOS. +- The calibration (branch vs branch) is `noise` on both fixtures at 1T and 4T. + The spread per side is 0.1-0.7 % at 1T. + +Before this round, arm64 used the portable header kernel, and +`_flow_trace_backend()` returned `"scalar"`. The branch (`126df39`) was parity +with `492ff9b`. All rows were `noise`: fixture3d was -0.8 % at 1T and -4 % at +4T, from the parallel finiteness scan. + +The NEON port is `src/cpp/flow/flow_density_neon.cxx`. It uses the same +packed-channel design as the SSE translation unit. + +- CMake builds it under the `BIOIMAGE_FLOW_NEON` option when the target is + `arm64` or `aarch64`. If `CMAKE_OSX_ARCHITECTURES` is set, it decides + instead of the host processor. +- It needs no runtime check, and `trace_backend()` returns `"neon"`. +- `BIOIMAGE_CPP_FLOW_FORCE_SCALAR=1` still selects the portable kernel. +- Intrinsics: `vld1q_f32`, `vfmaq_f32`, `vcvtq_s32_f32`, `vcvtq_f32_s32`, + `vgetq_lane_s32`, `vdupq_laneq_f32`, `vmulq_laneq_f32`, `vextq_f32` for + `movehl`, `vabsq_f32` and `vcgeq/vcltq/vcgtq`. The movemask test is + `vshrn_n_u32` with `vget_lane_u64`. + +Two semantic differences to SSE do not change the result: + +- `vcvtq_s32_f32` saturates and converts NaN to 0. `cvttps` returns + `INT_MIN`. The kernel converts only clipped or bounds-tested, finite values. +- `vmaxq/vminq` propagate NaN. `maxps/minps` return the second operand. The + midpoint clip uses `vmaxnmq/vminnmq`, which give the same result as SSE for + every finite input. + +Paired A/B of `492ff9b` (portable kernel) vs the NEON kernel (K = 16) on all +cases, with ABBA x4 and 3 inner calls. All rows are `significant`: + +``` +case threads previous neon change parity +fixture3d 1 1.5367 s 0.6515 s -57.6 % identical +fixture2d 1 0.1425 s 0.0765 s -46.3 % identical +stripes 1/4 1 1.0165 s 0.2754 s -72.9 % identical +stripes 3/4 1 1.7839 s 0.7719 s -56.7 % identical +random s=1 1 2.1132 s 0.9332 s -55.8 % identical +random s=10 1 1.9793 s 0.8541 s -56.9 % identical +random s=5 1 1.6117 s 0.7129 s -55.8 % identical +random s=20 1 1.6298 s 0.7461 s -54.2 % identical +random s=40 1 0.8707 s 0.4972 s -42.9 % identical +2D random 10 1 0.2370 s 0.1209 s -49.0 % differs (2D tree) +fixture3d 4 P 0.4271 s 0.1861 s -56.4 % identical +fixture2d 4 P 0.0368 s 0.0203 s -44.9 % identical +stripes 1/4 4 P 0.2635 s 0.0782 s -70.3 % identical +stripes 3/4 4 P 0.4540 s 0.2016 s -55.6 % identical +random s=1 4 P 0.5576 s 0.2437 s -56.3 % identical +random s=10 4 P 0.6323 s 0.2554 s -59.6 % identical +random s=5 4 P 0.4386 s 0.1931 s -56.0 % identical +random s=20 4 P 0.5436 s 0.2322 s -57.3 % identical +random s=40 4 P 0.2846 s 0.1517 s -46.7 % identical +2D random 10 4 P 0.0613 s 0.0318 s -48.1 % differs (2D tree) +``` + +The gain is larger than on x86 (-30..-40 %). The lane count is the main +reason. + +**Lane-count sweep.** Each value is the minimum at 1T from a paired run +against the previous best K: + +``` +K fixture3d fixture2d stripes 1/4 +2 1.416 s 0.226 s 0.897 s +3 0.988 s 0.153 s 0.983 s +4 0.770 s 0.118 s 0.894 s +6 0.666 s 0.083 s 0.627 s +8 0.658 s 0.078 s 0.483 s +10 0.656 s 0.077 s 0.399 s +12 0.657 s 0.077 s 0.344 s +16 0.652 s 0.077 s 0.279 s +24 0.657 s 0.076 s 0.278 s +32 0.656 s 0.077 s 0.281 s +``` + +- The fixtures do not improve after K = 8. +- K = 16 also wins on flows where few lanes of a group survive. Against + K = 8, it is -42 % on stripes 1/4, -4..-14 % on random s=10/5/20/40 and + -3 % on 2D random 10. The other rows are `noise`. +- K = 24 and K = 32 are parity with K = 16. The only exception is random s=40, + at -5 %. +- K = 16 is the default. `BIOIMAGE_FLOW_NEON_LANES` overrides it for sweeps. + +On Zen 3, a full FP register file limited the kernel, so K = 4 was best. The +M2 core has 32 vector registers and a much larger reorder window. It overlaps +more independent chains, so K = 4 is 14 % (3D) to 34 % (2D) slower here. + +**Codegen.** We emitted the assembly of the translation unit with the exact +build command. The chunk lambdas of `trace_all_packed<2,16>` and `<3,16>` make +no out-of-line call. The compiler inlines `sample_packed`, `step_packed` and +`trace_block_packed`. The 3D function has 255 `fmla`, which is 16 lanes x 15 +plus 15 for the K = 1 tail. The linker strips local symbols from the `.so` +(`-Wl,-x -Wl,-S`), so `otool -tV` on the `.so` cannot name them. Inspect the +object file or the `-S` output instead. + +**Contraction differs from GCC.** The first port used the unfused position +update of the SSE kernel, `pos + (dt * step)`. It matched the portable kernel +bitwise on the fixtures and on all 400 differential cases. It differed on the +chaotic random flows: on random s=10, 38 % of the voxels differed, with equal +sums. The cause is the default `-ffp-contract=on` of clang. It fuses +`proposed[axis] += dt * step[axis]` in the scalar kernel into one FMA. The GCC +build does not, because the unfused SSE kernel matches it bitwise. The NEON +kernel now fuses the update with `vfmaq_f32(pos, dt, step)`. Every 3D case is +then bitwise identical to the arm64 portable kernel, and the speed is parity +(-1 %, `noise`). The difference on 2D random 10 comes from the weight-tree +sampler, as on x86. + +**Correctness.** + +- `differential_check.py`: 400/400 bitwise identical for `492ff9b` vs NEON, + and 400/400 for NEON vs `FORCE_SCALAR=1`. +- `check_flow_density.py --dim both --repeats 3` is unchanged: 3D + `rel_diff = 0.0470`, `pearson = 0.9690`, 2D `0.0201` / `0.9975`. +- Full suite: 1443 passed, 21 skipped. The skipped tests need optional + packages (nifty, skfmm, pygeodesic) or the x86-only AVX2 filter backends. +- The parity test is now `test_forced_scalar_matches_simd_backend`. It + accepts `"fma"` and `"neon"`. + +**Unification (step 6): keep two translation units.** The two kernels have the +same structure. But a `detail/simd4.hxx` wrapper does not remove three +differences between the platforms: + +- the lane count (4 vs 16), +- the contraction of the position update, which matches the scalar kernel of + each compiler, +- the NaN semantics of the clip. + +A shared header can also expose the x86 translation unit to the inlining +budget of GCC again (see the codegen trap). Decide again after a measurement +of the x86 kernel with K > 4. Zen 3 measured K = 6 as parity, but nobody tried +K >= 8. + +**Open items.** + +- Measure K = 16 for the x86 kernel. Without AVX-512 there are only 16 XMM + registers, so spills are likely. +- Measure an E-core-aware thread count. We did not measure 8T, because 4 of + the 8 cores are efficiency cores. +- Repeat the table on AC power with Low Power Mode off. + ## Correctness - The current suite reports 1057 passed / 8 skipped. The 17 diff --git a/development/flow/_flow_cases.py b/development/flow/_flow_cases.py index 6e5b11a..9d63ab0 100644 --- a/development/flow/_flow_cases.py +++ b/development/flow/_flow_cases.py @@ -58,6 +58,11 @@ def preload_core(so_path: str | os.PathLike | None) -> None: def pin_cpus(cpus: str | None) -> list[int]: """Pin the current process to ``cpus`` ('36,37') before allocating data.""" + if not hasattr(os, "sched_setaffinity"): + # macOS has no affinity API, so the scheduler places the threads. + if cpus: + raise SystemExit("--cpu is not supported on this platform") + return list(range(os.cpu_count() or 1)) if not cpus: return sorted(os.sched_getaffinity(0)) selected = {int(c) for c in cpus.split(",") if c} diff --git a/include/bioimage_cpp/flow/flow_density.hxx b/include/bioimage_cpp/flow/flow_density.hxx index f743dea..bafa639 100644 --- a/include/bioimage_cpp/flow/flow_density.hxx +++ b/include/bioimage_cpp/flow/flow_density.hxx @@ -677,12 +677,67 @@ bool try_trace_all_fma( #endif +#if defined(BIOIMAGE_FLOW_NEON) + +// NEON+FMA is baseline on arm64, so the packed tracer needs no runtime check. +void trace_all_neon_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_neon_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_neon( + 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 (force_scalar_requested()) { + return false; + } + if constexpr (D == 2) { + trace_all_neon_2d(positions, channels, grid, mask, n_threads, n_iter, dt, tol); + } else { + trace_all_neon_3d(positions, channels, grid, mask, n_threads, n_iter, dt, tol); + } + return true; +} + +#endif + } // namespace detail // Name of the tracer the default RK2/convergence/mask path will use on this -// machine: "fma" (runtime-dispatched FMA translation unit) or "scalar" -// (portable kernel). Exposed to Python as `_core._flow_trace_backend`. +// machine: "fma" (runtime-dispatched x86 FMA translation unit), "neon" (arm64 +// translation unit) or "scalar" (portable kernel). Exposed to Python as +// `_core._flow_trace_backend`. inline const char *trace_backend() noexcept { +#if defined(BIOIMAGE_FLOW_NEON) + if (!detail::force_scalar_requested()) { + return "neon"; + } +#endif #if defined(BIOIMAGE_FLOW_FMA_DISPATCH) if (!detail::force_scalar_requested() && detail::runtime_fma_supported()) { return "fma"; @@ -788,6 +843,13 @@ void compute_flow_density( )) { break; } +#endif +#if defined(BIOIMAGE_FLOW_NEON) + if (detail::try_trace_all_neon( + 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 diff --git a/src/bindings/flow.cxx b/src/bindings/flow.cxx index 8108600..9ae4b44 100644 --- a/src/bindings/flow.cxx +++ b/src/bindings/flow.cxx @@ -159,7 +159,7 @@ void bind_flow(nb::module_ &m) { "_flow_trace_backend", &flow::trace_backend, "Return the internal flow tracing backend selected for the default path " - "('fma' or 'scalar')." + "('fma', 'neon' or 'scalar')." ); m.def( "_compute_flow_density_2d_float32", diff --git a/src/cpp/flow/flow_density_neon.cxx b/src/cpp/flow/flow_density_neon.cxx new file mode 100644 index 0000000..d96abfc --- /dev/null +++ b/src/cpp/flow/flow_density_neon.cxx @@ -0,0 +1,494 @@ +#include "bioimage_cpp/flow/flow_density.hxx" + +#include "bioimage_cpp/detail/force_inline.hxx" +#include "bioimage_cpp/detail/profile.hxx" +#include "bioimage_cpp/detail/threading.hxx" + +#if !defined(BIOIMAGE_FLOW_NEON) +#error "flow_density_neon.cxx must only be built with BIOIMAGE_FLOW_NEON" +#endif +#if !defined(__aarch64__) && !defined(_M_ARM64) +#error "flow_density_neon.cxx requires an arm64 target" +#endif + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +// NEON port of the packed-channel tracer in flow_density_fma.cxx. +// See development/flow/PERFORMANCE_NOTES.md for the design. + +namespace bioimage_cpp::flow::detail { +namespace { + +constexpr std::ptrdiff_t kPackedTailSlack = 4; // floats past the last voxel + +template +struct PackedGrid { + const float *data = nullptr; + float32x4_t upper_f; // shape-1 per axis (float), 0 in unused lanes + uint32x4_t lane_mask; // all-ones in the D used lanes + std::array pstride{}; // float-element strides of the padded packed grid + std::array mstride{}; // element strides of the mask grid +}; + +BIOIMAGE_FORCE_INLINE float32x4_t make_f32x4(const float a, const float b, const float c, const float d) { + const float values[4] = {a, b, c, d}; + return vld1q_f32(values); +} + +// Repack the flow channel-last, padded by one replicated voxel per axis. +// Returns false if the offsets do not fit int32 or the allocation fails. +template +bool make_packed_grid( + const std::array &channels, + const GridLayout &grid, + const std::size_t n_threads, + std::unique_ptr &storage, + PackedGrid &out +) { + static_assert(D == 2 || D == 3, "packed tracer supports 2D and 3D"); + constexpr std::ptrdiff_t int_max = std::numeric_limits::max(); + constexpr std::ptrdiff_t channels_per_voxel = static_cast(D); + + std::array padded{}; + std::ptrdiff_t n_padded = 1; + for (std::size_t axis = 0; axis < D; ++axis) { + padded[axis] = grid.shape[axis] + 1; + n_padded *= padded[axis]; + } + const std::ptrdiff_t total = channels_per_voxel * n_padded + kPackedTailSlack; + std::array pstride{}; + pstride[D - 1] = channels_per_voxel; + for (std::size_t axis = D - 1; axis > 0; --axis) { + pstride[axis - 1] = pstride[axis] * padded[axis]; + } + if (total > int_max) { + return false; + } + for (std::size_t axis = 0; axis < D; ++axis) { + if (grid.strides[axis] > int_max) { + return false; + } + } + try { + storage.reset(new float[static_cast(total)]); + } catch (const std::bad_alloc &) { + return false; + } + float *packed = storage.get(); + for (std::ptrdiff_t i = total - kPackedTailSlack; i < total; ++i) { + packed[i] = 0.0f; + } + + // Copy one source row and duplicate its last voxel. + const std::ptrdiff_t width = grid.shape[D - 1]; + const auto fill_row = [&](float *o, const std::ptrdiff_t src_offset) { + if constexpr (D == 3) { + const float *c0 = channels[0] + src_offset; + const float *c1 = channels[1] + src_offset; + const float *c2 = channels[2] + src_offset; + for (std::ptrdiff_t x = 0; x < width; ++x) { + o[3 * x] = c0[x]; + o[3 * x + 1] = c1[x]; + o[3 * x + 2] = c2[x]; + } + o[3 * width] = c0[width - 1]; + o[3 * width + 1] = c1[width - 1]; + o[3 * width + 2] = c2[width - 1]; + } else { + const float *c0 = channels[0] + src_offset; + const float *c1 = channels[1] + src_offset; + for (std::ptrdiff_t x = 0; x < width; ++x) { + o[2 * x] = c0[x]; + o[2 * x + 1] = c1[x]; + } + o[2 * width] = c0[width - 1]; + o[2 * width + 1] = c1[width - 1]; + } + }; + + if constexpr (D == 3) { + const std::ptrdiff_t n_planes = padded[0]; + const auto pack_threads = ::bioimage_cpp::detail::normalize_thread_count( + n_threads, static_cast(n_planes) + ); + ::bioimage_cpp::detail::parallel_for_chunks( + pack_threads, static_cast(n_planes), + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + for (std::size_t zp = begin; zp < end; ++zp) { + const auto z = std::min(static_cast(zp), grid.shape[0] - 1); + for (std::ptrdiff_t yp = 0; yp < padded[1]; ++yp) { + const std::ptrdiff_t y = std::min(yp, grid.shape[1] - 1); + fill_row( + packed + static_cast(zp) * pstride[0] + yp * pstride[1], + z * grid.strides[0] + y * grid.strides[1] + ); + } + } + } + ); + } else { + const std::ptrdiff_t n_rows = padded[0]; + const auto pack_threads = ::bioimage_cpp::detail::normalize_thread_count( + n_threads, static_cast(n_rows) + ); + ::bioimage_cpp::detail::parallel_for_chunks( + pack_threads, static_cast(n_rows), + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + for (std::size_t yp = begin; yp < end; ++yp) { + const auto y = std::min(static_cast(yp), grid.shape[0] - 1); + fill_row(packed + static_cast(yp) * pstride[0], y * grid.strides[0]); + } + } + ); + } + + out.data = packed; + for (std::size_t axis = 0; axis < D; ++axis) { + out.pstride[axis] = static_cast(pstride[axis]); + out.mstride[axis] = static_cast(grid.strides[axis]); + } + if constexpr (D == 3) { + out.upper_f = make_f32x4(grid.upper[0], grid.upper[1], grid.upper[2], 0.0f); + const std::uint32_t lanes[4] = {~0u, ~0u, ~0u, 0u}; + out.lane_mask = vld1q_u32(lanes); + } else { + out.upper_f = make_f32x4(grid.upper[0], grid.upper[1], 0.0f, 0.0f); + const std::uint32_t lanes[4] = {~0u, ~0u, 0u, 0u}; + out.lane_mask = vld1q_u32(lanes); + } + return true; +} + +// lo + f * (hi - lo) with one rounding, like the contracted scalar lerp. +BIOIMAGE_FORCE_INLINE float32x4_t lerp(const float32x4_t lo, const float32x4_t hi, const float32x4_t f) { + return vfmaq_f32(lo, f, vsubq_f32(hi, lo)); +} + +// Lower corner, fraction and flat offset. Truncation equals floor because p >= 0. +// vcvtq_s32_f32 differs from cvttps only for NaN and overflow, and p is clipped. +struct Stencil { + float32x4_t frac; + int base; +}; + +template +BIOIMAGE_FORCE_INLINE Stencil make_stencil(const PackedGrid &g, const float32x4_t p) { + const int32x4_t i0 = vcvtq_s32_f32(p); + const float32x4_t frac = vsubq_f32(p, vcvtq_f32_s32(i0)); + int base = vgetq_lane_s32(i0, 0) * g.pstride[0] + vgetq_lane_s32(i0, 1) * g.pstride[1]; + if constexpr (D == 3) { + base += vgetq_lane_s32(i0, 2) * g.pstride[2]; + } + return Stencil{frac, base}; +} + +// 3D: nested lerps in the order of the scalar kernel. Lane 3 holds finite junk. +BIOIMAGE_FORCE_INLINE float32x4_t sample_packed(const PackedGrid<3> &g, const float32x4_t p) { + const Stencil s = make_stencil<3>(g, p); + constexpr int dx = 3; + const int dy = g.pstride[1]; + const int dz = g.pstride[0]; + const float *p0 = g.data + s.base; + const float *p1 = p0 + dz; + const float32x4_t c000 = vld1q_f32(p0); + const float32x4_t c001 = vld1q_f32(p0 + dx); + const float32x4_t c010 = vld1q_f32(p0 + dy); + const float32x4_t c011 = vld1q_f32(p0 + dy + dx); + const float32x4_t c100 = vld1q_f32(p1); + const float32x4_t c101 = vld1q_f32(p1 + dx); + const float32x4_t c110 = vld1q_f32(p1 + dy); + const float32x4_t c111 = vld1q_f32(p1 + dy + dx); + const float32x4_t fz = vdupq_laneq_f32(s.frac, 0); + const float32x4_t fy = vdupq_laneq_f32(s.frac, 1); + const float32x4_t fx = vdupq_laneq_f32(s.frac, 2); + const float32x4_t c00 = lerp(c000, c001, fx); + const float32x4_t c01 = lerp(c010, c011, fx); + const float32x4_t c10 = lerp(c100, c101, fx); + const float32x4_t c11 = lerp(c110, c111, fx); + const float32x4_t c0 = lerp(c00, c01, fy); + const float32x4_t c1 = lerp(c10, c11, fy); + return lerp(c0, c1, fz); +} + +// 2D: weighted sum of the four corners, in the same order as the SSE kernel. +BIOIMAGE_FORCE_INLINE float32x4_t sample_packed(const PackedGrid<2> &g, const float32x4_t p) { + const Stencil s = make_stencil<2>(g, p); + const float *r0p = g.data + s.base; + // Row loads: [c0(x0), c1(x0), c0(x1), c1(x1)]. + const float32x4_t r0 = vld1q_f32(r0p); + const float32x4_t r1 = vld1q_f32(r0p + g.pstride[0]); + const float32x4_t f = s.frac; // [fy, fx, 0, 0] + const float32x4_t omf = vsubq_f32(vdupq_n_f32(1.0f), f); // [gy, gx, 1, 1] + const float32x4_t vx = vcombine_f32(vdup_laneq_f32(omf, 1), vdup_laneq_f32(f, 1)); // [gx, gx, fx, fx] + const float32x4_t w0 = vmulq_laneq_f32(vx, omf, 0); // row y0 + const float32x4_t w1 = vmulq_laneq_f32(vx, f, 0); // row y1 + const float32x4_t t = vfmaq_f32(vmulq_f32(r0, w0), r1, w1); // [x0 terms | x1 terms] + return vaddq_f32(t, vextq_f32(t, t, 2)); // lanes 0,1 = the sum +} + +struct StepConstants { + float32x4_t zero; + float32x4_t half; + float32x4_t half_dt; + float32x4_t dt; + float32x4_t tol; +}; + +inline StepConstants make_step_constants(const float dt, const float tol) { + return StepConstants{ + vdupq_n_f32(0.0f), + vdupq_n_f32(0.5f), + vdupq_n_f32(0.5f * dt), + vdupq_n_f32(dt), + vdupq_n_f32(tol), + }; +} + +// True if any lane of the comparison mask is set. +BIOIMAGE_FORCE_INLINE bool any_lane(const uint32x4_t m) { + return vget_lane_u64(vreinterpret_u64_u16(vshrn_n_u32(m, 16)), 0) != 0; +} + +constexpr int kStillAlive = -1; + +// One RK2 step of a single particle with convergence and mask termination. +// Returns kStillAlive, or the ActiveTraceStats exit reason if tracing stops. +template +BIOIMAGE_FORCE_INLINE int step_packed( + float32x4_t &pos, const PackedGrid &g, const StepConstants &c, const std::uint8_t *mask +) { + // The clip zeroes the junk lane of the first sample. The mask zeroes it in the second. + float32x4_t step = sample_packed(g, pos); + float32x4_t mid = vfmaq_f32(pos, c.half_dt, step); + // vmaxnmq/vminnmq match maxps/minps: a NaN lane clips to 0. + mid = vminnmq_f32(vmaxnmq_f32(mid, c.zero), g.upper_f); + step = vreinterpretq_f32_u32(vandq_u32(vreinterpretq_u32_f32(sample_packed(g, mid)), g.lane_mask)); + + // The particle converged if no lane has |dt * step| >= tol. + const float32x4_t disp = vmulq_f32(c.dt, step); + if (!any_lane(vcgeq_f32(vabsq_f32(disp), c.tol))) { + return ActiveTraceStats::Converged; + } + + // Fused, because clang fuses `proposed += dt * step` in the scalar kernel. + const float32x4_t proposed = vfmaq_f32(pos, c.dt, step); + const uint32x4_t outside = vorrq_u32(vcltq_f32(proposed, c.zero), vcgtq_f32(proposed, g.upper_f)); + if (any_lane(outside)) { + return ActiveTraceStats::LeftMask; + } + // Round half up to the mask voxel. Truncation equals floor because proposed >= 0. + const int32x4_t rounded = vcvtq_s32_f32(vaddq_f32(proposed, c.half)); + int flat = vgetq_lane_s32(rounded, 0) * g.mstride[0] + vgetq_lane_s32(rounded, 1) * g.mstride[1]; + if constexpr (D == 3) { + flat += vgetq_lane_s32(rounded, 2) * g.mstride[2]; + } + if (mask[flat] == 0) { + return ActiveTraceStats::LeftMask; + } + pos = proposed; + return kStillAlive; +} + +template +BIOIMAGE_FORCE_INLINE float32x4_t load_position(const std::array &p) { + if constexpr (D == 3) { + return make_f32x4(p[0], p[1], p[2], 0.0f); + } else { + return make_f32x4(p[0], p[1], 0.0f, 0.0f); + } +} + +template +BIOIMAGE_FORCE_INLINE void store_position(std::array &p, const float32x4_t v) { + float tmp[4]; + vst1q_f32(tmp, v); + for (std::size_t axis = 0; axis < D; ++axis) { + p[axis] = tmp[axis]; + } +} + +template +BIOIMAGE_FORCE_INLINE void for_each_lane(F &&f, std::index_sequence) { + (f(std::integral_constant{}), ...); +} + +// Trace K particles in lockstep. Compile-time lane indices keep the state in registers. +template +void trace_block_packed( + std::array *positions, + const PackedGrid &g, + const StepConstants &c, + const std::uint8_t *mask, + const std::size_t n_iter, + ActiveTraceStats &stats +) { + float32x4_t pos[K]; + bool alive[K]; + [[maybe_unused]] std::size_t lane_steps[K] = {}; + [[maybe_unused]] int lane_exit[K] = {}; + for_each_lane([&](auto k) { + pos[k] = load_position(positions[k]); + alive[k] = true; + if constexpr (ActiveTraceStats::enabled) { + lane_exit[k] = ActiveTraceStats::MaxIter; + } + }, std::make_index_sequence{}); + + for (std::size_t iter = 0; iter < n_iter; ++iter) { + bool any_alive = false; + for_each_lane([&](auto k) { + if (!alive[k]) { + return; + } + if constexpr (ActiveTraceStats::enabled) { + ++lane_steps[k]; + } + const int exit_reason = step_packed(pos[k], g, c, mask); + if (exit_reason != kStillAlive) { + alive[k] = false; + if constexpr (ActiveTraceStats::enabled) { + lane_exit[k] = exit_reason; + } + return; + } + any_alive = true; + }, std::make_index_sequence{}); + if (!any_alive) { + break; + } + } + + for_each_lane([&](auto k) { + store_position(positions[k], pos[k]); + if constexpr (ActiveTraceStats::enabled) { + stats.record(lane_steps[k], lane_exit[k]); + } + }, std::make_index_sequence{}); +} + +template +void trace_all_packed( + std::vector> &positions, + const PackedGrid &g, + const std::uint8_t *mask, + const std::size_t n_threads, + const std::size_t n_iter, + const float dt, + const float tol +) { + const StepConstants c = make_step_constants(dt, tol); + ActiveTraceStats null_stats{}; + std::vector per_thread_stats; + if constexpr (ActiveTraceStats::enabled) { + per_thread_stats.assign(n_threads, ActiveTraceStats(n_iter)); + } + ::bioimage_cpp::detail::parallel_for_chunks( + n_threads, positions.size(), + [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { + ActiveTraceStats *stats = &null_stats; + (void)thread_id; +#ifdef BIOIMAGE_PROFILE + stats = &per_thread_stats[thread_id]; + const auto chunk_start = std::chrono::steady_clock::now(); +#endif + std::size_t i = begin; + if constexpr (K > 1) { + for (; i + K <= end; i += K) { + trace_block_packed(&positions[i], g, c, mask, n_iter, *stats); + } + } + for (; i < end; ++i) { + trace_block_packed(&positions[i], g, c, mask, n_iter, *stats); + } +#ifdef BIOIMAGE_PROFILE + stats->seconds += std::chrono::duration( + std::chrono::steady_clock::now() - chunk_start + ).count(); +#endif + } + ); + if constexpr (ActiveTraceStats::enabled) { + ActiveTraceStats::report(per_thread_stats, n_iter, true); + } +} + +// K = 16 measured best on Apple M2. See PERFORMANCE_NOTES.md for the sweep. +#if defined(BIOIMAGE_FLOW_NEON_LANES) +constexpr std::size_t kPackedLanes = BIOIMAGE_FLOW_NEON_LANES; +#else +constexpr std::size_t kPackedLanes = 16; +#endif + +template +void trace_all_neon( + 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 +) { + BIOIMAGE_PROFILE_INIT(profiler); + { + std::unique_ptr storage; + PackedGrid packed{}; + bool ready = false; + { + BIOIMAGE_PROFILE_SCOPE(profiler, "neon_pack"); + ready = make_packed_grid(channels, grid, n_threads, storage, packed); + } + if (ready) { + BIOIMAGE_PROFILE_SCOPE(profiler, "neon_trace_packed"); + trace_all_packed(positions, packed, mask, n_threads, n_iter, dt, tol); + BIOIMAGE_PROFILE_REPORT_NAMED(profiler, "[bioimage profile: flow neon]"); + return; + } + } + // Fallback if the packed grid cannot be built. + { + BIOIMAGE_PROFILE_SCOPE(profiler, "neon_trace_scalar"); + trace_all(positions, channels, grid, mask, n_threads, n_iter, dt, tol); + } + BIOIMAGE_PROFILE_REPORT_NAMED(profiler, "[bioimage profile: flow neon]"); +} + +} // namespace + +void trace_all_neon_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_neon<2>(positions, channels, grid, mask, n_threads, n_iter, dt, tol); +} + +void trace_all_neon_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_neon<3>(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 384a689..90889d9 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -266,10 +266,11 @@ def test_large_finite_flow_is_accepted_multithreaded(): assert density.sum() == np.prod(shape) -def test_forced_scalar_matches_fma_backend(monkeypatch): +def test_forced_scalar_matches_simd_backend(monkeypatch): monkeypatch.delenv("BIOIMAGE_CPP_FLOW_FORCE_SCALAR", raising=False) - if bic._core._flow_trace_backend() != "fma": - pytest.skip("FMA flow backend is not available on this CPU/build") + backend = bic._core._flow_trace_backend() + if backend not in ("fma", "neon"): + pytest.skip("no SIMD flow backend is available on this CPU/build") rng = np.random.default_rng(2026) cases = [] @@ -291,7 +292,7 @@ def run_all(): assert bic._core._flow_trace_backend() == "scalar" scalar = run_all() monkeypatch.setenv("BIOIMAGE_CPP_FLOW_FORCE_SCALAR", "0") - assert bic._core._flow_trace_backend() == "fma" + assert bic._core._flow_trace_backend() == backend # The two backends may differ in floating-point contraction, so the bar is # particle conservation plus near-identical densities, not bitwise equality. From b1b37be42e63014737ced72e62f0b20b776b962a Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Thu, 5 Nov 2026 09:04:44 +0100 Subject: [PATCH 2/2] Correct the Apple Silicon section of the flow performance notes - fmla count: 272 (16 x 16 + 16) for the fused position update, not 255. - K = 4 vs K = 16: K = 4 is 18 % / 53 % slower, not 14 % / 34 %. - 4T rows are unpinned, not P-core only. - Open item: GCC on Linux aarch64 likely breaks bitwise parity. Co-Authored-By: Claude Opus 5.5 (1M context) --- development/flow/PERFORMANCE_NOTES.md | 38 ++++++++++++++++----------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/development/flow/PERFORMANCE_NOTES.md b/development/flow/PERFORMANCE_NOTES.md index e0cee1b..cad15ce 100644 --- a/development/flow/PERFORMANCE_NOTES.md +++ b/development/flow/PERFORMANCE_NOTES.md @@ -586,7 +586,8 @@ Two semantic differences to SSE do not change the result: every finite input. Paired A/B of `492ff9b` (portable kernel) vs the NEON kernel (K = 16) on all -cases, with ABBA x4 and 3 inner calls. All rows are `significant`: +cases, with ABBA x4 and 3 inner calls. All rows are `significant`. The 4T +rows are not pinned, so macOS can put threads on the efficiency cores: ``` case threads previous neon change parity @@ -600,16 +601,16 @@ random s=5 1 1.6117 s 0.7129 s -55.8 % identical random s=20 1 1.6298 s 0.7461 s -54.2 % identical random s=40 1 0.8707 s 0.4972 s -42.9 % identical 2D random 10 1 0.2370 s 0.1209 s -49.0 % differs (2D tree) -fixture3d 4 P 0.4271 s 0.1861 s -56.4 % identical -fixture2d 4 P 0.0368 s 0.0203 s -44.9 % identical -stripes 1/4 4 P 0.2635 s 0.0782 s -70.3 % identical -stripes 3/4 4 P 0.4540 s 0.2016 s -55.6 % identical -random s=1 4 P 0.5576 s 0.2437 s -56.3 % identical -random s=10 4 P 0.6323 s 0.2554 s -59.6 % identical -random s=5 4 P 0.4386 s 0.1931 s -56.0 % identical -random s=20 4 P 0.5436 s 0.2322 s -57.3 % identical -random s=40 4 P 0.2846 s 0.1517 s -46.7 % identical -2D random 10 4 P 0.0613 s 0.0318 s -48.1 % differs (2D tree) +fixture3d 4 0.4271 s 0.1861 s -56.4 % identical +fixture2d 4 0.0368 s 0.0203 s -44.9 % identical +stripes 1/4 4 0.2635 s 0.0782 s -70.3 % identical +stripes 3/4 4 0.4540 s 0.2016 s -55.6 % identical +random s=1 4 0.5576 s 0.2437 s -56.3 % identical +random s=10 4 0.6323 s 0.2554 s -59.6 % identical +random s=5 4 0.4386 s 0.1931 s -56.0 % identical +random s=20 4 0.5436 s 0.2322 s -57.3 % identical +random s=40 4 0.2846 s 0.1517 s -46.7 % identical +2D random 10 4 0.0613 s 0.0318 s -48.1 % differs (2D tree) ``` The gain is larger than on x86 (-30..-40 %). The lane count is the main @@ -642,14 +643,17 @@ K fixture3d fixture2d stripes 1/4 On Zen 3, a full FP register file limited the kernel, so K = 4 was best. The M2 core has 32 vector registers and a much larger reorder window. It overlaps -more independent chains, so K = 4 is 14 % (3D) to 34 % (2D) slower here. +more independent chains, so K = 4 is 18 % (3D) to 53 % (2D) slower than +K = 16 here. **Codegen.** We emitted the assembly of the translation unit with the exact build command. The chunk lambdas of `trace_all_packed<2,16>` and `<3,16>` make no out-of-line call. The compiler inlines `sample_packed`, `step_packed` and -`trace_block_packed`. The 3D function has 255 `fmla`, which is 16 lanes x 15 -plus 15 for the K = 1 tail. The linker strips local symbols from the `.so` -(`-Wl,-x -Wl,-S`), so `otool -tV` on the `.so` cannot name them. Inspect the +`trace_block_packed`. The 3D function has 272 `fmla`, which is 16 lanes x 16 +plus 16 for the K = 1 tail. Each lane-step has 7 lerps per sample, one +midpoint FMA and the fused position update. The 2D function has 68 `fmla` +(16 x 4 + 4). The unfused first port had 255 (16 x 15 + 15) in 3D. The +linker strips local symbols from the `.so` (`-Wl,-x -Wl,-S`), so `otool -tV` on the `.so` cannot name them. Inspect the object file or the `-S` output instead. **Contraction differs from GCC.** The first port used the unfused position @@ -693,6 +697,10 @@ K >= 8. - Measure K = 16 for the x86 kernel. Without AVX-512 there are only 16 XMM registers, so spills are likely. +- Measure parity with GCC on Linux aarch64 before those wheels ship. GCC + probably does not fuse the scalar position update, as on x86. NEON vs + `FORCE_SCALAR=1` is then not bitwise identical there. The parity test + accepts this, but the bitwise claims above apply to clang only. - Measure an E-core-aware thread count. We did not measure 8T, because 4 of the 8 cores are efficiency cores. - Repeat the table on AC power with Low Power Mode off.