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-11322-d49f7daf3cdb6c42ca5b0c6262e44cbc39e215e4 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 2f33b89afe0e0a361554d0db4daafd5018ab493d Author: mightsleep <[email protected]> AuthorDate: Tue Oct 13 21:02:05 2026 +0000 perf(arrow-buffer): faster portable fallback for bit_util::compress (#11322) # Which issue does this PR close? - Closes #11213. # Rationale for this change Without `pext` (every aarch64 build, and x86-64 builds without BMI2, the default target) `compress` walks the kept bits one at a time. With masks from independent rows the number of kept bits changes from word to word, so the loop exit mispredicts on almost every word, and a dense word takes up to 64 steps. # What changes are included in this PR? The fallback dispatches on the number of kept bits `k`: - `k <= 2`: the lowest two kept bits, no loop - `k <= 16`: the existing loop, still the cheapest when its branches are predictable (clustered or periodic rows) - `k >= 62`: the word with its at most two dropped bits removed, no loop - otherwise `compress_bytes`: constant time, no table. The parallel suffix network of HD 7-4 inside each byte (three rounds on all eight bytes at once), then the bytes joined at offsets from one multiply by `0x0101..01`. Builds with BMI2 enabled keep `pext` and are unchanged. # Are these changes tested? `test_compress_portable` checks every path against a reference: every mask with at most two kept or two dropped bits, random masks at every popcount. The existing test only draws uniform masks (about 32 kept bits), which reach ne path of four. Filter tests pass on x86-64, x86-64-v2 and with BMI2. ## Measurements main / this PR, above 1 is faster. GitHub runners: Neoverse-N2 (`ubuntu-24.04-arm`) and AMD EPYC 7763 (`ubuntu-latest`), both default target, plus x86-64-v2 (which has POPCNT). Each variant built once, three interleaved rounds. Between rounds a row moves by 0.6 % (median), at most 5.5 %. `filter_bits` batches from #11271 (512 batches of 8K rows): | mask | Neoverse-N2 | EPYC 7763 | EPYC 7763, v2 | |---|---|---|---| | random, kept 1/1024 | 0.99 | 1.01 | 0.95 | | random, kept 1/256 | 1.03 | 0.96 | 1.03 | | random, kept 1/64 | 1.21 | 1.20 | 1.27 | | random, kept 1/16 | 0.97 | 0.91 | 0.97 | | random, kept 1/4 | 1.21 | 1.12 | 1.19 | | random, kept 1/2 | 3.37 | 2.65 | 2.74 | | random, kept 3/4 | 4.68 | 3.54 | 3.64 | | random, kept 15/16 | 5.40 | 3.61 | 3.94 | | runs of 64, kept 1/8 | 2.11 | 1.77 | 1.87 | | runs of 64, kept 1/2 | 3.00 | 2.35 | 2.51 | | runs of 64, kept 7/8 | 4.75 | 3.46 | 3.78 | | runs of 512, kept 1/2 | 6.99 | 5.44 | 5.70 | The other callers of `compress`: | benchmark | Neoverse-N2 | EPYC 7763 | |---|---|---| | `filter_kernels`: filter context i32 w NULLs (kept 1/2) | 1.80 | 1.69 | | `filter_kernels`: filter context u8 w NULLs (kept 1/2) | 1.83 | 1.66 | | `filter_kernels`: filter context string dictionary w NULLs (kept 1/2) | 1.77 | 1.63 | | `filter_kernels`: other w NULLs cases | 0.96 to 1.78 | 0.96 to 1.64 | | `arrow_reader`: ListArray and struct (definition levels) | 0.99 to 1.11 | 0.98 to 1.07 | ## What gets slower, ideas welcome Sparse masks. Up to 9 % on EPYC 7763 (random, kept 1/16), 7 % with x86-64-v2 (`filter_bits` indices, kept 1/10), 5 % on Neoverse-N2 (indices, kept 1/1024). On a word with one or two kept bits the old loop did almost nothing, and it did it very well. The new code first counts the bits and branches on the count, and on the default x86-64 target that count is a software popcount, so every mispredicted dispatch also waits for it. What I tried, so nobody has to again: - Counting one word ahead, so the count is ready before its branch. `perf stat` had shown the same instructions and the same mispredictions as main, just about five cycles more per misprediction: the software popcount the branch waits on. The lookahead hid it, and Zen 5 hated it (EPYC 9V45, a large mask kept 1/256: 0.56 against main instead of 0.89). - Finding the `k <= 2` and `k >= 62` cases without a count, `m & (m - 1)` twice. Cheaper for those words, but every other word paid both tests before its count; words with exactly three kept bits went to 0.59 to 0.75. - The same tests plus the previous word's count as a hint for the rest: smaller gains everywhere, one slow case fixed and several new ones. On random rows the previous word is a poor hint. What I have not tried, in case it itches someone: - Pipeline by blocks instead of words: popcount four to eight words at once (SWAR works even on the SSE2 baseline) and dispatch on the block. A sparse block could then walk its set bits across words in one loop, and pay for the loop exit once per block, not once per word. - Anything that gets the sparse case back to the naive loop and then past it. The runners are shared VMs. The `macos-15` runner is not in the tables: its rounds differed by up to 30 % even on cases that never reach `compress`. # Are there any user-facing changes? No. `compress` keeps its signature and its results; only its speed without BMI2 changes. --------- Co-authored-by: Andrew Lamb <[email protected]> --- arrow-buffer/src/util/bit_util.rs | 182 +++++++++++++++++++++++++++++++++++--- 1 file changed, 168 insertions(+), 14 deletions(-) diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index e50e111bef..6cacd7c039 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -26,7 +26,8 @@ use crate::bit_chunk_iterator::BitChunks; /// Equivalent to the x86 BMI2 `PEXT` instruction. When compiled with the /// `bmi2` target feature enabled (for example `-C target-cpu=x86-64-v3`) /// this lowers to the hardware `pext` instruction; otherwise it falls back -/// to a portable scalar loop. +/// to a portable version that picks its method by the number of set bits in +/// the mask. /// /// # Functional Example /// @@ -51,8 +52,16 @@ use crate::bit_chunk_iterator::BitChunks; /// assert_eq!(compress(0b1011_0100, 0b0110_1101), 0b0000_1010); /// ``` // -// Replace with `value.compress(mask)` when `uint_gather_scatter_bits` is -// stabilised: <https://github.com/rust-lang/rust/issues/149069> +// When `uint_gather_scatter_bits` is stabilised +// (<https://github.com/rust-lang/rust/issues/149069>), measure +// `value.extract_bits(mask)` against this before switching. Its portable +// version runs six shift stages for every mask, with each stage's mask taken +// from prefix counts of zero bits (Hacker's Delight 7-4, see +// <https://github.com/rust-lang/rust/blob/main/library/core/src/num/imp/int_bits.rs>). +// The fallback here does less work on sparse and dense masks, and on x86 also +// in between (three stages inside each byte plus one multiply, fewer shifts in +// total); on big AArch64 cores the six-stage version is faster in that middle +// band. #[inline] pub fn compress(value: u64, mask: u64) -> u64 { #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))] @@ -64,6 +73,29 @@ pub fn compress(value: u64, mask: u64) -> u64 { #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] { + compress_portable(value, mask) + } +} + +/// [`compress`] without `pext`, by the number of kept bits `k`: +/// +/// * `k <= 2`: the two lowest kept bits, tested directly +/// * `k <= 16`: one step per kept bit +/// * `k >= 62`: `value` with its at most two dropped bits removed +/// * otherwise [`compress_bytes`] +/// +/// The two ends need no loop, so no loop exit to mispredict when `k` +/// varies from word to word. +// Always: called once a word, so a call would cost more than a sparse word +#[cfg(any(test, not(all(target_arch = "x86_64", target_feature = "bmi2"))))] +#[inline(always)] +fn compress_portable(value: u64, mask: u64) -> u64 { + let kept = mask.count_ones(); + if kept <= 2 { + let lowest = mask & mask.wrapping_neg(); + let second = mask ^ lowest; + u64::from(value & lowest != 0) | (u64::from(value & second != 0) << 1) + } else if kept <= 16 { let mut mask = mask; let mut result = 0_u64; let mut dest_bit = 1_u64; @@ -78,9 +110,78 @@ pub fn compress(value: u64, mask: u64) -> u64 { mask = rest; } result + } else if kept >= 62 { + if kept == 64 { + return value; + } + // Remove the highest dropped bit first, so the lower ones stay put. + // With nothing left to drop a step keeps every bit, and each step + // shifts a zero in at the top. + let mut value = value; + let mut dropped = !mask; + for _ in 0..2 { + let below = if dropped == 0 { + u64::MAX + } else { + (1 << dropped.ilog2()) - 1 + }; + value = (value & below) | ((value >> 1) & !below); + dropped &= below; + } + value + } else { + compress_bytes(value, mask) } } +/// `b` in every byte of a `u64` +#[cfg(any(test, not(all(target_arch = "x86_64", target_feature = "bmi2"))))] +const fn repeat_byte(b: u8) -> u64 { + b as u64 * 0x0101_0101_0101_0101 +} + +/// [`compress`] in constant time, a byte at a time and all bytes at once. +/// +/// Within a byte each kept bit moves down by the number of dropped bits +/// below it. Written in binary, that distance is at most 7, so three rounds +/// move by 1, 2 and 4 the bits whose distance has that bit set (the parallel +/// suffix network, Hacker's Delight 7-4); shifts are masked to stay within +/// the byte. Byte `j` then lands at the sum of the kept counts of bytes +/// `0..j`, all eight sums from one multiply by `0x0101..01`. +#[cfg(any(test, not(all(target_arch = "x86_64", target_feature = "bmi2"))))] +#[inline] +fn compress_bytes(value: u64, mask: u64) -> u64 { + let mut packed = value & mask; + let mut dropped = !mask; + let mut shift = 1; + while shift < 8 { + // The move mask of this round: the parity of the dropped bits + // below each bit, in strides of `shift` + let mut move_mask = dropped; + let mut step = shift; + while step < 8 { + move_mask ^= (move_mask << step) & repeat_byte(0xFF << step); + step <<= 1; + } + let moving = packed & move_mask; + packed ^= moving ^ ((moving >> shift) & repeat_byte(0xFF >> shift)); + dropped &= !move_mask; + dropped ^= (dropped >> shift) & repeat_byte(0xFF >> shift); + shift <<= 1; + } + // Kept bits per byte (SWAR popcount, HD 5-1) + let pairs = mask - ((mask >> 1) & repeat_byte(0x55)); + let nibbles = (pairs & repeat_byte(0x33)) + ((pairs >> 2) & repeat_byte(0x33)); + let kept = (nibbles + (nibbles >> 4)) & repeat_byte(0x0F); + // Where byte `j` lands: the kept bits of bytes 0..j, all eight at once + let dest = kept.wrapping_mul(repeat_byte(1)) << 8; + let mut result = 0; + for j in 0..8 { + result |= ((packed >> (8 * j)) & 0xFF) << ((dest >> (8 * j)) & 0xFF); + } + result +} + /// Parallel bit deposit: scatter the lowest `mask.count_ones()` bits of /// `value` into the set positions of `mask`, preserving their order. /// All other bits in the result are zero; excess input bits are ignored. @@ -1016,18 +1117,18 @@ mod tests { use rand::rngs::StdRng; use rand::{RngExt, SeedableRng}; + /// Reference for [`compress`]: gather the `mask`-selected bits of `value` + /// into contiguous low bits, least-significant first + fn compress_reference(value: u64, mask: u64) -> u64 { + (0..64) + .filter(|&i| (mask >> i) & 1 == 1) + .enumerate() + .map(|(dest, i)| ((value >> i) & 1) << dest) + .sum() + } + #[test] fn test_compress() { - // Reference: gather the `mask`-selected bits of `value` into - // contiguous low bits, least-significant first - fn reference(value: u64, mask: u64) -> u64 { - (0..64) - .filter(|&i| (mask >> i) & 1 == 1) - .enumerate() - .map(|(dest, i)| ((value >> i) & 1) << dest) - .sum() - } - assert_eq!(compress(0b1010, 0b1111), 0b1010); assert_eq!(compress(0b1010, 0b1010), 0b11); assert_eq!(compress(0b1010, 0b0101), 0); @@ -1042,12 +1143,65 @@ mod tests { let (value, mask): (u64, u64) = rng.random(); assert_eq!( compress(value, mask), - reference(value, mask), + compress_reference(value, mask), "value={value:#x} mask={mask:#x}" ); } } + #[test] + fn test_compress_portable() { + // Every path of the fallback, whatever this build compiles + // `compress` to: uniform random masks alone have about 32 kept bits + // and would only reach `compress_bytes` + let check = |value: u64, mask: u64| { + let want = compress_reference(value, mask); + let kept = mask.count_ones(); + assert_eq!( + compress_portable(value, mask), + want, + "value={value:#x} mask={mask:#x} kept={kept}" + ); + }; + let mut rng = StdRng::seed_from_u64(42); + + // Every mask with at most two kept or at most two dropped bits + let value: u64 = rng.random(); + check(value, 0); + check(value, u64::MAX); + for i in 0..64 { + check(value, 1 << i); + check(value, !(1 << i)); + for j in i + 1..64 { + let two = (1 << i) | (1 << j); + check(value, two); + check(value, !two); + } + } + + // Random masks at every density: ANDing random words thins them out, + // ORing fills them in + for _ in 0..4096 { + let value: u64 = rng.random(); + let mut sparse = u64::MAX; + let mut dense = 0; + for _ in 0..8 { + let word: u64 = rng.random(); + sparse &= word; + dense |= word; + check(value, sparse); + check(value, dense); + } + // every popcount from 0 to 64 + let kept = rng.random_range(0..=64); + let mut mask = 0_u64; + while mask.count_ones() < kept { + mask |= 1 << rng.random_range(0..64); + } + check(value, mask); + } + } + #[test] fn test_expand() { fn reference(values: u64, mask: u64) -> u64 {
