faster true count

robert3005 · robert3005 · commit b6712cf8d173 · 2026-03-12T14:23:14.000Z
diff --git a/vortex-buffer/benches/vortex_bitbuffer.rs b/vortex-buffer/benches/vortex_bitbuffer.rs
@@ -25,6 +25,24 @@ impl FromIterator<bool> for Arrow<BooleanBuffer> {
 }
 
 const INPUT_SIZE: &[usize] = &[128, 1024, 2048, 16_384, 65_536];
+const TRUE_COUNT_INPUT_SIZE: &[usize] = &[128, 1024, 16_384, 65_536, 1_048_576];
+const TRUE_COUNT_SLICE_CASES: &[(usize, usize)] = &[
+    (128, 1),
+    (128, 7),
+    (1_024, 1),
+    (1_024, 7),
+    (16_384, 1),
+    (16_384, 7),
+    (65_536, 1),
+    (65_536, 7),
+    (1_048_576, 1),
+    (1_048_576, 7),
+];
+
+#[inline]
+fn true_count_pattern(i: usize) -> bool {
+    (i.is_multiple_of(3)) ^ (i.is_multiple_of(11))
+}
 
 #[cfg(not(codspeed))]
 #[divan::bench(args = INPUT_SIZE)]
@@ -158,22 +176,46 @@ fn slice_arrow_buffer(bencher: Bencher, length: usize) {
         });
 }
 
-#[divan::bench(args = INPUT_SIZE)]
+#[divan::bench(args = TRUE_COUNT_INPUT_SIZE)]
 fn true_count_vortex_buffer(bencher: Bencher, length: usize) {
-    let buffer = BitBuffer::from_iter((0..length).map(|i| i % 2 == 0));
+    let buffer = BitBuffer::from_iter((0..length).map(true_count_pattern));
     bencher
         .with_inputs(|| &buffer)
         .bench_refs(|buffer| buffer.true_count())
 }
 
-#[divan::bench(args = INPUT_SIZE)]
+#[divan::bench(args = TRUE_COUNT_INPUT_SIZE)]
 fn true_count_arrow_buffer(bencher: Bencher, length: usize) {
-    let buffer = Arrow(BooleanBuffer::from_iter((0..length).map(|i| i % 2 == 0)));
+    let buffer = Arrow(BooleanBuffer::from_iter(
+        (0..length).map(true_count_pattern),
+    ));
     bencher
         .with_inputs(|| &buffer)
         .bench_refs(|buffer| buffer.0.count_set_bits());
 }
 
