[jvdd]

Problem

The SIMD functions leave the elements after the last full vector (up to one vector minus one) to a scalar loop. On short slices that loop is a large part of the call. Where it compiles to branches (for example argminmax for every type, and f32/f64 argmin/argmax before #133), its run time depends on the data: fast on sorted data, slow on random data.

This hurts most where the vectors are wide. With AVX-512 (the default dispatch on stable since #83), i16, u16 and f16 hold 32 lanes, so slices of 16-31 and 48-63 elements are 1.2-2.5x slower than with argminmax 0.6.3.

Changes

• The SIMD core (SIMDCore::_core_argminmax) takes any slice of at least one vector. If the length is not a multiple of the vector length, the core also compares the last vector of the slice, before its one horizontal reduction. That vector overlaps elements that the core already saw, and these keep their first index on ties. Its lanes get the indices of the vector after the last full one, and the core maps the result back after the reduction. This replaces the scalar loop on the remainder: task.rs no longer handles a remainder. (Only when one slice holds NaNs with different bit patterns can the unmasked nanarg* functions now return another of those NaNs, which the README limitations already allow.)
• The overflow-safe loop returns the core result directly when the slice fits in one chunk, which is nearly every call. Before, it merged that result with a start value: for f16 and bf16, that merge uses the half comparison, which branches on the sign bit. A last chunk shorter than one vector now starts in the previous chunk.
• The masked functions keep the scalar remainder: a second masked core call on the last vector was slower on almost all lengths.
• On AVX-512 CPUs, 16-bit slices shorter than 32 elements go to AVX2. At 16-31 elements, AVX2 runs one 16-element vector and the last vector, instead of running all of them in the scalar loop. The masked functions take the same route; there AVX2 runs one vector and the scalar remainder: 0.35-0.93x of main at 16-31 elements on the i7-1185G7 (all 180 cells faster), 0.81-1.08x on random lengths 1..=31. For the 32 and 64-bit types this is a trade-off, so they keep AVX-512: AVX2 was faster on random data (i32/u32 0.4x, f32 and f64 0.8x), but about 2x slower on sorted floats (median; single cells up to 3.3x), up to 1.3x slower on sorted i32/u32 and 1.5x slower on 64-bit integers on any data.
• split_array and its four-way matches are gone. The 8-bit SSE comment is updated to the new numbers.
• Tests: the NaN and infinity tests now also use lengths that end in a last vector. New cases put the min, max or NaN only in the last vector (in one of its own lanes, or only where it overlaps the vector before) and in a last overflow chunk shorter than one vector. u32 gets the overflow test (its NEON chunks hold 252 elements).

Rework (since ed58966)

The first version of this PR (c9f093a to ed58966) called the SIMD core a second time on the last vector when 8 or more elements remained (a second horizontal reduction plus a merge), and kept the scalar loop for 1-7 remaining elements (MIN_SIMD_REMAINDER). The rework (26c60bf, 14e21a2) folds that vector into the first core call and adds the one-chunk shortcut. MIN_SIMD_REMAINDER, find_final_index_min/max and the remainder blocks in task.rs are gone. In the library code (lib.rs, generic.rs, task.rs) the PR is now +127 −226 lines (first version: +129 −156); the tests grow by +93 −29.

Reworked PR vs the first version (median of the per-cell time ratios; in brackets the share of cells that are more than 5% faster / slower against both copies of the base):

new / first version	i7-1185G7 (AVX-512)	Ryzen 9 5950X (AVX2)	M5 (NEON)
Shorter than one vector	1.00 (17% faster, 7% slower)	1.01 (13% faster, 18% slower)	1.00 (5% faster, 9% slower)
Exact multiple of the vector	0.97 (34% faster, 1% slower)	0.96 (36% faster, 6% slower)	0.99 (24% faster, 3% slower)
Remainder 1-7	0.84 (82% faster, 1% slower)	0.81 (87% faster, 0% slower)	0.82 (63% faster, 6% slower)
Remainder 8 or more	0.70 (94% faster, 0% slower)	0.80 (86% faster, 1% slower)	0.77 (87% faster, 0% slower)
Random lengths 1..=16	0.98 (27% faster, 0% slower)	0.82 (64% faster, 3% slower)	0.69 (85% faster, 0% slower)
Random lengths 1..=64	0.67 (100% faster, 0% slower)	0.68 (91% faster, 0% slower)	0.73 (82% faster, 3% slower)
Random lengths 1..=256	0.65 (100% faster, 0% slower)	0.81 (88% faster, 0% slower)	0.88 (58% faster, 0% slower)
Long (4096, 65536)	1.00 (0% faster, 1% slower)	0.99 (17% faster, 11% slower)	1.00 (1% faster, 0% slower)

Faster: every slice with a remainder (one horizontal reduction instead of two, and no scalar loop), and short slices on f16/bf16 (no merge). Slower than the first version:

• Exact multiples of the vector, argmin/argmax on x86: the index mapping after the reduction costs 0.4-0.9 ns. The one-chunk shortcut hides it in most cells; it shows where the slice spans two overflow chunks, e.g. u8 at 256 and 512 elements on AVX-512: +3-8% (up to 0.8 ns at 256, 0.8-1.5 ns at 512; one cell +4 ns, 12%).
• 64-bit types with a remainder: the last vector waits for the running min/max, while the old scalar tail of 1-7 elements ran in parallel with the reduction. i64 argmin/argmax on NEON at 31 and 63 elements: +7-11% (about 1 ns); i64 argminmax on AVX-512 at 12 and 20 elements: +5-9% on sorted and reversed data, where the old scalar tail predicted well (random data: 0.72x).
• A nanarg* slice of one chunk whose first element is NaN is now scanned in full; before, the merge loop returned at once. The result is the same (unless the NaNs have different bit patterns, see above), and the run time no longer depends on where the NaN is.
• Slices shorter than one vector: the same scalar code, slower in 7-18% of the cells (median +0.25 ns on the i7, +0.6 ns on the Ryzen, +0.45 ns on the M5). The compiler places and allocates that code differently; see "Code placement" below.

Benchmarks

i7-1185G7 (AVX-512), Ryzen 9 5950X (AVX2; other users' jobs on other cores) and M5 (NEON). 11 types (f16, f32, f64, i16, i32, i64, u8, u16, u32, and f32/f64 nanarg*), argmin, argmax and argminmax, random, sorted, reversed and pairwise-swapped data, 25 lengths from 1 to 65536 and streams of random lengths 1..=k. Every call gets fresh data from an L2-resident pool; p25 of 11 rounds. main and the new head each have an identical copy as a control.

This PR vs main (median per group; share of cells more than 5% faster / slower against both copies):

new / main	i7-1185G7 (AVX-512)	Ryzen 9 5950X (AVX2)	M5 (NEON)
Shorter than one vector	0.91 (74% faster, 0% slower)	0.92 (54% faster, 1% slower)	0.93 (41% faster, 1% slower)
Exact multiple of the vector	0.96 (36% faster, 2% slower)	0.96 (30% faster, 3% slower)	0.99 (26% faster, 3% slower)
Remainder 1-7	0.68 (92% faster, 1% slower)	0.71 (93% faster, 0% slower)	0.77 (66% faster, 5% slower)
Remainder 8 or more	0.58 (95% faster, 0% slower)	0.65 (93% faster, 1% slower)	0.49 (85% faster, 0% slower)
Random lengths 1..=16	0.96 (42% faster, 0% slower)	0.68 (76% faster, 0% slower)	0.64 (82% faster, 0% slower)
Random lengths 1..=64	0.50 (100% faster, 0% slower)	0.56 (91% faster, 0% slower)	0.69 (82% faster, 0% slower)
Random lengths 1..=256	0.47 (100% faster, 0% slower)	0.67 (94% faster, 0% slower)	0.89 (52% faster, 0% slower)
Long (4096, 65536)	1.00 (0% faster, 1% slower)	1.00 (9% faster, 8% slower)	1.00 (0% faster, 0% slower)

Typical calls:

Call (random data), ns: main / first version / this PR	i7-1185G7	Ryzen 9 5950X	M5
f32 argmin, 16 elements (Lance 4-bit PQ)	9.2 / 8.6 / 8.9 (0.97x)	9.7 / 9.4 / 9.2 (0.95x)	4.5 / 4.5 / 4.3 (0.95x)
f32 argmin, 256 elements (Lance 8-bit PQ)	21.4 / 20.8 / 20.3 (0.95x)	26.5 / 26.0 / 26.5 (1.00x)	41.2 / 41.1 / 41.4 (1.01x)
f32 argmax, 127 elements	35.3 / 21.8 / 14.9 (0.42x)	31.8 / 30.8 / 17.7 (0.56x)	30.1 / 30.2 / 22.8 (0.76x)
f16 argmax, 63 elements	36.0 / 14.9 / 8.3 (0.23x)	39.5 / 24.4 / 16.6 (0.42x)	31.9 / 30.5 / 22.3 (0.70x)
i16 argminmax, 40 elements	24.1 / 15.3 / 9.2 (0.38x)	24.3 / 11.7 / 8.9 (0.37x)	8.5 / 8.5 / 8.7 (1.02x)
i64 argminmax, 12 elements	17.0 / 15.1 / 10.9 (0.64x)	8.5 / 8.6 / 7.7 (0.91x)	7.0 / 6.3 / 6.4 (0.92x)
f32 argminmax, random lengths 1..=256	54.5 / 42.0 / 25.7 (0.47x)	49.4 / 47.0 / 28.1 (0.57x)	44.2 / 44.1 / 38.7 (0.87x)
i16 argmin, random lengths 1..=64	20.8 / 14.9 / 11.2 (0.54x)	17.2 / 15.1 / 12.1 (0.70x)	9.1 / 8.3 / 6.1 (0.67x)
f32 argmax, 65536 elements	2829 / 2846 / 2828 (1.00x)	4062 / 4092 / 4098 (1.01x)	9943 / 9958 / 9937 (1.00x)

Regressions vs main

• Exact multiples, x86 argmin/argmax: the index mapping (see above). It shows where the one-chunk shortcut does not apply: u8 at 256-512 elements on AVX-512 +6-11% (1.1-2.0 ns; one cell +4.9 ns, 15%).
• 64-bit with a remainder: NEON i64 argmin/argmax at 31 and 63 elements +9-15% (about 1 ns), f64 +5-7%; AVX-512 i64 argminmax at 12 and 20 elements +7-10% (0.7-1 ns) on sorted and reversed data (random data: 0.64x).
• Single cells (less than 5% of the cells per CPU): Ryzen long i16/u16 arrays and f16 argminmax at 4096 (+8-20%), M5 u8 below 16 elements (up to 1.3x, 8 elements 6.1 to 8.1 ns), M5 i16/i32 at 40-48 elements (+6-13%). With every loop aligned to 64 bytes (both x86 CPUs), most of these cells are back at parity. A few long-array cells stay 6-10% slower (Ryzen u16 argminmax at 65536 and i16 at 4096, i7 u32 argmin at 65536); I could not explain them, and the medians of these groups stay at 1.00.
• With perf: start the scalar loops that ignore NaNs at the first non-NaN value #133 merged as well, f32 argmax/argmin on random lengths 1..=16 on AVX-512 were about 2 ns (1.18x) slower than with the first version of this PR: there LLVM now unrolls the unchanged scalar loop 4x, which mispredicts more when the length varies. At fixed lengths the unrolled loop is 0.94-0.98x.

Code placement

A change to these functions moves the code after it. In the study for this rework, cells that changed by ±5-23% while their code did not change went back to parity in a build that aligns every loop to 64 bytes (-C llvm-args=-align-loops=64). In a user's build, functions are aligned to 16 bytes only, so where a loop falls relative to the 64-byte lines depends on the rest of the binary, for main and for this PR alike.

Follow-up

The index mapping is the only cost that comes from the design. Giving the last vector its true indices (new_index + broadcast(rem)) instead removes it: in a prototype, exact multiples on x86 were 0.92-1.00x of the version without the mapping. That needs an index broadcast for each of the 56 SIMD implementations (for example in place of the INDEX_INCREMENT constant, about +110 lines), so it is a separate PR.

Checked

• Tests: cargo test --all-features and --no-default-features on an M5 (NEON). On the i7-1185G7, native tests (debug, release, no default features) and Intel SDE runs with -spr, -hsw, -nhm and an AVX-512F-only CPU (21 of 21 runs pass). On the Ryzen, release with all features and debug without default features. Clippy with -D warnings for x86_64, i686, aarch64 and armv7 (tests/argminmax_test.rs has a useless_conversion warning with clippy 1.95 on main as well).
• The results match the scalar implementation in every benchmark cell. Earlier review rounds ran differential tests against the scalar code: up to 555,000 arrays on NEON with simulated lane counts, and 615,000 arrays on the SSE, AVX2 and AVX-512 code of the i7-1185G7 (all 12 types, every length up to 300 and the f16/bf16 chunk boundaries, NaN, infinity and signed-zero patterns). Review round 6 (Codex and Opus, on the rework) found no correctness bug: 18.9 million (Codex) and 28.4 million (Opus) differential comparisons with the scalar code, with emulated lane counts 2-64 (8-bit index lanes, float indices, chunks of 1-9 vectors) and natively on NEON, gave no difference. Both found the same test gap: no test had its winner only in a last chunk shorter than one vector. 9ed1b66 adds such cases: a mutant that merges that chunk with the wrong start now fails 8 tests (it passed all tests before). Review round 7 (Codex and Opus, on the whole PR at ff79bc2) found no correctness bug either: 43.7 million differential comparisons with the scalar code (Codex, emulated lane counts 2-64 and native NEON) and 3.0 million arrays (Opus, natively on SSE4.2, AVX2 and AVX-512 of the i7-1185G7, and on NEON) gave no difference. Codex found that the default length of the overflow test overflowed a 32-bit usize (1 << 33 in a debug build on i686; CI only builds that target): it is now computed only when no length is given. Opus found that the masked 16-bit route to AVX2 was neither stated nor measured: it is now measured (see above).

[codspeed]

Merging this PR will improve performance by 2.05%

1 warning

⚠️ Base and head ran on different runner types

Some benchmarks with significant performance changes were compared across different runner types, so the changes might not come from the code. Run base and head on the same runner type, or use CodSpeed Macro Runners for more stable results.

Learn more about Macro Runners

Open the report in CodSpeed to investigate

⚡ 2 improved benchmarks
✅ 216 untouched benchmarks

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
⚡	sse_f32_argminmax_in	458.8 µs	449.4 µs	+2.09%
⚡	sse_f64_argminmax_in	915.2 µs	897.2 µs	+2.01%

Tip

Curious why performance improved? Comment @codspeedbot explain why performance improved on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing perf/simd-last-vector-remainder (5f5f4f7) with main (f68aba5)}