diff --git a/benches/bitpacking.rs b/benches/bitpacking.rs index 0701b32..e8ef08c 100644 --- a/benches/bitpacking.rs +++ b/benches/bitpacking.rs @@ -1,3 +1,4 @@ +use std::mem::MaybeUninit; use std::mem::size_of; use arrayref::{array_mut_ref, array_ref}; @@ -90,6 +91,145 @@ fn unpack_single_16_from_3(bencher: Bencher) { }); } +const MAX_BENCHMARK_INDICES: usize = 128; + +#[derive(Clone, Copy)] +enum IndexDistribution { + Uniform, + Clustered, + UnorderedWithDuplicates, +} + +fn benchmark_indices( + distribution: IndexDistribution, + 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 + }; + } + } + } + + indices +} + +macro_rules! unpack_indices_benchmarks { + ($module:ident, $type:ty, $width:expr, $distribution:expr) => { + mod $module { + use super::*; + + const WIDTH: usize = $width; + const PACKED_LENGTH: usize = 1024 * WIDTH / <$type>::T; + + fn fixture( + num_indices: usize, + ) -> ([$type; PACKED_LENGTH], [usize; MAX_BENCHMARK_INDICES]) { + let mask = ((1_u128 << WIDTH) - 1) as $type; + let values = + 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); + (packed, indices) + } + + #[divan::bench(args = [1, 8, 32, 128], 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]; + + bencher.bench_local(|| { + let packed = black_box(&packed); + let indices = black_box(&indices[..num_indices]); + 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); + } + black_box(&*output); + }); + } + + #[divan::bench(args = [1, 8, 32, 128], 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]; + + bencher.bench_local(|| { + let packed = black_box(&packed); + let indices = black_box(&indices[..num_indices]); + let output = black_box(&mut output[..num_indices]); + 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) + }); + } + black_box(&*output); + }); + } + + #[divan::bench(args = [1, 8, 32, 128], 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]; + + bencher.bench_local(|| { + let packed = black_box(&packed); + let indices = black_box(&indices[..num_indices]); + 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) }; + for (&index, value) in indices.iter().zip(output.iter_mut()) { + value.write(unpacked[index]); + } + black_box(&*output); + }); + } + } + }; +} + +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 +); + #[divan::bench(sample_count = 10000)] fn throughput_compress(bencher: Bencher) { const WIDTH: usize = 3; diff --git a/src/bitpacking.rs b/src/bitpacking.rs index 26fd861..7d13de4 100644 --- a/src/bitpacking.rs +++ b/src/bitpacking.rs @@ -1,4 +1,5 @@ use const_for::const_for; +use core::mem::MaybeUninit; use core::mem::size_of; use pastey::paste; @@ -54,6 +55,41 @@ pub trait BitPacking: FastLanes { /// /// These lengths are checked only with `debug_assert` (i.e., not checked on release builds). unsafe fn unchecked_unpack_single(width: usize, input: &[Self], index: usize) -> Self; + + /// Unpacks selected elements from a packed array of 1024 `W` bit elements. + /// + /// # Panics + /// + /// Panics if the output length differs from the index length or, when `W` is not zero, an index + /// is not less than 1024. + fn unpack_indices( + packed: &[Self; B], + indices: &[usize], + output: &mut [MaybeUninit], + ); + + /// Unpacks selected elements where `W` is known only at runtime. + /// + /// This method dispatches on `width` once for the complete index batch. + /// + /// # Safety + /// + /// - The input slice must contain exactly `1024 * W / T` elements, where `T` is the unpacked bit + /// width of `Self` and `W` is `width`. + /// - The `width` must be less than or equal to the unpacked bit width of `Self`. + /// + /// These conditions are checked only with `debug_assert`. + /// + /// # Panics + /// + /// Panics if the output length differs from the index length or, when `width` is not zero, an + /// index is not less than 1024. + unsafe fn unchecked_unpack_indices( + width: usize, + input: &[Self], + indices: &[usize], + output: &mut [MaybeUninit], + ); } macro_rules! impl_packing { @@ -117,7 +153,6 @@ 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 @@ -155,6 +190,7 @@ macro_rules! impl_packing { } /// Unpacks a single element at the provided index from a packed array of 1024 `W` bit elements. + #[inline] fn unpack_single(packed: &[Self; B], index: usize) -> Self { const { @@ -163,21 +199,11 @@ macro_rules! impl_packing { } if W == 0 { - // Special case for W=0, we just need to zero the output. return 0 as $T; } - // 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. + // Packing transposes each logical index into a FastLanes lane and row. + // Map the requested index back to those coordinates. assert!(index < 1024, "Index must be less than 1024, got {}", index); let (lane, row): (usize, usize) = { const LANES: [u8; 1024] = lanes_by_index::<$T>(); @@ -186,7 +212,6 @@ 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]; } @@ -197,16 +222,15 @@ macro_rules! impl_packing { let remaining_bits = <$T>::T - lo_shift; let lo = packed[<$T>::LANES * start_word + lane] >> lo_shift; - return if remaining_bits >= W { - // in this case we will mask out all bits of hi word + if remaining_bits >= W { 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] unsafe fn unchecked_unpack_single(width: usize, packed: &[Self], index: usize) -> Self { const T: usize = <$T>::T; @@ -218,13 +242,79 @@ macro_rules! impl_packing { match width { #(W => { const B: usize = 1024 * W / T; - return <$T>::unpack_single::(unsafe { crate::as_array_unchecked(packed) }, index); + return <$T>::unpack_single::( + unsafe { crate::as_array_unchecked(packed) }, + index, + ); },)* // seq_t has exclusive upper bound T => { const W: usize = T; const B: usize = 1024; - return <$T>::unpack_single::(unsafe { crate::as_array_unchecked(packed) }, index); + return <$T>::unpack_single::( + unsafe { crate::as_array_unchecked(packed) }, + index, + ); + }, + _ => unreachable!("Unsupported width: {}", width) + } + })) + } + + fn unpack_indices( + packed: &[Self; B], + indices: &[usize], + output: &mut [MaybeUninit], + ) { + const { + assert!(supported_bit_width(W, 8 * core::mem::size_of::<$T>())); + assert!(B == 1024 * W / Self::T); + } + + assert_eq!(indices.len(), output.len(), "Output length must equal index length"); + if W == 0 { + for value in output { + value.write(0 as Self); + } + return; + } + for (&index, value) in indices.iter().zip(output) { + value.write(<$T>::unpack_single::(packed, index)); + } + } + + unsafe fn unchecked_unpack_indices( + width: usize, + packed: &[Self], + indices: &[usize], + output: &mut [MaybeUninit], + ) { + const T: usize = <$T>::T; + debug_assert!(width <= T, "Width must be less than or equal to {}", T); + #[cfg(debug_assertions)] + { + let packed_len = 128 * width / size_of::(); + debug_assert_eq!(packed.len(), packed_len, "Input buffer must be of size {}", packed_len); + } + + paste!(seq_t!(W in $T { + match width { + #(W => { + const B: usize = 1024 * W / T; + return <$T>::unpack_indices::( + unsafe { crate::as_array_unchecked(packed) }, + indices, + output, + ); + },)* + T => { + const W: usize = T; + const B: usize = 1024; + return <$T>::unpack_indices::( + unsafe { crate::as_array_unchecked(packed) }, + indices, + output, + ); }, _ => unreachable!("Unsupported width: {}", width) } @@ -296,6 +386,87 @@ mod test { } } + fn assume_initialized(output: &[MaybeUninit]) -> Vec { + output + .iter() + .map(|value| { + // SAFETY: callers use this helper only after an unpack method initializes every + // output element. + unsafe { value.assume_init() } + }) + .collect() + } + + fn assert_u32_unpack_indices(indices: &[usize]) { + const WIDTH: usize = 13; + const PACKED_LENGTH: usize = 1024 * WIDTH / u32::T; + + let values = array::from_fn(|index| ((index as u32).wrapping_mul(17)) & 0x1fff); + let mut packed = [0; PACKED_LENGTH]; + BitPacking::pack::(&values, &mut packed); + let expected = indices + .iter() + .map(|&index| values[index]) + .collect::>(); + + let mut output = vec![MaybeUninit::uninit(); indices.len()]; + BitPacking::unpack_indices::(&packed, indices, &mut output); + assert_eq!(assume_initialized(&output), expected); + + let mut output = vec![MaybeUninit::uninit(); indices.len()]; + // SAFETY: `packed` contains exactly one packed FastLanes block. + unsafe { + BitPacking::unchecked_unpack_indices(WIDTH, &packed, indices, &mut output); + } + assert_eq!(assume_initialized(&output), expected); + } + + #[test] + fn test_unpack_indices_empty() { + assert_u32_unpack_indices(&[]); + } + + #[test] + fn test_unpack_indices_preserves_order_and_duplicates() { + assert_u32_unpack_indices(&[1023, 0, 17, 17, 511, 17, 511]); + } + + #[test] + fn test_unpack_indices_full_block() { + assert_u32_unpack_indices(&(0..1024).collect::>()); + } + + #[test] + #[should_panic(expected = "Output length must equal index length")] + fn test_unpack_indices_rejects_output_length_mismatch() { + let packed = [0_u32; 32]; + let mut output = [MaybeUninit::uninit(); 1]; + BitPacking::unpack_indices::<1, 32>(&packed, &[], &mut output); + } + + #[test] + #[should_panic(expected = "Output length must equal index length")] + fn test_unchecked_unpack_indices_rejects_output_length_mismatch() { + let packed = [0_u32; 32]; + let mut output = [MaybeUninit::uninit(); 1]; + // SAFETY: the packed input has the required length. The mismatched output is a documented + // panic condition, not a safety requirement. + unsafe { BitPacking::unchecked_unpack_indices(1, &packed, &[], &mut output) }; + } + + #[test] + fn test_unpack_indices_zero_width_ignores_indices() { + let indices = [usize::MAX]; + let mut output = [MaybeUninit::::uninit()]; + BitPacking::unpack_indices::<0, 0>(&[], &indices, &mut output); + assert_eq!(assume_initialized(&output), [0]); + + let mut output = [MaybeUninit::::uninit()]; + // SAFETY: the zero-width packed representation contains no elements. + unsafe { BitPacking::unchecked_unpack_indices(0, &[], &indices, &mut output) }; + assert_eq!(assume_initialized(&output), [0]); + } + fn assert_bitpack_roundtrip(tc: &TestCase) where T: BitPacking + Debug + Integer + 'static, @@ -393,6 +564,51 @@ mod test { } } + fn assert_bitpack_unpack_indices_matches_bulk(tc: &TestCase) + where + T: BitPacking + Debug + Integer + 'static, + { + let packed_source = tc.draw( + gs::vecs(gs::integers::()) + .min_size(BUFFER_SIZE) + .max_size(BUFFER_SIZE), + ); + let indices = tc.draw( + gs::vecs( + gs::integers::() + .min_value(0) + .max_value(BUFFER_SIZE - 1), + ) + .min_size(0) + .max_size(128), + ); + + for width in 0..=T::T { + let packed_length = (BUFFER_SIZE * width) / T::T; + let packed_input = &packed_source[..packed_length]; + let mut unpacked = vec![T::zero(); BUFFER_SIZE]; + // SAFETY: both buffers have the required lengths for `width`. + unsafe { T::unchecked_unpack(width, packed_input, &mut unpacked) }; + let expected = indices + .iter() + .map(|&index| unpacked[index]) + .collect::>(); + let mut output = vec![MaybeUninit::uninit(); indices.len()]; + + // SAFETY: `packed_input` contains exactly one packed FastLanes block. + unsafe { + T::unchecked_unpack_indices(width, packed_input, &indices, &mut output); + } + + assert_eq!( + assume_initialized(&output), + expected, + "indexed unpack failed for type={} width={width} indices={indices:?}", + core::any::type_name::(), + ); + } + } + fn reference_pack(width: usize, input: &[T]) -> Vec where T: BitPacking, @@ -454,5 +670,6 @@ mod test { bitpack_property_tests!(bitpack_roundtrip, 10 for u8, u16, u32, u64); bitpack_property_tests!(bitpack_repack_roundtrip, 10 for u8, u16, u32, u64); bitpack_property_tests!(bitpack_unpack_single_matches_bulk, 10 for u8, u16, u32, u64); + bitpack_property_tests!(bitpack_unpack_indices_matches_bulk, 10 for u8, u16, u32, u64); bitpack_property_tests!(bitpack_matches_reference, 10 for u8, u16, u32, u64); }