mbutrovich commented on code in PR #12024:
URL: https://github.com/apache/arrow-rs/pull/12024#discussion_r4446688807
##########
arrow-buffer/src/builder/boolean.rs:
##########
@@ -276,9 +276,51 @@ impl BooleanBufferBuilder {
self.advance(additional);
let offset = self.len() - additional;
- for (i, v) in slice.iter().enumerate() {
+
+ // Too short to fill a whole byte, so packing would only add overhead
+ if additional < 8 {
+ for (i, v) in slice.iter().enumerate() {
+ if *v {
+ unsafe { bit_util::set_bit_raw(self.buffer.as_mut_ptr(),
offset + i) }
+ }
+ }
+ } else {
+ self.pack_slice(offset, slice)
+ }
+ }
+
+ /// Writes `slice` into the already allocated bits starting at `offset`
+ ///
+ /// Kept out of line so that the short slice path of [`Self::append_slice`]
+ /// stays small when inlined
+ #[inline(never)]
+ fn pack_slice(&mut self, offset: usize, slice: &[bool]) {
+ let additional = slice.len();
+ let buffer = self.buffer.as_slice_mut();
+
+ // Set bits individually until the write position is byte aligned
+ let head = ((8 - offset % 8) % 8).min(additional);
+ let (head_slice, rest) = slice.split_at(head);
+ for (i, v) in head_slice.iter().enumerate() {
if *v {
- unsafe { bit_util::set_bit_raw(self.buffer.as_mut_ptr(),
offset + i) }
+ bit_util::set_bit(buffer, offset + i)
+ }
+ }
+
+ // Pack 8 bools into each byte without branching: a `bool` is a byte
+ // holding 0 or 1, and the multiply moves byte `i` to bit `56 + i`
+ let byte_offset = (offset + head) / 8;
+ let (chunks, remainder) = rest.as_chunks::<8>();
+ for (dst, chunk) in buffer[byte_offset..].iter_mut().zip(chunks) {
+ let bytes: [u8; 8] = std::array::from_fn(|i| chunk[i] as u8);
+ let packed =
u64::from_le_bytes(bytes).wrapping_mul(0x0102_0408_1020_4080) >> 56;
+ *dst = packed as u8;
+ }
+
+ let remainder_offset = self.len - remainder.len();
+ for (i, v) in remainder.iter().enumerate() {
+ if *v {
+ bit_util::set_bit(buffer, remainder_offset + i)
}
}
Review Comment:
Could we build the head, remainder, and short-slice bits in a register with
`(*v as u8) << i` instead of `if *v { set_bit(..) }`? Those three loops still
branch on every value. In the release build of `pack_slice` at this commit, the
head and remainder loops contain 14 branches on a loaded bool, on both aarch64
(`tbz`/`cbz w9`) and x86_64 (`cmp byte ptr [rdx + 1], 1` then `jne`). With the
suggestion below, each value is a load and a shifted OR, and the only branches
left compare against the length:
```asm
ldrb w9, [x2, #1]
orr w12, w12, w9, lsl #1 ; value 1 into bit 1
cmp x10, #2 ; x10 = head, stop after 2 values
b.eq LBB40_9
```
Each region then does one read-modify-write of its byte, and the short path
writes at most two bytes. This also covers @sunchao's case of eight values at
offset 3 without an alignment-aware cutoff. Its five head values and three
remainder values become two writes instead of eight branches.
I measured this with extra cases in `boolean_buffer_builder`. The 1-to-7
rows are 7168 appends cycling through 1 to 7 values from bit offset 3. The
8-value rows are 7168 appends of 8 values. All rows read consecutive fresh
values. These are us per iteration on an Apple M5 Max with thin LTO:
| case | main | this PR | suggestion |
|---|---|---|---|
| 1 to 7 values, random | 35.3 | 26.3 | 16.4 |
| 1 to 7 values, all false | 16.8 | 16.8 | 15.6 |
| 1 to 7 values, all true | 23.9 | 21.1 | 15.6 |
| 8 values at offset 0, random | 93.1 | 21.7 | 17.3 |
| 8 values at offset 0, all false | 24.3 | 21.1 | 17.0 |
| 8 values at offset 3, random | 89.6 | 45.5 | 21.0 |
| 8 values at offset 3, all false | 20.8 | 25.4 | 21.1 |
| 8 values at offset 3, all true | 31.1 | 31.6 | 21.2 |
| 64Ki, offset 0 or 3, all three inputs | 58 to 133 | 1.05 to 1.08 | 1.04 to
1.07 |
The PR and suggestion columns agree within about 1 us in a default release
build and in a build with function and block alignment fixed. The cost is in
aligned appends, where the head now ORs 0 into the first byte. The offset 0
rows measure that case, and they aren't slower. The suggestion also removes the
`unsafe` call to `set_bit_raw`.
```suggestion
// Too short to fill a whole byte, so packing would only add overhead
if additional < 8 {
if additional == 0 {
return;
}
// At most 7 bits starting at bit 7 end at bit 13, so they span
at most two bytes
let bits = (Self::pack_bits(slice) as u16) << (offset % 8);
let buffer = self.buffer.as_slice_mut();
buffer[offset / 8] |= bits as u8;
if let Some(next) = buffer.get_mut(offset / 8 + 1) {
*next |= (bits >> 8) as u8;
}
} else {
self.pack_slice(offset, slice)
}
}
/// Packs up to 7 bools into the low bits of a byte, `slice[i]` at bit
`i`,
/// without a branch per value
#[inline]
fn pack_bits(slice: &[bool]) -> u8 {
slice
.iter()
.enumerate()
.fold(0, |acc, (i, v)| acc | ((*v as u8) << i))
}
/// Writes `slice` into the already allocated bits starting at `offset`
///
/// Kept out of line so that the short slice path of
[`Self::append_slice`]
/// stays small when inlined
#[inline(never)]
fn pack_slice(&mut self, offset: usize, slice: &[bool]) {
let additional = slice.len();
let buffer = self.buffer.as_slice_mut();
// Fill the first byte up to the next byte boundary
let head = ((8 - offset % 8) % 8).min(additional);
let (head_slice, rest) = slice.split_at(head);
buffer[offset / 8] |= Self::pack_bits(head_slice) << (offset % 8);
// Pack 8 bools into each byte without branching: a `bool` is a byte
// holding 0 or 1, and the multiply moves byte `i` to bit `56 + i`
let byte_offset = (offset + head) / 8;
let (chunks, remainder) = rest.as_chunks::<8>();
for (dst, chunk) in buffer[byte_offset..].iter_mut().zip(chunks) {
let bytes: [u8; 8] = std::array::from_fn(|i| chunk[i] as u8);
let packed =
u64::from_le_bytes(bytes).wrapping_mul(0x0102_0408_1020_4080) >> 56;
*dst = packed as u8;
}
// The remainder starts on a byte boundary
if !remainder.is_empty() {
buffer[(self.len - remainder.len()) / 8] |=
Self::pack_bits(remainder);
}
```
--
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]
