I'll document the changes tomorrow - a bit too tired now :upside_down_face:

Long story short -> used more traits & structs to provide more flexibility (& less duplicate code) yay! :tada:

-- did not add support for aarch64/arm (as I currently do not have access to my ARM testing device)

Apparently we can indeed create pulls into pull requests.

Here we go. This should add some functions and types for the NEON/arm64 SIMD specialisations.

Currently the ArgMinMax trait is only implemented for ndarray::ArrayView1.
It might be interesting to implement the trait for

• Vec
• array?

Credits @minimalrust, see https://github.com/minimalrust/argmm/issues/8

Fixed the merge issues of #24

With a lot of patience I debugged & updated the NEON code on my raspberry pi 3 :upside_down_face:

Thx for the headstart @varon :)

This PR also fixes a bug in the tests in where we checked for the wrong target feature (and thus performed illegal instructions)

• Add rust toolchain file to default to nightly
• .gitignore improvements

This is a super-tiny PR - just some minor touchups during the initial onboarding of working with this code.

Adding the toolchain file solved a common gotcha, and the .gitignore adds support for working in CLion.

To make codspeed work I had to add "half" as default feature. I will revert this once https://github.com/CodSpeedHQ/codspeed-rust/issues/1 is resolved

• [x] use slice (&[T]) under the hood instead of ArrayView1
• [x] ndarray -> optional dependency
• extend trait implementation, resolves #9
  • [x] slice
  • [x] Vec
    • [x] test this
  • [x] ndarray
    • [x] test this
  • [x] arrow
    • [x] test this
• [x] update Readme

• [x] final rerun of benchmarks
  • [x] no regressions (i.e., same performance or even better :rocket:) for scalar implementation
  • [ ] no regressions (i.e., same performance or even better :rocket:) for SIMD implementation => I notices very few (minor) regressions for certain dtype / SIMD combinations & highlighted those in the benches/results file.. But, since the codspeed continuous monitoring only shows 2% regressions for the AVX2 u8 & i8 - (which is most likely not the most common datatypes for end users) - I am pro for merging this into main, given the significant enhanced flexibility that this PR adds :)

