From 844cc78ccdf9441f3ae7266cad5d5774c1e6e15a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 18 Sep 2026 15:11:39 -0400 Subject: [PATCH] refactor: make bitpacked CPU kernels consume chunk layouts Signed-off-by: "Matt Katz" --- .../bitpacking/array/bitpack_decompress.rs | 34 ++- .../fastlanes/src/bitpacking/array/mod.rs | 263 ++++++++++++++++-- .../src/bitpacking/array/unpack_iter.rs | 163 +++++++---- .../src/bitpacking/compute/between.rs | 5 +- .../src/bitpacking/compute/compare.rs | 2 +- .../src/bitpacking/compute/compare_fused.rs | 19 +- .../src/bitpacking/compute/filter.rs | 14 +- .../src/bitpacking/compute/is_constant.rs | 5 +- .../fastlanes/src/bitpacking/compute/mod.rs | 2 +- .../bitpacking/compute/stream_predicate.rs | 5 +- .../fastlanes/src/bitpacking/compute/take.rs | 9 +- encodings/fastlanes/src/bitpacking/mod.rs | 2 + .../src/bitpacking/vtable/operations.rs | 6 +- .../fastlanes/src/for/array/for_decompress.rs | 3 +- 14 files changed, 408 insertions(+), 124 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 0684aea5e6a..6911b44fe65 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -17,14 +17,13 @@ use vortex_array::match_each_integer_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::patches::Patches; use vortex_array::scalar::Scalar; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::FL_CHUNK_SIZE; -use crate::unpack_iter::BitPacked as BitPackedUnpack; -use crate::unpack_iter::BitUnpackedChunks; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedUnpack; +use crate::bitpacking::unpack_iter::BitUnpackedChunks; /// Unpacks a bit-packed array into a primitive array. pub fn unpack_array( @@ -124,8 +123,9 @@ where // SAFETY: `decode` writes a value to every slot in this range. let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) }; + let widths = array.chunk_widths(ctx)?; let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; - let mut chunks = array.unpacked_chunks::(&mut scratch)?; + let mut chunks = array.unpacked_chunks::(&widths, &mut scratch)?; decode(&mut chunks, uninit_slice, &map); if let Some(patches) = array.patches() { @@ -164,19 +164,28 @@ pub(crate) fn apply_patches_to_uninit_range, index: usize) -> Scalar { - let bit_width = array.bit_width() as usize; +pub fn unpack_single( + array: ArrayView<'_, BitPacked>, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { let ptype = array.dtype().as_ptype(); - // let packed = array.packed().into_primitive()?; let index_in_encoded = index + array.offset() as usize; + let chunk = index_in_encoded / FL_CHUNK_SIZE; + let index_in_chunk = index_in_encoded % FL_CHUNK_SIZE; + let (range, bit_width) = array.chunk_range(chunk, ctx)?; let scalar: Scalar = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |P| { + let packed_chunk = + &array.packed_slice::

()[range.start / size_of::

()..range.end / size_of::

()]; + // SAFETY: `packed_chunk` is exactly one packed block at `bit_width`, and the index is + // within the chunk. unsafe { - unpack_single_primitive::

