#simd #nan #ndarray #array #index #traits #max

argminmax

ArgMinMax (argmin & argmax in 1 function) with SIMD for floats and integers

11 unstable releases

0.6.2 Mar 4, 2024
0.6.1 Apr 14, 2023
0.5.0 Mar 23, 2023
0.4.0 Feb 13, 2023
0.1.1 Oct 2, 2022

#57 in Algorithms

Download history 9800/week @ 2023-12-23 16754/week @ 2023-12-30 18306/week @ 2024-01-06 16913/week @ 2024-01-13 16661/week @ 2024-01-20 17011/week @ 2024-01-27 16664/week @ 2024-02-03 18526/week @ 2024-02-10 17987/week @ 2024-02-17 14733/week @ 2024-02-24 15796/week @ 2024-03-02 16476/week @ 2024-03-09 15984/week @ 2024-03-16 18880/week @ 2024-03-23 18117/week @ 2024-03-30 14124/week @ 2024-04-06

69,161 downloads per month
Used in 108 crates (via polars-ops)

MIT license

475KB
9K SLoC

ArgMinMax

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 f162, f322, f643, i8, i16, i32, i64, u8, u16, u32, u64.

🤝 The trait is implemented for slice, Vec, 1D ndarray::ArrayBase4, apache arrow::PrimitiveArray5 and arrow2::PrimitiveArray6.

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 of NEON you should enable the (default) "nightly_simd" feature (requires nightly Rust).
2 for f16 you should enable the "half" feature.
3 for f32 and f64 you should enable the (default) "float" feature.
4 for ndarray::ArrayBase you should enable the "ndarray" feature.
5 for arrow::PrimitiveArray you should enable the "arrow" feature.
6 for arrow2::PrimitiveArray you should enable the "arrow2" feature.

Installing

Add the following to your Cargo.toml:

[dependencies]
argminmax = "0.6.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.

Features

  • [default] "nightly_simd": enables the use of non-stable SIMD intrinsics (AVX512 and most of NEON), which are only available on nightly Rust.
  • [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.

Acknowledgements

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

Dependencies

~0.1–10MB
~67K SLoC