Skip to content

jvdd/argminmax

Repository files navigation

ArgMinMax

Efficient argmin & argmax (in 1 function) with SIMD (SSE, AVX(2), AVX512, NEON) for f16, f32, f64, i8, i16, i32, i64, u8, u16, u32, u64.

πŸš€ The function is generic over the type of the array, so it can be used on &[T] or Vec<T> where T can be f161, f322, f642, i8, i16, i32, i64, u8, u16, u32, u64.

🀝 The trait is implemented for slice, Vec, 1D ndarray::ArrayBase3, and apache arrow::PrimitiveArray4.

⚑ 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 f16 you should enable the "half" feature.
2 for f32 and f64 you should enable the (default) "float" feature.
3 for ndarray::ArrayBase you should enable the "ndarray" feature.
4 for arrow::PrimitiveArray you should enable the "arrow" feature.

Installing

Add the following to your Cargo.toml:

[dependencies]
argminmax = "0.5"

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.

Features

  • [default] "float": support f32 and f64 argminmax (uses NaN-handling - see below).
  • "half": support f16 argminmax (through using the half crate).
  • "ndarray": add ArgMinMax trait to ndarray its Array1 & ArrayView1.
  • "arrow": add ArgMinMax trait to arrow its PrimitiveArray.

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 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.

❗ The "half" feature (f16 support) is not yet fully supported:

  • ReturnNaN variant
  • IgnoreNaN variant -> currently defaults to the ReturnNaN variant

Acknowledgements

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