Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
77 changes: 19 additions & 58 deletions benches/bitpacking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;

Expand All @@ -148,31 +120,33 @@ 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::<WIDTH, PACKED_LENGTH>(&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);
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);
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);
Expand All @@ -181,26 +155,27 @@ 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);
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) };
unsafe { BitPacking::unchecked_unpack(width, packed, unpacked) };
for (&index, value) in indices.iter().zip(output.iter_mut()) {
value.write(unpacked[index]);
}
Expand All @@ -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) {
Expand Down
45 changes: 34 additions & 11 deletions src/bitpacking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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>();
Expand All @@ -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];
}

Expand All @@ -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]
Expand Down Expand Up @@ -386,12 +400,21 @@ mod test {
}
}

fn initialized_output<T: Integer>(expected: &[T]) -> Vec<MaybeUninit<T>> {
expected
.iter()
.map(|&value| {
let initial = if value == T::MIN { T::MAX } else { T::MIN };
MaybeUninit::new(initial)
})
.collect()
}

fn assume_initialized<T: Copy>(output: &[MaybeUninit<T>]) -> Vec<T> {
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()
Expand All @@ -409,11 +432,11 @@ mod test {
.map(|&index| values[index])
.collect::<Vec<_>>();

let mut output = vec![MaybeUninit::uninit(); indices.len()];
let mut output = initialized_output(&expected);
BitPacking::unpack_indices::<WIDTH, PACKED_LENGTH>(&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);
Expand Down Expand Up @@ -457,11 +480,11 @@ mod test {
#[test]
fn test_unpack_indices_zero_width_ignores_indices() {
let indices = [usize::MAX];
let mut output = [MaybeUninit::<u32>::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::<u32>::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]);
Expand Down Expand Up @@ -593,7 +616,7 @@ mod test {
.iter()
.map(|&index| unpacked[index])
.collect::<Vec<_>>();
let mut output = vec![MaybeUninit::uninit(); indices.len()];
let mut output = initialized_output(&expected);

// SAFETY: `packed_input` contains exactly one packed FastLanes block.
unsafe {
Expand Down
Loading