vortex-data / vortex / 16051822264

88%

push

github

feat: migrate take_primitive_simd to stable AVX2 kernel (#3579)

An implementation of TakeKernel for PrimitiveArray that uses AVX2
explicit instructions, falling back to Scalar. For non-x86_64 platforms
the `portable_simd` implementation is still loaded if nightly compiler
is being used.

This is part of the #3546 series of PRs.

Additionally: fixed a soundness issue with lack of bounds checking of
indices for portable_simd impl. This requires us to do a full scan of
the indices upfront before running the kernel to avoid out of bounds
memory access.

## Implementation

The biggest source of added complexity in this PR is the new `avx2`
module which implements a take kernel for primitive indices/values that
uses the AVX2 GATHER operation.

Intel ISA provides 4 different gather intrinsics:

- `_mm256_i32gather_epi32` -> gathering 8x 32bit values with 32bit
indices
- `_mm_i32gather_epi64` -> gathering 4x 64bit values with 32bit indices
- `_mm_i64gather_epi32` -> gathering 8x 32bit values with 32bit indices
- `_mm256_i32gather_epi32` -> gathering 8x 32bit values with 32bit
indices

We implement a generic inner loop with a trait parameter `GatherFn<I,
V>`, and allow specialization to insert the proper loop logic for each
valid index/value type combination.

---------

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

171 of 202 new or added lines in 3 files covered. (84.65%)

43466 of 55812 relevant lines covered (77.88%)

55486.81 hits per line