Efficient argmin & argmax (in 1 function) with SIMD (SSE, AVX(2), AVX5121, NEON1) ⚡
🚀 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
2, f32
2, f64
3, i8
, i16
, i32
, i64
, u8
, u16
, u32
, u64
.
🤝 The trait is implemented for slice
, Vec
, 1D ndarray::ArrayBase
4, apache arrow::PrimitiveArray
5 and arrow2::PrimitiveArray
6.
⚡ 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 implementation contains no if checks, ensuring that the runtime of the function is independent of the input data its order (best-case = worst-case = average-case).
🪄 Efficient support for f16 and uints: through (bijective aka symmetric) bitwise operations, f16 (optional1) and uints are converted to ordered integers, allowing to use integer SIMD instructions.
1 for
AVX512
and most ofNEON
you should enable the (default)"nightly_simd"
feature (requires nightly Rust).
2 forf16
you should enable the"half"
feature.
3 forf32
andf64
you should enable the (default)"float"
feature.
4 forndarray::ArrayBase
you should enable the"ndarray"
feature.
5 forarrow::PrimitiveArray
you should enable the"arrow"
feature.
6 forarrow2::PrimitiveArray
you should enable the"arrow2"
feature.
Add the following to your Cargo.toml
:
[dependencies]
argminmax = "0.6.1"
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]);
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.
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 thanNaNArgMinMax
.
- [default] "nightly_simd": enables the use of non-stable SIMD intrinsics (
AVX512
and most ofNEON
), which are only available on nightly Rust. - [default] "float": support
f32
andf64
argminmax (uses NaN-handling - see below). - "half": support
f16
argminmax (through using thehalf
crate). - "ndarray": add
ArgMinMax
trait tondarray
itsArray1
&ArrayView1
. - "arrow": add
ArgMinMax
trait toarrow
itsPrimitiveArray
.
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:"
To run the tests use the following command:
cargo test --message-format=short --all-features
The library handles NaNs! 🚀
Some (minor) limitations:
ArgMinMax
its functions ignores NaN values.- ❗ When the array contains exclusively NaNs and/or infinities unexpected behaviour can occur (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.
Some parts of this library are inspired by the great work of minimalrust's argmm project.