This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11818-5649fe151886aa91b925e900438bdf5f2cfd520f in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
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();
