feat(bitpacking): Add batched index unpacking

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::<WIDTH, PACKED_LENGTH>(&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;