Skip to content

fix: validate unsigned AVX2 take indices #9036

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
connortsui20 wants to merge 3 commits into
base: develop
Choose a base branch
from
+92 −19
Open

fix: validate unsigned AVX2 take indices #9036

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2d6ec8f
fix: validate unsigned AVX2 take indices
connortsui20 Jul 28, 2026
96343dd
fix: use the AVX2 take valid mask
connortsui20 Jul 28, 2026
19d2840
fix: preserve full-width AVX2 take indices
connortsui20 Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 92 additions & 19 deletions vortex-array/src/arrays/primitive/compute/take/avx2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -172,12 +173,43 @@ trait GatherFn<Idx, Values> {
///
/// 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<Idx>,
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 }),+) => {
$(
Expand All @@ -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");
}
Expand All @@ -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 {

@robert3005 robert3005 Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this was non optional to avoid the branch here, you should lift it out

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::<SCALE>(zero_vec, src.cast(), indices_vec, invalid_mask) };
let values_vec = unsafe { $masked_gather::<SCALE>(zero_vec, src.cast(), indices_vec, valid_mask) };

// Write the vec out to dst.
unsafe { $store(dst.cast(), values_vec) };
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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::<Out>::with_capacity_aligned(indices_len, Alignment::of::<__m256i>());
let buf_uninit = buffer.spare_capacity_mut();
Expand Down Expand Up @@ -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]
Expand Down
Loading