This pull request updates the way we iterate over arrays in our code by replacing the use of the [exact_chunks] (https://docs.rs/ndarray/latest/ndarray/struct.ArrayBase.html#method.exact_chunks) method with a pointer that we offset ourselves. Our profiling showed that this change resulted in significant speed improvements, with up to a 50% speedup for 8-bit data types and noticeable speedups for other data types. This should result in a notable improvement in the overall performance of our code.

:sparkles: Contributions of this PR (focussing on performance) :rocket:

• [x] optimized iteration over array using (mutable) pointer that we offset ourselves
• [x] reorder SIMD inner loop
• [x] for 16-bit dtypes: _horiz_min and _horiz_max with SIMD
• [x] use NEON SIMD implementation on arm for 32-bit dtypes (as this is now 2-3x faster than scalar)
• [x] explore prefetching - as currently the bottleneck is loading the data => included in the code but commented out, as this has minor / no impact (only measurable positive impact on large arrays - 500M+ data points)

This PR was mainly driven by analyzing stacktraces of perf.

:eyes: short (& incomplete) summary of observed improvements

x86(_64)

• 8-bit dtypes: 80% faster :rocket: :rocket:
• 16-bit dtypes: 5-12% faster (12% for floats, 10% for ints, 5% for uints)
• 32-bit dtypes: 0-5% faster (5% for floats, 2% for ints, 0% for uints)
• 64-bit dtypes: 3-17% faster (3% for floats, 17% for ints & uints)

AArch64

• 8-bit dtypes: 34-42% faster :rocket: :rocket: (34% for ints and 42% for floats)
• 16-bit dtypes: 7% slower :cry: (=> bc of horiz SIMD? :thinking:)
• 32-bit dtypes: 0-20% faster (20% for floats, 0% for ints & uints)

ARM

• 8-bit: 400% faster :rocket: :rocket:
• 16-bit: 350-460% faster :rocket: :rocket: (350% for floats, 460% for ints & uints)
• 32-bit: 420-450% faster :rocket: :rocket: (420% for floats, 450% for ints & uints)

:sparkles: add support for int8

• [x] SIMD + SCALAR implementation
  • [x] alleviate bottleneck (horizontal argmin in final vector) -> algorithm: https://stackoverflow.com/a/9798369
    • [x] SSE
    • [x] AVX2
    • [x] AVX512(bw)
    • [x] ARM/AARCH64

:thinking: add support for uint8?
-> would probably require the XOR conversion ~~but would exit the overflow loop 2x fewer than int8 (guess this will be net 1.5x faster than int8 for longer arrays)~~ :sweat_smile: nvm still limited by the i8::MAX for overflow TODO: bench + horizontal SIMD implementation for neon

• [x] SSE
• [x] AVX2
• [x] AVX512(bw)
• [x] ARM/AARCH64

:snowflake: minor improvements

• [x] faster implementation of overflow loop
• [ ] ~~remove _get_min_max_index_value~~ => is for another refactoring :recycle: PR

Other stuff

• [x] :robot: rerun benchmarks!

• [x] :tada: add support for SIMD (NEON) on aarch64 cpu architectures
• [x] :recycle: benchmark code on raspberry pi 4 (which is aarch64)
• [x] :broom: fix some clippy warnings

Thx @matdemeue for providing me with a raspberry pi 4!! :hugs: (this allowed me to develop & bench the code!)

This PR will be squashed when merged - it includes #4

Contributions of this PR (in addition to #4):

• [x] minor refactoring of naming
• [ ] ~~lazy_static~~ => once_cell => will skip this for now

Hi @jvdd I am considering adding these to polars, but given that we define our own NaN handling, I only want to expose it for integer backed data.

Could the floating points be feature gated? Polars compile times are really strained so I'd rather not include code we don't need. Let me know what you think!

Handle NaNs, closes #16 Behavior changes in this PR:

• .argminmax returns the index of the first NaN value (instead of ignoring it) To realize this functionality, we use the transformation as detailed in #16 & explored in #18
• .nanargminmax (new function :tada:) ignores NaNs (while being even faster for floats) => Only "downside" - if data contains ONLY NaNs and/or +inf/-inf this will return 0 (I believe we can accept this unexpected for now - seems like a very uncommon use case)

Changing the "architecture":

• [x] SIMD ops in stand-alone trait (better w.r.t. coupling & cohesion)
  • [x] create SIMDOps trait & add additional traits (see sketch below)
  • [x] add auto-implement for SIMDCore traits
  • [x] implement the SIMDOps trait for the various data types
    • [x] x86 / x86_64
    • [x] arm / Aarch64
• [x] implement IgnoreNaN and ReturnNan variant for floats the SIMDInstructionSet structs
  • [x] IgnoreNaN
  • [x] ReturnNaN

Changing the default behavior

• [ ] switch from IgnoreNan to ReturnNan for argminmax
  • [x] simd_f16.rs should implement SIMDInstructionSet and not the IgnoreNan structs
  • [x] document header of IgnoreNan & ReturnNan implementations
  • [x] add tests for IgnoreNan
  • [x] add tests for ReturnNan
  • [ ] update benches (currently we only bench FloatIgnoreNaN) -> first assure that we currently have no regressions - is still the case! - if so, change the FloatIgnoreNan benches to FloatReturnNan (will most likely result in some regressions) and move the FloatIgnoreNan benches to dedicated bench file.
  • [x] update the ArgMinMax trait
• [x] add scalar support
  • [x] update tests with correct scalar implementation
  • [x] implement SCALARIgnoreNaN
  • [x] double check correct scalar support for f16

Minor TODOs during this (waaay too large) PR:

• [x] try to get rid of overhead for signed & unsigned integers (added to ignore NaNs in the first SIMD Vec, see commit edd083c5343e009c338b7a9f35cf68827adf90ab)
• [x] add tests for infinities, resolves #20
• [ ] :thinking: decide what we will default to: argminmax / nanargminmax?

Overview of the new "architecture"

• default SIMDInstructionSet struct (e.g., AVX2) its argminmax is "return NaN" (e.g. simd_f32.rs)
• "ignore NaN" is served in a distinct struct (e.g. AVX2FloatIgnoreNaN)

This PR adds some (failing) tests for NaN handling.

This does work and is standalone, but obviously the tests fail, so do be aware of that prior to merging.

This can either be used as the base branch for other PRs in this series, such as #18 or can be extended here to a complete implementation.

This PR aims at handling Nans, #16.

For the scalar implementation I check whether the scalar value is not equal to itself. This check will only be true if the scalar value is NAN, as the following is correct in Rust:

let v = f32::NAN;
assert!(v != v);

For the SIMD implementation I used a transformation similar to the one used in #1 - this transformation projects the NANs to integer values that are either higher / lower than the "real" floating point values. The transformation leverages the 2-complement https://observablehq.com/@rreusser/half-precision-floating-point-visualized

Some remarks:

• + inf & - inf will get projected as well
  • [x] validate if these are inbetween the NaNs & the "real" floating point values (pretty sure this is the case)
    => this is indeed the case - see plot :arrow_down:
• there are a plethora of ways to represent NaNs in floating point representation: exponent should be 11111 and the fraction should be non-zero. Thus the sign bit may be 1 or 0 -> resulting in half of the NaNs getting projected above and the other half below the "real" floating point values. Thus, only 1 of the 2 checks (either > or <) will fire & thus detect the NaN. This might be a problem when we want to implement argmin and argmax as separate functions..

Paths I looked into but did not seem worthwhile:

• I tried to handle NaNs with SIMD instructions, but the SIMD instruction to check whether there are a NaNs in the register turns out to be rather computationally expense

We want to properly handle NaNs.

What are NaNs?

NaNs are only present in float datatypes - and thus do not exist in (unsigend) integer datatypes.
For simplicity I will illustrate things for 16-bit numbers.

Some background :arrow_down:

So how do NaNs look like?

nans occur when exponent is all 1s and mantissa is not all 0s. (sign does not matter)

How does the current implementation cope with NaNs?

Necessary background info: every comparison (gt, lt, eq) with NaNs results in false :arrow_down:

let v = f32::NAN;
assert_eq!(v > 5.0, false);
assert_eq!(v < 5.0, false);
assert_eq!(v == 5.0, false);

=> Current implementation deals as follows with NaNs in the data:

• if there are NO NaNs in the first SIMD lane: current implementation will ignore the NaNs (as comparisons with NaNs result in false and thus no NaNs will be added in the accumulating SIMD vector)
  • this corresponds to the np.nanargmax / np.nanargmin behavior
• if there are ONE or MULTIPLE NaNs in first SIMD lane: current implementation will never update the NaNs in the accumulating SIMD vector (as comparisons will result in false and thus the NaNs will never be updated).

TODOs

• [ ] fix this issue with NaN values in the first SIMD vector -> if we can create/ensure a "NaN-free" initial SIMD vec, than no NaNs will be added in the inner SIMD loop => we have an implementation of np.nanargmin / np.nanargmax.

How should we handle NaNs

Ideally we support, just like numpy, two variants of the argmim / argmax algorithm:

1. ignoring NaNs: as the current version does (if we fix the TODO above) -> corresponds to np.nanargmin / np.nanargmax
2. returning first NaN index once there is one present in the data -> corresponds to np.argmin / np.argmax

we can serve both functionalities by adding a new function to the ArgMinMax trait and let it default for non-float datatypes to the current implementation.