[jvdd]

Problem

On the CPU, argmin, argmax, min and max (ReduceIndex in cpu_backend) use a scalar loop: one compare and one data-dependent branch per element. For f16 and bf16 each compare also converts both values to f32, which costs about 3.5 ns per element. These ops run in greedy and Gumbel sampling (argmax over the vocabulary), in candle_nn::ops::softmax (max over the last dim), and in codebook lookups such as snac's quantizer (argmin over 4096 distances).

The loop also gives a NaN result that depends on where the NaN is. It starts from the first element, and every compare with NaN is false. So a NaN in the first position always wins, and a NaN anywhere else is skipped:

Tensor::new(&[f32::NAN, 1., 2.], &Device::Cpu)?.argmin(0)?  // 0 (the NaN), min() is NaN
Tensor::new(&[3., f32::NAN, 1.], &Device::Cpu)?.argmin(0)?  // 2, min() is 1

Changes

• First commit, the NaN fix. A row that starts with NaN now starts from its first non-NaN value. NaN values are then skipped at every position, and a row of only NaN values gives index 0 (value NaN). This applies to all three loops: contiguous, strided and non-contiguous. The first example above now gives index 1 and min() 1. Rows without NaN give the same results as before. The check is one compare per row, outside the loop: a check inside the loop made the f32 loop up to 5x slower.
• Second commit, the speed-up. When a contiguous tensor is reduced over its last dim, each row goes to argmin or argmax from the argminmax crate. It selects SSE4.1, AVX2 or AVX-512 at run time on x86 and uses NEON on aarch64.
  • Types: u8, u32, i16, i32, i64, f16, bf16, f32 and f64. f8e4m3 keeps the loop.
  • Only rows of 128 elements or more, except for f16 and bf16, which win at every length (1.4-2x already at 2 elements). Below 128, argminmax's fixed cost per call and the scalar handling of the last partial vector make it slower than the loop for some lengths. For example, f32 on AVX2 runs at 0.78x for rows of 63 and 0.96x for rows of 79. The threshold is checked once per call, not per row.
  • The results are the same as with the fixed loop: NaN skipped, the lowest index wins on ties, and -0.0 and +0.0 compare equal.
  • The strided and non-contiguous paths do not change.
• New dependency: argminmax 0.7.1 (MIT), with its half feature. Its only dependencies are num-traits and half, which candle already uses. It has over 15 million downloads on crates.io, and Polars uses it for arg_min and arg_max. Disclosure: I maintain argminmax.
• New test arg_reduce_skips_nan (CPU only): rows of 6 (the loop) and 143 (argminmax, with the extremes after the last full SIMD vector), the strided and non-contiguous paths, f32, f64, f16 and bf16, a leading NaN followed by a second NaN, ties, and rows of only NaN values. It fails without the first commit.

Benchmarks

A candle-core example, not part of this PR (I can share it). The inputs are uniform in [-1, 1) from a fixed random stream, so both builds get the same data, and every call gets fresh data. The checksums of all outputs are equal between main and this PR, except for Gumbel sampling, which is random. Single thread (RAYON_NUM_THREADS=1), pinned to one core. Each cell is the median of 7 rounds, with main and this PR alternating. Built with the repository's .cargo/config.toml (target-cpu=native). argminmax picks its instruction set at run time, so builds without target-cpu=native take the same SIMD code.

