diff --git a/Cargo.lock b/Cargo.lock index 2448cb9aff8..68f2696d2f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11007,6 +11007,7 @@ dependencies = [ "prost 0.14.4", "rand 0.10.2", "rstest", + "seq-macro", "vortex-alp", "vortex-array", "vortex-buffer", diff --git a/Cargo.toml b/Cargo.toml index ae164db13c2..b2969c4c998 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -242,6 +242,7 @@ rstest = "0.26.1" rstest_reuse = "0.7.0" rustc-hash = "2.1.1" rustix = { version = "1.1", features = ["fs"] } +seq-macro = "0.3.6" serde = "1.0.221" serde_json = "1.0.138" serde_test = "1.0.176" diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..42023005d9a 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -26,6 +26,7 @@ lending-iterator = { workspace = true } num-traits = { workspace = true } prost = { workspace = true } rand = { workspace = true, optional = true } +seq-macro = { workspace = true } vortex-array = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index a393db6ecc8..8fac1aabb37 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use fastlanes::BitPacking; use itertools::Itertools; use num_traits::PrimInt; use vortex_array::ArrayView; @@ -30,6 +29,7 @@ use vortex_mask::Mask; use crate::BitPacked; use crate::BitPackedArray; use crate::bitpack_decompress; +use crate::bitpacking::array::kernels::BitPackedPhysical; pub fn bitpack_to_best_bit_width( array: &PrimitiveArray, @@ -143,10 +143,11 @@ pub unsafe fn bitpack_unchecked(parray: &PrimitiveArray, bit_width: u8) -> ByteB /// Bitpack a slice of primitives down to the given width. /// /// See `bitpack` for more caller information. -pub fn bitpack_primitive(array: &[T], bit_width: u8) -> Buffer { +pub fn bitpack_primitive(array: &[T], bit_width: u8) -> Buffer { if bit_width == 0 { return Buffer::::empty(); } + let pack = T::resolve_pack(bit_width); let bit_width = bit_width as usize; // How many fastlanes vectors we will process. @@ -164,14 +165,15 @@ pub fn bitpack_primitive(array: &[T], bit_width: u8 (0..num_full_chunks).for_each(|i| { let start_elem = i * 1024; let output_len = output.len(); + // SAFETY: The capacity holds every block, so the new slots exist; the input is exactly + // 1024 values and the output exactly one block, which `pack` fully initializes. unsafe { output.set_len(output_len + packed_len); - BitPacking::unchecked_pack( - bit_width, + pack( &array[start_elem..][..1024], &mut output[output_len..][..packed_len], ); - }; + } }); // Pad the last chunk with zeros to a full 1024 elements. @@ -181,14 +183,12 @@ pub fn bitpack_primitive(array: &[T], bit_width: u8 last_chunk[..last_chunk_size].copy_from_slice(&array[array.len() - last_chunk_size..]); let output_len = output.len(); + // SAFETY: The capacity holds every block, so the new slots exist; the input is exactly + // 1024 values and the output exactly one block, which `pack` fully initializes. unsafe { output.set_len(output_len + packed_len); - BitPacking::unchecked_pack( - bit_width, - &last_chunk, - &mut output[output_len..][..packed_len], - ); - }; + pack(&last_chunk, &mut output[output_len..][..packed_len]); + } } output.freeze() diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 692e7dcdd7f..efa79553ddd 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -3,7 +3,6 @@ use std::mem::MaybeUninit; -use fastlanes::BitPacking; use itertools::Itertools; use num_traits::AsPrimitive; use vortex_array::ArrayView; @@ -22,6 +21,8 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::bitpacking::array::kernels::BitPackedKernels; +use crate::bitpacking::array::kernels::BitPackedPhysical; use crate::unpack_iter::BitPacked as BitPackedUnpack; use crate::unpack_iter::BitUnpackedChunks; @@ -159,39 +160,38 @@ pub(crate) fn apply_patches_to_uninit_range, index: usize) -> Scalar { - let bit_width = array.bit_width() as usize; let ptype = array.dtype().as_ptype(); - // let packed = array.packed().into_primitive()?; let index_in_encoded = index + array.offset() as usize; let scalar: Scalar = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |P| { - unsafe { - unpack_single_primitive::

(array.packed_slice::

(), bit_width, index_in_encoded) - .into() - } + unpack_single_primitive::

( + array.kernels::

(), + array.packed_slice::

(), + index_in_encoded, + ) + .into() }); // Cast to fix signedness and nullability scalar.cast(array.dtype()).vortex_expect("cast failure") } -/// # Safety +/// Unpacks the value at `index_to_decode` from `packed`, a buffer of whole FastLanes blocks +/// packed at the bit width `kernels` were resolved for. /// -/// The caller must ensure the following invariants hold: -/// * `packed.len() == (length + 1023) / 1024 * 128 * bit_width` -/// * `index_to_decode < length` +/// # Panics /// -/// Where `length` is the length of the array/slice backed by `packed` -/// (but is not provided to this function). -pub unsafe fn unpack_single_primitive( - packed: &[T], - bit_width: usize, +/// If `index_to_decode` falls outside the blocks held by `packed`. +pub fn unpack_single_primitive( + kernels: BitPackedKernels

