Skip to content

Commit d40e0fd

Browse files
committed
perf: fuse bitpacked chunk encoding
Signed-off-by: "Matt Katz" <mhkatz97@gmail.com>
1 parent 60c53a1 commit d40e0fd

2 files changed

Lines changed: 264 additions & 1 deletion

File tree

‎encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs‎

Lines changed: 211 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
use fastlanes::BitPacking;
55
use itertools::Itertools;
66
use num_traits::PrimInt;
7+
use num_traits::Zero;
78
use vortex_array::ArrayView;
89
use vortex_array::ExecutionCtx;
910
use vortex_array::IntoArray;
@@ -14,10 +15,12 @@ use vortex_array::buffer::BufferHandle;
1415
use vortex_array::dtype::IntegerPType;
1516
use vortex_array::dtype::NativePType;
1617
use vortex_array::dtype::PType;
18+
use vortex_array::dtype::PhysicalPType;
1719
use vortex_array::match_each_integer_ptype;
1820
use vortex_array::match_each_unsigned_integer_ptype;
1921
use vortex_array::patches::Patches;
2022
use vortex_array::validity::Validity;
23+
use vortex_buffer::BitBuffer;
2124
use vortex_buffer::Buffer;
2225
use vortex_buffer::BufferMut;
2326
use vortex_buffer::ByteBuffer;
@@ -35,15 +38,222 @@ use crate::bitpack_decompress::count_exceptions;
3538
use crate::bitpacking::array::ChunkWidths;
3639
use crate::bitpacking::array::chunk_packed_bytes;
3740

38-
/// Choose a cost-model width for each chunk, then pack values and gather exceptions.
41+
/// Bit-pack an array choosing the cost-model-optimal width for every 1024-element chunk.
42+
///
43+
/// Each chunk is charged for its packed block plus the exceptions left behind, so a chunk of small
44+
/// values stays narrow no matter how wide its neighbours are.
45+
///
46+
/// Every chunk is processed in one go while it sits in L1: histogram, width choice, exception
47+
/// gathering and packing, so the values are streamed from memory once.
3948
pub fn bitpack_to_best_chunk_widths(
4049
array: &PrimitiveArray,
4150
ctx: &mut ExecutionCtx,
51+
) -> VortexResult<BitPackedArray> {
52+
ensure_non_negative(array, ctx)?;
53+
let validity = array.validity()?;
54+
let mask = validity.execute_mask(array.len(), ctx)?;
55+
let patch_validity = match validity {
56+
Validity::NonNullable => Validity::NonNullable,
57+
_ => Validity::AllValid,
58+
};
59+
60+
let len = array.len();
61+
let (widths, packed, patches) = match_each_integer_ptype!(array.ptype(), |T| {
62+
encode_chunks::<T>(array.as_slice::<T>(), &mask, patch_validity)?
63+
});
64+
65+
let offsets = widths.offsets_array();
66+
let bitpacked = BitPacked::try_new(
67+
BufferHandle::new_host(packed),
68+
array.ptype(),
69+
validity,
70+
patches,
71+
widths.into_array(),
72+
offsets,
73+
len,
74+
0,
75+
)?;
76+
bitpacked.statistics().inherit_from(array.statistics());
77+
Ok(bitpacked)
78+
}
79+
80+
/// Multi-pass reference for [`bitpack_to_best_chunk_widths`]: one pass to choose widths, one to
81+
/// pack, and one to gather exceptions. Kept to check the fused encoder against.
82+
#[cfg(test)]
83+
pub(crate) fn bitpack_to_best_chunk_widths_multipass(
84+
array: &PrimitiveArray,
85+
ctx: &mut ExecutionCtx,
4286
) -> VortexResult<BitPackedArray> {
4387
let plan = chunk_width_plan(array.as_view(), ctx)?;
4488
bitpack_encode_planned(array, plan, ctx)
4589
}
4690

