mirror of
https://github.com/quickwit-oss/tantivy.git
synced 2026-01-06 17:22:54 +00:00
Faster range (#1954)
* Faster range queries This PR does several changes - ip compact space now uses u32 - the bitunpacker now gets a get_batch function - we push down range filtering, removing GCD / shift in the bitpacking codec. - we rely on AVX2 routine to do the filtering. * Apply suggestions from code review * Apply suggestions from code review * CR comments
This commit is contained in:
Paul Masurel
GitHub
@@ -1,10 +1,14 @@
|
||||
use std::convert::TryInto;
|
||||
use std::io;
|
||||
use std::ops::{Range, RangeInclusive};
|
||||
|
||||
use bitpacking::{BitPacker as ExternalBitPackerTrait, BitPacker1x};
|
||||
|
||||
pub struct BitPacker {
|
||||
mini_buffer: u64,
|
||||
mini_buffer_written: usize,
|
||||
}
|
||||
|
||||
impl Default for BitPacker {
|
||||
fn default() -> Self {
|
||||
BitPacker::new()
|
||||
@@ -118,6 +122,125 @@ impl BitUnpacker {
|
||||
let val_shifted = val_unshifted_unmasked >> bit_shift;
|
||||
val_shifted & self.mask
|
||||
}
|
||||
|
||||
// Decodes the range of bitpacked `u32` values with idx
|
||||
// in [start_idx, start_idx + output.len()).
|
||||
//
|
||||
// #Panics
|
||||
//
|
||||
// This methods panics if `num_bits` is > 32.
|
||||
fn get_batch_u32s(&self, start_idx: u32, data: &[u8], output: &mut [u32]) {
|
||||
assert!(
|
||||
self.bit_width() <= 32,
|
||||
"Bitwidth must be <= 32 to use this method."
|
||||
);
|
||||
|
||||
let end_idx = start_idx + output.len() as u32;
|
||||
|
||||
let end_bit_read = end_idx * self.num_bits;
|
||||
let end_byte_read = (end_bit_read + 7) / 8;
|
||||
assert!(
|
||||
end_byte_read as usize <= data.len(),
|
||||
"Requested index is out of bounds."
|
||||
);
|
||||
|
||||
// Simple slow implementation of get_batch_u32s, to deal with our ramps.
|
||||
let get_batch_ramp = |start_idx: u32, output: &mut [u32]| {
|
||||
for (out, idx) in output.iter_mut().zip(start_idx..) {
|
||||
*out = self.get(idx, data) as u32;
|
||||
}
|
||||
};
|
||||
|
||||
// We use an unrolled routine to decode 32 values at once.
|
||||
// We therefore decompose our range of values to decode into three ranges:
|
||||
// - Entrance ramp: [start_idx, fast_track_start) (up to 31 values)
|
||||
// - Highway: [fast_track_start, fast_track_end) (a length multiple of 32s)
|
||||
// - Exit ramp: [fast_track_end, start_idx + output.len()) (up to 31 values)
|
||||
|
||||
// We want the start of the fast track to start align with bytes.
|
||||
// A sufficient condition is to start with an idx that is a multiple of 8,
|
||||
// so highway start is the closest multiple of 8 that is >= start_idx.
|
||||
let entrance_ramp_len = 8 - (start_idx % 8) % 8;
|
||||
|
||||
let highway_start: u32 = start_idx + entrance_ramp_len;
|
||||
|
||||
if highway_start + BitPacker1x::BLOCK_LEN as u32 > end_idx {
|
||||
// We don't have enough values to have even a single block of highway.
|
||||
// Let's just supply the values the simple way.
|
||||
get_batch_ramp(start_idx, output);
|
||||
return;
|
||||
}
|
||||
|
||||
let num_blocks: u32 = (end_idx - highway_start) / BitPacker1x::BLOCK_LEN as u32;
|
||||
|
||||
// Entrance ramp
|
||||
get_batch_ramp(start_idx, &mut output[..entrance_ramp_len as usize]);
|
||||
|
||||
// Highway
|
||||
let mut offset = (highway_start * self.num_bits) as usize / 8;
|
||||
let mut output_cursor = (highway_start - start_idx) as usize;
|
||||
for _ in 0..num_blocks {
|
||||
offset += BitPacker1x.decompress(
|
||||
&data[offset..],
|
||||
&mut output[output_cursor..],
|
||||
self.num_bits as u8,
|
||||
);
|
||||
output_cursor += 32;
|
||||
}
|
||||
|
||||
// Exit ramp
|
||||
let highway_end = highway_start + num_blocks * BitPacker1x::BLOCK_LEN as u32;
|
||||
get_batch_ramp(highway_end, &mut output[output_cursor..]);
|
||||
}
|
||||
|
||||
pub fn get_ids_for_value_range(
|
||||
&self,
|
||||
range: RangeInclusive<u64>,
|
||||
id_range: Range<u32>,
|
||||
data: &[u8],
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
if self.bit_width() > 32 {
|
||||
self.get_ids_for_value_range_slow(range, id_range, data, positions)
|
||||
} else {
|
||||
if *range.start() > u32::MAX as u64 {
|
||||
positions.clear();
|
||||
return;
|
||||
}
|
||||
let range_u32 = (*range.start() as u32)..=(*range.end()).min(u32::MAX as u64) as u32;
|
||||
self.get_ids_for_value_range_fast(range_u32, id_range, data, positions)
|
||||
}
|
||||
}
|
||||
|
||||
fn get_ids_for_value_range_slow(
|
||||
&self,
|
||||
range: RangeInclusive<u64>,
|
||||
id_range: Range<u32>,
|
||||
data: &[u8],
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
positions.clear();
|
||||
for i in id_range {
|
||||
// If we cared we could make this branchless, but the slow implementation should rarely
|
||||
// kick in.
|
||||
let val = self.get(i, data);
|
||||
if range.contains(&val) {
|
||||
positions.push(i);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn get_ids_for_value_range_fast(
|
||||
&self,
|
||||
value_range: RangeInclusive<u32>,
|
||||
id_range: Range<u32>,
|
||||
data: &[u8],
|
||||
positions: &mut Vec<u32>,
|
||||
) {
|
||||
positions.resize(id_range.len(), 0u32);
|
||||
self.get_batch_u32s(id_range.start, data, positions);
|
||||
crate::filter_vec::filter_vec_in_place(value_range, id_range.start, positions)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -200,4 +323,58 @@ mod test {
|
||||
test_bitpacker_aux(num_bits, &vals);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic]
|
||||
fn test_get_batch_panics_over_32_bits() {
|
||||
let bitunpacker = BitUnpacker::new(33);
|
||||
let mut output: [u32; 1] = [0u32];
|
||||
bitunpacker.get_batch_u32s(0, &[0, 0, 0, 0, 0, 0, 0, 0], &mut output[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_batch_limit() {
|
||||
let bitunpacker = BitUnpacker::new(1);
|
||||
let mut output: [u32; 3] = [0u32, 0u32, 0u32];
|
||||
bitunpacker.get_batch_u32s(8 * 4 - 3, &[0u8, 0u8, 0u8, 0u8], &mut output[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[should_panic]
|
||||
fn test_get_batch_panics_when_off_scope() {
|
||||
let bitunpacker = BitUnpacker::new(1);
|
||||
let mut output: [u32; 3] = [0u32, 0u32, 0u32];
|
||||
// We are missing exactly one bit.
|
||||
bitunpacker.get_batch_u32s(8 * 4 - 2, &[0u8, 0u8, 0u8, 0u8], &mut output[..]);
|
||||
}
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn test_get_batch_u32s_proptest(num_bits in 0u8..=32u8) {
|
||||
let mask =
|
||||
if num_bits == 32u8 {
|
||||
u32::MAX
|
||||
} else {
|
||||
(1u32 << num_bits) - 1
|
||||
};
|
||||
let mut buffer: Vec<u8> = Vec::new();
|
||||
let mut bitpacker = BitPacker::new();
|
||||
for val in 0..100 {
|
||||
bitpacker.write(val & mask as u64, num_bits, &mut buffer).unwrap();
|
||||
}
|
||||
bitpacker.flush(&mut buffer).unwrap();
|
||||
let bitunpacker = BitUnpacker::new(num_bits);
|
||||
let mut output: Vec<u32> = Vec::new();
|
||||
for len in [0, 1, 2, 32, 33, 34, 64] {
|
||||
for start_idx in 0u32..32u32 {
|
||||
output.resize(len as usize, 0);
|
||||
bitunpacker.get_batch_u32s(start_idx, &buffer, &mut output);
|
||||
for i in 0..len {
|
||||
let expected = (start_idx + i as u32) & mask;
|
||||
assert_eq!(output[i], expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
365
bitpacker/src/filter_vec/avx2.rs
Normal file
365
bitpacker/src/filter_vec/avx2.rs
Normal file
@@ -0,0 +1,365 @@
|
||||
//! SIMD filtering of a vector as described in the following blog post.
|
||||
//! https://quickwit.io/blog/filtering%20a%20vector%20with%20simd%20instructions%20avx-2%20and%20avx-512
|
||||
use std::arch::x86_64::{
|
||||
__m256i as DataType, _mm256_add_epi32 as op_add, _mm256_cmpgt_epi32 as op_greater,
|
||||
_mm256_lddqu_si256 as load_unaligned, _mm256_or_si256 as op_or, _mm256_set1_epi32 as set1,
|
||||
_mm256_storeu_si256 as store_unaligned, _mm256_xor_si256 as op_xor, *,
|
||||
};
|
||||
use std::ops::RangeInclusive;
|
||||
|
||||
const NUM_LANES: usize = 8;
|
||||
|
||||
const HIGHEST_BIT: u32 = 1 << 31;
|
||||
|
||||
#[inline]
|
||||
fn u32_to_i32(val: u32) -> i32 {
|
||||
(val ^ HIGHEST_BIT) as i32
|
||||
}
|
||||
|
||||
#[inline]
|
||||
unsafe fn u32_to_i32_avx2(vals_u32x8s: DataType) -> DataType {
|
||||
const HIGHEST_BIT_MASK: DataType = from_u32x8([HIGHEST_BIT; NUM_LANES]);
|
||||
op_xor(vals_u32x8s, HIGHEST_BIT_MASK)
|
||||
}
|
||||
|
||||
pub fn filter_vec_in_place(range: RangeInclusive<u32>, offset: u32, output: &mut Vec<u32>) {
|
||||
// We use a monotonic mapping from u32 to i32 to make the comparison possible in AVX2.
|
||||
let range_i32: RangeInclusive<i32> = u32_to_i32(*range.start())..=u32_to_i32(*range.end());
|
||||
let num_words = output.len() / NUM_LANES;
|
||||
let mut output_len = unsafe {
|
||||
filter_vec_avx2_aux(
|
||||
output.as_ptr() as *const __m256i,
|
||||
range_i32,
|
||||
output.as_mut_ptr(),
|
||||
offset,
|
||||
num_words,
|
||||
)
|
||||
};
|
||||
let reminder_start = num_words * NUM_LANES;
|
||||
for i in reminder_start..output.len() {
|
||||
let val = output[i];
|
||||
output[output_len] = offset + i as u32;
|
||||
output_len += if range.contains(&val) { 1 } else { 0 };
|
||||
}
|
||||
output.truncate(output_len);
|
||||
}
|
||||
|
||||
#[target_feature(enable = "avx2")]
|
||||
unsafe fn filter_vec_avx2_aux(
|
||||
mut input: *const __m256i,
|
||||
range: RangeInclusive<i32>,
|
||||
output: *mut u32,
|
||||
offset: u32,
|
||||
num_words: usize,
|
||||
) -> usize {
|
||||
let mut output_tail = output;
|
||||
let range_simd = set1(*range.start())..=set1(*range.end());
|
||||
let mut ids = from_u32x8([
|
||||
offset,
|
||||
offset + 1,
|
||||
offset + 2,
|
||||
offset + 3,
|
||||
offset + 4,
|
||||
offset + 5,
|
||||
offset + 6,
|
||||
offset + 7,
|
||||
]);
|
||||
const SHIFT: __m256i = from_u32x8([NUM_LANES as u32; NUM_LANES]);
|
||||
for _ in 0..num_words {
|
||||
let word = load_unaligned(input);
|
||||
let word = u32_to_i32_avx2(word);
|
||||
let keeper_bitset = compute_filter_bitset(word, range_simd.clone());
|
||||
let added_len = keeper_bitset.count_ones();
|
||||
let filtered_doc_ids = compact(ids, keeper_bitset);
|
||||
store_unaligned(output_tail as *mut __m256i, filtered_doc_ids);
|
||||
output_tail = output_tail.offset(added_len as isize);
|
||||
ids = op_add(ids, SHIFT);
|
||||
input = input.offset(1);
|
||||
}
|
||||
output_tail.offset_from(output) as usize
|
||||
}
|
||||
|
||||
#[inline]
|
||||
#[target_feature(enable = "avx2")]
|
||||
unsafe fn compact(data: DataType, mask: u8) -> DataType {
|
||||
let vperm_mask = MASK_TO_PERMUTATION[mask as usize];
|
||||
_mm256_permutevar8x32_epi32(data, vperm_mask)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
#[target_feature(enable = "avx2")]
|
||||
unsafe fn compute_filter_bitset(val: __m256i, range: std::ops::RangeInclusive<__m256i>) -> u8 {
|
||||
let too_low = op_greater(*range.start(), val);
|
||||
let too_high = op_greater(val, *range.end());
|
||||
let inside = op_or(too_low, too_high);
|
||||
255 - std::arch::x86_64::_mm256_movemask_ps(std::mem::transmute::<DataType, __m256>(inside))
|
||||
as u8
|
||||
}
|
||||
|
||||
union U8x32 {
|
||||
vector: DataType,
|
||||
vals: [u32; NUM_LANES],
|
||||
}
|
||||
|
||||
const fn from_u32x8(vals: [u32; NUM_LANES]) -> DataType {
|
||||
unsafe { U8x32 { vals }.vector }
|
||||
}
|
||||
|
||||
const MASK_TO_PERMUTATION: [DataType; 256] = [
|
||||
from_u32x8([0, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([3, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 0, 0, 0, 0]),
|
||||
from_u32x8([4, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 0, 0, 0, 0]),
|
||||
from_u32x8([3, 4, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 0, 0, 0]),
|
||||
from_u32x8([5, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 5, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 5, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 5, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([3, 5, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 5, 0, 0, 0]),
|
||||
from_u32x8([4, 5, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 5, 0, 0, 0]),
|
||||
from_u32x8([3, 4, 5, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 5, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 5, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 5, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 5, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 5, 0, 0]),
|
||||
from_u32x8([6, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 6, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 6, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 6, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([3, 6, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 6, 0, 0, 0]),
|
||||
from_u32x8([4, 6, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 6, 0, 0, 0]),
|
||||
from_u32x8([3, 4, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 6, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 6, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 6, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 6, 0, 0]),
|
||||
from_u32x8([5, 6, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 5, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 5, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 5, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([3, 5, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 5, 6, 0, 0]),
|
||||
from_u32x8([4, 5, 6, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 5, 6, 0, 0]),
|
||||
from_u32x8([3, 4, 5, 6, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 5, 6, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 5, 6, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 5, 6, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 5, 6, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 5, 6, 0]),
|
||||
from_u32x8([7, 0, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([3, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 7, 0, 0, 0]),
|
||||
from_u32x8([4, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 7, 0, 0, 0]),
|
||||
from_u32x8([3, 4, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 7, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 7, 0, 0]),
|
||||
from_u32x8([5, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 5, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 5, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 5, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([3, 5, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 5, 7, 0, 0]),
|
||||
from_u32x8([4, 5, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 5, 7, 0, 0]),
|
||||
from_u32x8([3, 4, 5, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 5, 7, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 5, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 5, 7, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 5, 7, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 5, 7, 0]),
|
||||
from_u32x8([6, 7, 0, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 6, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 6, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([2, 6, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([3, 6, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([2, 3, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 6, 7, 0, 0]),
|
||||
from_u32x8([4, 6, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([2, 4, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 6, 7, 0, 0]),
|
||||
from_u32x8([3, 4, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 6, 7, 0, 0]),
|
||||
from_u32x8([2, 3, 4, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 6, 7, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 6, 7, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 6, 7, 0]),
|
||||
from_u32x8([5, 6, 7, 0, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 5, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([1, 5, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([2, 5, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 2, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([3, 5, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 3, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([2, 3, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([1, 2, 3, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 5, 6, 7, 0]),
|
||||
from_u32x8([4, 5, 6, 7, 0, 0, 0, 0]),
|
||||
from_u32x8([0, 4, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([1, 4, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 1, 4, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([2, 4, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 2, 4, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([1, 2, 4, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([0, 1, 2, 4, 5, 6, 7, 0]),
|
||||
from_u32x8([3, 4, 5, 6, 7, 0, 0, 0]),
|
||||
from_u32x8([0, 3, 4, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([1, 3, 4, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([0, 1, 3, 4, 5, 6, 7, 0]),
|
||||
from_u32x8([2, 3, 4, 5, 6, 7, 0, 0]),
|
||||
from_u32x8([0, 2, 3, 4, 5, 6, 7, 0]),
|
||||
from_u32x8([1, 2, 3, 4, 5, 6, 7, 0]),
|
||||
from_u32x8([0, 1, 2, 3, 4, 5, 6, 7]),
|
||||
];
|
||||
167
bitpacker/src/filter_vec/mod.rs
Normal file
167
bitpacker/src/filter_vec/mod.rs
Normal file
@@ -0,0 +1,167 @@
|
||||
use std::ops::RangeInclusive;
|
||||
use std::sync::atomic::AtomicU8;
|
||||
|
||||
#[cfg(any(target_arch = "x86_64"))]
|
||||
mod avx2;
|
||||
|
||||
mod scalar;
|
||||
|
||||
#[derive(Clone, Copy, Eq, PartialEq, Debug)]
|
||||
#[repr(u8)]
|
||||
enum FilterImplPerInstructionSet {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
AVX2 = 0u8,
|
||||
Scalar = 1u8,
|
||||
}
|
||||
|
||||
impl FilterImplPerInstructionSet {
|
||||
#[inline]
|
||||
pub fn is_available(&self) -> bool {
|
||||
match *self {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
FilterImplPerInstructionSet::AVX2 => is_x86_feature_detected!("avx2"),
|
||||
FilterImplPerInstructionSet::Scalar => true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// List of available implementation in preferred order.
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
const IMPLS: [FilterImplPerInstructionSet; 2] = [
|
||||
FilterImplPerInstructionSet::AVX2,
|
||||
FilterImplPerInstructionSet::Scalar,
|
||||
];
|
||||
|
||||
impl FilterImplPerInstructionSet {
|
||||
#[inline]
|
||||
fn from(code: u8) -> FilterImplPerInstructionSet {
|
||||
if code == FilterImplPerInstructionSet::AVX2 as u8 {
|
||||
FilterImplPerInstructionSet::AVX2
|
||||
} else {
|
||||
FilterImplPerInstructionSet::Scalar
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn filter_vec_in_place(self, range: RangeInclusive<u32>, offset: u32, output: &mut Vec<u32>) {
|
||||
match self {
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
FilterImplPerInstructionSet::AVX2 => avx2::filter_vec_in_place(range, offset, output),
|
||||
FilterImplPerInstructionSet::Scalar => {
|
||||
scalar::filter_vec_in_place(range, offset, output)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
#[inline]
|
||||
fn get_best_available_instruction_set() -> FilterImplPerInstructionSet {
|
||||
use std::sync::atomic::Ordering;
|
||||
static INSTRUCTION_SET_BYTE: AtomicU8 = AtomicU8::new(u8::MAX);
|
||||
let instruction_set_byte: u8 = INSTRUCTION_SET_BYTE.load(Ordering::Relaxed);
|
||||
if instruction_set_byte == u8::MAX {
|
||||
// Let's initialize the instruction set and cache it.
|
||||
let instruction_set = IMPLS
|
||||
.into_iter()
|
||||
.find(FilterImplPerInstructionSet::is_available)
|
||||
.unwrap();
|
||||
INSTRUCTION_SET_BYTE.store(instruction_set as u8, Ordering::Relaxed);
|
||||
return instruction_set;
|
||||
}
|
||||
FilterImplPerInstructionSet::from(instruction_set_byte)
|
||||
}
|
||||
|
||||
#[cfg(not(target_arch = "x86_64"))]
|
||||
#[inline]
|
||||
const fn get_best_available_instruction_set() -> FilterImplPerInstructionSet {
|
||||
FilterImplPerInstructionSet::Scalar
|
||||
}
|
||||
|
||||
pub fn filter_vec_in_place(range: RangeInclusive<u32>, offset: u32, output: &mut Vec<u32>) {
|
||||
get_best_available_instruction_set().filter_vec_in_place(range, offset, output)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_get_best_available_instruction_set() {
|
||||
// This does not test much unfortunately.
|
||||
// We just make sure the function returns without crashing and returns the same result.
|
||||
let instruction_set = get_best_available_instruction_set();
|
||||
assert_eq!(get_best_available_instruction_set(), instruction_set);
|
||||
}
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
#[test]
|
||||
fn test_instruction_set_to_code_from_code() {
|
||||
for instruction_set in [
|
||||
FilterImplPerInstructionSet::AVX2,
|
||||
FilterImplPerInstructionSet::Scalar,
|
||||
] {
|
||||
let code = instruction_set as u8;
|
||||
assert_eq!(instruction_set, FilterImplPerInstructionSet::from(code));
|
||||
}
|
||||
}
|
||||
|
||||
fn test_filter_impl_empty_aux(filter_impl: FilterImplPerInstructionSet) {
|
||||
let mut output = vec![];
|
||||
filter_impl.filter_vec_in_place(0..=u32::MAX, 0, &mut output);
|
||||
assert_eq!(&output, &[]);
|
||||
}
|
||||
|
||||
fn test_filter_impl_simple_aux(filter_impl: FilterImplPerInstructionSet) {
|
||||
let mut output = vec![3, 2, 1, 5, 11, 2, 5, 10, 2];
|
||||
filter_impl.filter_vec_in_place(3..=10, 0, &mut output);
|
||||
assert_eq!(&output, &[0, 3, 6, 7]);
|
||||
}
|
||||
|
||||
fn test_filter_impl_simple_aux_shifted(filter_impl: FilterImplPerInstructionSet) {
|
||||
let mut output = vec![3, 2, 1, 5, 11, 2, 5, 10, 2];
|
||||
filter_impl.filter_vec_in_place(3..=10, 10, &mut output);
|
||||
assert_eq!(&output, &[10, 13, 16, 17]);
|
||||
}
|
||||
|
||||
fn test_filter_impl_simple_outside_i32_range(filter_impl: FilterImplPerInstructionSet) {
|
||||
let mut output = vec![u32::MAX, i32::MAX as u32 + 1, 0, 1, 3, 1, 1, 1, 1];
|
||||
filter_impl.filter_vec_in_place(1..=i32::MAX as u32 + 1u32, 0, &mut output);
|
||||
assert_eq!(&output, &[1, 3, 4, 5, 6, 7, 8]);
|
||||
}
|
||||
|
||||
fn test_filter_impl_test_suite(filter_impl: FilterImplPerInstructionSet) {
|
||||
test_filter_impl_empty_aux(filter_impl);
|
||||
test_filter_impl_simple_aux(filter_impl);
|
||||
test_filter_impl_simple_aux_shifted(filter_impl);
|
||||
test_filter_impl_simple_outside_i32_range(filter_impl);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_implementation_avx2() {
|
||||
if FilterImplPerInstructionSet::AVX2.is_available() {
|
||||
test_filter_impl_test_suite(FilterImplPerInstructionSet::AVX2);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filter_implementation_scalar() {
|
||||
test_filter_impl_test_suite(FilterImplPerInstructionSet::Scalar);
|
||||
}
|
||||
|
||||
proptest::proptest! {
|
||||
#[test]
|
||||
fn test_filter_impl_proptest(
|
||||
start in proptest::prelude::any::<u32>(),
|
||||
end in proptest::prelude::any::<u32>(),
|
||||
offset in 0u32..2u32,
|
||||
mut vals in proptest::collection::vec(0..u32::MAX, 0..30)) {
|
||||
if FilterImplPerInstructionSet::AVX2.is_available() {
|
||||
let mut vals_clone = vals.clone();
|
||||
FilterImplPerInstructionSet::AVX2.filter_vec_in_place(start..=end, offset, &mut vals);
|
||||
FilterImplPerInstructionSet::Scalar.filter_vec_in_place(start..=end, offset, &mut vals_clone);
|
||||
assert_eq!(&vals, &vals_clone);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
13
bitpacker/src/filter_vec/scalar.rs
Normal file
13
bitpacker/src/filter_vec/scalar.rs
Normal file
@@ -0,0 +1,13 @@
|
||||
use std::ops::RangeInclusive;
|
||||
|
||||
pub fn filter_vec_in_place(range: RangeInclusive<u32>, offset: u32, output: &mut Vec<u32>) {
|
||||
// We restrict the accepted boundary, because unsigned integers & SIMD don't
|
||||
// play well.
|
||||
let mut output_cursor = 0;
|
||||
for i in 0..output.len() {
|
||||
let val = output[i];
|
||||
output[output_cursor] = offset + i as u32;
|
||||
output_cursor += if range.contains(&val) { 1 } else { 0 };
|
||||
}
|
||||
output.truncate(output_cursor);
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
mod bitpacker;
|
||||
mod blocked_bitpacker;
|
||||
mod filter_vec;
|
||||
|
||||
use std::cmp::Ordering;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user