mbutrovich commented on code in PR #10136:
URL: https://github.com/apache/arrow-rs/pull/10136#discussion_r4094215179
##########
arrow-buffer/src/util/bit_chunk_iterator.rs:
##########
@@ -324,6 +324,20 @@ impl<'a> BitChunks<'a> {
ceil(self.chunk_len * 64 + self.remainder_len, 8)
}
+ /// Returns the `index`th complete chunk of 64 bits, the value
+ /// [`Self::iter`] yields at that position
+ ///
+ /// # Panics
+ ///
+ /// Panics if `index >= self.chunk_len()`
+ #[inline]
Review Comment:
Could `chunk` get a doc example? It's new public API. Something like this
passes as a doctest at the head commit:
```suggestion
/// Returns the `index`th complete chunk of 64 bits, the value
/// [`Self::iter`] yields at that position
///
/// # Panics
///
/// Panics if `index >= self.chunk_len()`
///
/// # Example
///
/// ```
/// # use arrow_buffer::Buffer;
/// let buffer = Buffer::from((0..17_u8).collect::<Vec<_>>());
/// // Start at bit 8, so chunk 1 holds bytes 9 through 16
/// let chunks = buffer.bit_chunks(8, 128);
/// assert_eq!(chunks.chunk_len(), 2);
/// assert_eq!(chunks.chunk(1), 0x100f0e0d0c0b0a09);
/// assert_eq!(chunks.iter().nth(1), Some(chunks.chunk(1)));
/// ```
```
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -19,6 +19,68 @@
use crate::bit_chunk_iterator::BitChunks;
+/// Parallel bit extract: for each set bit in `mask`, extract the
+/// corresponding bit from `value` and pack them contiguously into the low
+/// bits of the return value.
+///
+/// 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.
+///
+/// # Functional Example
+///
+/// Using 8 bits for brevity (the function operates on all 64). Each
+/// set bit in `mask` selects the bit at the same position in `value`; the
+/// selected bits are then shifted down so they are contiguous in the low
+/// bits of the result, in their original order:
+///
+/// ```text
+/// bit: 7 6 5 4 3 2 1 0
+/// value: a b c d e f g h
+/// mask: 0 1 1 0 1 1 0 1 set bits select b, c, e, f and h
+/// | | | | |
+/// v v v v v copy the relevant bits into result
+/// result: 0 0 0 b c e f h
+/// ```
+///
+/// # Code Example
+///
+/// ```
+/// # use arrow_buffer::bit_util::compress;
+/// 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>
+#[inline]
+pub fn compress(value: u64, mask: u64) -> u64 {
+ #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
+ {
+ // SAFETY: the `bmi2` target feature is statically enabled for this
+ // build, so the `pext` instruction is guaranteed to be available.
+ unsafe { std::arch::x86_64::_pext_u64(value, mask) }
+ }
+
+ #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
+ {
+ let mut mask = mask;
+ let mut result = 0_u64;
+ let mut dest_bit = 1_u64;
+ while mask != 0 {
+ // Clear the lowest set bit; the loop-carried dependency is only
+ // this two-operation chain, everything else hangs off it
+ let rest = mask & (mask - 1);
+ let lowest = mask ^ rest;
+ let keep = ((value & lowest) != 0) as u64;
+ result |= dest_bit & keep.wrapping_neg();
+ dest_bit <<= 1;
+ mask = rest;
+ }
+ result
+ }
Review Comment:
What do you think about having the fallback walk whichever set of bits is
smaller? The current loop runs once per set bit in `mask`, so at 9/10 density
it takes about 58 iterations per word. Each iteration is 9 instructions on
aarch64:
```asm
LBB1036_14:
mov x26, x25
sub x25, x25, #1
and x25, x25, x26 ; mask &= mask - 1
eor x26, x25, x26 ; lowest set bit
tst x26, x7
csel x26, xzr, x24, eq
orr x6, x26, x6
lsl x24, x24, #1
cbnz x25, LBB1036_14
```
When more than half the bits are set, removing the dropped bits instead
(highest first, so lower positions don't move) takes about 6. The dense loop in
the suggestion below compiles to 8 instructions per dropped bit:
```asm
LBB1036_16:
clz x25, x6
eor x25, x25, #0x3f ; ilog2
lsl x25, x5, x25 ; !below
bic x26, x7, x25 ; value & below
and x7, x25, x7, lsr #1 ; (value >> 1) & !below
orr x7, x26, x7
bics x6, x6, x25 ; dropped &= below
b.ne LBB1036_16
```
The compiler also reuses one `count_ones` for both the `kept > 32` check and
`Packer::push`, so the extra branch costs no extra popcount.
On an Apple M5 Max, `filter_bits` at 9/10 goes from 33.9 us to 3.3 us (-90%)
for the lazy, optimized, and sliced cases, and 1/2 improves by 7-11%. The 1/10
through 1/1024 cases move between -5% and +8%, mostly under 3%.
`test_compress`, the `arrow-select` filter tests, the `compress` doctest, and
clippy all pass. The Parquet caller of `compress` gets the same speedup. I
haven't checked whether the `dense` cutoff in `filter_bits` should move now
that dense masks are cheaper here.
```suggestion
{
let kept = mask.count_ones();
if kept > 32 {
// Dense mask: remove the dropped bits instead, highest first, so
// the loop runs once per clear bit of `mask`
let mut value = value;
let mut dropped = !mask;
while dropped != 0 {
let below = (1_u64 << dropped.ilog2()) - 1;
value = (value & below) | ((value >> 1) & !below);
dropped &= below;
}
return value & (u64::MAX >> (64 - kept));
}
let mut mask = mask;
let mut result = 0_u64;
let mut dest_bit = 1_u64;
while mask != 0 {
// Clear the lowest set bit; the loop-carried dependency is only
// this two-operation chain, everything else hangs off it
let rest = mask & (mask - 1);
let lowest = mask ^ rest;
let keep = ((value & lowest) != 0) as u64;
result |= dest_bit & keep.wrapping_neg();
dest_bit <<= 1;
mask = rest;
}
result
}
```
##########
arrow-select/src/filter.rs:
##########
@@ -719,6 +739,85 @@ fn filter_bits(buffer: &BooleanBuffer, predicate:
&FilterPredicate) -> Buffer {
}
}
+/// Filter the packed bitmask `buffer` with `predicate` by extracting the kept
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`)
+///
+/// Not inlined: within `filter_array` the packing state spills to the stack
+#[inline(never)]
+fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate)
-> Buffer {
+ /// Packs the bits extracted from successive words into the low `filled`
+ /// bits of `current`; once complete it is written at `idx` and restarts
+ /// from the bits that did not fit
+ struct Packer {
+ ptr: *mut u64,
+ idx: usize,
+ current: u64,
+ filled: u32,
+ }
+
+ impl Packer {
+ #[inline(always)]
+ fn push(&mut self, values: u64, mask: u64) {
+ let bits = bit_util::compress(values, mask);
+ self.current |= bits << self.filled;
+ let total = self.filled + mask.count_ones();
+ if total < 64 {
+ self.filled = total;
+ } else {
+ // SAFETY: `count` is the number of set bits in the filter, so
+ // at most `count / 64` words are ever completed and the
+ // buffer holds `count / 64 + 1`
+ unsafe { self.ptr.add(self.idx).write(self.current) };
+ self.idx += 1;
+ // `bits >> (64 - filled)`, written so that `filled == 0`
+ // shifts everything out
+ self.current = (bits >> 1) >> (63 - self.filled);
+ self.filled = total - 64;
+ }
+ }
+ }
+
+ assert!(buffer.len() >= predicate.filter.len());
+ let mask_chunks = predicate.filter.values().bit_chunks();
+ let value_chunks = BitChunks::new(buffer.values(), buffer.offset(),
predicate.filter.len());
+ // `count` is the filter's set bit count, which the buffer size and the
+ // raw writes below rely on, and both chunk views cover
+ // `predicate.filter.len()` bits, so indexing `value_chunks` by the
+ // position in `mask_chunks` stays in bounds
+ debug_assert_eq!(predicate.count, predicate.filter.true_count());
+ debug_assert_eq!(mask_chunks.chunk_len(), value_chunks.chunk_len());
+
+ // One word beyond the complete ones for the trailing partial word
+ let mut out: Vec<u64> = Vec::with_capacity(predicate.count / 64 + 1);
+ let mut packer = Packer {
+ ptr: out.as_mut_ptr(),
+ idx: 0,
+ current: 0,
+ filled: 0,
+ };
+
+ for (index, mask) in mask_chunks.iter().enumerate() {
+ // Words with no kept bits are skipped before the corresponding values
+ // are read, so only the mask is touched for them
+ if mask == 0 {
+ continue;
+ }
+ packer.push(value_chunks.chunk(index), mask);
+ }
+ packer.push(value_chunks.remainder_bits(), mask_chunks.remainder_bits());
+
+ // The trailing partial word; its bits above `filled` are zero
+ // SAFETY: `idx <= count / 64`, so this and every word below it is
+ // within the buffer and written
+ unsafe {
+ packer.ptr.add(packer.idx).write(packer.current);
+ out.set_len(packer.idx + 1);
+ }
Review Comment:
The raw writes in `Packer::push` and here depend on `idx <= count / 64`,
which only holds if `count` matches the filter's set bits. The
`debug_assert_eq!` above checks `count`, but nothing checks the resulting
index. Since `idx` only grows, one check before the final write catches an
overrun in any of the writes when the tests run:
```suggestion
// The trailing partial word; its bits above `filled` are zero
// SAFETY: `idx <= count / 64`, so this and every word below it is
// within the buffer and written
debug_assert!(packer.idx < out.capacity());
unsafe {
packer.ptr.add(packer.idx).write(packer.current);
out.set_len(packer.idx + 1);
}
```
--
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]
