alamb commented on code in PR #10136: URL: https://github.com/apache/arrow-rs/pull/10136#discussion_r4103536011
########## 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: Filed a follow on ticket to look into improving the performance of the fallback - https://github.com/apache/arrow-rs/issues/11213 -- 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]
