mbutrovich commented on code in PR #10136:
URL: https://github.com/apache/arrow-rs/pull/10136#discussion_r4085568404
##########
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]
+ pub fn chunk(&self, index: usize) -> u64 {
Review Comment:
`chunk` is new public API, so could it get its own unit test?
`test_filter_bits` covers it indirectly, but a direct test would pin down the
contract. For example, check that `chunk(i)` equals `iter().nth(i)` for every
`i`, for bit offsets 0 through 7 and lengths that do and don't end on a word
boundary, and add a `#[should_panic(expected = "chunk index out of bounds")]`
test for `index == chunk_len()`.
##########
arrow-buffer/src/util/bit_util.rs:
##########
@@ -19,6 +19,45 @@
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.
+//
+// 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"))]
Review Comment:
Have you measured this on AMD Zen 1 or Zen 2? On those cores `pext` is
microcoded, and its latency grows with the number of set bits in the mask
(Agner Fog's instruction tables list it in the tens to hundreds of cycles). Zen
3 made it fast. Zen 2 supports `x86-64-v3`, so a binary built with `-C
target-cpu=x86-64-v3` or `target-cpu=native` on those machines takes this path,
and at 1/2 density that's about 32 mask bits per word. The Parquet caller
already has this exposure, but the filter kernel runs on far more queries. If
Zen 2 turns out to be slower than the fallback, a doc note on `compress` saying
so would help people choosing target flags.
@andygrove, I think you have a Zen 1 or Zen 2 machine. Could you run `cargo
bench -p arrow-select --bench filter_bits` at this PR's head with and without
`RUSTFLAGS="-C target-cpu=native"` and post the numbers here?
##########
arrow-select/src/filter.rs:
##########
@@ -719,6 +739,86 @@ 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() {
+ // Skipping words with no kept bits before touching the values makes
+ // sparse filters cost a load and a test per word; at moderate
+ // densities the branch is never taken
+ if mask == 0 {
+ continue;
+ }
+ packer.push(value_chunks.chunk(index), mask);
Review Comment:
Did you compare this against zipping the two iterators,
`mask_chunks.iter().zip(value_chunks.iter())`, with the same `mask == 0` skip?
On aarch64 (M5 Max, scalar fallback) it's mixed. With `chunk(index)`, lazy 1/2
is 20.1 us vs 21.5 us for the zip, and sliced 1/2 is 21.0 us vs 23.3 us. At the
sparse end the zip is faster: lazy 1/256 is 805 ns vs 638 ns, and lazy 1/1024
is 560 ns vs 406 ns. That's the opposite of what the comment at lines 800-802
says about sparse filters. If `chunk` wins on x86 with `bmi2`, keeping it makes
sense, but could you share those numbers and adjust the comment to match? If
the zip is as good, it would also mean `BitChunks` doesn't need the new public
`chunk` method.
--
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]
