mbutrovich commented on code in PR #11918:
URL: https://github.com/apache/arrow-rs/pull/11918#discussion_r4395258002
##########
arrow-select/src/filter.rs:
##########
@@ -801,80 +801,150 @@ fn filter_bits_strategy(buffer: &BooleanBuffer,
predicate: &FilterPredicate) ->
}
/// 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
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`).
#[inline(never)]
-fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate)
-> Buffer {
+pub(crate) fn filter_bits_compress(
+ buffer: &BooleanBuffer,
+ predicate: &FilterPredicate,
+ mut append: impl FnMut(u64, usize),
+) {
+ filter_bits_compress_impl::<false>(buffer, None, predicate, |word, _,
count| {
+ append(word, count)
+ })
+}
+
+/// Filters two packed bitmaps using the same predicate and mask traversal.
+#[inline(never)]
+pub(crate) fn filter_bits_compress_pair(
+ first: &BooleanBuffer,
+ second: &BooleanBuffer,
+ predicate: &FilterPredicate,
+ mut append: impl FnMut(u64, u64, usize),
+) {
+ filter_bits_compress_impl::<true>(first, Some(second), predicate, |first,
second, count| {
+ append(first, second.unwrap(), count)
+ })
+}
+
+#[inline(always)]
+fn filter_bits_compress_impl<const PAIR: bool>(
+ first: &BooleanBuffer,
+ second: Option<&BooleanBuffer>,
+ predicate: &FilterPredicate,
+ mut append: impl FnMut(u64, Option<u64>, usize),
+) {
/// 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
+ /// bits of `current`.
+ #[derive(Default)]
struct Packer {
- ptr: *mut u64,
- idx: usize,
+ /// Bits waiting to be written.
current: u64,
+ /// Number of populated low bits in `current`.
filled: u32,
}
impl Packer {
#[inline(always)]
- fn push(&mut self, values: u64, mask: u64) {
- let bits = bit_util::compress(values, mask);
+ fn push(&mut self, bits: u64, count: u32) -> Option<u64> {
self.current |= bits << self.filled;
- let total = self.filled + mask.count_ones();
+ let total = self.filled + count;
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;
+ return None;
}
+
+ let complete = self.current;
+ // `bits >> (64 - filled)`, written so that `filled == 0`
+ // shifts everything out
+ self.current = (bits >> 1) >> (63 - self.filled);
+ self.filled = total - 64;
+ Some(complete)
Review Comment:
Does sharing this loop slow down the regular Boolean `filter`? The 10% and
80% `single_boolean` rows are above the 1/16 sparse cutoff
([coalesce.rs](https://github.com/apache/arrow-rs/blob/4d091cd4980c284256784f98b41eb264de1fb55a/arrow-select/src/coalesce.rs#L59-L63)),
so they measure `collect_filter_bits_compress`, not the fused path (#11921
says the same). On the bot, the 10% rows have been 6% to 8% slower since the
`Packer` refactor
([run](https://github.com/apache/arrow-rs/pull/11918#issuecomment-6428336979)).
As I read the codegen, returning `Option<u64>` lets LLVM turn the `total <
64` branch into selects, so the carry is computed on every non-zero mask word.
With thin LTO on aarch64, the path where no word completes goes from 5
instructions at the
[merge-base](https://github.com/apache/arrow-rs/blob/69be95c1d925bd876c71c11ec6cb2a646d54a1ad/arrow-select/src/filter.rs#L823-L838)
to 12. These are the added ones:
```text
lsr x10, x23, #1 // carry, now on every word
eor w11, w11, #0x3f
lsr x10, x10, x11
cmp w20, #0x3f
csel w11, w11, w20, hi // filled
csel x10, x10, x21, hi // current
cmp w20, #0x40
```
With `-C target-cpu=x86-64-v3`, `current` also moves to the stack, with a
load and a store on every non-zero mask word.
What do you think about one packer for both bitmaps? They keep the same
number of bits per word, so they can share `filled`, and completed words can be
emitted inside the branch as at the merge-base. This also drops `PAIR`, the
`Option` and the `unwrap`s:
```rust
#[inline(always)]
fn filter_bits_compress_impl<const N: usize>(
buffers: [&BooleanBuffer; N],
predicate: &FilterPredicate,
mut append: impl FnMut([u64; N], usize),
) {
/// Packs the bits extracted from successive words into the low `filled`
/// bits of each of `current`. Every bitmap keeps the same number of bits
/// per word, so they share `filled`.
struct Packer<const N: usize> {
current: [u64; N],
filled: u32,
}
impl<const N: usize> Packer<N> {
#[inline(always)]
fn push(
&mut self,
values: [u64; N],
mask: u64,
append: &mut impl FnMut([u64; N], usize),
) {
let mut bits = [0; N];
for (bits, values) in bits.iter_mut().zip(values) {
*bits = bit_util::compress(values, mask);
}
for (current, bits) in self.current.iter_mut().zip(bits) {
*current |= bits << self.filled;
}
let total = self.filled + mask.count_ones();
if total < 64 {
self.filled = total;
} else {
append(self.current, 64);
for (current, bits) in self.current.iter_mut().zip(bits) {
// `bits >> (64 - filled)`, written so that `filled == 0`
// shifts everything out
*current = (bits >> 1) >> (63 - self.filled);
}
self.filled = total - 64;
}
}
}
let len = predicate.filter.len();
for buffer in buffers {
assert!(buffer.len() >= len);
}
let mask_chunks = predicate.filter.values().bit_chunks();
let chunks = buffers.map(|buffer| BitChunks::new(buffer.values(),
buffer.offset(), len));
debug_assert_eq!(predicate.count, predicate.filter.true_count());
let mut packer = Packer { current: [0; N], 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(chunks.each_ref().map(|c| c.chunk(index)), mask, &mut
append);
}
let mask = mask_chunks.remainder_bits();
if mask != 0 {
packer.push(chunks.each_ref().map(|c| c.remainder_bits()), mask,
&mut append);
}
if packer.filled != 0 {
append(packer.current, packer.filled as usize);
}
}
```
`filter_bits_compress` and `filter_bits_compress_pair` then pass `[buffer]`
and `[first, second]`. Two details matter for the generated code. The compress
step needs to be a loop, because `values.map(...)` leaves `compress` out of
line on aarch64. And `filter_bits_compress` should be `#[inline(always)]`, so
the loop lands in `collect_filter_bits_compress` and `idx` stays in a register
instead of being stored through the closure on every completed word.
Here is the change against the merge-base on an Apple silicon Mac with thin
LTO, as the median of three interleaved runs. Each cell gives the default build
first, then the build with `-C llvm-args=-align-all-functions=7 -C
llvm-args=-align-all-nofallthru-blocks=7`:
| case | this PR | shared packer |
|---|---|---|
| no nulls, 1% selected | 11% / 14% faster | 25% / 25% faster |
| 10% nulls, 1% selected | 30% / 34% faster | 34% / 34% faster |
| no nulls, 10% selected | 3% faster / 4% slower | 4% / 2% faster |
| 10% nulls, 10% selected | 3% faster / 2% slower | 7% / 2% faster |
| no nulls, 80% selected | 1% / 7% slower | 5% / 5% slower |
| 10% nulls, 80% selected | 3% / 11% slower | 7% / 9% slower |
The 80% rows are still slower with both versions. A profile puts about 80%
of that case in the compress loop, and I couldn't find the cause in the
listings. Could you look into that case and rerun the bot, which runs on
Neoverse-V2?
Could you also add a Boolean data case to `filter_kernels`, in its own
commit first? Today it uses Boolean arrays only as filter masks.
##########
arrow-select/src/filter.rs:
##########
@@ -801,80 +801,150 @@ fn filter_bits_strategy(buffer: &BooleanBuffer,
predicate: &FilterPredicate) ->
}
/// 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
+/// bits of each 64-bit word with [`bit_util::compress`] (`pext`).
#[inline(never)]
-fn filter_bits_compress(buffer: &BooleanBuffer, predicate: &FilterPredicate)
-> Buffer {
+pub(crate) fn filter_bits_compress(
+ buffer: &BooleanBuffer,
+ predicate: &FilterPredicate,
+ mut append: impl FnMut(u64, usize),
+) {
+ filter_bits_compress_impl::<false>(buffer, None, predicate, |word, _,
count| {
+ append(word, count)
+ })
+}
+
+/// Filters two packed bitmaps using the same predicate and mask traversal.
+#[inline(never)]
+pub(crate) fn filter_bits_compress_pair(
+ first: &BooleanBuffer,
+ second: &BooleanBuffer,
+ predicate: &FilterPredicate,
+ mut append: impl FnMut(u64, u64, usize),
+) {
+ filter_bits_compress_impl::<true>(first, Some(second), predicate, |first,
second, count| {
+ append(first, second.unwrap(), count)
+ })
+}
+
+#[inline(always)]
+fn filter_bits_compress_impl<const PAIR: bool>(
+ first: &BooleanBuffer,
+ second: Option<&BooleanBuffer>,
+ predicate: &FilterPredicate,
+ mut append: impl FnMut(u64, Option<u64>, usize),
+) {
/// 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
+ /// bits of `current`.
+ #[derive(Default)]
struct Packer {
- ptr: *mut u64,
- idx: usize,
+ /// Bits waiting to be written.
current: u64,
+ /// Number of populated low bits in `current`.
filled: u32,
}
impl Packer {
#[inline(always)]
- fn push(&mut self, values: u64, mask: u64) {
- let bits = bit_util::compress(values, mask);
+ fn push(&mut self, bits: u64, count: u32) -> Option<u64> {
self.current |= bits << self.filled;
- let total = self.filled + mask.count_ones();
+ let total = self.filled + count;
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;
+ return None;
}
+
+ let complete = self.current;
+ // `bits >> (64 - filled)`, written so that `filled == 0`
+ // shifts everything out
+ self.current = (bits >> 1) >> (63 - self.filled);
+ self.filled = total - 64;
+ Some(complete)
+ }
+
+ #[inline]
+ fn finish(self) -> Option<(u64, usize)> {
+ (self.filled != 0).then_some((self.current, self.filled as usize))
}
}
- assert!(buffer.len() >= predicate.filter.len());
+ assert!(first.len() >= predicate.filter.len());
+ if PAIR {
+ assert!(second.unwrap().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
+ let first_chunks = BitChunks::new(first.values(), first.offset(),
predicate.filter.len());
+ let second_chunks = second
+ .map(|buffer| BitChunks::new(buffer.values(), buffer.offset(),
predicate.filter.len()));
debug_assert_eq!(predicate.count, predicate.filter.true_count());
- debug_assert_eq!(mask_chunks.chunk_len(), value_chunks.chunk_len());
+ debug_assert_eq!(mask_chunks.chunk_len(), first_chunks.chunk_len());
+ if PAIR {
+ debug_assert_eq!(
+ mask_chunks.chunk_len(),
+ second_chunks.as_ref().unwrap().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,
- };
+ let mut first_packer = Packer::default();
+ let mut second_packer = Packer::default();
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);
+ let count = mask.count_ones();
+ let first =
first_packer.push(bit_util::compress(first_chunks.chunk(index), mask), count);
+ let second = PAIR.then(|| {
+ second_packer.push(
+
bit_util::compress(second_chunks.as_ref().unwrap().chunk(index), mask),
+ count,
+ )
+ });
+ if let Some(first) = first {
+ append(first, second.flatten(), 64);
+ }
}
- 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
- debug_assert!(packer.idx < out.capacity());
- unsafe {
- packer.ptr.add(packer.idx).write(packer.current);
- out.set_len(packer.idx + 1);
+ let mask = mask_chunks.remainder_bits();
+ if mask != 0 {
+ let count = mask.count_ones();
+ let first = first_packer.push(
+ bit_util::compress(first_chunks.remainder_bits(), mask),
+ count,
+ );
+ let second = PAIR.then(|| {
+ second_packer.push(
+
bit_util::compress(second_chunks.as_ref().unwrap().remainder_bits(), mask),
+ count,
+ )
+ });
+ if let Some(first) = first {
+ append(first, second.flatten(), 64);
+ }
+ }
+ if let Some((first, count)) = first_packer.finish() {
+ let second = PAIR.then(|| second_packer.finish().unwrap().0);
+ append(first, second, count);
}
+}
+
+/// Collects compressed bits into a new buffer.
+///
+/// Not inlined: within `filter_array` the packing state spills to the stack.
+#[inline(never)]
+fn collect_filter_bits_compress(buffer: &BooleanBuffer, predicate:
&FilterPredicate) -> Buffer {
+ // One word beyond the complete ones for the trailing partial word
+ let mut out: Vec<u64> = Vec::with_capacity(predicate.count / 64 + 1);
+ let ptr = out.as_mut_ptr();
+ let mut idx = 0;
+ filter_bits_compress(buffer, predicate, |word, _| {
Review Comment:
The merge-base had `debug_assert!(packer.idx < out.capacity())` ahead of its
last unchecked write, and the closure now has no check before
`ptr.add(idx).write(word)`. Could we add one so debug builds and Miri catch a
wrong `predicate.count`?
```suggestion
filter_bits_compress(buffer, predicate, |word, _| {
debug_assert!(idx <= predicate.count / 64);
```
I know a debug_assert gets compiled out anyway, but it's a nice documented
invariant.
--
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]
