diff --git a/encodings/fastlanes/src/bit_transpose/validity.rs b/encodings/fastlanes/src/bit_transpose/validity.rs index 826e1a2ad35..a6b2ec43cbc 100644 --- a/encodings/fastlanes/src/bit_transpose/validity.rs +++ b/encodings/fastlanes/src/bit_transpose/validity.rs @@ -107,8 +107,10 @@ fn bits_op_with_copy( let output_len = bytes.len().next_multiple_of(128); let mut output = ByteBufferMut::with_capacity(output_len); let (input_chunks, input_trailer) = bytes.as_chunks::<128>(); - // We can ignore the spare trailer capacity that can be an artifact of allocator as we requested 128 multiple chunks - let (output_chunks, _) = output.spare_capacity_mut().as_chunks_mut::<128>(); + // Bound to the requested `output_len`: `spare_capacity_mut` may expose extra over-aligned + // capacity, which would otherwise split into spurious trailing chunks and make `last_mut` + // below target a chunk past the data we actually initialize. + let (output_chunks, _) = output.spare_capacity_mut()[..output_len].as_chunks_mut::<128>(); for (input, output) in input_chunks.iter().zip(output_chunks.iter_mut()) { op(input, unsafe { @@ -189,4 +191,25 @@ mod tests { let roundtripped = untranspose_bitbuffer(transposed); assert_eq!(bits, roundtripped.slice(0..original_len)); } + + /// Regression: the copy path split the over-aligned spare capacity into extra 128-byte + /// chunks and wrote the padded remainder via `last_mut()`, which landed past the requested + /// length and left the real trailing chunk uninitialized. Whether the surplus capacity + /// produced an extra chunk depended on the allocation address, so we repeat across sizes to + /// defeat that luck; the fix bounds the spare slice so the result no longer depends on it. + #[test] + fn transpose_copy_path_survives_overallocation() { + for original_len in [129, 500, 1500, 9999] { + let bits = make_validity_bits(original_len); + for _ in 0..64 { + let transposed = transpose_bitbuffer(bits.clone()); + let roundtripped = untranspose_bitbuffer(transposed); + assert_eq!( + roundtripped.slice(0..original_len), + bits, + "len={original_len}" + ); + } + } + } } diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index fe2567e63c7..7dcbeb44950 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -70,7 +70,9 @@ where // Allocate a result array. let mut output = BufferMut::with_capacity(deltas.len()); - let (output_chunks, _) = output.spare_capacity_mut().as_chunks_mut::<1024>(); + // Bound to the requested length: `spare_capacity_mut` may expose extra over-aligned capacity + // beyond `deltas.len()`, which would desync the `zip_eq` with `chunks` below and panic. + let (output_chunks, _) = output.spare_capacity_mut()[..deltas.len()].as_chunks_mut::<1024>(); // Loop over all the chunks let mut transposed: [T; 1024] = [T::default(); 1024]; diff --git a/vortex-array/src/builders/primitive.rs b/vortex-array/src/builders/primitive.rs index 3738cb65682..870ad2275bd 100644 --- a/vortex-array/src/builders/primitive.rs +++ b/vortex-array/src/builders/primitive.rs @@ -594,12 +594,10 @@ mod tests { /// Test that creating an uninit range exceeding capacity panics. #[test] - #[should_panic( - expected = "uninit_range of len 10 exceeds builder with length 0 and capacity 6" - )] + #[should_panic(expected = "uninit_range of len 261 exceeds builder with length 0 and capacity")] fn test_uninit_range_exceeds_capacity_panics() { let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 5); - let _range = builder.uninit_range(10); + let _range = builder.uninit_range(261); } /// Test that `copy_from_slice` debug asserts on out-of-bounds access. diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs index bfb56a27670..c517972fb11 100644 --- a/vortex-array/src/builders/varbinview.rs +++ b/vortex-array/src/builders/varbinview.rs @@ -6,6 +6,7 @@ use std::ops::Range; use std::sync::Arc; use itertools::Itertools; +use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; @@ -41,7 +42,7 @@ pub struct VarBinViewBuilder { views_builder: BufferMut, nulls: LazyBitBufferBuilder, completed: CompletedBuffers, - in_progress: ByteBufferMut, + in_progress: Option, growth_strategy: BufferGrowthStrategy, compaction_threshold: f64, } @@ -83,10 +84,14 @@ impl VarBinViewBuilder { "VarBinViewBuilder DType must be Utf8 or Binary." ); Self { - views_builder: BufferMut::::with_capacity(capacity), + views_builder: BufferMut::with_capacity_preferred_aligned( + capacity, + Alignment::of::(), + None, + ), nulls: LazyBitBufferBuilder::new(capacity), completed, - in_progress: ByteBufferMut::empty(), + in_progress: None, dtype, growth_strategy, compaction_threshold, @@ -129,15 +134,14 @@ impl VarBinViewBuilder { } fn flush_in_progress(&mut self) { - if self.in_progress.is_empty() { + let Some(block) = self.in_progress.take() else { return; - } - let block = std::mem::take(&mut self.in_progress).freeze(); + }; assert!(block.len() < u32::MAX as usize, "Block too large"); let initial_len = self.completed.len(); - self.completed.push(block); + self.completed.push(block.freeze()); assert_eq!( self.completed.len(), initial_len + 1, @@ -145,23 +149,41 @@ impl VarBinViewBuilder { ); } + fn init_in_progress(&mut self, min_len: usize) { + let next_buffer_size = self.growth_strategy.next_size() as usize; + let to_reserve = next_buffer_size.max(min_len); + self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned( + to_reserve, + Alignment::of::(), + None, + )); + } + /// append a non inlined value to self.in_progress. fn append_value_to_buffer(&mut self, value: &[u8]) -> (u32, u32) { assert!( value.len() > BinaryView::MAX_INLINED_SIZE, "must inline small strings" ); - let required_cap = self.in_progress.len() + value.len(); - if self.in_progress.capacity() < required_cap { - self.flush_in_progress(); - let next_buffer_size = self.growth_strategy.next_size() as usize; - let to_reserve = next_buffer_size.max(value.len()); - self.in_progress.reserve(to_reserve); - } + + if let Some(in_progress) = &mut self.in_progress { + let required_cap = in_progress.len() + value.len(); + if in_progress.capacity() < required_cap { + self.flush_in_progress(); + self.init_in_progress(value.len()); + } + } else { + self.init_in_progress(value.len()) + }; + + let in_progress = self + .in_progress + .as_mut() + .vortex_expect("in_progress just set"); let buffer_idx = self.completed.len(); - let offset = u32::try_from(self.in_progress.len()).vortex_expect("too many buffers"); - self.in_progress.extend_from_slice(value); + let offset = u32::try_from(in_progress.len()).vortex_expect("too many buffers"); + in_progress.extend_from_slice(value); (buffer_idx, offset) } @@ -173,7 +195,7 @@ impl VarBinViewBuilder { /// Returns true if a non-empty in-progress buffer is staged (and would /// become a completed buffer on the next flush), false otherwise. pub fn in_progress(&self) -> bool { - !self.in_progress.is_empty() + self.in_progress.is_some() } /// Pushes buffers and pre-adjusted views into the builder. diff --git a/vortex-buffer/src/alignment.rs b/vortex-buffer/src/alignment.rs index 43a4c29c8b7..1c7bbe24da3 100644 --- a/vortex-buffer/src/alignment.rs +++ b/vortex-buffer/src/alignment.rs @@ -19,6 +19,15 @@ impl Alignment { /// Default alignment for device-to-host buffer copies. pub const HOST_COPY: Self = Alignment::new(256); + /// Default alignment for all buffers. + /// + /// Chosen to be larger than any SIMD register (e.g. AVX-512's 64-byte + /// registers) so that buffers can be processed with vectorized loads/stores + /// without alignment fixups, and to match the alignment guarantees of the + /// CUDA allocator (256 bytes) so host buffers can be copied to/from device + /// memory without re-alignment. + pub const DEFAULT_ALIGNMENT: Self = Alignment::new(256); + /// Create a new alignment. /// /// ## Panics diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs index 08640526c22..98398214108 100644 --- a/vortex-buffer/src/buffer.rs +++ b/vortex-buffer/src/buffer.rs @@ -95,8 +95,25 @@ impl Buffer { } /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. + /// + /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned pub fn copy_from_aligned(values: impl AsRef<[T]>, alignment: Alignment) -> Self { - BufferMut::copy_from_aligned(values, alignment).freeze() + Self::copy_from_preferred_aligned(values, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Returns a new `Buffer` copied from the provided slice and with the requested alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn copy_from_preferred_aligned( + values: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + BufferMut::copy_from_preferred_aligned(values, alignment, preferred_alignment).freeze() } /// Create a new zeroed `Buffer` with the given value. @@ -104,9 +121,26 @@ impl Buffer { Self::zeroed_aligned(len, Alignment::of::()) } - /// Create a new zeroed `Buffer` with the given value. + /// Create a new zeroed `Buffer` with the requested alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. + /// + /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - BufferMut::zeroed_aligned(len, alignment).freeze() + Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a new zeroed `Buffer` with the requested alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn zeroed_preferred_aligned( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + BufferMut::zeroed_preferred_aligned(len, alignment, preferred_alignment).freeze() } /// Create a new empty `ByteBuffer` with the provided alignment. @@ -803,6 +837,29 @@ mod test { assert_eq!(buf.as_slice(), &[0; LEN]); } + #[test] + fn copy_from_over_aligns_to_default() { + let values = [1u32, 2, 3]; + let buf = Buffer::::copy_from(values); + + // The buffer reports the scalar type's alignment, ... + assert_eq!(buf.alignment(), Alignment::of::()); + // ... but the underlying allocation is over-aligned to DEFAULT_ALIGNMENT. + assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT)); + assert_eq!(buf.as_slice(), &values); + } + + #[test] + fn zeroed_over_aligns_to_default() { + const LEN: usize = 17; + + let buf = Buffer::::zeroed(LEN); + + assert_eq!(buf.alignment(), Alignment::of::()); + assert!(buf.is_aligned(Alignment::DEFAULT_ALIGNMENT)); + assert_eq!(buf.as_slice(), &[0; LEN]); + } + #[test] fn from_vec() { let vec = vec![1, 2, 3, 4, 5]; diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs index 6b3885e18af..a7edfb00293 100644 --- a/vortex-buffer/src/buffer_mut.rs +++ b/vortex-buffer/src/buffer_mut.rs @@ -3,6 +3,7 @@ use core::mem::MaybeUninit; use std::any::type_name; +use std::cmp::max; use std::fmt::Debug; use std::fmt::Formatter; use std::io::Write; @@ -38,7 +39,33 @@ impl BufferMut { } /// Create a new `BufferMut` with the requested alignment and capacity. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`with_capacity_preferred_aligned`] to control the over-alignment. + /// + /// [`with_capacity_preferred_aligned`]: Self::with_capacity_preferred_aligned pub fn with_capacity_aligned(capacity: usize, alignment: Alignment) -> Self { + Self::with_capacity_preferred_aligned( + capacity, + alignment, + Some(Alignment::DEFAULT_ALIGNMENT), + ) + } + + /// Create a new `BufferMut` with the requested alignment and capacity. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn with_capacity_preferred_aligned( + capacity: usize, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + let actual = max( + alignment, + preferred_alignment.unwrap_or(Alignment::of::()), + ); + if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!( "Alignment {} must align to the scalar type's alignment {}", @@ -47,8 +74,8 @@ impl BufferMut { ); } - let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *alignment); - bytes.align_empty(alignment); + let mut bytes = BytesMut::with_capacity((capacity * size_of::()) + *actual); + bytes.align_empty(actual); Self { bytes, @@ -63,10 +90,29 @@ impl BufferMut { Self::zeroed_aligned(len, Alignment::of::()) } - /// Create a new zeroed `BufferMut`. + /// Create a new zeroed `BufferMut` with the requested alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`zeroed_preferred_aligned`] to control the over-alignment. + /// + /// [`zeroed_preferred_aligned`]: Self::zeroed_preferred_aligned pub fn zeroed_aligned(len: usize, alignment: Alignment) -> Self { - let mut bytes = BytesMut::zeroed((len * size_of::()) + *alignment); - bytes.advance(bytes.as_ptr().align_offset(*alignment)); + Self::zeroed_preferred_aligned(len, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a new zeroed `BufferMut` with the requested alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn zeroed_preferred_aligned( + len: usize, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + let preferred_alignment = preferred_alignment.unwrap_or(Alignment::of::()); + let actual_alignment = max(preferred_alignment, alignment); + let mut bytes = BytesMut::zeroed((len * size_of::()) + *actual_alignment); + bytes.advance(bytes.as_ptr().align_offset(*actual_alignment)); unsafe { bytes.set_len(len * size_of::()) }; let actual_len = bytes.len().checked_div(size_of::()).unwrap_or(0); Self { @@ -83,8 +129,24 @@ impl BufferMut { } /// Create a new empty `BufferMut` with the provided alignment. + /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`empty_preferred_aligned`] to control the over-alignment. + /// + /// [`empty_preferred_aligned`]: Self::empty_preferred_aligned pub fn empty_aligned(alignment: Alignment) -> Self { - BufferMut::with_capacity_aligned(0, alignment) + Self::empty_preferred_aligned(alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a new empty `BufferMut` with the provided alignment. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + pub fn empty_preferred_aligned( + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { + BufferMut::with_capacity_preferred_aligned(0, alignment, preferred_alignment) } /// Create a new full `BufferMut` with the given value. @@ -104,15 +166,37 @@ impl BufferMut { /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. /// + /// The allocation is over-aligned to [`Alignment::DEFAULT_ALIGNMENT`] when that is larger than + /// `alignment`. Use [`copy_from_preferred_aligned`] to control the over-alignment. + /// + /// [`copy_from_preferred_aligned`]: Self::copy_from_preferred_aligned + /// /// ## Panics /// /// Panics when the requested alignment isn't itself aligned to type T. pub fn copy_from_aligned(other: impl AsRef<[T]>, alignment: Alignment) -> Self { + Self::copy_from_preferred_aligned(other, alignment, Some(Alignment::DEFAULT_ALIGNMENT)) + } + + /// Create a mutable scalar buffer with the alignment by copying the contents of the slice. + /// + /// The buffer reports `alignment`, but the underlying allocation is over-aligned to the larger + /// of `alignment` and `preferred_alignment`. + /// + /// ## Panics + /// + /// Panics when the requested alignment isn't itself aligned to type T. + pub fn copy_from_preferred_aligned( + other: impl AsRef<[T]>, + alignment: Alignment, + preferred_alignment: Option, + ) -> Self { if !alignment.is_aligned_to(Alignment::of::()) { vortex_panic!("Given alignment is not aligned to type T") } let other = other.as_ref(); - let mut buffer = Self::with_capacity_aligned(other.len(), alignment); + let mut buffer = + Self::with_capacity_preferred_aligned(other.len(), alignment, preferred_alignment); buffer.extend_from_slice(other); debug_assert_eq!(buffer.alignment(), alignment); buffer @@ -214,6 +298,10 @@ impl BufferMut { /// reading from a file) before marking the data as initialized using the /// [`set_len`] method. /// + /// Note that the returned slice may be larger than the capacity requested at + /// construction, since the underlying allocation can be rounded up (e.g. to + /// satisfy alignment requirements). + /// /// [`set_len`]: BufferMut::set_len /// [`Vec::spare_capacity_mut`]: Vec::spare_capacity_mut ///