, + packed: &[P], index_to_decode: usize, -) -> T { +) -> P { let chunk_index = index_to_decode / 1024; let index_in_chunk = index_to_decode % 1024; - let elems_per_chunk: usize = 128 * bit_width / size_of::(); + let elems_per_chunk = kernels.packed_block_len(); - let packed_chunk = &packed[chunk_index * elems_per_chunk..][0..elems_per_chunk]; - unsafe { BitPacking::unchecked_unpack_single(bit_width, packed_chunk, index_in_chunk) } + let packed_chunk = &packed[chunk_index * elems_per_chunk..][..elems_per_chunk]; + // SAFETY: `packed_chunk` is exactly one block at the width `kernels` were resolved for. + unsafe { (kernels.unpack_single)(packed_chunk, index_in_chunk) } } pub fn count_exceptions(bit_width: u8, bit_width_freq: &[usize]) -> usize { diff --git a/encodings/fastlanes/src/bitpacking/array/kernels.rs b/encodings/fastlanes/src/bitpacking/array/kernels.rs new file mode 100644 index 00000000000..f098edf0f69 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/array/kernels.rs @@ -0,0 +1,432 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Width-resolved FastLanes kernels for bit-packed arrays. +//! +//! The FastLanes kernels are generic over the packed bit width `W`, so the runtime-width +//! `unchecked_*` entry points of the `fastlanes` crate dispatch on the width with a `match` on +//! every call. A [`BitPackedArray`](crate::BitPackedArray) knows its width but is type erased, so +//! it cannot name the instantiation statically. Instead, the array resolves function pointers to +//! the concrete instantiations once, when it is constructed, and hands them out as +//! [`BitPackedKernels`]. The decoding paths then call the resolved kernels block after block +//! without re-dispatching. + +use std::mem; + +use fastlanes::BitPacking; +use fastlanes::BitPackingCompare; +use fastlanes::FastLanesComparable; +use fastlanes::FoR; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; + +/// Packs one FastLanes block of 1024 values. +/// +/// # Safety +/// +/// `input` must hold exactly 1024 elements and `output` exactly `128 * bit_width / size_of::

()`. +/// The lengths are checked only with `debug_assert`. +pub type PackFn

= unsafe fn(input: &[P], output: &mut [P]); + +/// Unpacks one FastLanes block of 1024 values. +/// +/// # Safety +/// +/// `packed` must hold exactly [`BitPackedKernels::packed_block_len`] elements and `output` +/// exactly 1024. The lengths are checked only with `debug_assert`. +pub type UnpackFn

= unsafe fn(packed: &[P], output: &mut [P]); + +/// Unpacks the value at `index` of one packed FastLanes block. +/// +/// # Safety +/// +/// `packed` must hold exactly [`BitPackedKernels::packed_block_len`] elements. The length is +/// checked only with `debug_assert`. An `index` of 1024 or more panics. +pub type UnpackSingleFn

= unsafe fn(packed: &[P], index: usize) -> P; + +/// Unpacks one FastLanes block and wrapping-adds `reference` to every value. +/// +/// # Safety +/// +/// `packed` must hold exactly [`BitPackedKernels::packed_block_len`] elements and `output` +/// exactly 1024. The lengths are checked only with `debug_assert`. +pub type UnforPackFn

= unsafe fn(packed: &[P], reference: P, output: &mut [P]); + +/// Unpacks one FastLanes block, comparing each value against `rhs` with `cmp` and writing the +/// results as a lane-major 1024-bit mask. See [`BitPackingCompare::unpack_cmp`] for the layout. +/// +/// # Safety +/// +/// `packed` must hold exactly `128 * bit_width / size_of::

()` elements. The length is checked +/// only with `debug_assert`. +pub type UnpackCmpFn = unsafe fn(packed: &[P], output: &mut [u64; 16], cmp: F, rhs: V); + +/// FastLanes kernels resolved for one bit width of physical type `P`. +/// +/// Obtained from [`BitPackedData::kernels`](crate::BitPackedData::kernels). Each kernel is the +/// const-width instantiation for the array's bit width, so calling it does not dispatch on the +/// width. +#[derive(Clone, Copy, Debug)] +pub struct BitPackedKernels

{ + /// The unsigned [`PType`] of `P`, kept so the type-erased form can check it before + /// re-typing the pointers. + ptype: PType, + bit_width: u8, + /// Unpacks one packed block into 1024 values. + pub unpack: UnpackFn

, + /// Unpacks a single value of one packed block. + pub unpack_single: UnpackSingleFn

, + /// Unpacks one packed block, adding a frame-of-reference value. + pub unfor_pack: UnforPackFn

, +} + +impl

BitPackedKernels

{ + /// The bit width the kernels were resolved for. + #[inline] + pub fn bit_width(&self) -> u8 { + self.bit_width + } +} + +impl BitPackedKernels

{ + /// The number of `P` elements holding one packed block of 1024 values. + #[inline] + pub fn packed_block_len(&self) -> usize { + 128 * self.bit_width as usize / size_of::

() + } +} + +/// [`BitPackedKernels`] with the physical type erased, as stored by +/// [`BitPackedData`](crate::BitPackedData). +/// +/// The function pointers are those of a `BitPackedKernels

` for the recorded `ptype`, cast to +/// a placeholder element type. [`Self::typed`] casts them back; they are never called in the +/// erased form. +pub type ResolvedKernels = BitPackedKernels<()>; + +impl ResolvedKernels { + /// Resolves the kernels for an array of `ptype` packed to `bit_width` bits. + /// + /// Signed types resolve to their unsigned counterpart, which is what the packed buffer holds. + pub fn try_new(ptype: PType, bit_width: u8) -> VortexResult { + vortex_ensure!(ptype.is_int(), MismatchedTypes: "integer", ptype); + vortex_ensure!( + bit_width as usize <= ptype.bit_width(), + "Unsupported bit width {bit_width} for {ptype}" + ); + Ok(match_each_unsigned_integer_ptype!( + ptype.to_unsigned(), + |P| { P::resolve_kernels(bit_width).erase() } + )) + } + + /// The kernels typed for `P`, which must be the physical type they were resolved for. + /// + /// # Panics + /// + /// If `P` is not the physical type the kernels were resolved for. + #[inline] + pub fn typed(self) -> BitPackedKernels

{ + assert!( + self.ptype == P::PTYPE, + "BitPacked kernels were resolved for a different physical type" + ); + // SAFETY: `ptype` records the `P` these pointers were erased from (see `erase`), and + // transmuting a function pointer back to its original signature is lossless. + unsafe { + BitPackedKernels { + ptype: self.ptype, + bit_width: self.bit_width, + unpack: mem::transmute::, UnpackFn

>(self.unpack), + unpack_single: mem::transmute::, UnpackSingleFn

>( + self.unpack_single, + ), + unfor_pack: mem::transmute::, UnforPackFn

>(self.unfor_pack), + } + } + } +} + +impl BitPackedKernels

{ + /// Erases `P` from the function pointers; [`ResolvedKernels::typed`] restores it. + fn erase(self) -> ResolvedKernels { + // SAFETY: Function pointers of every signature share one layout, and the erased pointers + // are only ever called after `typed` casts them back to this signature, which `ptype` + // enforces. + unsafe { + BitPackedKernels { + ptype: self.ptype, + bit_width: self.bit_width, + unpack: mem::transmute::, UnpackFn<()>>(self.unpack), + unpack_single: mem::transmute::, UnpackSingleFn<()>>( + self.unpack_single, + ), + unfor_pack: mem::transmute::, UnforPackFn<()>>(self.unfor_pack), + } + } + } +} + +/// The physical storage types of a bit-packed array, i.e. the unsigned integers the FastLanes +/// kernels are implemented for. Signed arrays are packed as their unsigned counterpart. +pub trait BitPackedPhysical: NativePType + BitPacking + BitPackingCompare + FoR { + /// Resolves the kernels for `bit_width`, which must not exceed the width of `Self`. + fn resolve_kernels(bit_width: u8) -> BitPackedKernels; + + /// Resolves the pack kernel for `bit_width`, which must not exceed the width of `Self`. + /// + /// Packing happens before an array exists to cache kernels on, so callers resolve this once + /// per buffer instead. + fn resolve_pack(bit_width: u8) -> PackFn; + + /// Resolves the fused unpack-and-compare kernel for `bit_width`, which must not exceed the + /// width of `Self`. + fn resolve_unpack_cmp(bit_width: u8) -> UnpackCmpFn + where + V: FastLanesComparable, + F: Fn(V, V) -> bool; +} + +unsafe fn pack(input: &[P], output: &mut [P]) { + // SAFETY: The caller upholds the `PackFn` length contract. + unsafe { P::pack::(as_block(input), as_block_mut(output)) } +} + +unsafe fn unpack(packed: &[P], output: &mut [P]) { + // SAFETY: The caller upholds the `UnpackFn` length contract. + unsafe { P::unpack::(as_block(packed), as_block_mut(output)) } +} + +unsafe fn unpack_single( + packed: &[P], + index: usize, +) -> P { + // SAFETY: The caller upholds the `UnpackSingleFn` length contract. + P::unpack_single::(unsafe { as_block(packed) }, index) +} + +unsafe fn unfor_pack( + packed: &[P], + reference: P, + output: &mut [P], +) { + // SAFETY: The caller upholds the `UnforPackFn` length contract. + unsafe { P::unfor_pack::(as_block(packed), reference, as_block_mut(output)) } +} + +unsafe fn unpack_cmp( + packed: &[P], + output: &mut [u64; 16], + cmp: F, + rhs: V, +) where + V: FastLanesComparable, + F: Fn(V, V) -> bool, +{ + // SAFETY: The caller upholds the `UnpackCmpFn` length contract. + P::unpack_cmp::(unsafe { as_block(packed) }, output, cmp, rhs); +} + +/// Reinterprets `slice` as a block of exactly `N` elements. +/// +/// # Safety +/// +/// `slice.len()` must be `N`. This is checked only with `debug_assert`. +unsafe fn as_block(slice: &[P]) -> &[P; N] { + debug_assert_eq!(slice.len(), N); + // SAFETY: The caller guarantees `N` elements, and `[P; N]` has the alignment of `P`. + unsafe { &*slice.as_ptr().cast::<[P; N]>() } +} + +/// Reinterprets `slice` as a mutable block of exactly `N` elements. +/// +/// # Safety +/// +/// `slice.len()` must be `N`. This is checked only with `debug_assert`. +unsafe fn as_block_mut(slice: &mut [P]) -> &mut [P; N] { + debug_assert_eq!(slice.len(), N); + // SAFETY: The caller guarantees `N` elements, and `[P; N]` has the alignment of `P`. + unsafe { &mut *slice.as_mut_ptr().cast::<[P; N]>() } +} + +macro_rules! impl_bitpacked_physical { + ($P:ty, $bits:literal) => { + impl BitPackedPhysical for $P { + fn resolve_kernels(bit_width: u8) -> BitPackedKernels { + seq_macro::seq!(W in 0..=$bits { + match bit_width { + #(W => BitPackedKernels { + ptype: <$P as NativePType>::PTYPE, + bit_width, + unpack: unpack::<$P, W, { 1024 * W / $bits }>, + unpack_single: unpack_single::<$P, W, { 1024 * W / $bits }>, + unfor_pack: unfor_pack::<$P, W, { 1024 * W / $bits }>, + },)* + _ => vortex_panic!( + "Unsupported bit width {bit_width} for {}", + <$P as NativePType>::PTYPE + ), + } + }) + } + + fn resolve_pack(bit_width: u8) -> PackFn { + seq_macro::seq!(W in 0..=$bits { + match bit_width { + #(W => pack::<$P, W, { 1024 * W / $bits }>,)* + _ => vortex_panic!( + "Unsupported bit width {bit_width} for {}", + <$P as NativePType>::PTYPE + ), + } + }) + } + + fn resolve_unpack_cmp(bit_width: u8) -> UnpackCmpFn + where + V: FastLanesComparable, + F: Fn(V, V) -> bool, + { + seq_macro::seq!(W in 0..=$bits { + match bit_width { + #(W => unpack_cmp::<$P, W, { 1024 * W / $bits }, V, F>,)* + _ => vortex_panic!( + "Unsupported bit width {bit_width} for {}", + <$P as NativePType>::PTYPE + ), + } + }) + } + } + }; +} + +impl_bitpacked_physical!(u8, 8); +impl_bitpacked_physical!(u16, 16); +impl_bitpacked_physical!(u32, 32); +impl_bitpacked_physical!(u64, 64); + +#[cfg(test)] +mod tests { + use num_traits::WrappingAdd; + use rstest::rstest; + + use super::*; + + /// Every width of every physical type resolves to kernels that agree with the runtime-width + /// FastLanes entry points, both directly and after a round trip through the erased form. + fn assert_kernels_match_fastlanes

() + where + P: BitPackedPhysical + WrappingAdd + FastLanesComparable, + { + let values: [P; 1024] = std::array::from_fn(|i| P::from(i % 251).unwrap()); + for bit_width in 0..=(8 * size_of::

() as u8) { + let kernels = ResolvedKernels::try_new(P::PTYPE, bit_width) + .unwrap() + .typed::

(); + assert_eq!(kernels.bit_width(), bit_width); + assert_eq!( + kernels.unpack as usize, + P::resolve_kernels(bit_width).unpack as usize, + "erased round trip at width {bit_width}" + ); + + let block_len = kernels.packed_block_len(); + let mut expected_packed = vec![P::zero(); block_len]; + // SAFETY: `expected_packed` holds exactly one block at `bit_width` and `values` 1024 + // values. + unsafe { P::unchecked_pack(bit_width as usize, &values, &mut expected_packed) }; + + let mut packed = vec![P::zero(); block_len]; + // SAFETY: `values` holds 1024 values and `packed` exactly one block at `bit_width`. + unsafe { P::resolve_pack(bit_width)(&values, &mut packed) }; + assert_eq!(packed, expected_packed, "pack at width {bit_width}"); + + let mut expected = [P::zero(); 1024]; + // SAFETY: `packed` holds exactly one block at `bit_width` and `expected` 1024 values. + unsafe { P::unchecked_unpack(bit_width as usize, &packed, &mut expected) }; + + // SAFETY: `packed` holds exactly one block at `bit_width` and `unpacked` 1024 values. + let mut unpacked = [P::zero(); 1024]; + unsafe { (kernels.unpack)(&packed, &mut unpacked) }; + assert_eq!(unpacked, expected, "unpack at width {bit_width}"); + + for index in [0, 1, 511, 1023] { + assert_eq!( + // SAFETY: `packed` holds exactly one block at `bit_width`. + unsafe { (kernels.unpack_single)(&packed, index) }, + expected[index], + "unpack_single at width {bit_width} index {index}" + ); + } + + let reference = P::from(7).unwrap(); + let mut unfor = [P::zero(); 1024]; + // SAFETY: `packed` holds exactly one block at `bit_width` and `unfor` 1024 values. + unsafe { (kernels.unfor_pack)(&packed, reference, &mut unfor) }; + for (got, want) in unfor.iter().zip(expected) { + assert_eq!( + *got, + want.wrapping_add(&reference), + "unfor_pack at {bit_width}" + ); + } + + let rhs = P::from(100).unwrap(); + let mut mask = [0u64; 16]; + // SAFETY: `packed` holds exactly one block at `bit_width`. + unsafe { + P::resolve_unpack_cmp::(bit_width)(&packed, &mut mask, |a, b| a < b, rhs) + }; + let mut expected_mask = [0u64; 16]; + // SAFETY: `packed` holds exactly one block at `bit_width`. + unsafe { + P::unchecked_unpack_cmp::( + bit_width as usize, + &packed, + &mut expected_mask, + |a, b| a < b, + rhs, + ); + } + assert_eq!(mask, expected_mask, "unpack_cmp at width {bit_width}"); + } + } + + #[rstest] + #[case::u8(assert_kernels_match_fastlanes::)] + #[case::u16(assert_kernels_match_fastlanes::)] + #[case::u32(assert_kernels_match_fastlanes::)] + #[case::u64(assert_kernels_match_fastlanes::)] + fn kernels_match_fastlanes(#[case] check: fn()) { + check(); + } + + #[test] + fn signed_resolves_to_unsigned_kernels() { + let resolved = ResolvedKernels::try_new(PType::I16, 3).unwrap(); + assert_eq!( + resolved.typed::().unpack as usize, + u16::resolve_kernels(3).unpack as usize + ); + } + + #[test] + fn rejects_width_beyond_type() { + assert!(ResolvedKernels::try_new(PType::U8, 9).is_err()); + assert!(ResolvedKernels::try_new(PType::F32, 3).is_err()); + assert!(ResolvedKernels::try_new(PType::U8, 8).is_ok()); + } + + #[test] + #[should_panic(expected = "resolved for a different physical type")] + fn typed_rejects_other_types() { + ResolvedKernels::try_new(PType::U16, 3) + .unwrap() + .typed::(); + } +} diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 03fa3ed7f4c..0bab8c37d7d 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -26,10 +26,14 @@ use vortex_error::vortex_err; pub mod bitpack_compress; pub mod bitpack_decompress; +pub mod kernels; pub mod unpack_iter; use crate::BitPackedArray; use crate::bitpack_compress::bitpack_encode; +use crate::bitpacking::array::kernels::BitPackedKernels; +use crate::bitpacking::array::kernels::BitPackedPhysical; +use crate::bitpacking::array::kernels::ResolvedKernels; use crate::unpack_iter::BitPacked as BitPackedIter; use crate::unpack_iter::BitUnpackedChunks; @@ -69,15 +73,23 @@ 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, + /// FastLanes kernels for the physical type and bit width of this array, resolved at + /// construction so that decoding never dispatches on the runtime bit width. Also the only + /// record of the bit width. + kernels: ResolvedKernels, } 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, + "bit_width: {}, offset: {}", + self.bit_width(), + self.offset + ) } } @@ -116,7 +128,7 @@ impl BitPackedData { /// /// # Validation /// - /// * The `ptype` must be an integer + /// * The `ptype` must be an integer and `bit_width` must not exceed its width /// * `validity` must have `length` len /// * Any patches must have any `array_len` equal to `length` /// * The `packed` buffer must be exactly sized to hold `length` values of `bit_width` rounded @@ -126,10 +138,10 @@ impl BitPackedData { pub fn try_new( packed: BufferHandle, patches: Option, + ptype: PType, bit_width: u8, offset: u16, ) -> VortexResult { - vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}"); vortex_ensure!( offset < 1024, "Offset must be less than the full block i.e., 1024, got {offset}" @@ -137,9 +149,9 @@ impl BitPackedData { Ok(Self { offset, - bit_width, packed, patches_data: patches.as_ref().map(PatchesData::from_patches), + kernels: ResolvedKernels::try_new(ptype, bit_width)?, }) } @@ -239,7 +251,20 @@ impl BitPackedData { /// Bit-width of the packed values #[inline] pub fn bit_width(&self) -> u8 { - self.bit_width + self.kernels.bit_width() + } + + /// The FastLanes kernels for this array's bit width, resolved when the array was built. + /// + /// `P` must be the physical type of the array, i.e. the unsigned counterpart of its + /// [`PType`]. + /// + /// # Panics + /// + /// If `P` is not the array's physical type. + #[inline] + pub fn kernels(&self) -> BitPackedKernels

{ + self.kernels.typed::

() } #[inline] @@ -319,6 +344,12 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { fn unpacked_chunks(&self) -> VortexResult> { BitPackedData::unpacked_chunks::(self, self.as_ref().dtype(), self.as_ref().len()) } + + /// The FastLanes kernels for this array's bit width, see [`BitPackedData::kernels`]. + #[inline] + fn kernels(&self) -> BitPackedKernels

{ + BitPackedData::kernels::

(self) + } } impl> BitPackedArrayExt for T {} diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 3c77e146ad0..65532b51409 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -5,7 +5,6 @@ use std::mem; use std::mem::MaybeUninit; use std::ops::Range; -use fastlanes::BitPacking; use lending_iterator::gat; use lending_iterator::prelude::Item; #[gat(Item)] @@ -16,34 +15,34 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::BitPackedData; +use crate::bitpacking::array::kernels::BitPackedPhysical; +use crate::bitpacking::array::kernels::UnpackFn; const CHUNK_SIZE: usize = 1024; /// Strategy trait for fastlanes unpacking operations pub trait UnpackStrategy { - /// Unpack a chunk of packed data into the destination buffer + /// Unpack a chunk of packed data into the destination buffer. /// /// # Safety - /// - `chunk` must contain exactly `elems_per_chunk` elements - /// - `dst` must have exactly CHUNK_SIZE capacity - unsafe fn unpack_chunk(&self, bit_width: usize, chunk: &[T::Physical], dst: &mut [T::Physical]); + /// + /// - `chunk` must contain exactly one packed block (`128 * bit_width / size_of::()` + /// elements) + /// - `dst` must contain exactly `CHUNK_SIZE` elements + unsafe fn unpack_chunk(&self, chunk: &[T::Physical], dst: &mut [T::Physical]); } -/// BitPacking strategy - uses plain bitpacking without reference value -pub struct BitPackingStrategy; +/// BitPacking strategy - plain bitpacking without reference value, using the unpack kernel +/// resolved for the array's bit width. +pub struct BitPackingStrategy

{ + unpack: UnpackFn

, +} -impl> UnpackStrategy for BitPackingStrategy { +impl UnpackStrategy for BitPackingStrategy { #[inline(always)] - unsafe fn unpack_chunk( - &self, - bit_width: usize, - chunk: &[T::Physical], - dst: &mut [T::Physical], - ) { - // SAFETY: Caller must ensure [`BitPacking::unchecked_unpack`] safety requirements hold. - unsafe { - BitPacking::unchecked_unpack(bit_width, chunk, dst); - } + unsafe fn unpack_chunk(&self, chunk: &[T::Physical], dst: &mut [T::Physical]) { + // SAFETY: The caller upholds the `unpack_chunk` length contract, which is `UnpackFn`'s. + unsafe { (self.unpack)(chunk, dst) } } } @@ -91,12 +90,15 @@ pub struct UnpackedChunks> { buffer: [MaybeUninit; CHUNK_SIZE], } -pub type BitUnpackedChunks = UnpackedChunks; +pub type BitUnpackedChunks = + UnpackedChunks::Physical>>; impl BitUnpackedChunks { pub fn try_new(array: &BitPackedData, len: usize) -> VortexResult { Self::try_new_with_strategy( - BitPackingStrategy, + BitPackingStrategy { + unpack: array.kernels::().unpack, + }, array.packed().clone().unwrap_host(), array.bit_width() as usize, array.offset() as usize, @@ -111,7 +113,7 @@ impl BitUnpackedChunks { BitUnpackIterator::new( buffer_as_slice(&self.packed), &mut self.buffer, - self.bit_width, + self.strategy.unpack, elems_per_chunk, self.num_chunks - last_chunk_is_sliced, first_chunk_is_sliced, @@ -173,9 +175,9 @@ impl> UnpackedChunks { }; // SAFETY: // 1. chunk is elems_per_chunk. - // 2. buffer is exactly CHUNK_SIZE. + // 2. buffer is exactly CHUNK_SIZE, and `unpack_chunk` initializes all of it. unsafe { - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(chunk, dst); mem::transmute(&mut self.buffer[self.offset..][..header_end_slice]) } }) @@ -241,7 +243,7 @@ impl> UnpackedChunks { let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; unsafe { let dst: &mut [T::Physical] = mem::transmute(&mut self.buffer[..]); - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(chunk, dst); let unpacked: &mut [T] = mem::transmute(&mut self.buffer[..]); f(unpacked, local_idx..local_idx + CHUNK_SIZE); } @@ -307,7 +309,7 @@ impl> UnpackedChunks { 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(chunk, dst); } local_idx += CHUNK_SIZE; } @@ -328,9 +330,9 @@ impl> UnpackedChunks { let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; // SAFETY: // 1. chunk is elems_per_chunk. - // 2. buffer is exactly CHUNK_SIZE. + // 2. buffer is exactly CHUNK_SIZE, and `unpack_chunk` initializes all of it. unsafe { - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(chunk, dst); mem::transmute(&mut self.buffer[..self.last_chunk_length]) } }) @@ -349,7 +351,7 @@ impl> UnpackedChunks { pub struct BitUnpackIterator<'a, T: BitPacked + 'a> { packed: &'a [T::Physical], buffer: &'a mut [MaybeUninit; CHUNK_SIZE], - bit_width: usize, + unpack: UnpackFn, elems_per_chunk: usize, num_chunks: usize, idx: usize, @@ -359,7 +361,7 @@ impl<'a, T: BitPacked> BitUnpackIterator<'a, T> { pub fn new( packed: &'a [T::Physical], buffer: &'a mut [MaybeUninit; CHUNK_SIZE], - bit_width: usize, + unpack: UnpackFn, elems_per_chunk: usize, num_chunks: usize, first_chunk_is_sliced: bool, @@ -367,7 +369,7 @@ impl<'a, T: BitPacked> BitUnpackIterator<'a, T> { Self { packed, buffer, - bit_width, + unpack, elems_per_chunk, num_chunks, idx: if first_chunk_is_sliced { 1 } else { 0 }, @@ -390,10 +392,11 @@ impl<'a, T: BitPacked + 'a> LendingIterator for BitUnpackIterator<'a, T> { let chunk = &self.packed[self.idx * self.elems_per_chunk..][..self.elems_per_chunk]; let dst: &mut [MaybeUninit] = self.buffer; + // SAFETY: &[MaybeUninit] and &[T::Physical] have the same layout; `chunk` is exactly + // `elems_per_chunk` and `dst` exactly CHUNK_SIZE. unsafe { let dst: &mut [T::Physical] = mem::transmute(dst); - - BitPacking::unchecked_unpack(self.bit_width, chunk, dst); + (self.unpack)(chunk, dst); } self.idx += 1; // SAFETY: The buffer has the appropriate lifetime, the iterator signature doesn't account for it @@ -418,7 +421,7 @@ fn write_map(src: &[T], dst: &mut [MaybeUninit], f: &mut impl FnM } } -pub trait BitPacked: PhysicalPType {} +pub trait BitPacked: PhysicalPType {} impl BitPacked for i8 {} impl BitPacked for i16 {} diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index c9d6b815b0d..95cd018b180 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -11,8 +11,6 @@ //! [`BitPackedArray`]: crate::BitPackedArray //! [`BitBuffer`]: vortex_buffer::BitBuffer -use fastlanes::BitPacking; -use fastlanes::BitPackingCompare; use fastlanes::FastLanesComparable; use vortex_array::ArrayRef; use vortex_array::ArrayView; @@ -78,7 +76,6 @@ where T: NativePType + BitPackedIter + FastLanesComparable::Physical>, - ::Physical: BitPacking + NativePType + BitPackingCompare, { match operator { CompareOperator::Eq => { diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index f27384b898c..ca4fcd44b18 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -5,15 +5,16 @@ //! //! Where [`super::stream_predicate`] unpacks a full 1024-element FastLanes block into a scratch //! buffer and *then* folds a predicate over it, this path hands the comparison down into the -//! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which compares each value against -//! the constant *as it is unpacked*, accumulating the boolean results straight into a 1024-bit -//! mask (`[u64; 16]`) in transposed FastLanes lane order - one register-resident word per lane, no -//! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`transpose_bits`] per block then rotates -//! that mask into logical row order. +//! FastLanes [`fastlanes::BitPackingCompare::unpack_cmp`] kernel, which compares each value +//! against the constant *as it is unpacked*, accumulating the boolean results straight into a +//! 1024-bit mask (`[u64; 16]`) in transposed FastLanes lane order - one register-resident word per +//! lane, no `[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 the regular [`crate::unpack_iter::BitUnpackedChunks`] //! iterator (via [`crate::unpack_iter::BitUnpackedChunks::for_each_packed_chunk`]) rather than a -//! bespoke chunk loop, so chunk sizing and bounds live in one place. +//! bespoke chunk loop, so chunk sizing and bounds live in one place. The kernel instantiation for +//! the array's bit width is resolved once up front, so no block re-dispatches on the width. //! //! Slicing is handled by working in *padded* coordinates: bit `offset + i` holds element `i`. The //! output buffer is over-allocated to whole 1024-bit blocks, so every block - the sliced first @@ -26,8 +27,6 @@ //! [`BitPackedArray`]: crate::BitPackedArray //! [`BitBuffer`]: vortex_buffer::BitBuffer -use fastlanes::BitPacking; -use fastlanes::BitPackingCompare; use fastlanes::FastLanesComparable; use fastlanes::transpose_bits; use num_traits::AsPrimitive; @@ -49,6 +48,7 @@ use vortex_error::VortexResult; use super::stream_predicate::stream_predicate; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::bitpacking::array::kernels::BitPackedPhysical; use crate::unpack_iter::BitPacked as BitPackedIter; const CHUNK_SIZE: usize = 1024; @@ -74,11 +74,10 @@ where T: NativePType + BitPackedIter + FastLanesComparable::Physical>, - ::Physical: BitPacking + NativePType + BitPackingCompare, F: Fn(T, T) -> bool + Copy, { let len = array.len(); - let bit_width = array.bit_width() as usize; + let bit_width = array.bit_width(); let offset = array.offset() as usize; // A degenerate width has no packed payload for the fused kernel to consume; defer to the scalar @@ -93,6 +92,10 @@ where let mut words: BufferMut = BufferMut::zeroed(num_chunks * WORDS_PER_CHUNK); let chunks = array.unpacked_chunks::()?; + let unpack_cmp = <::Physical as BitPackedPhysical>::resolve_unpack_cmp::< + T, + F, + >(bit_width); { let words = words.as_mut_slice(); let mut lane_major = [0u64; WORDS_PER_CHUNK]; @@ -101,18 +104,10 @@ where let out = words[range.start / U64_BITS..] .first_chunk_mut::() .vortex_expect("over-allocated buffer holds a full block per chunk"); - // 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. - unsafe { - <::Physical as BitPackingCompare>::unchecked_unpack_cmp::( - bit_width, - packed_chunk, - &mut lane_major, - cmp, - rhs, - ); - } + // SAFETY: `packed_chunk` holds exactly one block at the array's bit width, which is + // the width `unpack_cmp` was resolved for. The kernel assigns every word in + // `lane_major`, so its previous contents are irrelevant. + unsafe { unpack_cmp(packed_chunk, &mut lane_major, cmp, rhs) }; transpose_bits::<::Physical>(&lane_major, out); }); } diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 0b1b9422f86..2f77a72b9ab 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -3,14 +3,12 @@ use std::mem::MaybeUninit; -use fastlanes::BitPacking; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::filter::FilterKernel; -use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::dtype::UnsignedPType; use vortex_array::match_each_unsigned_integer_ptype; @@ -26,6 +24,7 @@ use super::take::UNPACK_CHUNK_THRESHOLD; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitPackedData; +use crate::bitpacking::array::kernels::BitPackedPhysical; /// 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. @@ -106,7 +105,7 @@ impl FilterKernel for BitPacked { /// elements is relatively slow. /// /// Returns a tuple of (values buffer, validity mask). -fn filter_primitive_without_patches( +fn filter_primitive_without_patches( array: ArrayView<'_, BitPacked>, selection: &MaskValuesRef, ) -> VortexResult<(Buffer, Validity)> { @@ -118,12 +117,12 @@ fn filter_primitive_without_patches( Ok((values.freeze(), validity)) } -fn filter_with_indices( +fn filter_with_indices( array: &BitPackedData, indices: &[usize], ) -> BufferMut { let offset = array.offset() as usize; - let bit_width = array.bit_width() as usize; + let kernels = array.kernels::(); let mut values = BufferMut::with_capacity(indices.len()); // Some re-usable memory to store per-chunk indices. @@ -131,7 +130,7 @@ fn filter_with_indices( let packed_bytes = array.packed_slice::(); // Group the indices by the FastLanes chunk they belong to. - let chunk_size = 128 * bit_width / size_of::(); + let chunk_size = kernels.packed_block_len(); chunked_indices( indices.iter().copied(), @@ -141,21 +140,22 @@ fn filter_with_indices( if indices_within_chunk.len() == 1024 { // Unpack the entire chunk. + // SAFETY: The capacity holds every index, so the 1024 new slots exist; `packed` + // is exactly one block and `unpack` initializes all 1024 slots before they are + // read. unsafe { let values_len = values.len(); values.set_len(values_len + 1024); - BitPacking::unchecked_unpack( - bit_width, - packed, - &mut values.as_mut_slice()[values_len..], - ); + (kernels.unpack)(packed, &mut values.as_mut_slice()[values_len..]); } } else if indices_within_chunk.len() > UNPACK_CHUNK_THRESHOLD { // Unpack into a temporary chunk and then copy the values. + // SAFETY: &[MaybeUninit] and &[T] have the same layout; `packed` is exactly + // one block and `unpacked` exactly 1024 values. unsafe { let dst: &mut [MaybeUninit] = &mut unpacked; let dst: &mut [T] = std::mem::transmute(dst); - BitPacking::unchecked_unpack(bit_width, packed, dst); + (kernels.unpack)(packed, dst); } values.extend_trusted( indices_within_chunk @@ -164,9 +164,12 @@ fn filter_with_indices( ); } else { // Otherwise, unpack each element individually. - values.extend_trusted(indices_within_chunk.iter().map(|&idx| unsafe { - BitPacking::unchecked_unpack_single(bit_width, packed, idx) - })); + // SAFETY: `packed` is exactly one block and every index is within the chunk. + values.extend_trusted( + indices_within_chunk + .iter() + .map(|&idx| unsafe { (kernels.unpack_single)(packed, idx) }), + ); } }, ); diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 86e97623cf6..ec6ad18a3d4 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -4,7 +4,6 @@ use std::mem; use std::mem::MaybeUninit; -use fastlanes::BitPacking; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -12,7 +11,6 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; use vortex_array::dtype::IntegerPType; -use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_unsigned_integer_ptype; @@ -25,7 +23,7 @@ use vortex_error::VortexResult; use super::chunked_indices; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::bitpack_decompress; +use crate::bitpacking::array::kernels::BitPackedPhysical; // 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 @@ -70,7 +68,7 @@ impl TakeExecute for BitPacked { } } -fn take_primitive( +fn take_primitive( array: ArrayView<'_, BitPacked>, indices: &PrimitiveArray, taken_validity: Validity, @@ -81,7 +79,7 @@ fn take_primitive( } let offset = array.offset() as usize; - let bit_width = array.bit_width() as usize; + let kernels = array.kernels::(); let packed = array.packed_slice::(); @@ -93,7 +91,7 @@ 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::(); + let chunk_len = kernels.packed_block_len(); chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; @@ -104,10 +102,12 @@ fn take_primitive( // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets for offset_chunk in offset_chunks { if !have_unpacked { + // SAFETY: &[MaybeUninit] and &[T] have the same layout; `packed` is exactly + // one block and `unpacked` exactly 1024 values. unsafe { let dst: &mut [MaybeUninit] = &mut unpacked; let dst: &mut [T] = mem::transmute(dst); - BitPacking::unchecked_unpack(bit_width, packed, dst); + (kernels.unpack)(packed, dst); } have_unpacked = true; } @@ -128,9 +128,8 @@ fn take_primitive( // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place, // so we need to unpack each one individually for &index in remainder { - output.push(unsafe { - bitpack_decompress::unpack_single_primitive::(packed, bit_width, index) - }); + // SAFETY: `packed` is exactly one block and `index` is within the chunk. + output.push(unsafe { (kernels.unpack_single)(packed, index) }); } } } diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..1dedcc106ef 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -9,6 +9,7 @@ pub use array::BitPackedDataParts; pub use array::BitPackedSlots; pub use array::bitpack_compress; pub use array::bitpack_decompress; +pub use array::kernels; pub use array::unpack_iter; pub(crate) mod compute; diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index a621d085514..df57bbab18b 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -68,7 +68,7 @@ impl ArrayPlugin for BitPackedPatchedPlugin { let packed = bitpacked.packed().clone(); let ptype = bitpacked.dtype().as_ptype(); let validity = bitpacked.validity()?; - let bw = bitpacked.bit_width; + let bw = bitpacked.bit_width(); let len = bitpacked.len(); let offset = bitpacked.offset(); diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..b8ebee7a26f 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -75,7 +75,7 @@ pub struct BitPackedMetadata { impl ArrayHash for BitPackedData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.offset.hash(state); - self.bit_width.hash(state); + self.bit_width().hash(state); self.packed.array_hash(state, accuracy); self.patches_data.hash(state); } @@ -84,7 +84,7 @@ impl ArrayHash for BitPackedData { impl ArrayEq for BitPackedData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { self.offset == other.offset - && self.bit_width == other.bit_width + && self.bit_width() == other.bit_width() && self.packed.array_eq(&other.packed, accuracy) && self.patches_data == other.patches_data } @@ -118,7 +118,7 @@ impl VTable for BitPacked { dtype.as_ptype(), &validity, patches.as_ref(), - data.bit_width, + data.bit_width(), len, data.offset, ) @@ -239,6 +239,7 @@ impl VTable for BitPacked { let data = BitPackedData::try_new( packed, patches, + dtype.as_ptype(), u8::try_from(metadata.bit_width).map_err(|_| { vortex_err!( "BitPackedMetadata bit_width {} does not fit in u8", @@ -319,7 +320,7 @@ impl BitPacked { s.push(validity_to_child(&validity, len)); s }; - let data = BitPackedData::try_new(packed, patches, bit_width, offset)?; + let data = BitPackedData::try_new(packed, patches, ptype, bit_width, offset)?; Array::try_from_parts(ArrayParts::new(BitPacked, dtype, len, data).with_slots(slots)) } @@ -330,7 +331,7 @@ impl BitPacked { let data = array.into_data(); BitPackedDataParts { offset: data.offset, - bit_width: data.bit_width, + bit_width: data.bit_width(), len, packed: data.packed, patches, diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index cb94618c071..1fedbe494ed 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use fastlanes::FoR; use num_traits::PrimInt; use num_traits::WrappingAdd; use vortex_array::ArrayView; @@ -21,28 +20,26 @@ use crate::BitPacked; use crate::BitPackedArrayExt; use crate::FoRArray; use crate::bitpack_decompress; +use crate::bitpacking::kernels::BitPackedPhysical; +use crate::bitpacking::kernels::UnforPackFn; use crate::r#for::array::FoRArrayExt; use crate::r#for::array::FoRArraySlotsExt; use crate::unpack_iter::UnpackStrategy; use crate::unpack_iter::UnpackedChunks; -/// FoR unpacking strategy that applies a reference value during unpacking. +/// FoR unpacking strategy that applies a reference value during unpacking, using the fused +/// kernel resolved for the bit-packed child's width. struct FoRStrategy { reference: T, + unfor_pack: UnforPackFn, } -impl + FoR> UnpackStrategy for FoRStrategy { +impl + BitPackedPhysical> UnpackStrategy for FoRStrategy { #[inline(always)] - unsafe fn unpack_chunk( - &self, - bit_width: usize, - chunk: &[T::Physical], - dst: &mut [T::Physical], - ) { - // SAFETY: Caller ensures chunk and dst have correct sizes. - unsafe { - FoR::unchecked_unfor_pack(bit_width, chunk, self.reference, dst); - } + unsafe fn unpack_chunk(&self, chunk: &[T::Physical], dst: &mut [T::Physical]) { + // SAFETY: The caller upholds the `unpack_chunk` length contract, which is + // `UnforPackFn`'s. + unsafe { (self.unfor_pack)(chunk, self.reference, dst) } } } @@ -80,7 +77,7 @@ pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult + UnsignedPType + FoR + WrappingAdd, + T: PhysicalPType + UnsignedPType + BitPackedPhysical + WrappingAdd, >( for_: &FoRArray, bp: ArrayView<'_, BitPacked>, @@ -92,7 +89,10 @@ pub(crate) fn fused_decompress< .as_::() .vortex_expect("cannot be null"); - let strategy = FoRStrategy { reference: ref_ }; + let strategy = FoRStrategy { + reference: ref_, + unfor_pack: bp.kernels::().unfor_pack, + }; // Create [`UnpackedChunks`] with FoR strategy. let mut unpacked = UnpackedChunks::try_new_with_strategy(