ArgMinMax

Efficient argmin & argmax (in 1 function) with SIMD (SSE, AVX(2), AVX512, NEON¹) ⚡

🚀 The functions are generic over the type of the array, so it can be used on &[T] or Vec<T> where T can be f16², bf16², f32³, f64³, i8, i16, i32, i64, i128⁶, u8, u16, u32, u64, u128⁶.

🤝 The trait is implemented for slice, Vec, 1D ndarray::ArrayBase⁴ and apache arrow::PrimitiveArray⁵.

⚡ Runtime CPU feature detection is used to select the most efficient implementation for the current CPU. This means that the same binary can be used on different CPUs without recompilation.

👀 The SIMD inner loop contains no if checks, so its runtime does not depend on the values of the input or on their order (best-case = worst-case = average-case).

🪄 Efficient support for f16, bf16 and uints: through bitwise operations, f16 and bf16 (optional²) and uints are converted to ordered integers, allowing to use integer SIMD instructions.

¹ for NEON on 32-bit ARM you should enable the "nightly_simd" feature (requires nightly Rust).
² for f16 and bf16 you should enable the "half" feature.
³ for f32 and f64 you should enable the (default) "float" feature.
⁴ for ndarray::ArrayBase you should enable the "ndarray" feature.
⁵ for arrow::PrimitiveArray you should enable the "arrow" feature.
⁶ i128 and u128 (e.g., arrow Decimal128 arrays) use the scalar implementation.

Installing

Add the following to your Cargo.toml:

[dependencies]
argminmax = "0.7.1"

Example usage

use argminmax::ArgMinMax;  // import trait

let arr: Vec<i32> = (0..200_000).collect();  // create a vector

let (min, max) = arr.argminmax();  // apply extension

println!("min: {}, max: {}", min, max);
println!("arr[min]: {}, arr[max]: {}", arr[min], arr[max]);

Traits

ArgMinMax

Implemented for ints, uints, and floats (if "float" feature enabled).

Provides the following functions:

• argminmax: returns the index of the minimum and maximum element in the array.

When dealing with NaNs, ArgMinMax its functions ignore NaNs. For more info see Limitations.

NaNArgMinMax

Implemented for floats (if "float" feature enabled).

Provides the following functions:

• nanargminmax: returns the index of the minimum and maximum element in the array.

When dealing with NaNs, NaNArgMinMax its functions return the first NaN its index. For more info see Limitations.

Tip 💡: if you know that there are no NaNs in your the array, we advise you to use ArgMinMax as this should be 5-30% faster than NaNArgMinMax.

ArgMinMaxMasked & NaNArgMinMaxMasked

The same as ArgMinMax & NaNArgMinMax, but skipping the null elements that an Arrow validity bitmap (and its bit offset) marks. Their functions (argminmax_masked, nanargminmax_masked, ...) return None when there are no valid elements. When all valid values are NaN, the index of the first valid value is returned (see Limitations).

Tip 💡: when an array has no nulls, ArgMinMax & NaNArgMinMax are faster: call the masked functions only when null_count() > 0. This check is cheap: arrow and Polars keep the null count (they count the nulls once, e.g., when the bitmap is built or sliced; a 10M-element bitmap takes ~30 µs, 1% of an argminmax of its f32 values). Check null_count() > 0, not whether there is a validity bitmap: an array can have a validity bitmap without nulls.

use argminmax::ArgMinMaxMasked;

let arr: Vec<i32> = vec![-5, 3, 9, 4];
let validity: [u8; 1] = [0b0110]; // only the elements 1 and 2 are valid

assert_eq!(arr.argminmax_masked(&validity, 0), Some((1, 2)));

Features

• [default] "float": support f32 and f64 argminmax (uses NaN-handling - see below).
• "nightly_simd": enables the use of non-stable SIMD intrinsics (NEON on 32-bit ARM), which are only available on nightly Rust.
• "half": support f16 and bf16 argminmax (through using the half crate).
• "ndarray": add ArgMinMax trait to ndarray its Array1 & ArrayView1.
• "arrow": add ArgMinMax trait to arrow its PrimitiveArray (skipping the nulls).
  • ❗ The functions panic when all values are null (as for an empty array). To get an Option instead, check the null count first: (array.null_count() < array.len()).then(|| array.argmin()) (arrow keeps the null count, so this check is O(1)).

Benchmarks

Benchmarks on my laptop (AMD Ryzen 7 4800U, 1.8 GHz, 16GB RAM) using criterion show that the function is 3-20x faster than the scalar implementation (depending of data type).

See /benches/results.

Run the benchmarks yourself with the following command:

cargo bench --quiet --message-format=short --features half | grep "time:"

Tests

To run the tests use the following command:

cargo test --message-format=short --all-features

Limitations

The library handles NaNs! 🚀

Some (minor) limitations:

• ArgMinMax its functions ignores NaN values.
  • When the array contains only NaNs, index 0 is returned.
• NaNArgMinMax its functions returns the first NaN its index (if any present).
  • ❗ When multiple bit-representations for NaNs are used, no guarantee is made that the first NaN is returned.

---

Acknowledgements

Some parts of this library are inspired by the great work of minimalrust's argmm project.