From d85faa01f25da3ec435dd7664db64984ebe10838 Mon Sep 17 00:00:00 2001 From: Will Manning Date: Tue, 1 Sep 2026 16:49:04 -0400 Subject: [PATCH] test(bitpacking): Correct indexed unpack benchmarks --- benches/bitpacking.rs | 77 +++++++++++-------------------------------- src/bitpacking.rs | 45 ++++++++++++++++++------- 2 files changed, 53 insertions(+), 69 deletions(-) diff --git a/benches/bitpacking.rs b/benches/bitpacking.rs index e8ef08c..97aed98 100644 --- a/benches/bitpacking.rs +++ b/benches/bitpacking.rs @@ -91,49 +91,21 @@ fn unpack_single_16_from_3(bencher: Bencher) { }); } -const MAX_BENCHMARK_INDICES: usize = 128; +const MAX_BENCHMARK_INDICES: usize = 192; -#[derive(Clone, Copy)] -enum IndexDistribution { - Uniform, - Clustered, - UnorderedWithDuplicates, -} - -fn benchmark_indices( - distribution: IndexDistribution, - num_indices: usize, -) -> [usize; MAX_BENCHMARK_INDICES] { +fn benchmark_indices(num_indices: usize) -> [usize; MAX_BENCHMARK_INDICES] { assert!(num_indices <= MAX_BENCHMARK_INDICES); let mut indices = [0; MAX_BENCHMARK_INDICES]; - match distribution { - IndexDistribution::Uniform => { - for (position, index) in indices[..num_indices].iter_mut().enumerate() { - *index = position * 1024 / num_indices; - } - } - IndexDistribution::Clustered => { - for (position, index) in indices[..num_indices].iter_mut().enumerate() { - *index = 448 + position; - } - } - IndexDistribution::UnorderedWithDuplicates => { - for position in 0..num_indices { - indices[position] = if position > 0 && position % 4 == 0 { - indices[position - 1] - } else { - (position * 541 + 17) % 1024 - }; - } - } + for (position, index) in indices[..num_indices].iter_mut().enumerate() { + *index = position * 1024 / num_indices; } indices } macro_rules! unpack_indices_benchmarks { - ($module:ident, $type:ty, $width:expr, $distribution:expr) => { + ($module:ident, $type:ty, $width:expr, [$($num_indices:expr),+ $(,)?]) => { mod $module { use super::*; @@ -148,14 +120,15 @@ macro_rules! unpack_indices_benchmarks { std::array::from_fn(|index| ((index as $type).wrapping_mul(17)) & mask); let mut packed = [0; PACKED_LENGTH]; BitPacking::pack::(&values, &mut packed); - let indices = benchmark_indices($distribution, num_indices); + let indices = benchmark_indices(num_indices); (packed, indices) } - #[divan::bench(args = [1, 8, 32, 128], sample_size = 10_000)] + #[divan::bench(args = [$($num_indices),+], sample_size = 10_000)] fn batched(bencher: Bencher, num_indices: usize) { let (packed, indices) = fixture(num_indices); let mut output = [MaybeUninit::<$type>::uninit(); MAX_BENCHMARK_INDICES]; + let width = black_box(WIDTH); bencher.bench_local(|| { let packed = black_box(&packed); @@ -163,16 +136,17 @@ macro_rules! unpack_indices_benchmarks { let output = black_box(&mut output[..num_indices]); // SAFETY: `packed` contains exactly one packed FastLanes block. unsafe { - BitPacking::unchecked_unpack_indices(WIDTH, packed, indices, output); + BitPacking::unchecked_unpack_indices(width, packed, indices, output); } black_box(&*output); }); } - #[divan::bench(args = [1, 8, 32, 128], sample_size = 10_000)] + #[divan::bench(args = [$($num_indices),+], sample_size = 10_000)] fn repeated_single(bencher: Bencher, num_indices: usize) { let (packed, indices) = fixture(num_indices); let mut output = [MaybeUninit::<$type>::uninit(); MAX_BENCHMARK_INDICES]; + let width = black_box(WIDTH); bencher.bench_local(|| { let packed = black_box(&packed); @@ -181,18 +155,19 @@ macro_rules! unpack_indices_benchmarks { for (&index, value) in indices.iter().zip(output.iter_mut()) { // SAFETY: `packed` contains exactly one packed FastLanes block. value.write(unsafe { - BitPacking::unchecked_unpack_single(WIDTH, packed, index) + BitPacking::unchecked_unpack_single(width, packed, index) }); } black_box(&*output); }); } - #[divan::bench(args = [1, 8, 32, 128], sample_size = 10_000)] + #[divan::bench(args = [$($num_indices),+], sample_size = 10_000)] fn full_unpack_then_gather(bencher: Bencher, num_indices: usize) { let (packed, indices) = fixture(num_indices); let mut unpacked = [0; 1024]; let mut output = [MaybeUninit::<$type>::uninit(); MAX_BENCHMARK_INDICES]; + let width = black_box(WIDTH); bencher.bench_local(|| { let packed = black_box(&packed); @@ -200,7 +175,7 @@ macro_rules! unpack_indices_benchmarks { let unpacked = black_box(&mut unpacked); let output = black_box(&mut output[..num_indices]); // SAFETY: both buffers have the required lengths for `WIDTH`. - unsafe { BitPacking::unchecked_unpack(WIDTH, packed, unpacked) }; + unsafe { BitPacking::unchecked_unpack(width, packed, unpacked) }; for (&index, value) in indices.iter().zip(output.iter_mut()) { value.write(unpacked[index]); } @@ -211,24 +186,10 @@ macro_rules! unpack_indices_benchmarks { }; } -unpack_indices_benchmarks!( - unpack_indices_u16_width3_uniform, - u16, - 3, - IndexDistribution::Uniform -); -unpack_indices_benchmarks!( - unpack_indices_u32_width16_clustered, - u32, - 16, - IndexDistribution::Clustered -); -unpack_indices_benchmarks!( - unpack_indices_u64_width63_unordered_duplicates, - u64, - 63, - IndexDistribution::UnorderedWithDuplicates -); +unpack_indices_benchmarks!(unpack_indices_u8_width7, u8, 7, [1, 8, 16, 24]); +unpack_indices_benchmarks!(unpack_indices_u16_width3, u16, 3, [1, 8, 32, 48]); +unpack_indices_benchmarks!(unpack_indices_u32_width16, u32, 16, [1, 8, 64, 80]); +unpack_indices_benchmarks!(unpack_indices_u64_width63, u64, 63, [1, 8, 160, 192]); #[divan::bench(sample_count = 10000)] fn throughput_compress(bencher: Bencher) { diff --git a/src/bitpacking.rs b/src/bitpacking.rs index 7d13de4..abeab41 100644 --- a/src/bitpacking.rs +++ b/src/bitpacking.rs @@ -153,6 +153,7 @@ macro_rules! impl_packing { assert!(B == 1024 * W / Self::T); } + for lane in 0..Self::LANES { unpack!($T, W, input, lane, |$idx, $elem| { output[$idx] = $elem @@ -199,11 +200,21 @@ macro_rules! impl_packing { } if W == 0 { + // Special case for W=0, we just need to zero the output. return 0 as $T; } - // Packing transposes each logical index into a FastLanes lane and row. - // Map the requested index back to those coordinates. + // We can think of the input array as effectively a row-major, left-to-right + // 2-D array of with `Self::LANES` columns and `Self::T` rows. + // + // Meanwhile, we can think of the packed array as either: + // 1. `Self::T` rows of W-bit elements, with `Self::LANES` columns + // 2. `W` rows of `Self::T`-bit words, with `Self::LANES` columns + // + // Bitpacking involves a transposition of the input array ordering, such that + // decompression can be fused efficiently with encodings like delta and RLE. + // + // First step, we need to get the lane and row for interpretation #1 above. assert!(index < 1024, "Index must be less than 1024, got {}", index); let (lane, row): (usize, usize) = { const LANES: [u8; 1024] = lanes_by_index::<$T>(); @@ -212,6 +223,7 @@ macro_rules! impl_packing { }; if W == <$T>::T { + // Special case for W==T, we can just read the value directly return packed[<$T>::LANES * row + lane]; } @@ -222,12 +234,14 @@ macro_rules! impl_packing { let remaining_bits = <$T>::T - lo_shift; let lo = packed[<$T>::LANES * start_word + lane] >> lo_shift; - if remaining_bits >= W { + return if remaining_bits >= W { + // in this case we will mask out all bits of hi word lo & mask } else { + // guaranteed that lo_shift > 0 and thus remaining_bits < T let hi = packed[<$T>::LANES * (start_word + 1) + lane] << remaining_bits; (lo | hi) & mask - } + }; } #[inline] @@ -386,12 +400,21 @@ mod test { } } + fn initialized_output(expected: &[T]) -> Vec> { + expected + .iter() + .map(|&value| { + let initial = if value == T::MIN { T::MAX } else { T::MIN }; + MaybeUninit::new(initial) + }) + .collect() + } + fn assume_initialized(output: &[MaybeUninit]) -> Vec { output .iter() .map(|value| { - // SAFETY: callers use this helper only after an unpack method initializes every - // output element. + // SAFETY: every caller initializes each element before the unpack method. unsafe { value.assume_init() } }) .collect() @@ -409,11 +432,11 @@ mod test { .map(|&index| values[index]) .collect::>(); - let mut output = vec![MaybeUninit::uninit(); indices.len()]; + let mut output = initialized_output(&expected); BitPacking::unpack_indices::(&packed, indices, &mut output); assert_eq!(assume_initialized(&output), expected); - let mut output = vec![MaybeUninit::uninit(); indices.len()]; + let mut output = initialized_output(&expected); // SAFETY: `packed` contains exactly one packed FastLanes block. unsafe { BitPacking::unchecked_unpack_indices(WIDTH, &packed, indices, &mut output); @@ -457,11 +480,11 @@ mod test { #[test] fn test_unpack_indices_zero_width_ignores_indices() { let indices = [usize::MAX]; - let mut output = [MaybeUninit::::uninit()]; + let mut output = [MaybeUninit::new(1_u32)]; BitPacking::unpack_indices::<0, 0>(&[], &indices, &mut output); assert_eq!(assume_initialized(&output), [0]); - let mut output = [MaybeUninit::::uninit()]; + let mut output = [MaybeUninit::new(1_u32)]; // SAFETY: the zero-width packed representation contains no elements. unsafe { BitPacking::unchecked_unpack_indices(0, &[], &indices, &mut output) }; assert_eq!(assume_initialized(&output), [0]); @@ -593,7 +616,7 @@ mod test { .iter() .map(|&index| unpacked[index]) .collect::>(); - let mut output = vec![MaybeUninit::uninit(); indices.len()]; + let mut output = initialized_output(&expected); // SAFETY: `packed_input` contains exactly one packed FastLanes block. unsafe {