This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 003ab783ee perf(arrow-buffer): faster count_set_bits on aarch64 by
summing byte popcounts before widening (#11818)
003ab783ee is described below
commit 003ab783eeaa054acd33de94dce46284fb462498
Author: Matt Butrovich <[email protected]>
AuthorDate: Mon Nov 2 20:03:07 2026 +0000
perf(arrow-buffer): faster count_set_bits on aarch64 by summing byte
popcounts before widening (#11818)
# Which issue does this PR close?
- Closes #11776.
# Rationale for this change
See #11776.
# What changes are included in this PR?
On aarch64, `UnalignedBitChunk::count_ones` sums byte popcounts in `u8`
lanes before widening them, as the issue proposes. The words after the
last 64-byte block are counted with `u64::count_ones`. Other targets
keep the `u64::count_ones` sum for all words. On all targets,
`count_ones` now counts the prefix and suffix words separately from the
aligned middle.
# Are these changes tested?
Adds `test_unaligned_bit_chunk_count_ones` in
`arrow-buffer/src/util/bit_chunk_iterator.rs`. It covers sizes around
the 64-byte block and the 2047-block flush of the `u16` lanes, with
uniform and seeded random bytes.
The `true_count` cases of the `boolean_array` bench measure the change.
This PR adds 448-bit and 960-bit cases, which end in 7 words after the
last 64-byte block. On the bot's default runner (Google Axion, Neoverse
V2), `true_count` without nulls is 1.87x faster at 65536 bits
([results](https://github.com/apache/arrow-rs/pull/11818#issuecomment-6343044150)).
A run with `-C llvm-args=-align-all-functions=6 -C
llvm-args=-align-all-nofallthru-blocks=6` gave the same result, so code
layout doesn't explain it
([results](https://github.com/apache/arrow-rs/pull/11818#issuecomment-6344008687)).
<img width="1600" height="532" alt="count-set-bits-bot"
src="https://github.com/user-attachments/assets/0ad08c1c-1718-47cc-a9f0-3f47e710ead8";
/>
On amd64 (Xeon Platinum 8581C), the 65536-bit cases are unchanged within
noise and the generated x86_64 inner loop is the same before and after
([first](https://github.com/apache/arrow-rs/pull/11818#issuecomment-6343046830),
[second](https://github.com/apache/arrow-rs/pull/11818#issuecomment-6343510237)).
The 64-bit and 1024-bit cases are faster on both architectures in the
bot's build (thin LTO, 16 codegen units), where `main` copies the
`UnalignedBitChunk` into an out-of-line `Chain::fold`. With
`codegen-units = 1`, as DataFusion and Comet build, `main` inlines that
fold, so I don't expect consumers to see most of that gain.
# Are there any significant user-facing changes?
No.
## AI Disclosure
I wrote this PR with help from Claude but understand and support the
core approach, implementation, and test coverage.
---
arrow-array/benches/boolean_array.rs | 2 +-
arrow-buffer/src/util/bit_chunk_iterator.rs | 76 ++++++++++++++++++++++++++++-
2 files changed, 76 insertions(+), 2 deletions(-)
diff --git a/arrow-array/benches/boolean_array.rs
b/arrow-array/benches/boolean_array.rs
index f64440a795..f7d31f75cc 100644
--- a/arrow-array/benches/boolean_array.rs
+++ b/arrow-array/benches/boolean_array.rs
@@ -20,7 +20,7 @@ use criterion::*;
use std::hint;
fn criterion_benchmark(c: &mut Criterion) {
- for len in [64, 1024, 65536] {
+ for len in [64, 448, 960, 1024, 65536] {
// All true (no nulls)
let all_true = BooleanArray::from(vec![true; len]);
c.bench_function(&format!("true_count(all_true, {len})"), |b| {
diff --git a/arrow-buffer/src/util/bit_chunk_iterator.rs
b/arrow-buffer/src/util/bit_chunk_iterator.rs
index 15bdafc589..27bad6fdc1 100644
--- a/arrow-buffer/src/util/bit_chunk_iterator.rs
+++ b/arrow-buffer/src/util/bit_chunk_iterator.rs
@@ -17,6 +17,8 @@
//! Types for iterating over bitmasks in 64-bit chunks
+#[cfg(target_arch = "aarch64")]
+use crate::ToByteSlice;
use crate::util::bit_util::ceil;
use std::fmt::Debug;
@@ -168,10 +170,47 @@ impl<'a> UnalignedBitChunk<'a> {
/// Counts the number of ones
pub fn count_ones(&self) -> usize {
- self.iter().map(|x| x.count_ones() as usize).sum()
+ let edges = self.prefix.into_iter().chain(self.suffix);
+ edges.map(|x| x.count_ones() as usize).sum::<usize>() +
count_ones_words(self.chunks)
}
}
+#[cfg(not(target_arch = "aarch64"))]
+fn count_ones_words(words: &[u64]) -> usize {
+ words.iter().map(|x| x.count_ones() as usize).sum()
+}
+
+/// Summing four byte popcounts in `u8` lanes before widening lets LLVM emit
one
+/// widening per 64-byte block. A `u64::count_ones` sum widens every 16 bytes.
On x86
+/// the `u64` sum measured faster (#11776), so this is only used on aarch64.
+#[cfg(target_arch = "aarch64")]
+fn count_ones_words(words: &[u64]) -> usize {
+ // A block is 8 words (64 bytes), the unit that is widened once.
+ let (blocks, tail) = words.as_chunks::<8>();
+ let mut total = 0;
+ // `lane_counts` counts in u16 lanes so that one widening per block
suffices. Each
+ // lane gains at most 32 (four bytes of eight bits) per block, so flush
the lanes
+ // (add them into `total` and reset them) before they can overflow.
+ const BLOCKS_PER_FLUSH: usize = (u16::MAX / 32) as usize;
+ for flush_group in blocks.chunks(BLOCKS_PER_FLUSH) {
+ let mut lane_counts = [0u16; 16];
+ for block in flush_group {
+ let bytes = block.as_slice().to_byte_slice();
+ // Lane `lane` sums byte `lane` of each of the block's four
16-byte vectors.
+ for lane in 0..16 {
+ let count = bytes[lane].count_ones() as u8
+ + bytes[16 + lane].count_ones() as u8
+ + bytes[32 + lane].count_ones() as u8
+ + bytes[48 + lane].count_ones() as u8;
+ lane_counts[lane] += count as u16;
+ }
+ }
+ total += lane_counts.iter().copied().map(usize::from).sum::<usize>();
+ }
+ // Counting the tail in words avoids widening each byte popcount to u64.
+ total + tail.iter().map(|w| w.count_ones() as usize).sum::<usize>()
+}
+
impl<'a> IntoIterator for &UnalignedBitChunk<'a> {
type Item = u64;
type IntoIter = UnalignedBitChunkIterator<'a>;
@@ -464,6 +503,7 @@ mod tests {
use crate::buffer::Buffer;
use crate::util::bit_chunk_iterator::UnalignedBitChunk;
+ use crate::util::bit_util::get_bit;
#[test]
fn test_iter_aligned() {
@@ -822,6 +862,40 @@ mod tests {
assert_eq!(unaligned.trailing_padding(), 62);
}
+ #[test]
+ fn test_unaligned_bit_chunk_count_ones() {
+ // Sizes in bytes around the 64-byte block and the 2047-block flush in
+ // `count_ones_words`. With all-ones input, this fails if the flush
comes too
+ // late and a u16 lane counter overflows. Random input (`None`) fails
if a
+ // lane counts the wrong bytes, which uniform input hides.
+ let sizes = [
+ 0, 1, 7, 8, 63, 64, 65, 72, 136, 131_008, 131_072, 131_136,
131_207,
+ ];
+ let mut rng = StdRng::seed_from_u64(42);
+ for size in sizes {
+ for byte in [Some(0xFF_u8), Some(0xA5), Some(0x00), None] {
+ let bytes = match byte {
+ Some(byte) => vec![byte; size],
+ None => (0..size).map(|_| rng.random()).collect(),
+ };
+ let buffer = Buffer::from(bytes);
+ let bits = size * 8;
+ for offset in [0, 3.min(bits), 64.min(bits)] {
+ let len = bits - offset;
+ let unaligned = UnalignedBitChunk::new(buffer.as_slice(),
offset, len);
+ let expected = (offset..bits)
+ .filter(|i| get_bit(buffer.as_slice(), *i))
+ .count();
+ assert_eq!(
+ unaligned.count_ones(),
+ expected,
+ "{size} {byte:#x?} {offset}"
+ );
+ }
+ }
+ }
+ }
+
#[test]
fn fuzz_unaligned_bit_chunk_iterator() {
let mut rng = rng();
