alamb opened a new issue, #11213: URL: https://github.com/apache/arrow-rs/issues/11213
### Is your feature request related to a problem or challenge? @devanbenz / @Rich-T-kid and others updated the filter kernel to be more efficient - https://github.com/apache/arrow-rs/pull/10136 If the CPU has the [`pext`](https://en.wikipedia.org/wiki/X86_Bit_manipulation_instruction_set#Parallel_bit_deposit_and_extract) instruction, the `compress` function uses that (and is quite speedy on intel CPUs) ```rust #[inline] pub fn compress(value: u64, mask: u64) -> u64 { #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))] { unsafe { std::arch::x86_64::_pext_u64(value, mask) } } #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] { // <--------- Fallback code HERE } } ``` However, as @mbutrovich and others pointed out , the fallback seem like it could be improved as it operates one bit at a time ### Describe the solution you'd like Find a better / faster implementation for `compress` on hardware that doesn't support `pext` ### Describe alternatives you've considered @mbutrovich suggests https://github.com/apache/arrow-rs/pull/10136/changes#r4094221099 > 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: > > 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: > > 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. > > Suggested change > { > 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 > } > { > 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 > } I had claude code suggest a table based approach from hackers delight: https://github.com/apache/arrow-rs/pull/10136/changes#r4087490049 > Claude flagged this as not very efficient as it is still doing bit at a time > > Here is some version it wrote that operates using a precomputed table (citing the [Hacker's Delight](https://www.amazon.com/dp/0321842685?lv=shuf&channelId=500&plpRedirect=mhFallback), one of my favorite books -- I think [@XiangpengHao](https://github.com/XiangpengHao) has a copy!) > > Perhaps we can flag this as a potential improvement for follow on work in another ticket > > { > // Nibble at a time through a 256 byte table: 16 steps per word > // regardless of how many bits `mask` selects > let mut result = 0_u64; > let mut shift = 0_u32; > let mut i = 0; > while i < 16 { > let m = ((mask >> (i * 4)) & 0xF) as usize; > let v = ((value >> (i * 4)) & 0xF) as usize; > result |= (COMPRESS_NIBBLE[m][v] as u64) << shift; > shift += m.count_ones(); > i += 1; > } > result > } > > /// `COMPRESS_NIBBLE[mask][value]` is `compress(value, mask)` for 4-bit inputs > #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] > static COMPRESS_NIBBLE: [[u8; 16]; 16] = { > let mut t = [[0u8; 16]; 16]; > let mut m = 0; > while m < 16 { > let mut v = 0; > while v < 16 { > let mut r = 0u8; > let mut d = 0; > let mut b = 0; > while b < 4 { > if m >> b & 1 == 1 { > r |= ((v >> b & 1) as u8) << d; > d += 1; > } > b += 1; > } > t[m][v] = r; > v += 1; > } > m += 1; > } > t > }; ### Additional context _No response_ -- 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]