(array.packed_slice::

(), bit_width, index_in_encoded) - .into() + BitPacking::unchecked_unpack_single(bit_width as usize, packed_chunk, index_in_chunk) } + .into() }); // Cast to fix signedness and nullability - scalar.cast(array.dtype()).vortex_expect("cast failure") + scalar.cast(array.dtype()) } /// # Safety @@ -227,12 +236,13 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_session::VortexSession; use super::*; use crate::BitPackedArray; use crate::BitPackedData; - use crate::bitpack_compress::bitpack_encode; + use crate::bitpacking::bitpack_compress::bitpack_encode; fn encode(array: &PrimitiveArray, bit_width: u8) -> BitPackedArray { bitpack_encode(array, bit_width, None, &mut SESSION.create_execution_ctx()).unwrap() @@ -259,7 +269,7 @@ mod tests { .iter() .enumerate() .for_each(|(i, v)| { - let scalar: u16 = (&unpack_single(compressed.as_view(), i)) + let scalar: u16 = (&unpack_single(compressed.as_view(), i, &mut ctx).unwrap()) .try_into() .unwrap(); assert_eq!(scalar, *v); diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 05cbee8b3ef..8002320d434 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -3,11 +3,15 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; use std::mem::MaybeUninit; +use std::ops::Range; use fastlanes::BitPacking; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::Primitive; @@ -21,6 +25,8 @@ use vortex_array::patches::Patches; use vortex_array::patches::PatchesData; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -35,6 +41,183 @@ use crate::bitpack_compress::bitpack_encode; use crate::unpack_iter::BitPacked as BitPackedIter; use crate::unpack_iter::BitUnpackedChunks; +/// Bytes occupied by one packed FastLanes chunk of `bit_width`-bit values. +#[inline] +pub const fn chunk_packed_bytes(bit_width: u8) -> usize { + (FL_CHUNK_SIZE / 8) * bit_width as usize +} + +/// Chunk widths and byte offsets used while encoding or executing bit-packed data. +/// Execution borrows the materialized children; only encoding computes prefix sums. +#[derive(Clone, Debug)] +pub struct ChunkWidths { + widths: Widths, + byte_offsets: Buffer, + max_width: u8, +} + +#[derive(Clone, Debug)] +enum Widths { + Uniform { width: u8, len: usize }, + PerChunk(Buffer), +} + +impl ChunkWidths { + /// Build an encoding plan, computing byte offsets from one width per chunk. + pub fn new(widths: Buffer) -> Self { + let mut byte_offsets = BufferMut::::with_capacity(widths.len() + 1); + let mut total = 0u64; + byte_offsets.push(0); + for &width in widths.iter() { + total += chunk_packed_bytes(width) as u64; + byte_offsets.push(total); + } + Self::from_buffers(Widths::PerChunk(widths), byte_offsets.freeze()) + } + + /// `num_chunks` chunks all packed at `bit_width`. + pub fn uniform(bit_width: u8, num_chunks: usize) -> Self { + Self::from_buffers( + Widths::Uniform { + width: bit_width, + len: num_chunks, + }, + Buffer::from_iter((0..=num_chunks).map(|i| (i * chunk_packed_bytes(bit_width)) as u64)), + ) + } + + fn from_buffers(widths: Widths, byte_offsets: Buffer) -> Self { + let max_width = match &widths { + Widths::Uniform { width, .. } => *width, + Widths::PerChunk(widths) => widths.iter().copied().max().unwrap_or(0), + }; + Self { + widths, + byte_offsets, + max_width, + } + } + + /// Number of chunks. + #[inline] + pub fn len(&self) -> usize { + match &self.widths { + Widths::Uniform { len, .. } => *len, + Widths::PerChunk(widths) => widths.len(), + } + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The bit width of `chunk`. + #[inline] + pub fn width(&self, chunk: usize) -> u8 { + match &self.widths { + Widths::Uniform { width, .. } => *width, + Widths::PerChunk(widths) => widths[chunk], + } + } + + /// The widest chunk width. + #[inline] + pub fn max_width(&self) -> u8 { + self.max_width + } + + /// The single width shared by every chunk, if they all agree. + pub fn uniform_width(&self) -> Option { + if self.is_empty() { + return None; + } + match &self.widths { + Widths::Uniform { width, .. } => Some(*width), + Widths::PerChunk(widths) => { + let first = widths[0]; + widths.iter().all(|&w| w == first).then_some(first) + } + } + } + + /// Whether every chunk shares one width. An empty array counts as uniform. + pub fn is_uniform(&self) -> bool { + self.is_empty() || self.uniform_width().is_some() + } + + /// Materialize the widths as one byte per chunk. + pub fn as_buffer(&self) -> Buffer { + match &self.widths { + Widths::Uniform { width, len } => Buffer::from_iter(std::iter::repeat_n(*width, *len)), + Widths::PerChunk(widths) => widths.clone(), + } + } + + /// The offsets child, including the trailing boundary. Slices may start at a nonzero offset. + pub fn offsets_array(&self) -> ArrayRef { + self.byte_offsets.clone().into_array() + } + + /// Byte offset relative to the packed buffer. Passing the chunk count yields the total size. + #[inline] + pub fn byte_offset(&self, chunk: usize) -> usize { + (self.byte_offsets[chunk] - self.byte_offsets[0]) as usize + } + + /// Total packed bytes. + #[inline] + pub fn packed_bytes(&self) -> usize { + self.byte_offset(self.len()) + } + + /// Restrict to chunks without copying or rebasing the offset buffer. + pub fn slice(&self, chunks: Range) -> Self { + let widths = match &self.widths { + Widths::Uniform { width, .. } => Widths::Uniform { + width: *width, + len: chunks.len(), + }, + Widths::PerChunk(widths) => Widths::PerChunk(widths.slice(chunks.clone())), + }; + Self::from_buffers( + widths, + self.byte_offsets.slice(chunks.start..chunks.end + 1), + ) + } +} + +impl PartialEq for ChunkWidths { + fn eq(&self, other: &Self) -> bool { + self.len() == other.len() && (0..self.len()).all(|i| self.width(i) == other.width(i)) + } +} + +impl Eq for ChunkWidths {} + +impl Hash for ChunkWidths { + fn hash(&self, state: &mut H) { + self.len().hash(state); + for i in 0..self.len() { + self.width(i).hash(state); + } + } +} + +impl Display for ChunkWidths { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self.uniform_width() { + Some(w) => write!(f, "bit_width: {w}"), + None => write!( + f, + "bit_widths: {} chunks, max {}", + self.len(), + self.max_width + ), + } + } +} + #[array_slots(crate::BitPacked)] pub struct BitPackedSlots { /// The indices of exception values that don't fit in the bit-packed representation. @@ -224,27 +407,28 @@ impl BitPackedData { unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) } } - /// Accessor for bit unpacked chunks - pub fn unpacked_chunks<'a, T: BitPackedIter>( - &'a self, - dtype: &DType, - len: usize, - scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], - ) -> VortexResult> { - assert_eq!( - T::PTYPE, - self.ptype(dtype), - "Requested type doesn't match the array ptype" - ); - BitUnpackedChunks::try_new(self, len, scratch) - } - /// Bit-width of the packed values #[inline] pub fn bit_width(&self) -> u8 { self.bit_width } + /// Access a chunk using the layout prepared for this operation. + #[inline] + pub(crate) fn packed_chunk( + &self, + widths: &ChunkWidths, + chunk: usize, + ) -> (&[T], usize) { + let bit_width = widths.width(chunk); + let start = widths.byte_offset(chunk) / size_of::(); + let len = chunk_packed_bytes(bit_width) / size_of::(); + ( + &self.packed_slice::()[start..][..len], + bit_width as usize, + ) + } + #[inline] pub fn offset(&self) -> u16 { self.offset @@ -283,6 +467,24 @@ impl BitPackedData { } pub trait BitPackedArrayExt: BitPackedArraySlotsExt { + /// Prepare a chunk layout for a bulk operation. + fn chunk_widths(&self, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(ChunkWidths::uniform( + self.bit_width(), + (self.as_ref().len() + self.offset() as usize).div_ceil(FL_CHUNK_SIZE), + )) + } + + /// Locate one chunk without allocating a layout for the entire array. + fn chunk_range( + &self, + chunk: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult<(Range, u8)> { + let len = chunk_packed_bytes(self.bit_width()); + Ok((chunk * len..(chunk + 1) * len, self.bit_width())) + } + #[inline] fn packed(&self) -> &BufferHandle { BitPackedData::packed(self) @@ -321,14 +523,14 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { #[inline] fn unpacked_chunks<'a, T: BitPackedIter>( &'a self, + widths: &'a ChunkWidths, scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], ) -> VortexResult> { - BitPackedData::unpacked_chunks::( - self, - self.as_ref().dtype(), - self.as_ref().len(), - scratch, - ) + vortex_ensure!( + T::PTYPE == self.as_ref().dtype().as_ptype(), + "Requested unpack type does not match array dtype" + ); + BitUnpackedChunks::try_new(self, self.as_ref().len(), widths, scratch) } } @@ -343,8 +545,10 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_buffer::Buffer; + use vortex_buffer::buffer; use vortex_session::VortexSession; + use super::ChunkWidths; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; @@ -407,4 +611,21 @@ mod test { &mut ctx ); } + #[test] + fn chunk_widths_offsets() { + assert_eq!(ChunkWidths::uniform(3, 3).uniform_width(), Some(3)); + assert_eq!(ChunkWidths::new(Buffer::::empty()).packed_bytes(), 0); + + let widths = ChunkWidths::new(buffer![3u8, 0, 16]); + assert_eq!(widths.uniform_width(), None); + assert_eq!(widths.len(), 3); + assert_eq!(widths.max_width(), 16); + assert_eq!(widths.width(1), 0); + assert_eq!(widths.byte_offset(0), 0); + assert_eq!(widths.byte_offset(1), 128 * 3); + assert_eq!(widths.byte_offset(2), 128 * 3); + assert_eq!(widths.packed_bytes(), 128 * 19); + assert_eq!(widths.slice(1..3), ChunkWidths::new(buffer![0u8, 16])); + assert_eq!(widths.slice(0..1).uniform_width(), Some(3)); + } } diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 4877fa9c57f..0f9ae7d6e58 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -16,6 +16,8 @@ use vortex_error::vortex_ensure; use crate::BitPackedData; use crate::FL_CHUNK_SIZE; +use crate::bitpacking::array::ChunkWidths; +use crate::bitpacking::array::chunk_packed_bytes; const CHUNK_SIZE: usize = FL_CHUNK_SIZE; @@ -48,6 +50,16 @@ impl> UnpackStrategy for BitPackingStr } } +/// The packed FastLanes block of `chunk` and its bit width. +#[allow(clippy::inline_always)] +#[inline(always)] +fn packed_chunk<'a, P>(packed: &'a [P], widths: &ChunkWidths, chunk: usize) -> (&'a [P], usize) { + let bit_width = widths.width(chunk); + let start = widths.byte_offset(chunk) / size_of::

