diff --git a/vortex-array/src/arrays/primitive/compute/take/avx2.rs b/vortex-array/src/arrays/primitive/compute/take/avx2.rs index fb64e6748b0..5378321b005 100644 --- a/vortex-array/src/arrays/primitive/compute/take/avx2.rs +++ b/vortex-array/src/arrays/primitive/compute/take/avx2.rs @@ -14,7 +14,6 @@ use std::arch::x86_64::_mm_setzero_si128; use std::arch::x86_64::_mm_shuffle_epi32; use std::arch::x86_64::_mm_storeu_si128; use std::arch::x86_64::_mm_unpacklo_epi64; -use std::arch::x86_64::_mm256_andnot_si256; use std::arch::x86_64::_mm256_cmpgt_epi32; use std::arch::x86_64::_mm256_cmpgt_epi64; use std::arch::x86_64::_mm256_cvtepu8_epi32; @@ -27,10 +26,12 @@ use std::arch::x86_64::_mm256_loadu_si256; use std::arch::x86_64::_mm256_mask_i32gather_epi32; use std::arch::x86_64::_mm256_mask_i64gather_epi32; use std::arch::x86_64::_mm256_mask_i64gather_epi64; +use std::arch::x86_64::_mm256_movemask_epi8; use std::arch::x86_64::_mm256_set1_epi32; use std::arch::x86_64::_mm256_set1_epi64x; use std::arch::x86_64::_mm256_setzero_si256; use std::arch::x86_64::_mm256_storeu_si256; +use std::arch::x86_64::_mm256_xor_si256; use std::convert::identity; use vortex_buffer::Alignment; @@ -172,12 +173,43 @@ trait GatherFn { /// /// This function can read up to `STRIDE` elements through `indices`, and read/write up to /// `WIDTH` elements through `src` and `dst` respectively. - unsafe fn gather(indices: *const Idx, max_idx: Idx, src: *const Values, dst: *mut Values); + unsafe fn gather( + indices: *const Idx, + max_idx: Option, + src: *const Values, + dst: *mut Values, + ); } /// AVX2 version of [`GatherFn`] defined for 32- and 64-bit value types. enum AVX2Gather {} +macro_rules! cmpgt_epu32 { + ($lhs:expr, $rhs:expr) => {{ + // AVX2 only supplies a signed integer comparison. XORing each lane with its sign bit + // maps the unsigned ordering into the signed ordering without changing the relative + // order. + let sign_bit = _mm256_set1_epi32(i32::MIN); + _mm256_cmpgt_epi32( + _mm256_xor_si256($lhs, sign_bit), + _mm256_xor_si256($rhs, sign_bit), + ) + }}; +} + +macro_rules! cmpgt_epu64 { + ($lhs:expr, $rhs:expr) => {{ + // AVX2 only supplies a signed integer comparison. XORing each lane with its sign bit + // maps the unsigned ordering into the signed ordering without changing the relative + // order. + let sign_bit = _mm256_set1_epi64x(i64::MIN); + _mm256_cmpgt_epi64( + _mm256_xor_si256($lhs, sign_bit), + _mm256_xor_si256($rhs, sign_bit), + ) + }}; +} + macro_rules! impl_gather { ($idx:ty, $({$value:ty => load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal }),+) => { $( @@ -191,7 +223,7 @@ macro_rules! impl_gather { #[allow(unused_unsafe, clippy::cast_possible_truncation)] #[inline(always)] - unsafe fn gather(indices: *const $idx, max_idx: $idx, src: *const $value, dst: *mut $value) { + unsafe fn gather(indices: *const $idx, max_idx: Option<$idx>, src: *const $value, dst: *mut $value) { const { assert!($WIDTH <= $STRIDE, "dst cannot advance by more than the stride"); } @@ -202,19 +234,30 @@ macro_rules! impl_gather { // Extend indices to fill vector register. let indices_vec = unsafe { $extend(indices_vec) }; - // Create a vec of the max idx. - let max_idx_vec = unsafe { $splat(max_idx as _) }; - // Create a mask for valid indices (where the max_idx > provided index). - let invalid_mask = unsafe { _mm256_andnot_si256($mask_indices(indices_vec, max_idx_vec), $splat(-1)) }; - let invalid_mask = { - let $mask_var = invalid_mask; + // A representable length is an exclusive upper bound. If it does not fit in + // the index type, every possible index value is in-bounds. + let valid_mask = if let Some(max_idx) = max_idx { + let max_idx_vec = unsafe { $splat(max_idx as _) }; + // Passing the valid mask to the gather masks every invalid lane before it + // can access `src`. + unsafe { $mask_indices!(max_idx_vec, indices_vec) } + } else { + unsafe { $splat(-1) } + }; + assert_eq!( + unsafe { _mm256_movemask_epi8(valid_mask) }, + -1, + "take index out of bounds" + ); + let valid_mask = { + let $mask_var = valid_mask; $mask_cvt }; let zero_vec = unsafe { $zero_vec() }; // Gather the values into new vector register, for masked positions // it substitutes zero instead of accessing the src. - let values_vec = unsafe { $masked_gather::(zero_vec, src.cast(), indices_vec, invalid_mask) }; + let values_vec = unsafe { $masked_gather::(zero_vec, src.cast(), indices_vec, valid_mask) }; // Write the vec out to dst. unsafe { $store(dst.cast(), values_vec) }; @@ -231,7 +274,7 @@ impl_gather!(u8, extend: _mm256_cvtepu8_epi32, splat: _mm256_set1_epi32, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, + mask_indices: cmpgt_epu32, mask_cvt: |x| { x }, gather: _mm256_mask_i32gather_epi32, store: _mm256_storeu_si256, @@ -244,7 +287,7 @@ impl_gather!(u8, extend: _mm256_cvtepu8_epi64, splat: _mm256_set1_epi64x, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, + mask_indices: cmpgt_epu64, mask_cvt: |x| { x }, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, @@ -260,7 +303,7 @@ impl_gather!(u16, extend: _mm256_cvtepu16_epi32, splat: _mm256_set1_epi32, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, + mask_indices: cmpgt_epu32, mask_cvt: |x| { x }, gather: _mm256_mask_i32gather_epi32, store: _mm256_storeu_si256, @@ -273,7 +316,7 @@ impl_gather!(u16, extend: _mm256_cvtepu16_epi64, splat: _mm256_set1_epi64x, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, + mask_indices: cmpgt_epu64, mask_cvt: |x| { x }, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, @@ -289,7 +332,7 @@ impl_gather!(u32, extend: identity, splat: _mm256_set1_epi32, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, + mask_indices: cmpgt_epu32, mask_cvt: |x| { x }, gather: _mm256_mask_i32gather_epi32, store: _mm256_storeu_si256, @@ -302,7 +345,7 @@ impl_gather!(u32, extend: _mm256_cvtepu32_epi64, splat: _mm256_set1_epi64x, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, + mask_indices: cmpgt_epu64, mask_cvt: |x| { x }, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, @@ -317,7 +360,7 @@ impl_gather!(u64, extend: identity, splat: _mm256_set1_epi64x, zero_vec: _mm_setzero_si128, - mask_indices: _mm256_cmpgt_epi64, + mask_indices: cmpgt_epu64, mask_cvt: |m| { unsafe { let lo_bits = _mm256_extracti128_si256::<0>(m); // lower half @@ -338,7 +381,7 @@ impl_gather!(u64, extend: identity, splat: _mm256_set1_epi64x, zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, + mask_indices: cmpgt_epu64, mask_cvt: |x| { x }, gather: _mm256_mask_i64gather_epi64, store: _mm256_storeu_si256, @@ -370,7 +413,7 @@ where ); let indices_len = indices.len(); - let max_index = Idx::from(values.len()).unwrap_or_else(|| Idx::max_value()); + let max_index = Idx::from(values.len()); let mut buffer = BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); let buf_uninit = buffer.spare_capacity_mut(); @@ -501,6 +544,36 @@ mod avx2_tests { assert_eq!(&vec![65535; indices.len()], result.as_slice()); } + #[test] + #[should_panic(expected = "take index out of bounds")] + fn test_avx2_take_u32_max_index_u32_lane() { + let values = vec![0u32; 8]; + // The first eight indices execute in the SIMD loop; the scalar remainder is valid. + let indices = vec![0, u32::MAX, 2, 3, 4, 5, 6, 7, 0]; + + unsafe { take_avx2(&values, &indices) }; + } + + #[test] + #[should_panic(expected = "take index out of bounds")] + fn test_avx2_take_u64_max_index_u32_lane() { + let values = vec![0u32; 8]; + // The first four indices execute in the SIMD loop; the scalar remainder is valid. + let indices = vec![0, u64::MAX, 2, 3, 0]; + + unsafe { take_avx2(&values, &indices) }; + } + + #[test] + #[should_panic(expected = "take index out of bounds")] + fn test_avx2_take_u64_max_index_u64_lane() { + let values = vec![0u64; 8]; + // The first four indices execute in the SIMD loop; the scalar remainder is valid. + let indices = vec![0, u64::MAX, 2, 3, 0]; + + unsafe { take_avx2(&values, &indices) }; + } + /// A `[u8; 4]` is a 4-byte `Copy` POD that is not a `NativePType`. This proves the kernel /// gathers an arbitrary 4-byte value type through the `u32` SIMD lane. #[test]