From e53eab23a05c9934ada6d981f4ca08ced76df870 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 4 Jun 2026 17:20:30 +0000 Subject: [PATCH] perf(mask): SIMD popcount for valid_counts_for_indices `Mask::valid_counts_for_indices` walked the validity bit buffer one bit at a time via `BitIterator::next`, accumulating a running valid count. The pco/zstd slice decoders call it with `[slice_start, slice_stop]`, so the cost was a bit-by-bit scan of the entire prefix up to `slice_stop`. Replace the per-bit walk with an incremental SIMD popcount over each gap between consecutive indices. Total scanned bits are unchanged, but the work now goes through the vectorized `count_ones` path. Add `BitBuffer::count_range(start, end)`, which counts set bits in a range directly over the backing buffer without cloning a sliced `BitBuffer`, so the many-indices case has no per-gap allocation overhead. Benchmark (vortex-mask/benches/valid_counts.rs, median): slice_bounds 16384 8.7us -> 51ns (~170x) slice_bounds 262144 138us -> 310ns (~445x) slice_bounds 1M 554us -> 1.86us (~300x) many_indices 16384 10.1us -> 1.68us (~6x) many_indices 1M 613us -> 3.33us (~185x) Signed-off-by: Joe Isaacs Signed-off-by: Robert Kruszewski --- AGENTS.md | 44 ++++++++++ encodings/runend/src/compute/filter.rs | 22 +++-- vortex-buffer/src/bit/buf.rs | 110 +++++++++++++++++++++++++ vortex-duckdb/src/exporter/bool.rs | 17 ++-- vortex-mask/benches/mask_iteration.rs | 12 +++ vortex-mask/src/lib.rs | 24 +++--- 6 files changed, 204 insertions(+), 25 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index e5c3d0cc13b..30ddeeb9515 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -151,6 +151,47 @@ Notes: self-explanatory code. - Keep public APIs small and consistent with neighboring crates. +## Performance: avoid hidden-cost accessors in hot loops + +The most common performance trap in this codebase is calling a *per-element accessor that +hides non-trivial work* inside an `O(n)` loop, instead of doing the work once over the whole +chunk. The call site looks like a cheap getter, but each call re-pays a cost that is constant +(or amortizable) across the loop, making the loop `O(n · k)`. + +Watch for these accessors used inside `for i in 0..n { ... }`: + +| Per-element accessor (in a loop) | Hidden cost | Bulk replacement | +| --- | --- | --- | +| `Validity::is_valid(i)` / `is_null(i)` | for array-backed validity, **allocates an `ExecutionCtx` and runs a scalar lookup per call** | `validity.execute_mask(len, ctx)?` once, then `Mask::value(i)` | +| `array.scalar_at(i)` / `array.execute_scalar(i, ctx)` | per-element execution through the compute stack | canonicalize once (`execute::` / `as_slice`) then index | +| `BitBuffer::value(i)` / `Mask::value(i)` accumulated into a count | recomputes the byte address each call; defeats popcount | `true_count()`, `BitBuffer::count_range(start, end)`, `set_indices()` | +| `BitIterator::next()` to accumulate a running rank/prefix count | bit-at-a-time | `count_range` over each gap (SIMD popcount) | +| re-deriving a value inside the loop (e.g. `self.validity()?` each iteration) | re-runs the derivation `n` times | hoist it above the loop | + +Decide per site — bulk is not always the answer: + +- **Sequential / contiguous access** over an accessor that hides amortizable work → hoist and + go bulk (materialize once, then index or iterate the chunk). +- **Gather over arbitrary indices** → you cannot amortize a per-element *decode*, but you can + still materialize the backing buffer once (e.g. `execute_mask`) and then do cheap `O(1)` + random reads, avoiding per-call context/allocation. +- **The accessor is already genuinely `O(1)`** (e.g. reading an already-materialized `Mask`/ + slice, or a native bitmap) → leave it; bulk would not help. + +Even after materializing into a `Mask`/`BitBuffer`, **do not loop with `value(i)` per +element** to act on each set bit — the per-element branch dominates. Iterate a `u64` word +at a time with all-set / all-unset fast paths: use [`BitBuffer::for_each_set_index`]. It +beats `for i in 0..len { if buf.value(i) {..} }` by 2-45x (more at low density) and beats +collecting `set_indices()` by ~2x at mid/high density, while self-adapting from sparse to +dense — see `vortex-mask/benches/mask_iteration.rs`. Reach for the cached `indices()` / +`slices()` representations when you need them more than once; otherwise `for_each_set_index` +needs no materialization. + +When you touch such a loop, back the change with a benchmark (see +`vortex-array/benches/validity_is_valid.rs` for the `is_valid` case, +`vortex-mask/benches/valid_counts.rs` for the popcount case, and +`vortex-mask/benches/mask_iteration.rs` for the per-element-vs-word-vs-sparse comparison). + ## Tests - Strongly consider `rstest` cases when parameterizing repetitive test logic. @@ -179,6 +220,9 @@ Check new and modified lines against this list before finishing: - Updating expected test output to match buggy behavior without independently verifying the intended semantics. - Silently reducing the scope of an approved plan when implementation is harder than expected. +- Calling a hidden-cost per-element accessor (`Validity::is_valid`, `scalar_at`, `BitBuffer:: + value` accumulation) inside a hot loop instead of materializing once — see + "Performance: avoid hidden-cost accessors in hot loops". ## Summaries diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index c460dd6552d..6d8d4aab6ea 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/encodings/runend/src/compute/filter.rs @@ -5,6 +5,7 @@ use std::cmp::min; use std::ops::AddAssign; use num_traits::AsPrimitive; +use num_traits::NumCast; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -87,16 +88,21 @@ pub fn filter_run_end_primitive + AsPrim let mut count = R::zero(); let new_mask: Mask = BitBuffer::collect_bool(run_ends.len(), |i| { - let mut keep = false; let end = min(run_ends[i].as_() - offset, length); - // Safety: predicate must be the same length as the array the ends have been taken from - for pred in (start..end).map(|i| unsafe { - mask.value_unchecked(i.try_into().vortex_expect("index must fit in usize")) - }) { - count += >::from(pred); - keep |= pred - } + // SIMD popcount of the predicate bits in this run. The range matches the + // bit-by-bit `value_unchecked` read it replaces, so `end <= mask.len()`. + let start_usize = start + .try_into() + .vortex_expect("run start index must fit in usize"); + let end_usize = end + .try_into() + .vortex_expect("run end index must fit in usize"); + let run_trues = mask.count_range(start_usize, end_usize); + count += ::from(run_trues) + .vortex_expect("run popcount must fit in run-end native type"); + let keep = run_trues > 0; + // this is to avoid branching new_run_ends[j] = count; j += keep as usize; diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index a8bd9e1929c..d229329fcd2 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -379,6 +379,20 @@ impl BitBuffer { count_ones(self.buffer.as_slice(), self.offset, self.len) } + /// Get the number of set bits in the bit range `[start, end)`. + /// + /// Unlike `self.slice(start..end).true_count()`, this counts directly over the + /// existing backing buffer without allocating or cloning a new [`BitBuffer`], + /// making it cheap to call repeatedly over many small ranges. + /// + /// Panics if `start > end` or `end > len`. + #[inline] + pub fn count_range(&self, start: usize, end: usize) -> usize { + assert!(start <= end, "start {start} exceeds end {end}"); + assert!(end <= self.len, "end {end} exceeds len {}", self.len); + count_ones(self.buffer.as_slice(), self.offset + start, end - start) + } + /// Returns the position of the `nth` set bit (0-indexed). /// /// This is the "select" operation on a bitmap: given a rank `nth`, find @@ -410,6 +424,33 @@ impl BitBuffer { BitSliceIterator::new(self.buffer.as_slice(), self.offset, self.len) } + /// Invoke `f(index)` for every set bit, in ascending order, processing a `u64` + /// word at a time. + /// + /// This is the fast way to "do something for each set bit": it skips all-zero + /// words, fast-paths all-one words, and walks the remaining bits with + /// `trailing_zeros`. Prefer it over `for i in 0..len { if buf.value(i) { f(i) } }` + /// (which pays a branch per element) and over collecting [`Self::set_indices`] + /// (whose per-`next` iterator state does not inline as well). + #[inline] + pub fn for_each_set_index(&self, mut f: F) { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k); + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize); + w &= w - 1; + } + } + base += 64; + } + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -857,6 +898,75 @@ mod tests { } } + #[rstest] + #[case(0, 0)] + #[case(0, 64)] + #[case(5, 70)] + #[case(64, 130)] + #[case(0, 200)] + fn test_count_range(#[case] start: usize, #[case] end: usize) { + let len = 200; + let buf = BitBuffer::collect_bool(len, |i| i % 3 == 0); + let expected = (start..end).filter(|i| i % 3 == 0).count(); + assert_eq!(buf.count_range(start, end), expected); + // Must agree with slicing then counting. + assert_eq!( + buf.count_range(start, end), + buf.slice(start..end).true_count() + ); + } + + #[rstest] + #[case(3)] + #[case(7)] + fn test_count_range_with_offset(#[case] offset: usize) { + let len = 150; + let buf = BitBuffer::collect_bool(offset + len, |i| i % 2 == 0); + let view = BitBuffer::new_with_offset(buf.inner().clone(), len, offset); + for (start, end) in [(0, len), (10, 100), (1, 2), (63, 129)] { + let expected = (offset + start..offset + end) + .filter(|i| i % 2 == 0) + .count(); + assert_eq!(view.count_range(start, end), expected, "[{start}, {end})"); + } + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(63)] + #[case(64)] + #[case(65)] + #[case(200)] + #[case(1000)] + fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); + let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); + buf.for_each_set_index(|i| got.push(i)); + assert_eq!(got, expected); + } + + #[rstest] + #[case(3, 200)] + #[case(7, 130)] + fn test_for_each_set_index_with_offset(#[case] offset: usize, #[case] len: usize) { + let base = BitBuffer::collect_bool(offset + len, |i| i % 3 == 0); + let view = BitBuffer::new_with_offset(base.inner().clone(), len, offset); + let expected: Vec = view.set_indices().collect(); + let mut got = Vec::new(); + view.for_each_set_index(|i| got.push(i)); + assert_eq!(got, expected); + } + + #[test] + fn test_for_each_set_index_all_set() { + let buf = BitBuffer::new_set(130); + let mut got = Vec::new(); + buf.for_each_set_index(|i| got.push(i)); + assert_eq!(got, (0..130).collect::>()); + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value diff --git a/vortex-duckdb/src/exporter/bool.rs b/vortex-duckdb/src/exporter/bool.rs index 4d0a2b29883..fc1dd3e10c3 100644 --- a/vortex-duckdb/src/exporter/bool.rs +++ b/vortex-duckdb/src/exporter/bool.rs @@ -46,13 +46,16 @@ impl ColumnExporter for BoolExporter { ) -> VortexResult<()> { // DuckDB uses byte bools, not bit bools. // maybe we can convert into these from a compressed array sometimes?. - unsafe { vector.as_slice_mut(len) }.copy_from_slice( - &self - .bit_buffer - .slice(offset..(offset + len)) - .iter() - .collect::>(), - ); + // Unpack the bits directly into the destination byte slice, avoiding the + // throwaway `Vec` allocation and the extra copy pass. The sliced + // iterator already accounts for the buffer's bit offset. + let dst = unsafe { vector.as_slice_mut(len) }; + for (slot, bit) in dst + .iter_mut() + .zip(self.bit_buffer.slice(offset..(offset + len)).iter()) + { + *slot = bit; + } Ok(()) } diff --git a/vortex-mask/benches/mask_iteration.rs b/vortex-mask/benches/mask_iteration.rs index 50f5c321303..84ca37991ef 100644 --- a/vortex-mask/benches/mask_iteration.rs +++ b/vortex-mask/benches/mask_iteration.rs @@ -110,6 +110,18 @@ fn word_trailing_zeros(bencher: Bencher, (len, density): (usize, f64)) { }); } +#[divan::bench(args = ARGS)] +fn for_each_set_index(bencher: Bencher, (len, density): (usize, f64)) { + let (buf, values) = make(len, density); + bencher + .with_inputs(|| (&buf, &values)) + .bench_refs(|(buf, values)| { + let mut acc = 0u64; + buf.for_each_set_index(|i| acc = acc.wrapping_add(values[i])); + acc + }); +} + #[divan::bench(args = ARGS)] fn set_slices(bencher: Bencher, (len, density): (usize, f64)) { let (buf, values) = make(len, density); diff --git a/vortex-mask/src/lib.rs b/vortex-mask/src/lib.rs index 401077f0b34..1f41c367947 100644 --- a/vortex-mask/src/lib.rs +++ b/vortex-mask/src/lib.rs @@ -628,23 +628,24 @@ impl Mask { /// Given monotonically increasing `indices` in [0, n_rows], returns the /// count of valid elements up to each index. /// - /// This is O(n_rows). + /// This is O(n_rows), but the per-gap counts are computed with a SIMD + /// popcount over the underlying bit buffer rather than walking bit-by-bit. pub fn valid_counts_for_indices(&self, indices: &[usize]) -> Vec { match self { Self::AllTrue(_) => indices.to_vec(), Self::AllFalse(_) => vec![0; indices.len()], Self::Values(values) => { - let mut bool_iter = values.bit_buffer().iter(); + let buffer = values.bit_buffer(); let mut valid_counts = Vec::with_capacity(indices.len()); let mut valid_count = 0; - let mut idx = 0; + let mut prev = 0; for &next_idx in indices { - while idx < next_idx { - idx += 1; - valid_count += bool_iter - .next() - .unwrap_or_else(|| vortex_panic!("Row indices exceed array length")) - as usize; + assert!(next_idx <= buffer.len(), "Row indices exceed array length"); + // `indices` is monotonically increasing, so each gap is counted once; + // the total work across all gaps scans the prefix `[0, last_idx)` once. + if next_idx > prev { + valid_count += buffer.count_range(prev, next_idx); + prev = next_idx; } valid_counts.push(valid_count); } @@ -779,7 +780,10 @@ impl MaskValues { } let mut indices = Vec::with_capacity(self.true_count); - indices.extend(self.buffer.set_indices()); + // Word-at-a-time set-bit walk; faster than collecting `set_indices()`, + // whose per-`next` iterator state inlines less well (see + // `vortex-mask/benches/mask_iteration.rs`). + self.buffer.for_each_set_index(|i| indices.push(i)); debug_assert!(indices.is_sorted()); assert_eq!(indices.len(), self.true_count); indices