+#[divan::bench(args = TRUE_COUNT_SLICE_CASES)]
+fn true_count_vortex_buffer_sliced(bencher: Bencher, (length, offset): (usize, usize)) {
+    let buffer = BitBuffer::from_iter((0..length + offset).map(true_count_pattern));
+    let sliced = buffer.slice(offset..offset + length);
+
+    bencher
+        .with_inputs(|| &sliced)
+        .bench_refs(|buffer| buffer.true_count())
+}
+
+#[divan::bench(args = TRUE_COUNT_SLICE_CASES)]
+fn true_count_arrow_buffer_sliced(bencher: Bencher, (length, offset): (usize, usize)) {
+    let buffer = Arrow(BooleanBuffer::from_iter(
+        (0..length + offset).map(true_count_pattern),
+    ));
+    let sliced = Arrow(buffer.0.slice(offset, length));
+
+    bencher
+        .with_inputs(|| &sliced)
+        .bench_refs(|buffer| buffer.0.count_set_bits());
+}
+
 #[divan::bench(args = INPUT_SIZE)]
 fn bitwise_and_vortex_buffer(bencher: Bencher, length: usize) {
     let a = BitBuffer::from_iter((0..length).map(|i| i % 2 == 0));
diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs
@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+use std::arch::is_x86_feature_detected;
 use std::fmt::Display;
 use std::fmt::Formatter;
 use std::fmt::Result as FmtResult;
@@ -316,7 +318,7 @@ impl BitBuffer {
 
     /// Get the number of set bits in the buffer.
     pub fn true_count(&self) -> usize {
-        self.unaligned_chunks().count_ones()
+        true_count_impl(self.buffer.as_slice(), self.offset, self.len)
     }
 
     /// Get the number of unset bits in the buffer.
@@ -356,6 +358,198 @@ impl BitBuffer {
     }
 }
 
+#[inline]
+fn true_count_impl(bytes: &[u8], offset: usize, len: usize) -> usize {
+    let Some((head, middle, tail)) = byte_aligned_region(bytes, offset, len) else {
+        return 0;
+    };
+
+    let mut count = head.map_or(0, |v| v.count_ones() as usize);
+
+    if !middle.is_empty() {
+        count += count_aligned_bytes(middle);
+    }
+
+    count + tail.map_or(0, |v| v.count_ones() as usize)
+}
+
+#[inline]
+fn byte_aligned_region(
+    bytes: &[u8],
+    offset: usize,
+    len: usize,
+) -> Option<(Option<u8>, &[u8], Option<u8>)> {
+    if len == 0 {
+        return None;
+    }
+
+    let start_byte = offset / 8;
+    let start_bit = offset % 8;
+    let end_bit = offset + len;
+    let end_byte = end_bit / 8;
+    let head = (start_bit != 0).then(|| {
+        let head_len = (8 - start_bit).min(len);
+        mask_partial_byte(bytes[start_byte], start_bit, head_len)
+    });
+
+    let middle_start = start_byte + usize::from(start_bit != 0);
+    let middle_end = end_byte;
+    let middle = if middle_start < middle_end {
+        &bytes[middle_start..middle_end]
+    } else {
+        &[]
+    };
+
+    let consumed = if start_bit != 0 {
+        (8 - start_bit).min(len)
+    } else {
+        0
+    } + middle.len() * 8;
+    let tail_len = len - consumed;
+    let tail = (tail_len != 0).then(|| mask_partial_byte(bytes[middle_end], 0, tail_len));
+
+    Some((head, middle, tail))
+}
+
+#[inline]
+fn mask_partial_byte(byte: u8, bit_offset: usize, bit_len: usize) -> u8 {
+    debug_assert!(bit_offset < 8);
+    debug_assert!(bit_len <= 8 - bit_offset);
+
+    let shifted = byte >> bit_offset;
+    let mask = if bit_len == 8 {
+        u8::MAX
+    } else {
+        (1u8 << bit_len) - 1
+    };
+
+    shifted & mask
+}
+
+#[inline]
+fn count_aligned_bytes(bytes: &[u8]) -> usize {
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    {
+        if bytes.len() >= 64
+            && is_x86_feature_detected!("avx512f")
+            && is_x86_feature_detected!("avx512vpopcntdq")
+        {
+            // SAFETY: Runtime detection guarantees the required target features.
+            return unsafe { count_aligned_bytes_avx512(bytes) };
+        }
+
+        if bytes.len() >= 32 && is_x86_feature_detected!("avx2") {
+            // SAFETY: Runtime detection guarantees the required target features.
+            return unsafe { count_aligned_bytes_avx2(bytes) };
+        }
+    }
+
+    count_aligned_bytes_scalar(bytes)
+}
+
+#[inline]
+fn count_aligned_bytes_scalar(bytes: &[u8]) -> usize {
+    let (words, tail) = bytes.as_chunks::<8>();
+    let mut count = words
+        .iter()
+        .map(|word| u64::from_le_bytes(*word).count_ones() as usize)
+        .sum::<usize>();
+
+    count += tail
+        .iter()
+        .map(|byte| byte.count_ones() as usize)
+        .sum::<usize>();
+
+    count
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+unsafe fn count_aligned_bytes_avx2(bytes: &[u8]) -> usize {
+    use std::arch::x86_64::__m256i;
+    use std::arch::x86_64::_mm256_add_epi8;
+    use std::arch::x86_64::_mm256_add_epi64;
+    use std::arch::x86_64::_mm256_and_si256;
+    use std::arch::x86_64::_mm256_loadu_si256;
+    use std::arch::x86_64::_mm256_sad_epu8;
+    use std::arch::x86_64::_mm256_set1_epi8;
+    use std::arch::x86_64::_mm256_setr_epi8;
+    use std::arch::x86_64::_mm256_setzero_si256;
+    use std::arch::x86_64::_mm256_shuffle_epi8;
+    use std::arch::x86_64::_mm256_srli_epi16;
+    use std::arch::x86_64::_mm256_storeu_si256;
+
+    #[inline]
+    unsafe fn byte_popcount(chunk: __m256i, mask: __m256i, lookup: __m256i) -> __m256i {
+        let lo = unsafe { _mm256_and_si256(chunk, mask) };
+        let hi = unsafe { _mm256_and_si256(_mm256_srli_epi16(chunk, 4), mask) };
+        unsafe {
+            _mm256_add_epi8(
+                _mm256_shuffle_epi8(lookup, lo),
+                _mm256_shuffle_epi8(lookup, hi),
+            )
+        }
+    }
+
+    let lookup = _mm256_setr_epi8(
+        0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3,
+        3, 4,
+    );
+    let mask = _mm256_set1_epi8(0x0f);
+    let zero = _mm256_setzero_si256();
+    let mut accum = _mm256_setzero_si256();
+    let mut index = 0;
+
+    while index + 128 <= bytes.len() {
+        for lane in 0..4 {
+            let ptr = unsafe { bytes.as_ptr().add(index + lane * 32) }.cast::<__m256i>();
+            let chunk = unsafe { _mm256_loadu_si256(ptr) };
+            let counts = unsafe { byte_popcount(chunk, mask, lookup) };
+            accum = _mm256_add_epi64(accum, _mm256_sad_epu8(counts, zero));
+        }
+        index += 128;
+    }
+
+    while index + 32 <= bytes.len() {
+        let ptr = unsafe { bytes.as_ptr().add(index) }.cast::<__m256i>();
+        let chunk = unsafe { _mm256_loadu_si256(ptr) };
+        let counts = unsafe { byte_popcount(chunk, mask, lookup) };
+        accum = _mm256_add_epi64(accum, _mm256_sad_epu8(counts, zero));
+        index += 32;
+    }
+
+    let mut lanes = [0u64; 4];
+    unsafe { _mm256_storeu_si256(lanes.as_mut_ptr().cast::<__m256i>(), accum) };
+
+    lanes.iter().sum::<u64>() as usize + count_aligned_bytes_scalar(&bytes[index..])
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx512f,avx512vpopcntdq")]
+unsafe fn count_aligned_bytes_avx512(bytes: &[u8]) -> usize {
+    use std::arch::x86_64::__m512i;
+    use std::arch::x86_64::_mm512_add_epi64;
+    use std::arch::x86_64::_mm512_loadu_si512;
+    use std::arch::x86_64::_mm512_popcnt_epi64;
+    use std::arch::x86_64::_mm512_setzero_si512;
+    use std::arch::x86_64::_mm512_storeu_si512;
+
+    let mut accum = _mm512_setzero_si512();
+    let mut index = 0;
+
+    while index + 64 <= bytes.len() {
+        let ptr = unsafe { bytes.as_ptr().add(index) }.cast::<__m512i>();
+        let chunk = unsafe { _mm512_loadu_si512(ptr) };
+        accum = _mm512_add_epi64(accum, _mm512_popcnt_epi64(chunk));
+        index += 64;
+    }
+
+    let mut lanes = [0u64; 8];
+    unsafe { _mm512_storeu_si512(lanes.as_mut_ptr().cast::<__m512i>(), accum) };
+
+    lanes.iter().sum::<u64>() as usize + count_aligned_bytes_scalar(&bytes[index..])
+}
+
 // Conversions
 
 impl BitBuffer {
@@ -784,4 +978,33 @@ mod tests {
             assert_eq!(mapped.value(i), expected, "Mismatch at index {}", i);
         }
     }
+
+    #[rstest]
+    fn test_true_count_matches_iteration_for_slices(
+        #[values(
+            0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
+            23, 24, 25, 26, 27, 28, 29, 30
+        )]
+        offset: usize,
+        #[values(
+            0usize, 1, 2, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 255, 256, 257, 513
+        )]
+        slice_len: usize,
+    ) {
+        let len = 513;
+        let buf = BitBuffer::collect_bool(len + 31, |i| (i % 3 == 0) ^ (i % 11 == 0));
+
+        if offset + slice_len > buf.len() {
+            return;
+        }
+
+        let sliced = buf.slice(offset..offset + slice_len);
+        let expected = sliced.iter().filter(|bit| *bit).count();
+
+        assert_eq!(
+            sliced.true_count(),
+            expected,
+            "offset={offset} len={slice_len}"
+        );
+    }
 }