91+
/// Histogram, choose a width, gather exceptions and pack, one 1024-element chunk at a time.
92+
///
93+
/// Patch indices use the narrowest unsigned type that can address every value.
94+
fn encode_chunks<T>(
95+
values: &[T],
96+
mask: &Mask,
97+
patch_validity: Validity,
98+
) -> VortexResult<(ChunkWidths, ByteBuffer, Option<Patches>)>
99+
where
100+
T: NativePType + PrimInt + PhysicalPType,
101+
T::Physical: BitPacking + NativePType,
102+
{
103+
let len = values.len();
104+
if len < u8::MAX as usize {
105+
encode_chunks_indexed::<T, u8>(values, mask, patch_validity)
106+
} else if len < u16::MAX as usize {
107+
encode_chunks_indexed::<T, u16>(values, mask, patch_validity)
108+
} else if len < u32::MAX as usize {
109+
encode_chunks_indexed::<T, u32>(values, mask, patch_validity)
110+
} else {
111+
encode_chunks_indexed::<T, u64>(values, mask, patch_validity)
112+
}
113+
}
114+
115+
fn encode_chunks_indexed<T, P>(
116+
values: &[T],
117+
mask: &Mask,
118+
patch_validity: Validity,
119+
) -> VortexResult<(ChunkWidths, ByteBuffer, Option<Patches>)>
120+
where
121+
T: NativePType + PrimInt + PhysicalPType,
122+
T::Physical: BitPacking + NativePType,
123+
P: IntegerPType,
124+
{
125+
let bits = T::PTYPE.bit_width();
126+
let bytes_per_exception = bytes_per_exception(T::PTYPE);
127+
let num_chunks = values.len().div_ceil(FL_CHUNK_SIZE);
128+
129+
let mut widths = BufferMut::<u8>::with_capacity(num_chunks);
130+
let mut chunk_offsets = BufferMut::<u64>::with_capacity(num_chunks);
131+
let mut indices = BufferMut::<P>::empty();
132+
let mut patch_values = BufferMut::<T>::empty();
133+
134+
let validity = match mask.bit_buffer() {
135+
AllOr::All => None,
136+
AllOr::Some(bits) => Some(bits),
137+
// Every value is null: every chunk is zero-width and there is nothing to pack.
138+
AllOr::None => {
139+
widths.extend_trusted(std::iter::repeat_n(0u8, num_chunks));
140+
chunk_offsets.extend_trusted(std::iter::repeat_n(0u64, num_chunks));
141+
return Ok((ChunkWidths::new(widths.freeze()), ByteBuffer::empty(), None));
142+
}
143+
};
144+
145+
// Every chunk packs into a whole block, so the padded size bounds the output; a short trailing
146+
// chunk can pack to more than its raw size. The buffer is shrunk to its exact size at the end.
147+
let mut packed = BufferMut::<T::Physical>::with_capacity(num_chunks * FL_CHUNK_SIZE);
148+
let mut histogram = vec![0usize; bits + 1];
149+
// Zero-padded copy of the trailing partial chunk.
150+
let mut padded = [T::Physical::zero(); FL_CHUNK_SIZE];
151+
152+
for (chunk_idx, chunk) in values.chunks(FL_CHUNK_SIZE).enumerate() {
153+
let base = chunk_idx * FL_CHUNK_SIZE;
154+
let chunk_validity = validity.map(|v| v.slice(base..base + chunk.len()));
155+
156+
histogram.fill(0);
157+
for_each_valid_width(chunk, chunk_validity.as_ref(), |_, _, width| {
158+
histogram[width] += 1;
159+
});
160+
161+
let bit_width = best_chunk_width(&histogram, bytes_per_exception);
162+
widths.push(bit_width);
163+
chunk_offsets.push(patch_values.len() as u64);
164+
165+
// The chunk is still in L1, so a second walk over it is cheaper than remembering widths.
166+
if count_exceptions(bit_width, &histogram) > 0 {
167+
for_each_valid_width(chunk, chunk_validity.as_ref(), |i, value, width| {
168+
if width > bit_width as usize {
169+
indices.push(P::from(base + i).vortex_expect("cast index from usize"));
170+
patch_values.push(value);
171+
}
172+
});
173+
}
174+
175+
if bit_width > 0 {
176+
let input: &[T::Physical] = if chunk.len() == FL_CHUNK_SIZE {
177+
as_physical(chunk)
178+
} else {
179+
padded[..chunk.len()].copy_from_slice(as_physical(chunk));
180+
&padded
181+
};
182+
let packed_len = chunk_packed_bytes(bit_width) / size_of::<T::Physical>();
183+
let start = packed.len();
184+
// SAFETY: `input` holds exactly 1024 values and the output window is exactly one
185+
// packed block at `bit_width`, within the raw-size capacity reserved above.
186+
unsafe {
187+
packed.set_len(start + packed_len);
188+
BitPacking::unchecked_pack(bit_width as usize, input, &mut packed[start..]);
189+
}
190+
}
191+
}
192+
193+
let packed = if packed.len() < packed.capacity() {
194+
let mut exact = BufferMut::<T::Physical>::with_capacity(packed.len());
195+
exact.extend_from_slice(&packed);
196+
exact.freeze()
197+
} else {
198+
packed.freeze()
199+
};
200+
201+
let patches = if indices.is_empty() {
202+
None
203+
} else {
204+
Some(Patches::new(
205+
values.len(),
206+
0,
207+
indices.into_array(),
208+
PrimitiveArray::new(patch_values, patch_validity).into_array(),
209+
Some(chunk_offsets.into_array()),
210+
)?)
211+
};
212+
213+
Ok((
214+
ChunkWidths::new(widths.freeze()),
215+
packed.into_byte_buffer(),
216+
patches,
217+
))
218+
}
219+
220+
/// Call `f(index, value, bit_width)` for every value of `chunk`; nulls report a width of zero.
221+
#[inline]
222+
fn for_each_valid_width<T: NativePType + PrimInt>(
223+
chunk: &[T],
224+
validity: Option<&BitBuffer>,
225+
mut f: impl FnMut(usize, T, usize),
226+
) {
227+
let bits = T::PTYPE.bit_width();
228+
match validity {
229+
None => {
230+
for (i, &v) in chunk.iter().enumerate() {
231+
f(i, v, bits - PrimInt::leading_zeros(v) as usize);
232+
}
233+
}
234+
Some(validity) => {
235+
for ((i, &v), valid) in chunk.iter().enumerate().zip(validity.iter()) {
236+
let width = if valid {
237+
bits - PrimInt::leading_zeros(v) as usize
238+
} else {
239+
0
240+
};
241+
f(i, v, width);
242+
}
243+
}
244+
}
245+
}
246+
247+
/// View signed or unsigned values as their unsigned physical twin, which FastLanes packs.
248+
fn as_physical<T: PhysicalPType>(values: &[T]) -> &[T::Physical] {
249+
const {
250+
assert!(size_of::<T>() == size_of::<T::Physical>());
251+
assert!(align_of::<T>() == align_of::<T::Physical>());
252+
}
253+
// SAFETY: `Physical` is the same-width unsigned integer, so the layouts match exactly.
254+
unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) }
255+
}
256+
47257
/// The cost-model-optimal bit width of every 1024-element chunk of `array`.
48258
pub fn best_chunk_widths(
49259
array: ArrayView<'_, Primitive>,

‎encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ use crate::bitpacked_v2_id;
6565
use crate::bitpacking::bitpack_compress::bitpack_encode_with_widths;
6666
use crate::bitpacking::bitpack_compress::bitpack_to_best_bit_width;
6767
use crate::bitpacking::bitpack_compress::bitpack_to_best_chunk_widths;
68+
use crate::bitpacking::bitpack_compress::bitpack_to_best_chunk_widths_multipass;
6869
use crate::bitpacking::plugin::BitPackedV2Metadata;
6970
use crate::bitpacking::vtable::BitPackedMetadata;
7071

@@ -833,3 +834,55 @@ fn conformance(#[case] array: BitPackedArray) {
833834
test_cast_conformance(&array, &mut ctx);
834835
test_binary_numeric_array(&array, &mut ctx);
835836
}
837+
838+
/// The fused single-walk encoder must produce exactly what the multi-pass one does.
839+
#[rstest]
840+
#[case::varied(PrimitiveArray::from_iter(varied(100)))]
841+
#[case::varied_exact(PrimitiveArray::from_iter(varied(0)))]
842+
#[case::tiny(PrimitiveArray::from_iter([5u32, 1 << 20, 7]))]
843+
#[case::nullable_signed(PrimitiveArray::new(
844+
Buffer::from_iter(varied(50).into_iter().map(|v| v as i32)),
845+
Validity::from_iter((0..4 * FL_CHUNK_SIZE + 50).map(|i| i % 7 != 0)),
846+
))]
847+
#[case::all_null(PrimitiveArray::new(Buffer::from_iter(varied(9)), Validity::AllInvalid))]
848+
#[case::short_and_wide(short_and_wide())]
849+
fn fused_matches_multipass(#[case] array: PrimitiveArray) -> VortexResult<()> {
850+
let mut ctx = SESSION.create_execution_ctx();
851+
let fused = bitpack_to_best_chunk_widths(&array, &mut ctx)?;
852+
let multipass = bitpack_to_best_chunk_widths_multipass(&array, &mut ctx)?;
853+
assert_eq!(
854+
fused.chunk_widths(&mut SESSION.create_execution_ctx())?,
855+
multipass.chunk_widths(&mut SESSION.create_execution_ctx())?
856+
);
857+
assert_eq!(fused.packed().as_host(), multipass.packed().as_host());
858+
assert_eq!(
859+
fused.patches().map(|p| p.num_patches()),
860+
multipass.patches().map(|p| p.num_patches())
861+
);
862+
assert_eq!(fused.nbytes(), multipass.nbytes());
863+
assert_arrays_eq!(fused, array, &mut ctx);
864+
Ok(())
865+
}
866+
867+
/// 200 u8 values needing 7 bits: the single padded block (896 bytes) is larger than the raw
868+
/// array (200 bytes), which an encoder sizing its output by the raw length overflows.
869+
fn short_and_wide() -> PrimitiveArray {
870+
PrimitiveArray::from_iter((0..200u8).map(|i| i.wrapping_mul(97) % 128))
871+
}
872+
873+
#[test]
874+
fn short_chunk_packs_wider_than_raw() -> VortexResult<()> {
875+
let mut ctx = SESSION.create_execution_ctx();
876+
let array = short_and_wide();
877+
let packed = bitpack_to_best_chunk_widths(&array, &mut ctx)?;
878+
assert_eq!(
879+
packed
880+
.chunk_widths(&mut SESSION.create_execution_ctx())?
881+
.as_buffer()
882+
.as_slice(),
883+
&[7]
884+
);
885+
assert!(packed.packed().len() > array.nbytes() as usize);
886+
assert_arrays_eq!(packed, array, &mut ctx);
887+
Ok(())
888+
}

0 commit comments

Comments
 (0)