perf: const-generic bit unpacking with wide u64 reads
Specialize extract_bits_msb and unpack_bits via const generics so the compiler knows BPS at compile time. Replace the byte-by-byte accumulation loop with a single u64::from_be_bytes read + shift + mask. Benchmarks (1024x1024 gray, uncompressed): 10-bit: 1.32ms → 886µs (1.49x faster) 12-bit: 1.34ms → 871µs (1.54x faster) 14-bit: 1.49ms → 963µs (1.55x faster) 24-bit: 1.84ms → 776µs (2.37x faster)
This commit is contained in:
1 parent
fcdb86ca7f
commit
3434d0220b
1 file changed
+58
-20
+58
-20
@@ -1329,52 +1329,90 @@ fn unpack_bits(
|
|||||||
bits_per_sample: u8,
|
bits_per_sample: u8,
|
||||||
num_samples: usize,
|
num_samples: usize,
|
||||||
output_width: u8,
|
output_width: u8,
|
||||||
|
) {
|
||||||
|
// Dispatch to const-generic specializations so the compiler can unroll the
|
||||||
|
// byte-gathering loop and precompute shift/mask constants.
|
||||||
|
macro_rules! dispatch {
|
||||||
|
($($bps:literal),+) => {
|
||||||
|
match bits_per_sample {
|
||||||
|
$( $bps => unpack_bits_const::<$bps>(src, dst, num_samples, output_width), )+
|
||||||
|
_ => unpack_bits_const::<0>(src, dst, num_samples, output_width),
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
// Cover every bit depth from 9..=32 that could reach this path, plus 0 as fallback.
|
||||||
|
dispatch!(
|
||||||
|
9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
|
||||||
|
32
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inner loop with const BPS so the compiler can optimize the bit extraction.
|
||||||
|
/// BPS=0 is the fallback for unexpected bit depths (uses runtime value).
|
||||||
|
#[inline(always)]
|
||||||
|
fn unpack_bits_const<const BPS: u8>(
|
||||||
|
src: &[u8],
|
||||||
|
dst: &mut [u8],
|
||||||
|
num_samples: usize,
|
||||||
|
output_width: u8,
|
||||||
) {
|
) {
|
||||||
debug_assert!(output_width == 16 || output_width == 32);
|
debug_assert!(output_width == 16 || output_width == 32);
|
||||||
let bps = u32::from(bits_per_sample);
|
|
||||||
|
|
||||||
match output_width {
|
match output_width {
|
||||||
16 => {
|
16 => {
|
||||||
debug_assert!(dst.len() >= num_samples * 2);
|
debug_assert!(dst.len() >= num_samples * 2);
|
||||||
let mut bit_offset: usize = 0;
|
let mut bit_offset: usize = 0;
|
||||||
for i in 0..num_samples {
|
for i in 0..num_samples {
|
||||||
let val = extract_bits_msb(src, bit_offset, bps);
|
let val = if BPS > 0 {
|
||||||
|
extract_bits_msb_const::<BPS>(src, bit_offset)
|
||||||
|
} else {
|
||||||
|
extract_bits_msb_const::<16>(src, bit_offset) // unreachable fallback
|
||||||
|
};
|
||||||
dst[i * 2..i * 2 + 2].copy_from_slice(&(val as u16).to_ne_bytes());
|
dst[i * 2..i * 2 + 2].copy_from_slice(&(val as u16).to_ne_bytes());
|
||||||
bit_offset += bps as usize;
|
bit_offset += if BPS > 0 { BPS as usize } else { 16 };
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
32 => {
|
32 => {
|
||||||
debug_assert!(dst.len() >= num_samples * 4);
|
debug_assert!(dst.len() >= num_samples * 4);
|
||||||
let mut bit_offset: usize = 0;
|
let mut bit_offset: usize = 0;
|
||||||
for i in 0..num_samples {
|
for i in 0..num_samples {
|
||||||
let val = extract_bits_msb(src, bit_offset, bps);
|
let val = if BPS > 0 {
|
||||||
|
extract_bits_msb_const::<BPS>(src, bit_offset)
|
||||||
|
} else {
|
||||||
|
extract_bits_msb_const::<32>(src, bit_offset) // unreachable fallback
|
||||||
|
};
|
||||||
dst[i * 4..i * 4 + 4].copy_from_slice(&val.to_ne_bytes());
|
dst[i * 4..i * 4 + 4].copy_from_slice(&val.to_ne_bytes());
|
||||||
bit_offset += bps as usize;
|
bit_offset += if BPS > 0 { BPS as usize } else { 32 };
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
_ => unreachable!(),
|
_ => unreachable!(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Extract a value of `num_bits` bits starting at `bit_offset` from an MSB-first packed byte
|
/// Extract a value of `BPS` bits starting at `bit_offset` from an MSB-first packed byte stream.
|
||||||
/// stream. Returns the value as a u32 (right-justified, zero-extended).
|
/// Uses a single wide read + shift/mask for branchless extraction.
|
||||||
#[inline]
|
#[inline(always)]
|
||||||
fn extract_bits_msb(src: &[u8], bit_offset: usize, num_bits: u32) -> u32 {
|
fn extract_bits_msb_const<const BPS: u8>(src: &[u8], bit_offset: usize) -> u32 {
|
||||||
debug_assert!(num_bits <= 32);
|
|
||||||
|
|
||||||
let byte_idx = bit_offset / 8;
|
let byte_idx = bit_offset / 8;
|
||||||
let bit_idx = bit_offset % 8;
|
let bit_idx = bit_offset % 8;
|
||||||
|
let mask: u64 = (1u64 << BPS) - 1;
|
||||||
|
|
||||||
// Read enough bytes to cover the value (up to 5 bytes for 32 bits crossing a byte boundary).
|
// Read a u64 from as many bytes as available, big-endian, without going out of bounds.
|
||||||
let mut acc: u64 = 0;
|
// For BPS <= 32 (+ up to 7 bits offset = 39 bits max), we need at most 5 bytes,
|
||||||
let bytes_needed = (bit_idx + num_bits as usize).div_ceil(8);
|
// but reading 8 at once is faster and the branch is only taken near row ends.
|
||||||
for i in 0..bytes_needed {
|
let remaining = src.len() - byte_idx;
|
||||||
acc = (acc << 8) | u64::from(src[byte_idx + i]);
|
let acc = if remaining >= 8 {
|
||||||
}
|
u64::from_be_bytes(src[byte_idx..byte_idx + 8].try_into().unwrap())
|
||||||
|
} else {
|
||||||
|
// Near the end of the buffer — read available bytes into the high bits of a u64.
|
||||||
|
let mut buf = [0u8; 8];
|
||||||
|
buf[..remaining].copy_from_slice(&src[byte_idx..]);
|
||||||
|
u64::from_be_bytes(buf)
|
||||||
|
};
|
||||||
|
|
||||||
// Shift right to align the value to bit 0 and mask.
|
// The value starts at bit_idx from the MSB of our u64. Shift it to the LSB and mask.
|
||||||
let shift = bytes_needed * 8 - bit_idx - num_bits as usize;
|
let shift = 64 - bit_idx - BPS as usize;
|
||||||
((acc >> shift) & ((1u64 << num_bits) - 1)) as u32
|
((acc >> shift) & mask) as u32
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ReadoutLayout {
|
impl ReadoutLayout {
|
||||||
|
|||||||
Reference in new issue
Block a user