mbutrovich opened a new issue, #11776: URL: https://github.com/apache/arrow-rs/issues/11776
### Is your feature request related to a problem or challenge? I was inspired while reviewing @neilconway's https://github.com/apache/datafusion/pull/26507, which speeds up repeated counting of set bits in DataFusion's list functions. `BooleanBuffer::count_set_bits` and `Buffer::count_set_bits_offset` are on many hot paths. Every `NullBuffer::new` computes its null count with them ([null.rs](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/buffer/null.rs#L41-L44)). `BooleanArray::true_count` uses them when there are no nulls ([boolean_array.rs](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-array/src/array/boolean_array.rs#L175-L187)), and the filter kernel calls `true_count` for the selected-row count ([filter.rs](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-select/src/filter.rs#L345)). #11206 added APIs so that callers who already know the count can skip computing it. Both end in `UnalignedBitChunk::count_ones` ([bit_chunk_iterator.rs](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-buffer/src/util/bit_chunk_iterator.rs#L169-L172)), which sums `u64::count_ones` over the words. LLVM vectorizes this loop. On aarch64 it widens the byte counts every 16 bytes, which takes four vector instructions per 16 bytes (rustc 1.99, `aarch64-apple-darwin`, release profile): ```asm cnt.16b v5, v5 ; popcount of each byte movi.2d v17, #0 udot.4s v17, v2, v5 ; sum bytes into u32 lanes uadalp.2d v0, v17 ; widen into u64 accumulators ``` ### Describe the solution you'd like Summing four byte popcounts in `u8` lanes before widening needs one widening per 64 bytes. This can be written in portable Rust, without intrinsics or `unsafe`: ```rust /// Counts the number of ones pub fn count_ones(&self) -> usize { let edges = self.prefix.into_iter().chain(self.suffix); edges.map(|x| x.count_ones() as usize).sum::<usize>() + count_ones_bytes(self.chunks.to_byte_slice()) } /// Counts the set bits in `bytes`. /// /// Summing four byte popcounts in `u8` lanes before widening lets LLVM emit /// one widening per 64 bytes. A `u64::count_ones` sum widens every 16 bytes. fn count_ones_bytes(bytes: &[u8]) -> usize { let (chunks, rem) = bytes.as_chunks::<64>(); let mut total = 0; // Each u16 lane gains at most 32 per chunk, so flush before 65535 / 32 chunks. for block in chunks.chunks(2000) { let mut acc = [0u16; 16]; for c in block { for i in 0..16 { let s = c[i].count_ones() as u8 + c[16 + i].count_ones() as u8 + c[32 + i].count_ones() as u8 + c[48 + i].count_ones() as u8; acc[i] += s as u16; } } total += acc.iter().map(|&a| a as usize).sum::<usize>(); } total + rem.iter().map(|b| b.count_ones() as usize).sum::<usize>() } ``` On aarch64 the inner loop becomes nine vector instructions per 64 bytes, against sixteen today: ```asm cnt.16b v2, v2 cnt.16b v3, v3 add.16b v2, v3, v2 cnt.16b v3, v3 cnt.16b v4, v4 add.16b v3, v3, v4 add.16b v2, v2, v3 ; four byte counts summed in u8 lanes uaddw2.8h v1, v1, v2 ; one widening per 64 bytes uaddw.8h v0, v0, v2 ``` These are the `true_count` cases of the existing `boolean_array` bench, on an Apple M5 Max with rustc 1.99.1 and the default bench profile, against 01a9917: | case | main | patched | change | |---|---|---|---| | `true_count(all_true, 64)` | 4.52 ns | 4.70 ns | +4.1% | | `true_count(all_false, 64)` | 4.49 ns | 4.73 ns | +5.4% | | `true_count(mixed_early, 64)` | 4.67 ns | 4.76 ns | +2.0% | | `true_count(all_true, 1024)` | 5.09 ns | 4.75 ns | -6.6% | | `true_count(all_false, 1024)` | 5.06 ns | 4.75 ns | -6.1% | | `true_count(mixed_early, 1024)` | 5.10 ns | 4.75 ns | -6.9% | | `true_count(all_true, 65536)` | 96.98 ns | 70.08 ns | -27.7% | | `true_count(all_false, 65536)` | 97.03 ns | 70.00 ns | -27.9% | | `true_count(mixed_early, 65536)` | 97.15 ns | 70.05 ns | -27.9% | The 64-bit cases have no full 64-byte chunk, so the new loop never runs for them. An earlier run of the same change (using `chunks_exact` in place of `as_chunks`) measured them at -0.9% to -3.8%, so I read those differences as noise. In a standalone benchmark, 16 KB of words went from 189 ns to 137 ns (-27%), and 8 MB went from 98.9 us to 76.5 us (-23%). The change is not a win on every target. With `-C target-cpu=znver4`, today's loop already compiles to `vpopcntq` on 512-bit registers (256 bytes per four instructions). The byte version compiles to `vpopcntb` on 128-bit registers plus widening, which I expect to be slower there. With `x86-64-v3`, both versions use the `vpshufb` nibble lookup. Today's loop does it on 256-bit registers (four `vpshufb` and two `vpsadbw` per 64 bytes). The byte version does it on 128-bit registers (eight `vpshufb` per 64 bytes), which may also be slower. I have not measured either x86 case. So I'd suggest selecting the byte version with `cfg(target_arch = "aarch64")`, which keeps it portable Rust, and enabling it elsewhere only after x86 benchmarks. Does a target-specific selection like that seem acceptable? ### Describe alternatives you've considered - NEON intrinsics (`vcntq_u8`, then `vpadalq_u8` into two `u16` accumulators) were faster still. That version took `true_count(all_true, 65536)` to 61.5 ns, against 70 ns for the portable version. It needs `unsafe` and `std::arch`, like the `pext` path in `bit_util::compress` (#10136). A second accumulator in the portable version did not close the gap. - For x86 with AVX2, the Harley-Seal method from [Faster Population Counts Using AVX2 Instructions](https://arxiv.org/abs/1611.07612v9) counts more than one vector per widening step. The authors measured it at about twice the speed of `popcnt` on large arrays on Haswell. LLVM's current `x86-64-v3` loop uses the nibble lookup from the same paper but not Harley-Seal. I have not measured it. ### Additional context `BooleanArray::true_count` with nulls takes a different path: it zips two `iter_padded` iterators and counts `a & b` per word ([boolean_array.rs](https://github.com/apache/arrow-rs/blob/01a9917658fecd639991e92a5a5810be710b539c/arrow-array/src/array/boolean_array.rs#L178-L183)). In the same bench, `true_count(nulls_all_true, 65536)` takes 417 ns against 97 ns without nulls. This issue does not cover that path. I'm copying people who have worked on bitmap and filter performance or on target-specific code recently, since they may have a view on the approach: @Jefffrey @devanbenz @Rich-T-kid @mightsleep @sdf-jkl @adriangb @neilconway @rluvaton @haohuaijin -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
