diff --git a/vortex-buffer/src/bit/take.rs b/vortex-buffer/src/bit/take.rs index 7affc34822b..c9ea1f54e17 100644 --- a/vortex-buffer/src/bit/take.rs +++ b/vortex-buffer/src/bit/take.rs @@ -79,9 +79,14 @@ where *bools.get_unchecked(idx) }) } else { + #[cfg(target_arch = "x86_64")] + if let Some(taken) = avx2::take(bits, indices, None) { + return taken; + } + let ptr = bits.inner().as_ptr(); let offset = bits.offset(); - BitBuffer::collect_bool(indices.len(), |i| unsafe { + BitBuffer::collect_bool_multiversioned(indices.len(), |i| unsafe { // SAFETY: we're already iterating over every index, a bound // is excessive let idx = indices.get_unchecked(i).as_(); @@ -91,6 +96,11 @@ where }; }; + #[cfg(target_arch = "x86_64")] + if let Some(taken) = avx2::take(bits, indices, Some(validity)) { + return taken; + } + let ptr = validity.inner().as_ptr(); if bits.len() <= COLLECT_TO_VEC { let offset = validity.offset(); @@ -132,8 +142,14 @@ fn first_unset(bits: BitBufferView<'_>) -> Option { (target < bits.len()).then_some(target) } +#[cfg(target_arch = "x86_64")] +#[path = "take_avx2.rs"] +mod avx2; + #[cfg(test)] mod tests { + use num_traits::AsPrimitive; + use num_traits::Zero; use rand::RngExt; use rand::SeedableRng; use rand::rngs::StdRng; @@ -152,34 +168,32 @@ mod tests { (0..len).map(|_| rng.random_range(0..bound)).collect() } - #[rstest] - #[case(300)] - #[case(5000)] - fn gathers(#[case] bits_len: usize) { + fn gathers(bits_len: usize, garbage: [I; 2]) + where + I: AsPrimitive + TryFrom + Ord + Zero, + >::Error: std::fmt::Debug, + { let bits = random_bits(bits_len + 11); let view = bits.as_view().slice(11..); - let len = 400; - let indices: Vec = random_indices(len, bits_len) + let len = 403; + let indices: Vec = random_indices(len, bits_len) .into_iter() - .map(|idx| i64::try_from(idx).unwrap()) + .map(|idx| I::try_from(idx).unwrap()) .collect(); let taken = take_bits(view, &indices, None); assert_eq!(taken.len(), len); for (i, &idx) in indices.iter().enumerate() { - assert_eq!( - taken.value(i), - bits.value(usize::try_from(idx).unwrap() + 11) - ); + assert_eq!(taken.value(i), bits.value(idx.as_() + 11)); } let validity = BitBuffer::collect_bool(len, |i| i % 3 != 0); - let with_garbage: Vec = indices + let with_garbage: Vec = indices .iter() .enumerate() .map(|(i, &idx)| match i % 6 { - 0 => i64::MAX, - 3 => -1, + 0 => garbage[0], + 3 => garbage[1], _ => idx, }) .collect(); @@ -188,14 +202,32 @@ mod tests { assert_eq!(taken.len(), len); for (i, &idx) in with_garbage.iter().enumerate() { if validity.value(i) { - assert_eq!( - taken.value(i), - bits.value(usize::try_from(idx).unwrap() + 11) - ); + assert_eq!(taken.value(i), bits.value(idx.as_() + 11)); } } } + #[rstest] + #[case(300)] + #[case(5000)] + fn gathers_i64(#[case] bits_len: usize) { + gathers::(bits_len, [i64::MAX, -1]); + } + + #[rstest] + #[case(300)] + #[case(5000)] + fn gathers_u16(#[case] bits_len: usize) { + gathers::(bits_len, [u16::MAX, u16::MAX - 1]); + } + + #[rstest] + #[case(300)] + #[case(5000)] + fn gathers_u32(#[case] bits_len: usize) { + gathers::(bits_len, [u32::MAX, u32::MAX - 1]); + } + #[rstest] #[case(BitBuffer::new_set(100), |_: usize| true)] #[case(BitBuffer::new_unset(100), |_: usize| false)] @@ -232,4 +264,18 @@ mod tests { let taken = take_bits(bits.as_view(), &[7u32, 8, 9], Some(validity.as_view())); assert_eq!(taken, BitBuffer::new_unset(3)); } + + #[test] + #[should_panic(expected = "out of bounds")] + fn out_of_bounds() { + let bits = BitBuffer::collect_bool(10_000, |i| i % 2 == 0); + let view = bits.as_view().slice(3..); + let mut indices: Vec = random_indices(403, 9_000) + .into_iter() + .map(|idx| u32::try_from(idx).unwrap()) + .collect(); + indices[80] = u32::MAX; + let validity = BitBuffer::collect_bool(403, |i| i != 200); + let _taken = take_bits(view, &indices, Some(validity.as_view())); + } } diff --git a/vortex-buffer/src/bit/take_avx2.rs b/vortex-buffer/src/bit/take_avx2.rs new file mode 100644 index 00000000000..08758c9a2d8 --- /dev/null +++ b/vortex-buffer/src/bit/take_avx2.rs @@ -0,0 +1,252 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::__m128i; +use std::arch::x86_64::__m256; +use std::arch::x86_64::__m256i; +use std::arch::x86_64::_mm_loadu_si128; +use std::arch::x86_64::_mm256_add_epi32; +use std::arch::x86_64::_mm256_and_si256; +use std::arch::x86_64::_mm256_castsi256_ps; +use std::arch::x86_64::_mm256_cmpeq_epi32; +use std::arch::x86_64::_mm256_cvtepu16_epi32; +use std::arch::x86_64::_mm256_i32gather_epi32; +use std::arch::x86_64::_mm256_loadu_si256; +use std::arch::x86_64::_mm256_mask_i32gather_epi32; +use std::arch::x86_64::_mm256_max_epu32; +use std::arch::x86_64::_mm256_min_epu32; +use std::arch::x86_64::_mm256_movemask_ps; +use std::arch::x86_64::_mm256_set1_epi32; +use std::arch::x86_64::_mm256_setr_epi32; +use std::arch::x86_64::_mm256_setzero_si256; +use std::arch::x86_64::_mm256_slli_epi32; +use std::arch::x86_64::_mm256_srli_epi32; +use std::arch::x86_64::_mm256_srlv_epi32; +use std::slice::from_raw_parts; + +use num_traits::AsPrimitive; + +use crate::BitBuffer; +use crate::BitBufferView; +use crate::BufferMut; +use crate::bit::get_bit; +use crate::bit::get_bit_unchecked; + +/// Caller must verify indices.len() == validity.len() +pub(super) fn take>( + bits: BitBufferView<'_>, + indices: &[I], + validity: Option>, +) -> Option { + if !is_x86_feature_detected!("avx2") { + return None; + } + if bits.len() >= (1 << 31) - 8 { + return None; + } + let last_dword = (bits.offset() + bits.len() - 1) / 32; + if (last_dword + 1) * 4 > bits.inner().len() { + return None; + } + + let size = size_of::(); + if size == size_of::() { + // SAFETY: if no validity is provided, caller has verified the indices + // are in bounds of "bits". If validity is provided, we clamp every + // index so it never produces an out of bound reads and panic if any + // clamped index isn't equal to original index. + let indices = unsafe { from_raw_parts(indices.as_ptr().cast::(), indices.len()) }; + let indices_ptr: *const u16 = indices.as_ptr(); + unsafe { + let group = |group_idx| -> __m256i { + let group_ptr: *const __m128i = indices_ptr.add(group_idx * 8).cast(); + // copy 128 bits into __m128i + let vector: __m128i = _mm_loadu_si128(group_ptr); + // zero-extend unsigned 16-bit integers in __m128i to signed + // 32-bit integers in __m256i + _mm256_cvtepu16_epi32(vector) + }; + return Some(take_group(bits, indices, validity, group)); + } + } + if size == size_of::() { + // SAFETY: see u16 case above + let indices = unsafe { from_raw_parts(indices.as_ptr().cast::(), indices.len()) }; + let indices_ptr: *const u32 = indices.as_ptr(); + unsafe { + let group = |group_idx| -> __m256i { + let group_ptr: *const __m256i = indices_ptr.add(group_idx * 8).cast(); + // copy 256 bits into __m256i + _mm256_loadu_si256(group_ptr) + }; + return Some(take_group(bits, indices, validity, group)); + } + } + None +} + +/// By the time we call this function we already handle cases of u16 and u32 +/// separately as there's different packing code. Now we also need to branch +/// on validity. take(), the first function, is called from two places, one +/// with validity and one without. Consequently, there are different gather +/// instructions depending on whether validity is present. +#[target_feature(enable = "avx2")] +unsafe fn take_group>( + bits: BitBufferView<'_>, + indices: &[I], + validity: Option>, + load_group: impl Fn(usize) -> __m256i, +) -> BitBuffer { + let inner = bits.inner(); + let base = inner.as_ptr(); + let zero: __m256i = _mm256_setzero_si256(); + + let Some(validity) = validity else { + // If validity isn't present, caller has verified "indices" are valid + // offsets into "bits" (see min/max code in take.rs), so we can use + // unchecked access. + let gather = |_group_idx: usize, indices: __m256i| -> (__m256i, __m256i) { + let gathered = unsafe { _mm256_i32gather_epi32::<4>(base.cast(), indices) }; + (gathered, zero) + }; + // inside take_lanes we do min(x, max_index) to check out of bound + // reads. Without validity there can't be out of bound reads so we + // pass u32::MAX. LLVM therefore reduces min(x, MAX) to x. + return unsafe { take_lanes(bits, indices, load_group, gather, usize::MAX, None) }; + }; + + // If validity is present, we can't trust "indices" are valid offsets. + // There also may be garbage under a NULL index, and we need to avoid + // reading bits's offset at this garbage. + let valid_bytes = validity.inner(); + let validity_offset = validity.offset(); + let validity_shift = validity_offset % 8; + let validity_first_byte = validity_offset / 8; + let has_non_byte_shift: usize = (validity_shift != 0).as_(); + + let lane_bits: __m256i = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128); + + let gather = |group_idx: usize, indices: __m256i| -> (__m256i, __m256i) { + let byte = validity_first_byte + group_idx; + let next_byte = byte + has_non_byte_shift; + + // We've verified in take.rs validity read is in bounds + let lo = unsafe { *valid_bytes.get_unchecked(byte) as u16 }; + let hi = unsafe { *valid_bytes.get_unchecked(next_byte) as u16 }; + + let window = lo | (hi << 8); + let valid = i32::from(window >> validity_shift); + + // mask = "valid" in every element of vector + let mask = _mm256_set1_epi32(valid); + // selected = mask & lane_bits + let selected = _mm256_and_si256(mask, lane_bits); + // lanes = selected == lane_bits + let lanes = _mm256_cmpeq_epi32(selected, lane_bits); + + let gathered = + unsafe { _mm256_mask_i32gather_epi32::<4>(zero, base.cast(), indices, lanes) }; + (gathered, lanes) + }; + + unsafe { + take_lanes( + bits, + indices, + load_group, + gather, + bits.len() - 1, + Some((valid_bytes, validity_offset)), + ) + } +} + +#[target_feature(enable = "avx2")] +unsafe fn take_lanes( + bits: BitBufferView<'_>, + indices: &[I], + load_group: impl Fn(usize) -> __m256i, + gather: impl Fn(usize, __m256i) -> (__m256i, __m256i), + max_index: usize, + validity: Option<(&[u8], usize)>, +) -> BitBuffer +where + I: AsPrimitive, +{ + let total = indices.len(); + let groups = total / 8; + + let bit_offset: __m256i = _mm256_set1_epi32(bits.offset().as_()); + let low_bits: __m256i = _mm256_set1_epi32(31); + let max_index: __m256i = _mm256_set1_epi32(max_index.as_()); + let mut max_taken_index: __m256i = _mm256_setzero_si256(); + + let mut out = BufferMut::::with_capacity(total.div_ceil(8)); + for group in 0..groups { + let original_indices: __m256i = load_group(group); + // clamped = min(original_indices, max_index) + // "clamped" never reads beyond bits's buffer end + let clamped = _mm256_min_epu32(original_indices, max_index); + // clamped_shifted = clamped + bit_offset + let clamped_shifted = _mm256_add_epi32(clamped, bit_offset); + + // gather_indices = clamped_shifted / 32 + // Gather can read only 32-bit words + let gather_indices = _mm256_srli_epi32::<5>(clamped_shifted); + let (gathered, lanes) = gather(group, gather_indices); + // taken_indices = original_indices & lanes + let taken_indices = _mm256_and_si256(original_indices, lanes); + // max_taken_index = max(max_taken_index, taken_indices) + max_taken_index = _mm256_max_epu32(max_taken_index, taken_indices); + + // shift = clamped_shifted % 32 + // Bit we need is now in position 0 + let shift = _mm256_and_si256(clamped_shifted, low_bits); + // shifted = gathered >> shift + let shifted = _mm256_srlv_epi32(gathered, shift); + // top = shifted << 31 + // Bit we need is now in position 31 + let top: __m256i = _mm256_slli_epi32::<31>(shifted); + // AVX2 doesn't have integral MSB for 32bit lanes, cast to floats + let top: __m256 = _mm256_castsi256_ps(top); + // per-element 1 if MSB is set, 0 otherwise + let msb: u8 = _mm256_movemask_ps(top).as_(); + + // SAFETY: out has sufficient capacity + unsafe { out.push_unchecked(msb) }; + } + + // max_clamped = min(max_taken_index, max_index) + let max_clamped = _mm256_min_epu32(max_taken_index, max_index); + // in_bounds = max_clamped == max_taken_index + let in_bounds: __m256i = _mm256_cmpeq_epi32(max_clamped, max_taken_index); + // AVX2 doesn't have integral MSB for 32bit lanes, cast to floats + let in_bounds: __m256 = _mm256_castsi256_ps(in_bounds); + // per-element 1 if MSB is set, 0 otherwise + let msb = _mm256_movemask_ps(in_bounds); + assert_eq!(msb, 0xFF, "index out of bounds"); + + let tail = total % 8; + if tail == 0 { + return BitBuffer::new(out.freeze(), total); + } + + let start = groups * 8; + let inner = bits.inner(); + let mut byte = 0u8; + for bit in 0..tail { + let pos = start + bit; + let keep = validity.map_or(usize::MAX, |(valid_bytes, validity_offset)| { + // SAFETY: validity has same length as indices + (unsafe { get_bit_unchecked(valid_bytes.as_ptr(), validity_offset + pos) } as usize) + .wrapping_neg() + }); + // SAFETY: pos stays within indices + let idx = unsafe { indices.get_unchecked(pos) }.as_() & keep; + let value = get_bit(inner, bits.offset() + idx); + byte |= (value as u8) << bit; + } + // SAFETY: out has sufficient capacity + unsafe { out.push_unchecked(byte) } + BitBuffer::new(out.freeze(), total) +}