Scan M4 decimation buckets with SIMD argmin/argmax

Merged
JV
4d ago
+257-27

Problem

M4 decimation (m4_range in src/kernels.rs) handles one row at a time. For every row it computes the bucket from x, compares y with the running min and max, and flushes when the bucket changes. So it reads x and y in full, and the bucket state machine keeps it scalar (simd.rs notes that a two-phase SIMD version was slower). M4 runs for the first payload of every line or area trace over DECIMATION_THRESHOLD (10k rows) and for every zoom (decimate_view). On one core, 1M rows into 2048 buckets take about 3 ms.

Changes

  • x is ascending, so the rows of a bucket are contiguous. The new m4_range_buckets handles one bucket at a time: m4_bucket_end probes x to find where the bucket stops (two probes when the buckets are equally wide, else a galloping search), and m4_bucket scans the bucket's y once. x is only probed, not read in full.
  • A bucket of 32 rows or more goes to the new simd::try_argminmax, which calls the argminmax crate. The crate selects SSE4.1, AVX2 or AVX-512 at run time on x86, uses NEON on aarch64, and uses a scalar loop on other targets (wasm32, armv7). Shorter buckets use a scalar loop.
  • A window whose buckets average fewer than 8 rows (M4_BUCKET_WALK_MIN_ROWS, deep zooms) keeps the per-row loop. That loop is the old m4_range, unchanged and now called m4_range_rows. It is #[inline(never)] so that its machine code stays the same as on main: when it was inlined next to the bucket walk, it lost registers and was 8% slower on the i7.
  • The rows do not change. NaN and ±∞ are still skipped (§19): argminmax skips NaN but not ±∞, so a bucket with a ±∞ minimum or maximum takes the scalar loop. Ties keep the first row, as before. XY_SIMD=0 also turns argminmax off.
  • New test fuzz_m4_bucket_walk_matches_per_row_loop: on 2000 hostile cases (uniform, duplicate-heavy and irregular x; NaN, ±∞, ±0, huge and repeated y), the bucket walk must give the same rows as the per-row loop. It fails when the ±∞ check, the bucket-end search or the tie rule is broken.
  • spec/design/rust-engine.md (the dependency list and §3.4), the simd.rs doc and the comment in Cargo.toml describe the change. The news fragment is news/535.performance.md.
  • New dependency: argminmax 0.7.1 (MIT), with default-features = false, features = ["float"]. It adds num-traits and, at build time, autocfg. It has over 16 million downloads on crates.io. Polars uses it for arg_min and arg_max, and tsdownsample uses it for the same kind of M4 and MinMax downsampling. Disclosure: I maintain argminmax.

Benchmarks

Machines: i7-1185G7 (Tiger Lake, AVX-512, 4 cores, Rust 1.95) and Ryzen 9 5950X (Zen 3, AVX2, 16 cores, Rust 1.94), both on Ubuntu, with the release profile of the repository (baseline x86-64, no target-cpu). main (d95f84b6) and this PR alternate. Each cell is the median of 5 rounds. "1 core" is pinned with taskset, so par_threads is 1, as on CodSpeed. "All cores" is not pinned (8 and 18 threads).

benchmarks/bench_native.py --sizes 1e5,1e6,1e7 (the kernel bench of CI, y = sin(i·1e-3), 2048 buckets):

rowsi7, 1 corei7, all coresRyzen, 1 coreRyzen, all cores
M4 (full)100k294 → 861 Mpt/s, 2.9x306 → 881, 2.9x262 → 586, 2.2x256 → 591, 2.3x
M4 (full)1M351 → 1922, 5.5x1067 → 3587, 3.4x318 → 2315, 7.3x1658 → 2238, 1.35x
M4 (full)10M358 → 1857, 5.2x1289 → 2501, 1.9x335 → 2545, 7.6x2448 → 7454, 3.0x
zoom re-decimate (0.5 rows per bucket)100k25 → 24 µs23 → 24 µs64 → 64 µs64 → 65 µs
zoom re-decimate (5 rows per bucket)1M51 → 54 µs ¹57 → 60 µs ¹99 → 99 µs101 → 101 µs
zoom re-decimate (49 rows per bucket)10M337 → 127 µs, 2.7x357 → 185 µs, 1.9x397 → 174 µs, 2.3x413 → 178 µs, 2.3x

¹ The bench takes the best of 3 calls. With 2000 calls of the same zoom (this window runs on one thread), the i7 gives 51.0–51.9 µs on main and 50.8–51.2 µs with this PR.

The other kernels in this bench (encode, zone maps, histogram, bin_2d, normalize, range scan) do not change; between the two builds they move 0.90–1.12x, which is code placement.

Through the Python API, 1 core, x = arange(n), random-walk y (as in bench_line.py), 2048 buckets:

rowsi7Ryzen
first payload: xy.chart(xy.line(x=x, y=y)).figure().build_payload_split(2048)100k1.10 → 0.86 ms, 1.28x1.57 → 1.16 ms, 1.35x
first payload1M7.55 → 5.19 ms, 1.46x8.48 → 5.26 ms, 1.61x
first payload10M74.9 → 51.0 ms, 1.47x79.5 → 51.9 ms, 1.53x
fig.decimate_view over 10%1M516 → 248 µs, 2.1x797 → 418 µs, 1.9x
fig.decimate_view over 10%10M3.30 → 0.50 ms, 6.6x3.95 → 0.66 ms, 5.9x
fig.decimate_view over 1% (5 rows per bucket)1M153 → 152 µs307 → 317 µs ²
control: density scatter payload (no M4)1M1.00x1.00x

² Within the spread of that cell (±6%). The Rust microbench below gives 1.00–1.02x there.

With normal y (as in test_codspeed_kernels.py), the first payload gains less (1.16–1.28x). With all cores, the first payload of 1M rows on the Ryzen is 1.10x faster (1.23x at 10M). On the i7 laptop with all 8 threads, the first payload varied between processes with this PR: 2.3–4.2 ms for the random walk at 1M, against 2.8–2.9 ms on main. It also varied with argminmax forced to AVX2 (2.3–3.2 ms), and the M4 call alone was 3–5x faster in every run. main's build did not vary with the PR's library preloaded. I could not find the cause; the Ryzen and single-core runs do not show it.

Rows per bucket decide the gain. The M4 kernel alone (a Rust copy of this PR's code), 1M rows, 2048 buckets, a window that holds the given rows per bucket, sin, normal and random-walk y, 1 core:

rows per bucket0.5–48163264488 (full range)
i7, time vs main1.00–1.010.80–0.920.60–0.700.28–0.300.18–0.200.20–0.22
Ryzen, time vs main1.00–1.020.70–0.880.61–0.920.23–0.360.15–0.200.11–0.13

Without the per-row loop for sparse windows, the bucket walk was 1.1–1.7x slower at 0.5–2 rows per bucket, and up to 1.3x slower on the i7 at 4–5 rows. The Ryzen gained at 4 rows, the i7 did not, so the threshold is 8.

Why a dependency, and not a two-pass scan? The same bucket walk with xy's own AVX2 min/max needs no new crate and gives the same rows:

let (min, max) = simd::try_min_max(ys).unwrap_or_else(|| min_max_scalar(ys));
let min_i = ys.iter().position(|&v| v == min);
let max_i = ys.iter().position(|&v| v == max);

It is 1.6–2.2x faster than main, against 3.0–9.1x with argminmax: M4 over 1M rows (same microbench, normal, random-walk and irregular-x data) takes 0.55–0.63 (i7) and 0.45–0.54 (Ryzen) of main's time with it, against 0.20–0.33 and 0.11–0.16 with argminmax. Its two position scans are scalar and read most of the bucket again, while argminmax keeps the indices in SIMD registers in one pass.

Notes

  • Dependency form. rust-engine.md prefers optional crates behind a feature, with the lean build as default. I made argminmax a plain dependency: every build, including the release wheels (hatch_build.py), runs cargo build --release with default features, so a crate that is off by default would not reach users. If you prefer a default-on argminmax feature (lean build with --no-default-features), I can change that. Either way, the dev sandbox needs the crate vendored or cached.
  • The C ABI does not change, so ABI_VERSION stays the same.
  • #424 changes m4_indices and the thread gates. This PR changes only the serial m4_range they call, so the two do not conflict.
  • Not measured: speed on aarch64 (the tests pass there), the Pyodide wheel (argminmax builds for wasm32 and uses its scalar loop there) and Windows.

Checked

  • cargo test (139 passed) and cargo clippy --all-targets -- -D warnings on macOS aarch64, on the i7 (AVX-512) and on the Ryzen (AVX2). The M4 tests also under Intel SDE as Haswell (AVX2), Nehalem (SSE4.2) and Merom (no SSE4.1, so argminmax's scalar loop), and with XY_SIMD=0.
  • cargo check --release --lib for i686, armv7, aarch64-linux and wasm32-unknown-unknown.
  • pytest on macOS aarch64: 4096 passed. On the Ryzen: 4089 passed; 6 tests fail there with main's library too, because they need Node 20 or newer, Chromium or a git checkout. test_legend_best_live.py::test_update_payload_keeps_auto_legend_chrome_and_scores_after_animation (a browser probe) also failed on macOS in some runs, on main and with this PR.
  • make news-check and pre-commit run --all-files: ruff check and ruff format pass. The docs/app codespell hook needs that app's environment, so I ran codespell on the changed files instead.

Summary by CodeRabbit

  • Performance
    • Line and area charts with large datasets, including charts with 1 million points, are prepared and re-decimated after zooming more quickly.
    • Decimation is faster for dense chart data, reducing wait times when viewing or exploring large charts.
  • Reliability
    • Decimation continues to produce the same chart rows, preserving displayed results.