Case	i7-1185G7 (AVX-512)	Ryzen 9 5950X (AVX2)
argmax f32, 1 row of 32 000	28.0 → 2.4 µs, 11.9x	10.5 → 1.8 µs, 6.0x
argmax f32, 1 row of 151 936	137 → 13.4 µs, 10.2x	49.2 → 7.9 µs, 6.2x
argmax f32, 1 row of 262 144	241 → 40.4 µs, 6.0x	84.7 → 13.7 µs, 6.2x
argmax bf16, 1 row of 151 936	529 → 6.8 µs, 77x	430 → 6.7 µs, 64x
argmin [1, 1024, 1024] over dim 2 (the arg_reduce bench), f32	948 → 182 µs, 5.2x	339 → 67 µs, 5.1x
the same, f16 / bf16	3.66 / 3.62 ms → 71 / 71 µs, 51x	3.27 / 3.17 ms → 52 / 52 µs, 63x / 61x
max f32 [1, 1024, 1024] over dim 2	887 → 178 µs, 5.0x	221 → 64 µs, 3.4x
argmin f32 [512, 4096] over dim 1	1.92 ms → 430 µs, 4.5x	686 → 199 µs, 3.5x
argmax f32, rows of 256	960 → 186 µs, 5.2x	335 → 106 µs, 3.2x
argmax, rows of 143, the 7 other types	1.6-5.1x	1.4-4.9x
argmax, rows of 143, f16 / bf16	17x	18x
argmax, rows of 127 (the loop), all 9 types	0.98-1.00x	0.99-1.00x
argmax f32, rows of 4 / 10 / 64 (the loop)	0.92x / 0.92x / 0.98x	1.05x / 0.94x / 1.00x
max f32, rows of 4 (the loop)	1.05x	0.78x
argmin f32 [1, 1024, 1024] over dim 1 (strided, the loop)	1.04x	1.01x

The cases with many rows (rows of 4 to 256) reduce about one million elements per call.

The rows below 128 run the loop. A row of 4 takes 2-3 ns, and at that size code layout alone moves the result by ±10-20%. In a variant whose short-row loop was textually identical to main, the rows-of-4 cases still moved: one went from 0.78x to 1.00x, another from 1.05x to 0.91x.

The ops are 5-12x faster for f32 and 50-77x for f16 and bf16. In the callers above, the gain is smaller:

Code path	i7-1185G7	Ryzen 9 5950X
Greedy sampling: Qwen2.5-0.5B lm_head (896 → 151 936, f32) + argmax	31.19 → 30.96 ms, 1.01x	18.77 → 18.76 ms, 1.00x
candle_nn::sampling::gumbel_softmax over 151 936 logits	1.32 → 1.22 ms, 1.08x	1.07 → 1.03 ms, 1.04x
candle_nn::ops::softmax over the last dim (bert, clip and stable diffusion attention), 12 heads, key length 128 / 512	1.10x / 1.15x	1.03x / 1.03x
The same, key length 8 / 16	0.99x / 1.00x	1.01x / 1.00x
snac VectorQuantize: distance + argmin, 512 frames, codebook 4096 × 8	7.31 → 5.71 ms, 1.28x	3.58 → 3.00 ms, 1.20x

Greedy sampling saves about 0.1 ms per token, which is within the noise of the matmul. candle_nn::ops::softmax_last_dim has its own kernel and does not change.

The i7 columns are from this PR's code. The Ryzen columns are from a clean run of the second commit before the NaN fix was added; the NaN fix only adds the first-value check in the loop. A run of the final code on the Ryzen gave ratios within a few percent of these, but other jobs were running on the machine (load 6-13), so I did not use it.

Notes

• The threshold of 128 can drop with the next argminmax version. Pending argminmax changes make short slices 2-4x faster: perf: fold the last vector into the SIMD core, AVX2 for short 16-bit slices jvdd/argminmax#131 (the last partial vector in SIMD) and a follow-up (the scalar loop and slices shorter than 8 elements). My first measurements put the break-even with the loop at about 32 elements on both CPUs. I will send a follow-up that lowers the threshold after that release.
• Skip NaN or return NaN? I chose to skip NaN because the old loop already skipped NaN everywhere except in the first position, so this changes the fewest results. PyTorch's argmax and max return the NaN instead. If you prefer that, argminmax has nanargmin and nanargmax (the first NaN wins), and the fix in the loops is just as small. The CUDA and Metal kernels have no NaN handling either, and this PR does not change them.
• Checked: cargo fmt --all -- --check and cargo clippy -p candle-core --tests --examples --benches -- -D warnings. cargo test -p candle-core on aarch64 (NEON), and on x86 with AVX-512 and with AVX2. The lib and tensor_tests binaries also under Intel SDE as Haswell (AVX2) and Nehalem (SSE4.2).