Skip to content

refactor: make bitpacked CPU kernels consume stored chunk layouts #9938

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mhk197 wants to merge 1 commit into from
Closed

refactor: make bitpacked CPU kernels consume stored chunk layouts #9938

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 22 additions & 12 deletions encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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_layout(ctx)?;
let mut scratch = [const { MaybeUninit::<F>::uninit() }; FL_CHUNK_SIZE];
let mut chunks = array.unpacked_chunks::<F>(&mut scratch)?;
let mut chunks = array.unpacked_chunks::<F>(&widths, &mut scratch)?;
decode(&mut chunks, uninit_slice, &map);

if let Some(patches) = array.patches() {
Expand Down Expand Up @@ -164,19 +164,28 @@ pub(crate) fn apply_patches_to_uninit_range<S: NativePType, T: NativePType, F: F
Ok(())
}

pub fn unpack_single(array: ArrayView<'_, BitPacked>, index: usize) -> Scalar {
let bit_width = array.bit_width() as usize;
pub fn unpack_single(
array: ArrayView<'_, BitPacked>,
index: usize,
ctx: &mut ExecutionCtx,
) -> VortexResult<Scalar> {
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::<P>()[range.start / size_of::<P>()..range.end / size_of::<P>()];
// SAFETY: `packed_chunk` is exactly one packed block at `bit_width`, and the index is
// within the chunk.
unsafe {
unpack_single_primitive::<P>(array.packed_slice::<P>(), 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
Expand Down Expand Up @@ -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()
Expand All @@ -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);
Expand Down
148 changes: 93 additions & 55 deletions encodings/fastlanes/src/bitpacking/array/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,6 @@ pub(crate) fn materialized_layout(offsets: &ArrayRef) -> VortexResult<Option<Chu

pub struct BitPackedDataParts {
pub offset: u16,
pub bit_width: u8,
pub chunk_offsets: ArrayRef,
pub len: usize,
pub packed: BufferHandle,
Expand All @@ -263,15 +262,14 @@ pub struct BitPackedData {
/// The offset within the first block (created with a slice).
/// 0 <= offset < 1024
pub(super) offset: u16,
pub(super) bit_width: u8,
pub(super) packed: BufferHandle,
/// Patch metadata for reconstructing Patches from slots.
pub(super) patches_data: Option<PatchesData>,
}

impl Display for BitPackedData {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "bit_width: {}, offset: {}", self.bit_width, self.offset)
write!(f, "offset: {}", self.offset)
}
}

Expand Down Expand Up @@ -302,25 +300,23 @@ impl BitPackedData {
/// * Any patches must have any `array_len` equal to `length`
/// * The offsets child must hold `num_chunks + 1` non-nullable `u64` byte boundaries.
///
/// Offset differences must still imply the scalar `bit_width`, and the child must be
/// materialized until the kernels are migrated. The packed buffer must match its span.
/// Adjacent boundaries must differ by a multiple of 128 bytes, implying a width no greater
/// than `ptype`. The final boundary minus the origin must equal the packed buffer size.
/// Compressed offsets are checked at execution time, before unpacking.
///
/// Any violation of these preconditions will result in an error.
pub fn try_new(
packed: BufferHandle,
patches: Option<Patches>,
bit_width: u8,
offset: u16,
) -> VortexResult<Self> {
vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}");
vortex_ensure!(
(offset as usize) < FL_CHUNK_SIZE,
"Offset must be less than the full block i.e., {FL_CHUNK_SIZE}, got {offset}"
);

Ok(Self {
offset,
bit_width,
packed,
patches_data: patches.as_ref().map(PatchesData::from_patches),
})
Expand Down Expand Up @@ -359,15 +355,10 @@ impl BitPackedData {
num_chunks + 1,
offsets.len()
);
let widths = materialized_layout(offsets)?.ok_or_else(|| {
vortex_err!("BitPacked chunk layout must be materialized while kernels use bit_width")
})?;
vortex_ensure!(
widths.is_empty() || widths.uniform_width() == Some(self.bit_width),
"Chunk offsets must imply bit_width {}",
self.bit_width
);
Self::validate_layout(&self.packed, ptype, &widths)?;
// Compressed offsets are checked once materialized, before any unchecked unpacking.
if let Some(widths) = materialized_layout(offsets)? {
Self::validate_layout(&self.packed, ptype, &widths)?;
}
Ok(())
}

Expand Down Expand Up @@ -432,25 +423,20 @@ 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<T>; FL_CHUNK_SIZE],
) -> VortexResult<BitUnpackedChunks<'a, T>> {
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
/// The packed FastLanes block of `chunk` as `T` words, along with that chunk's bit width.
#[inline]
pub fn bit_width(&self) -> u8 {
self.bit_width
pub(crate) fn packed_chunk<T: NativePType + BitPacking>(
&self,
widths: &ChunkLayout,
chunk: usize,
) -> (&[T], usize) {
let bit_width = widths.width(chunk);
let start = widths.byte_offset(chunk) / size_of::<T>();
let len = chunk_packed_bytes(bit_width) / size_of::<T>();
(
&self.packed_slice::<T>()[start..][..len],
bit_width as usize,
)
}

#[inline]
Expand Down Expand Up @@ -480,14 +466,6 @@ impl BitPackedData {
.map_err(|a| vortex_err!(InvalidArgument: "Bitpacking can only encode primitive arrays, got {}", a.encoding_id()))?;
bitpack_encode(&parray, bit_width, None, ctx)
}

/// Calculate the maximum value that **can** be contained by this array, given its bit-width.
///
/// Note that this value need not actually be present in the array.
#[inline]
pub fn max_packed_value(&self) -> usize {
(1 << self.bit_width()) - 1
}
}

pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
Expand All @@ -496,11 +474,6 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
BitPackedData::packed(self)
}

#[inline]
fn bit_width(&self) -> u8 {
BitPackedData::bit_width(self)
}

/// Materialize and validate offsets once for bulk access, without computing prefix sums.
fn chunk_layout(&self, ctx: &mut ExecutionCtx) -> VortexResult<ChunkLayout> {
let layout = match materialized_layout(self.chunk_offsets())? {
Expand All @@ -516,6 +489,71 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
Ok(layout)
}

/// Read and validate offsets only when the child is already materialized.
fn materialized_chunk_layout(&self) -> VortexResult<Option<ChunkLayout>> {
let widths = materialized_layout(self.chunk_offsets())?;
if let Some(widths) = &widths {
BitPackedData::validate_layout(
self.packed(),
self.as_ref().dtype().as_ptype(),
widths,
)?;
}
Ok(widths)
}

/// Read one byte boundary without executing the entire offsets child.
fn chunk_byte_offset(&self, boundary: usize, ctx: &mut ExecutionCtx) -> VortexResult<u64> {
vortex_ensure!(
boundary < self.chunk_offsets().len(),
"Chunk boundary out of bounds"
);
if let Some(offsets) = self
.chunk_offsets()
.as_opt::<Primitive>()
.filter(|a| a.buffer_handle().is_on_host())
{
Ok(offsets.as_slice::<u64>()[boundary])
} else {
u64::try_from(&self.chunk_offsets().execute_scalar(boundary, ctx)?)
}
}

/// Locate and validate one chunk using scalar child access, without materializing the offsets.
fn chunk_range(
&self,
chunk: usize,
ctx: &mut ExecutionCtx,
) -> VortexResult<(Range<usize>, u8)> {
vortex_ensure!(
chunk < self.chunk_offsets().len() - 1,
"Chunk index out of bounds"
);
let base = self.chunk_byte_offset(0, ctx)?;
let start = if chunk == 0 {
base
} else {
self.chunk_byte_offset(chunk, ctx)?
};
let end = self.chunk_byte_offset(chunk + 1, ctx)?;
let width = width_from_offsets(start, end)?;
vortex_ensure!(
width as usize <= self.as_ref().dtype().as_ptype().bit_width(),
"Unsupported bit width {width}"
);
let start = start
.checked_sub(base)
.ok_or_else(|| vortex_err!("Chunk offset precedes buffer origin"))?;
let end = end
.checked_sub(base)
.ok_or_else(|| vortex_err!("Chunk offset precedes buffer origin"))?;
vortex_ensure!(
start % (FL_CHUNK_SIZE / 8) as u64 == 0 && end <= self.packed().len() as u64,
"Chunk offsets are unaligned or exceed the packed buffer"
);
Ok((usize::try_from(start)?..usize::try_from(end)?, width))
}

#[inline]
fn offset(&self) -> u16 {
BitPackedData::offset(self)
Expand All @@ -541,17 +579,17 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt {
BitPackedData::packed_slice::<T>(self)
}

#[inline]
/// Iterate packed chunks using boundaries materialized for this operation.
fn unpacked_chunks<'a, T: BitPackedIter>(
&'a self,
widths: &'a ChunkLayout,
scratch: &'a mut [MaybeUninit<T>; FL_CHUNK_SIZE],
) -> VortexResult<BitUnpackedChunks<'a, T>> {
BitPackedData::unpacked_chunks::<T>(
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)
}
}

Expand Down
Loading
Loading