[jvdd]

Problem

The SIMD functions run a scalar loop on slices shorter than one vector, and on the elements after the last full vector. For the floats that ignore NaNs (ArgMinMax on f32 and f64), that loop tests a first_non_nan_update flag at every value. Because of that flag, LLVM compiles the update of the minimum (maximum) as a branch instead of a conditional move. On random data that branch mispredicts: about 2 ns per element, against 0.3-0.6 ns for a plain if v > max { max = v; index = i } loop. On an i7-1185G7 (AVX-512, 16 f32 lanes), argmax on 2-15 random f32 values takes 9-30 ns; such a plain loop takes 2-8 ns.

This is why candle (huggingface/candle#4051) calls argminmax only for rows of at least 128 elements.

Changes

• The ignore-NaN scalar loops start at the first non-NaN value (index 0 if all values are NaN). From there NaNs are never lower or higher, so the loop needs neither a NaN check nor the flag. SCALARInit::_init_min, _init_max and _allow_first_non_nan_update become one _start_index (0 for Int and FloatReturnNaN). The f16/bf16 scalar code already works this way.
• f32 and f64 argmin/argmax now compile without branches in the loop (AVX2 f32 argmax: 109 → 53 jumps). argminmax keeps its if .. else if ..: two separate ifs made integer and i128 argminmax slower.
• The results do not change: the first index wins on ties, NaNs are skipped, and an all-NaN slice returns 0. The integer code keeps the same jumps and cmovs; the return-NaN code keeps its structure (a few jumps fewer).

Checked

• Tests: cargo test (all features and no default features) on an M5 (NEON). On an i7-1185G7: native debug, release and no default features, and Intel SDE with -spr, -hsw, -nhm and an AVX-512F-only CPU. Native on a Ryzen 9 5950X. Clippy as in CI, also for x86_64, i686, armv7 and wasm32.
• Benchmarks of this commit against main on the i7-1185G7 (AVX-512) and the Ryzen 9 5950X (AVX2):
  • Coverage: f32, f64 and i32; argmin, argmax and argminmax; every length from 1 to 72, plus 127, 128, 135, 143, 1000 and 4096.
  • Data: random, sorted, reversed, sorted with random neighbour swaps, and random lengths 1..=k. Each call gets fresh data.
  • Controls: an identical copy of main (and, for f64 on the i7, of this commit) to catch code-layout effects. One AVX-512 f64 argmin build was 3.7 ns slower per call at every length from 8; the effect followed whichever copy sat at that address.
  • Shorter than one vector, random data: f32 0.30x (i7) and 0.43x (Ryzen), f64 0.38x and 0.60x. With neighbour swaps: 0.30-0.64x.
  • Remainders of 1-7 elements, random data: 0.60-0.86x. On AVX-512, 8-15 remaining f32 elements (main runs them in the scalar loop too): 0.53x.
  • Random lengths 1..=k: f32 0.54-0.59x, f64 0.52-0.67x.
  • Sorted data: f32 0.86-1.06x (median per group; single lengths up to 1.11x, e.g. 15 remaining elements on AVX-512: 16.5 → 18.3 ns). f64 shorter than one vector: about 1.07x (0.2-0.5 ns). Single elements: 0.2-0.5 ns slower on any data, because of the scan for the first non-NaN value.
  • argminmax: 0.84-1.0x. Integers: unchanged. Long slices: unchanged, or faster where main leaves 8-15 elements to the scalar loop (1000 f32 on AVX-512: 0.8x, 69 → 58 ns).
  • Every benchmark cell returns the same indices as a plain loop.
• Known cost: when no SIMD is used for the whole slice, the branch-free loop runs at a fixed speed, set by the chain from maxss through the compare to the index cmov. That applies to the public SCALAR type, to wasm32 and ARMv7 without nightly_simd, and to f32 on x86 without SSE4.1 and f64 without AVX. 4096 f32 or f64 values, argmin/argmax:
  • i7: about 2x slower on random and sorted data (1.6-1.9 → 3.6-3.9 µs).
  • Ryzen: 0.75-0.8x (1.8-2.3 → 1.4-1.8 µs).
  • The old branchy loop was fast while its branch predicted well, and about 5x slower when it did not. On sorted data with random neighbour swaps, argmax took 8.3 µs on the i7 (now 3.6 µs) and 9.1 µs on the Ryzen (now 1.7 µs).
  • So on this path, f32/f64 argmin/argmax no longer depend on the data (after any leading NaNs, which the scan for the first non-NaN value reads first), but on Intel they are slower for random data. argminmax keeps its if .. else if .. branches. CodSpeed's scalar benchmarks may show this.
• Also measured on top of perf: fold the last vector into the SIMD core, AVX2 for short 16-bit slices #131, which this PR does not depend on. In candle's own harness (candle#4051 without its row threshold), the row length from which argminmax beats candle's loop at every length drops from 87 to 30 for f32, and from 56 to 39 for f64, on both CPUs.

[codspeed]

Merging this PR will regress 2 benchmarks

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 8 improved benchmarks
❌ 2 regressed benchmarks
✅ 208 untouched benchmarks

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
❌	scalar_f32_argminmax_masked	1.1 ms	1.2 ms	-5.79%
❌	scalar_f64_argminmax_masked	1.5 ms	1.5 ms	-4.14%
⚡	scalar_f32_argminmax_rn	959.7 µs	905.8 µs	+5.95%
⚡	scalar_f32_argminmax_in	848.7 µs	801.2 µs	+5.93%
⚡	scalar_f32_argmax_in	770 µs	737.4 µs	+4.42%
⚡	scalar_f32_argmin_in	770 µs	737.4 µs	+4.42%
⚡	scalar_f64_argminmax_in	1.2 ms	1.2 ms	+3.87%
⚡	scalar_f64_argminmax_rn	1.3 ms	1.3 ms	+3.71%
⚡	scalar_f64_argmin_in	1.1 ms	1.1 ms	+2.27%
⚡	scalar_f64_argmax_in	1.1 ms	1.1 ms	+2.27%

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing perf/scalar-first-non-nan (fea95ff) with main (f68aba5)}