(); + let len = chunk_packed_bytes(bit_width) / size_of::

(); + (&packed[start..][..len], bit_width as usize) +} + /// Accessor to unpacked chunks of bitpacked arrays /// /// The usual pattern of usage should follow @@ -63,23 +75,22 @@ impl> UnpackStrategy for BitPackingStr /// use vortex_buffer::buffer; /// use vortex_fastlanes::BitPackedData; /// use vortex_fastlanes::BitPackedArrayExt; +/// use vortex_fastlanes::FL_CHUNK_SIZE; /// use vortex_fastlanes::unpack_iter::BitUnpackedChunks; /// /// let mut ctx = vortex_array::array_session().create_execution_ctx(); /// let array = BitPackedData::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).unwrap(); -/// let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; -/// let mut unpacked_chunks: BitUnpackedChunks = -/// array.unpacked_chunks(&mut scratch).unwrap(); +/// let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; +/// let widths = array.chunk_widths(&mut ctx).unwrap(); +/// let mut unpacked_chunks: BitUnpackedChunks = array.unpacked_chunks(&widths, &mut scratch).unwrap(); /// /// if let Some(header) = unpacked_chunks.initial() { /// // handle partial initial chunk /// } /// -/// { -/// let mut chunks_iter = unpacked_chunks.full_chunks(); -/// while let Some(chunk) = chunks_iter.next() { -/// // handle full bitpacked chunks of 1024 elements -/// } +/// let mut chunks_iter = unpacked_chunks.full_chunks(); +/// while let Some(chunk) = chunks_iter.next() { +/// // handle full bitpacked chunks of 1024 elements /// } /// /// if let Some(trailer) = unpacked_chunks.trailer() { @@ -88,7 +99,7 @@ impl> UnpackStrategy for BitPackingStr /// ``` pub struct UnpackedChunks<'a, T: PhysicalPType, S: UnpackStrategy> { strategy: S, - bit_width: usize, + widths: &'a ChunkWidths, offset: usize, len: usize, num_chunks: usize, @@ -104,12 +115,13 @@ impl<'a, T: BitPacked> BitUnpackedChunks<'a, T> { pub fn try_new( array: &'a BitPackedData, len: usize, + widths: &'a ChunkWidths, scratch: &'a mut [MaybeUninit; CHUNK_SIZE], ) -> VortexResult { Self::try_new_with_strategy( BitPackingStrategy, array.packed_slice::(), - array.bit_width() as usize, + widths, array.offset() as usize, len, scratch, @@ -117,14 +129,12 @@ impl<'a, T: BitPacked> BitUnpackedChunks<'a, T> { } pub fn full_chunks(&mut self) -> BitUnpackIterator<'_, T> { - let elems_per_chunk = self.elems_per_chunk(); let last_chunk_is_sliced = self.last_chunk_is_sliced() as usize; let first_chunk_is_sliced = self.first_chunk_is_sliced(); BitUnpackIterator::new( self.packed, + self.widths, self.scratch, - self.bit_width, - elems_per_chunk, self.num_chunks - last_chunk_is_sliced, first_chunk_is_sliced, ) @@ -135,16 +145,16 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { pub fn try_new_with_strategy( strategy: S, packed: &'a [T::Physical], - bit_width: usize, + widths: &'a ChunkWidths, offset: usize, len: usize, scratch: &'a mut [MaybeUninit; CHUNK_SIZE], ) -> VortexResult { let (num_chunks, last_chunk_length) = - validate_packed::(packed.len(), bit_width, offset, len)?; + validate_packed::(packed.len(), widths, offset, len)?; Ok(Self { strategy, - bit_width, + widths, offset, len, num_chunks, @@ -156,14 +166,14 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { #[allow(clippy::inline_always)] #[inline(always)] - const fn elems_per_chunk(&self) -> usize { - 128 * self.bit_width / size_of::() + fn chunk(&self, chunk: usize) -> (&'a [T::Physical], usize) { + packed_chunk(self.packed, self.widths, chunk) } /// Access first chunk of the array if the last chunk has fewer than 1024 due to slicing pub fn initial(&mut self) -> Option<&mut [T]> { (self.first_chunk_is_sliced() || self.num_chunks == 1).then(|| { - let chunk: &[T::Physical] = &self.packed[..self.elems_per_chunk()]; + let (chunk, bit_width) = self.chunk(0); let dst: &mut [MaybeUninit] = self.scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; @@ -173,10 +183,10 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { CHUNK_SIZE - self.offset }; // SAFETY: - // 1. chunk is elems_per_chunk. + // 1. chunk holds exactly one packed block at bit_width. // 2. buffer is exactly CHUNK_SIZE. unsafe { - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width, chunk, dst); mem::transmute(&mut self.scratch[self.offset..][..header_end_slice]) } }) @@ -230,13 +240,18 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { } if self.num_chunks > 1 { - let packed_slice = self.packed; - let elems_per_chunk = self.elems_per_chunk(); - for i in self.full_chunks_range() { - let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; + let range = self.full_chunks_range(); + let packed = self.packed; + let widths: &'a ChunkWidths = self.widths; + let mut start = widths.byte_offset(range.start) / size_of::(); + for chunk_idx in range { + let bit_width = widths.width(chunk_idx); + let len = chunk_packed_bytes(bit_width) / size_of::(); + let chunk = &packed[start..start + len]; + start += len; unsafe { let dst: &mut [T::Physical] = mem::transmute(&mut self.scratch[..]); - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width as usize, chunk, dst); let unpacked: &mut [T] = mem::transmute(&mut self.scratch[..]); f(unpacked, local_idx..local_idx + CHUNK_SIZE); } @@ -265,16 +280,19 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { let mut local_idx = start_idx; - let packed_slice = self.packed; - let elems_per_chunk = self.elems_per_chunk(); - for i in self.full_chunks_range() { - let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; + let range = self.full_chunks_range(); + let mut start = self.widths.byte_offset(range.start) / size_of::(); + for chunk_idx in range { + let bit_width = self.widths.width(chunk_idx); + let len = chunk_packed_bytes(bit_width) / size_of::(); + let chunk = &self.packed[start..start + len]; + start += len; unsafe { let uninit_dst = &mut output[local_idx..local_idx + CHUNK_SIZE]; // SAFETY: &[T] and &[MaybeUninit] have the same layout. let dst: &mut [T::Physical] = mem::transmute(uninit_dst); - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width as usize, chunk, dst); } local_idx += CHUNK_SIZE; } @@ -289,15 +307,14 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { /// Access last chunk of the array if the last chunk has fewer than 1024 due to slicing pub fn trailer(&mut self) -> Option<&mut [T]> { (self.last_chunk_is_sliced() && self.num_chunks > 1).then(|| { - let chunk: &[T::Physical] = &self.packed - [(self.num_chunks - 1) * self.elems_per_chunk()..][..self.elems_per_chunk()]; + let (chunk, bit_width) = self.chunk(self.num_chunks - 1); let dst: &mut [MaybeUninit] = self.scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; // SAFETY: - // 1. chunk is elems_per_chunk. + // 1. chunk holds exactly one packed block at bit_width. // 2. buffer is exactly CHUNK_SIZE. unsafe { - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width, chunk, dst); mem::transmute(&mut self.scratch[..self.last_chunk_length]) } }) @@ -312,33 +329,49 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { } } -/// Walk every packed chunk in array order without allocating an unpack scratch buffer. +/// Walk every *packed* chunk in array order, yielding the raw packed FastLanes block, its bit +/// width, and the padded bit range it covers, without unpacking it. +/// +/// Unlike [`UnpackedChunks::for_each_unpacked_chunk`], this does not fill a scratch buffer: it +/// hands the still-packed block to the callback so fused kernels (e.g. compare) can unpack and +/// consume it in a single pass. +/// +/// The yielded range is in *padded* coordinates: block `c` covers +/// `[c * 1024, min((c + 1) * 1024, offset + len))`, so it includes the leading `offset` rows +/// that slicing skips. Block starts are therefore always 1024-aligned regardless of `offset`. +/// Callers must account for the array's `offset` when mapping a block's rows back to logical +/// output positions (e.g. by viewing the output buffer at a bit offset of `offset`). pub(crate) fn for_each_packed_chunk( packed: &[T::Physical], - bit_width: usize, + widths: &ChunkWidths, offset: usize, len: usize, mut f: F, ) -> VortexResult<()> where T: PhysicalPType, - F: FnMut(&[T::Physical], Range), + F: FnMut(&[T::Physical], usize, Range), { - let (num_chunks, _) = validate_packed::(packed.len(), bit_width, offset, len)?; - let elems_per_chunk = 128 * bit_width / size_of::(); + validate_packed::(packed.len(), widths, offset, len)?; let padded_len = offset + len; - for chunk in 0..num_chunks { - let packed_chunk = &packed[chunk * elems_per_chunk..][..elems_per_chunk]; - let start = chunk * CHUNK_SIZE; - let end = (start + CHUNK_SIZE).min(padded_len); - f(packed_chunk, start..end); + let mut start = 0; + for chunk in 0..widths.len() { + let bit_width = widths.width(chunk); + let packed_len = chunk_packed_bytes(bit_width) / size_of::(); + let packed_chunk = &packed[start..start + packed_len]; + start += packed_len; + let row_start = chunk * CHUNK_SIZE; + let row_end = (row_start + CHUNK_SIZE).min(padded_len); + f(packed_chunk, bit_width as usize, row_start..row_end); } Ok(()) } +/// Check that `packed_len` words of `T::Physical` hold exactly the chunks described by `widths` +/// for `offset + len` padded elements, returning the chunk count and the trailing chunk's length. fn validate_packed( packed_len: usize, - bit_width: usize, + widths: &ChunkWidths, offset: usize, len: usize, ) -> VortexResult<(usize, usize)> { @@ -346,12 +379,20 @@ fn validate_packed( offset < CHUNK_SIZE, "Invalid bit-packed offset {offset}, expected < {CHUNK_SIZE}" ); - let elems_per_chunk = 128 * bit_width / size_of::(); let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); vortex_ensure!( - packed_len == num_chunks * elems_per_chunk, - "Invalid packed length: got {packed_len}, expected {}", - num_chunks * elems_per_chunk + widths.len() == num_chunks, + "Invalid chunk widths: got {}, expected {num_chunks}", + widths.len() + ); + vortex_ensure!( + widths.max_width() as usize <= size_of::() * 8, + "Chunk width exceeds unpacked type width" + ); + let expected = widths.packed_bytes() / size_of::(); + vortex_ensure!( + packed_len == expected, + "Invalid packed length: got {packed_len}, expected {expected}" ); Ok((num_chunks, (offset + len) % CHUNK_SIZE)) } @@ -359,29 +400,30 @@ fn validate_packed( /// Iterator over full chunks of bitpacked array that yields unpacked chunks one at a time pub struct BitUnpackIterator<'a, T: BitPacked + 'a> { packed: &'a [T::Physical], + widths: &'a ChunkWidths, buffer: &'a mut [MaybeUninit; CHUNK_SIZE], - bit_width: usize, - elems_per_chunk: usize, num_chunks: usize, idx: usize, + /// Word offset of chunk `idx` within `packed`. + start: usize, } impl<'a, T: BitPacked> BitUnpackIterator<'a, T> { pub fn new( packed: &'a [T::Physical], + widths: &'a ChunkWidths, buffer: &'a mut [MaybeUninit; CHUNK_SIZE], - bit_width: usize, - elems_per_chunk: usize, num_chunks: usize, first_chunk_is_sliced: bool, ) -> Self { + let idx = if first_chunk_is_sliced { 1 } else { 0 }; Self { packed, + widths, buffer, - bit_width, - elems_per_chunk, num_chunks, - idx: if first_chunk_is_sliced { 1 } else { 0 }, + idx, + start: widths.byte_offset(idx) / size_of::(), } } } @@ -398,15 +440,18 @@ impl<'a, T: BitPacked + 'a> LendingIterator for BitUnpackIterator<'a, T> { return None; } - let chunk = &self.packed[self.idx * self.elems_per_chunk..][..self.elems_per_chunk]; + let bit_width = self.widths.width(self.idx); + let len = chunk_packed_bytes(bit_width) / size_of::(); + let chunk = &self.packed[self.start..self.start + len]; let dst: &mut [MaybeUninit] = self.buffer; unsafe { let dst: &mut [T::Physical] = mem::transmute(dst); - BitPacking::unchecked_unpack(self.bit_width, chunk, dst); + BitPacking::unchecked_unpack(bit_width as usize, chunk, dst); } self.idx += 1; + self.start += len; // SAFETY: The buffer has the appropriate lifetime, the iterator signature doesn't account for it Some(unsafe { mem::transmute::<&mut [MaybeUninit; 1024], &mut [T; 1024]>(self.buffer) }) } diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index 8010fa208d2..97d3444b18c 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -23,6 +23,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::bitpacking::compute::stream_predicate::stream_predicate; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; impl BetweenKernel for BitPacked { fn between( @@ -74,7 +75,7 @@ fn between_constant_typed( ctx: &mut ExecutionCtx, ) -> VortexResult where - T: NativePType + Copy + crate::unpack_iter::BitPacked, + T: NativePType + Copy + BitPackedIter, { // Branch on strictness once at the top so each call into `between_impl` monomorphises // a single tight predicate — same shape as `Primitive::between` in `vortex-array`. @@ -128,7 +129,7 @@ fn between_impl( ctx: &mut ExecutionCtx, ) -> VortexResult where - T: NativePType + Copy + crate::unpack_iter::BitPacked, + T: NativePType + Copy + BitPackedIter, Lo: Fn(T, T) -> bool, Up: Fn(T, T) -> bool, { diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index c9d6b815b0d..807b6f21b8f 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -28,7 +28,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::bitpacking::compute::compare_fused::stream_compare_fused; -use crate::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; impl CompareKernel for BitPacked { fn compare( diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index 1259ed815fe..8296bf723b1 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -11,7 +11,7 @@ //! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`transpose_bits`] per block then rotates //! that mask into logical row order. //! -//! The packed blocks are walked through [`crate::unpack_iter::for_each_packed_chunk`], so chunk +//! The packed blocks are walked through [`crate::bitpacking::unpack_iter::for_each_packed_chunk`], so chunk //! sizing and bounds live in one place without allocating an unpack scratch buffer. //! //! Slicing is handled by working in *padded* coordinates: bit `offset + i` holds element `i`. The @@ -48,8 +48,8 @@ use vortex_error::VortexResult; use super::stream_predicate::stream_predicate; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::unpack_iter::BitPacked as BitPackedIter; -use crate::unpack_iter::for_each_packed_chunk; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::for_each_packed_chunk; const CHUNK_SIZE: usize = 1024; const U64_BITS: usize = u64::BITS as usize; @@ -78,12 +78,12 @@ where F: Fn(T, T) -> bool + Copy, { let len = array.len(); - let bit_width = array.bit_width() as usize; + let widths = array.chunk_widths(ctx)?; let offset = array.offset() as usize; // A degenerate width has no packed payload for the fused kernel to consume; defer to the scalar // streaming predicate, which handles every layout (including the empty array). - if len == 0 || bit_width == 0 { + if len == 0 || widths.max_width() == 0 { return stream_predicate::(array, nullability, move |v| cmp(v, rhs), ctx); } @@ -97,14 +97,19 @@ where let mut lane_major = [0u64; WORDS_PER_CHUNK]; for_each_packed_chunk::( array.packed_slice::<::Physical>(), - bit_width, + &widths, offset, len, - |packed_chunk, range| { + |packed_chunk, bit_width, range| { // Block starts are always 1024-aligned (padded coords), so the slot is a full block. let out = words[range.start / U64_BITS..] .first_chunk_mut::() .vortex_expect("over-allocated buffer holds a full block per chunk"); + // A zero-width chunk has no packed payload: every value in it is zero. + if bit_width == 0 { + out.fill(if cmp(T::default(), rhs) { u64::MAX } else { 0 }); + return; + } // SAFETY: `packed_chunk` holds exactly `128 * bit_width / size_of::()` packed // elements and `bit_width <= U::T`, satisfying `unchecked_unpack_cmp`'s contract. The // kernel assigns every word in `transposed`, so its previous contents are irrelevant. diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 0b1b9422f86..5d7c9a4e41f 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -26,6 +26,7 @@ use super::take::UNPACK_CHUNK_THRESHOLD; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitPackedData; +use crate::ChunkWidths; /// The threshold over which it is faster to fully unpack the entire [`BitPackedArray`](crate::BitPackedArray) and then /// filter the result than to unpack only specific bitpacked values into the output buffer. @@ -65,7 +66,7 @@ impl FilterKernel for BitPacked { // Filter and patch using the correct unsigned type for FastLanes, then cast to signed if needed. let primitive = match_each_unsigned_integer_ptype!(array.dtype().as_ptype().to_unsigned(), |U| { - let (buffer, validity) = filter_primitive_without_patches::(array, values)?; + let (buffer, validity) = filter_primitive_without_patches::(array, values, ctx)?; // reinterpret_cast for signed types. let primitive = PrimitiveArray::new(buffer, validity); if array.dtype().as_ptype().is_signed_int() { @@ -109,8 +110,10 @@ impl FilterKernel for BitPacked { fn filter_primitive_without_patches( array: ArrayView<'_, BitPacked>, selection: &MaskValuesRef, + ctx: &mut ExecutionCtx, ) -> VortexResult<(Buffer, Validity)> { - let values = filter_with_indices(array.data(), selection.indices()); + let widths = array.chunk_widths(ctx)?; + let values = filter_with_indices(array.data(), &widths, selection.indices()); let validity = array .validity()? .filter(&Mask::Values(MaskValuesRef::clone(selection)))?; @@ -120,24 +123,21 @@ fn filter_primitive_without_patches( fn filter_with_indices( array: &BitPackedData, + widths: &ChunkWidths, indices: &[usize], ) -> BufferMut { let offset = array.offset() as usize; - let bit_width = array.bit_width() as usize; let mut values = BufferMut::with_capacity(indices.len()); // Some re-usable memory to store per-chunk indices. let mut unpacked = [const { MaybeUninit::::uninit() }; 1024]; - let packed_bytes = array.packed_slice::(); // Group the indices by the FastLanes chunk they belong to. - let chunk_size = 128 * bit_width / size_of::(); - chunked_indices( indices.iter().copied(), offset, |chunk_idx, indices_within_chunk| { - let packed = &packed_bytes[chunk_idx * chunk_size..][..chunk_size]; + let (packed, bit_width) = array.packed_chunk::(widths, chunk_idx); if indices_within_chunk.len() == 1024 { // Unpack the entire chunk. diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 0ab01a635ba..2250f1c47ab 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -23,7 +23,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::unpack_iter::BitPacked as BitPackedUnpack; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedUnpack; /// BitPacked-specific is_constant kernel with SIMD support. #[derive(Debug)] @@ -56,8 +56,9 @@ fn bitpacked_is_constant( array: ArrayView<'_, BitPacked>, ctx: &mut ExecutionCtx, ) -> VortexResult { + let widths = array.chunk_widths(ctx)?; let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; - let mut bit_unpack_iterator = array.unpacked_chunks::(&mut scratch)?; + let mut bit_unpack_iterator = array.unpacked_chunks::(&widths, &mut scratch)?; let patches = array .patches() .map(|p| -> VortexResult<_> { diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 38f86f781bb..4a06d67b9d2 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -53,7 +53,7 @@ mod tests { use vortex_array::compute::conformance::consistency::test_array_consistency; use crate::BitPackedArray; - use crate::bitpack_compress::bitpack_encode; + use crate::bitpacking::bitpack_compress::bitpack_encode; use crate::bitpacking::compute::chunked_indices; fn bp(array: &PrimitiveArray, bit_width: u8) -> BitPackedArray { diff --git a/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs b/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs index 9154ca736c1..1943ab7b43b 100644 --- a/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs +++ b/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs @@ -35,7 +35,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::FL_CHUNK_SIZE; -use crate::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; /// Stream `predicate` over the unpacked values of a [`BitPackedArray`](crate::BitPackedArray), one FastLanes /// block at a time, producing a [`BoolArray`]. @@ -53,8 +53,9 @@ where let mut words: BufferMut = BufferMut::zeroed(len.div_ceil(u64::BITS as usize)); if len > 0 { + let widths = array.chunk_widths(ctx)?; let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; - let mut chunks = array.unpacked_chunks::(&mut scratch)?; + let mut chunks = array.unpacked_chunks::(&widths, &mut scratch)?; let words = words.as_mut_slice(); if let Some(p) = array.patches() { diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 86e97623cf6..e916db2d100 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -25,7 +25,7 @@ use vortex_error::VortexResult; use super::chunked_indices; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::bitpack_decompress; +use crate::bitpacking::bitpack_decompress; // TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. /// assuming the buffer is already allocated (which will happen at most once) then unpacking @@ -80,10 +80,8 @@ fn take_primitive( return Ok(PrimitiveArray::new(Buffer::::empty(), taken_validity)); } + let widths = array.chunk_widths(ctx)?; let offset = array.offset() as usize; - let bit_width = array.bit_width() as usize; - - let packed = array.packed_slice::(); // Group indices by 1024-element chunk, *without* allocating on the heap let indices_iter = indices.as_slice::().iter().map(|i| { @@ -93,10 +91,9 @@ fn take_primitive( let mut output = BufferMut::::with_capacity(indices.len()); let mut unpacked = [const { MaybeUninit::uninit() }; 1024]; - let chunk_len = 128 * bit_width / size_of::(); chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { - let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; + let (packed, bit_width) = array.data().packed_chunk::(&widths, chunk_idx); let mut have_unpacked = false; let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index cf556df2780..ba87ae92f1b 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -7,8 +7,10 @@ pub use array::BitPackedArraySlotsExt; pub use array::BitPackedData; pub use array::BitPackedDataParts; pub use array::BitPackedSlots; +pub use array::ChunkWidths; pub use array::bitpack_compress; pub use array::bitpack_decompress; +pub use array::chunk_packed_bytes; pub use array::unpack_iter; pub(crate) mod compute; diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 2816407ac03..0fe41c2ee3a 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -8,15 +8,15 @@ use vortex_array::vtable::OperationsVTable; use vortex_error::VortexResult; use crate::BitPacked; -use crate::bitpack_decompress; use crate::bitpacking::array::BitPackedArrayExt; +use crate::bitpacking::bitpack_decompress; impl OperationsVTable for BitPacked { type ProbeState = (); fn scalar_at( array: ArrayView<'_, BitPacked>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { Ok( if let Some(patches) = array.patches() @@ -24,7 +24,7 @@ impl OperationsVTable for BitPacked { { patch } else { - bitpack_decompress::unpack_single(array, index) + bitpack_decompress::unpack_single(array, index, ctx)? }, ) } diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index e41eb27a5be..fca8198ce1e 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -96,6 +96,7 @@ pub(crate) fn fused_decompress< .as_::() .vortex_expect("cannot be null"); + let widths = bp.chunk_widths(ctx)?; let strategy = FoRStrategy { reference: ref_ }; let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; @@ -103,7 +104,7 @@ pub(crate) fn fused_decompress< let mut unpacked = UnpackedChunks::try_new_with_strategy( strategy, bp.packed_slice::(), - bp.bit_width() as usize, + &widths, bp.offset() as usize, bp.len(), &mut scratch,