diff --git a/vortex-cuda/benches/arrow_validity_cuda.rs b/vortex-cuda/benches/arrow_validity_cuda.rs index d1142e361f4..01420717560 100644 --- a/vortex-cuda/benches/arrow_validity_cuda.rs +++ b/vortex-cuda/benches/arrow_validity_cuda.rs @@ -41,7 +41,7 @@ const INPUT_OFFSET: usize = 5; const ARROW_OFFSET: usize = 3; const EXPORT_BENCH_SIZES: &[(usize, &str)] = &[(100_000_000, "100M")]; -fn validity_bitmap_byte_len(len: usize, bit_offset: usize) -> usize { +fn arrow_bitmap_byte_len(len: usize, bit_offset: usize) -> usize { (bit_offset + len).div_ceil(8) } @@ -98,7 +98,7 @@ fn benchmark_arrow_validity_export(c: &mut Criterion) { [("device_bitmap", 0), ("device_bitmap_repack", INPUT_OFFSET)] { group.throughput(Throughput::Bytes( - validity_bitmap_byte_len(len, validity_offset) as u64, + arrow_bitmap_byte_len(len, validity_offset) as u64, )); group.bench_with_input( BenchmarkId::new(format!("cuda/arrow_validity/export/{case}"), len_label), @@ -148,7 +148,7 @@ fn benchmark_arrow_validity_repack(c: &mut Criterion) { for &(len, len_label) in bench_config::BENCH_SIZES { group.throughput(Throughput::Bytes( - validity_bitmap_byte_len(len, INPUT_OFFSET) as u64, + arrow_bitmap_byte_len(len, INPUT_OFFSET) as u64 )); group.bench_with_input( BenchmarkId::new("cuda/arrow_validity/repack", len_label), @@ -165,7 +165,7 @@ fn benchmark_arrow_validity_repack(c: &mut Criterion) { .vortex_expect("failed to create validity fixture"); for _ in 0..iters { - let output = test_harness::repack_arrow_validity_buffer( + let output = test_harness::repack_arrow_bitmap( &input_buffer, input_offset, len, @@ -190,7 +190,7 @@ fn benchmark_arrow_validity_count_nulls(c: &mut Criterion) { for &(len, len_label) in bench_config::BENCH_SIZES { group.throughput(Throughput::Bytes( - validity_bitmap_byte_len(len, ARROW_OFFSET) as u64, + arrow_bitmap_byte_len(len, ARROW_OFFSET) as u64 )); group.bench_with_input( BenchmarkId::new("cuda/arrow_validity/count_nulls", len_label), diff --git a/vortex-cuda/kernels/src/arrow_validity.cu b/vortex-cuda/kernels/src/arrow_validity.cu index d139fa66f25..c048d3cb687 100644 --- a/vortex-cuda/kernels/src/arrow_validity.cu +++ b/vortex-cuda/kernels/src/arrow_validity.cu @@ -18,18 +18,21 @@ __device__ uint64_t load_input_word(const uint8_t *const input, int64_t word_idx if (byte_idx >= input_bytes) { return 0; } - if (byte_idx + sizeof(uint64_t) <= input_bytes) { - return reinterpret_cast(input)[word_idx]; + const uint64_t available_bytes = input_bytes - byte_idx; + // The host aligns input down within its allocation and adjusts the bit offset. + if (available_bytes >= sizeof(uint64_t)) { + return reinterpret_cast(input + byte_idx)[0]; } - // Trailing partial word: assemble byte-by-byte to avoid reading past the buffer. + // Only the final partial word needs byte loads; never overread the logical tail. + const uint64_t min_bytes = min(static_cast(sizeof(uint64_t)), available_bytes); uint64_t word = 0; - for (uint64_t i = byte_idx; i < input_bytes; i++) { - word |= static_cast(input[i]) << ((i - byte_idx) * 8); + for (uint64_t i = 0; i < min_bytes; i++) { + word |= static_cast(input[byte_idx + i]) << (i * 8); } return word; } -// Build one output word for sliced validity. The row bits are the same, but +// Build one output word for a sliced bitmap. The row bits are the same, but // row 0 may live at a different bit position in the source and Arrow bitmaps. // For example, `input_offset = 5` and `arrow_offset = 0` shifts row0 from bit 5 // in the input bitmap to bit 0 in the Arrow bitmap. @@ -39,7 +42,7 @@ __device__ uint64_t load_input_word(const uint8_t *const input, int64_t word_idx // Arrow bitmap: [ row0 ][ row1 ][ row2 ].... // ^ arrow_offset // -// Padding bits are cleared so word-sized validity readers can safely over-read. +// Padding bits are cleared for word-sized bitmap readers. __device__ uint64_t repack_word(const uint8_t *const input, uint64_t word_idx, int64_t shift, @@ -137,7 +140,7 @@ __device__ uint64_t block_sum_to_thread_zero(uint64_t value, uint64_t *const war } // namespace -// Repack sliced validity when the source bitmap offset does not match the +// Repack a sliced bitmap when the source bit offset does not match the // Arrow array offset. Each thread writes independent output words. // // thread 0 -> output word 0, word N, ... diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index e872cfdab51..12e19544615 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -18,6 +18,7 @@ use vortex::array::Canonical; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::arrays::Bool; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::Dict; use vortex::array::arrays::DictArray; @@ -44,6 +45,7 @@ use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::struct_::StructDataParts; use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; +use vortex::array::buffer::DeviceBuffer; use vortex::array::builtins::ArrayBuiltins; use vortex::array::match_each_decimal_value_type; use vortex::array::validity::Validity; @@ -80,11 +82,17 @@ use crate::arrow::cuda_decimal_value_type; use crate::arrow::list_view::export_device_list_view; use crate::cub::exclusive_sum_i32; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; +use crate::device_buffer::cuda_aligned_bitmap_view; use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; use crate::kernel::decode_onpair_varbin; +use crate::stream::zero_padding; + +const BITMAP_THREADS_PER_BLOCK: u32 = 256; +// Exercise grid-stride wrapping on small bitmaps in tests. +const MAX_BITMAP_BLOCKS: u32 = if cfg!(test) { 2 } else { 4096 }; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -317,24 +325,7 @@ fn export_canonical( let buffer = ctx.ensure_on_device(buffer).await?; export_fixed_size(buffer, len, 0, validity_buffer, null_count, ctx) } - Canonical::Bool(bool_array) => { - let len = bool_array.len(); - let validity = bool_array.validity()?; - let BoolDataParts { bits, meta } = bool_array.into_data().into_parts(len); - - let (validity_buffer, null_count) = - export_arrow_validity_buffer(validity, len, meta.offset(), ctx).await?; - - let bits = ctx.ensure_on_device(bits).await?; - export_fixed_size( - bits, - meta.len(), - meta.offset(), - validity_buffer, - null_count, - ctx, - ) - } + Canonical::Bool(bool_array) => export_bool(bool_array, ctx).await, Canonical::List(listview) => export_list_view(listview, ctx).await, Canonical::FixedSizeList(fixed_size_list) => { export_fixed_size_list(fixed_size_list, ctx).await @@ -350,6 +341,26 @@ fn export_canonical( }) } +async fn export_bool( + array: BoolArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let len = array.len(); + let validity = array.validity()?; + let BoolDataParts { bits, meta } = array.into_data().into_parts(len); + let (validity_buffer, null_count) = + export_arrow_validity_buffer(validity, len, meta.offset(), ctx).await?; + + let bits = ctx.ensure_on_device(bits).await?; + // cuDF reads bool values as padded mask words. Values and validity share Arrow's offset. + let bits = if len == 0 { + bits + } else { + export_arrow_bitmap(&bits, meta.offset(), len, meta.offset(), ctx)? + }; + export_fixed_size(bits, len, meta.offset(), validity_buffer, null_count, ctx) +} + /// Export a Vortex dictionary array as an Arrow Device dictionary array. /// /// Owns the codes buffers and recursively exported dictionary values. @@ -865,7 +876,7 @@ pub(super) async fn export_arrow_validity_buffer( // This only marks every row null via buffer 0, the validity bitmap. Validity::AllInvalid => Ok(( Some(device_zeroed_byte_buffer( - validity_bitmap_byte_len(len, arrow_offset)?, + arrow_bitmap_byte_len(len, arrow_offset)?, ctx, )?), i64::try_from(len)?, @@ -879,18 +890,7 @@ pub(super) async fn export_arrow_validity_buffer( })?; let BoolDataParts { bits, meta } = array.into_data().into_parts(len); let bitmap = ctx.ensure_on_device(bits).await?; - let bitmap = - match export_arrow_validity_bitmap(&bitmap, meta.offset(), len, arrow_offset, ctx)? - { - Some(bitmap) => bitmap, - None => repack_arrow_validity_buffer( - &bitmap, - meta.offset(), - len, - arrow_offset, - ctx, - )?, - }; + let bitmap = export_arrow_bitmap(&bitmap, meta.offset(), len, arrow_offset, ctx)?; // Keep nullable exports self-describing for consumers that require exact null counts. let null_count = count_arrow_validity_nulls(&bitmap, len, arrow_offset, ctx)?; Ok((Some(bitmap), null_count)) @@ -898,15 +898,15 @@ pub(super) async fn export_arrow_validity_buffer( } } -/// Return the byte length needed for `len` validity bits at the given bit offset. -fn validity_bitmap_byte_len(len: usize, arrow_offset: usize) -> VortexResult { +/// Byte length of a bitmap with `len` bits at the given offset. +fn arrow_bitmap_byte_len(len: usize, arrow_offset: usize) -> VortexResult { Ok(len .checked_add(arrow_offset) - .ok_or_else(|| vortex_err!("Arrow validity bit length overflows usize"))? + .ok_or_else(|| vortex_err!("Arrow bitmap bit length overflows usize"))? .div_ceil(8)) } -/// Allocate a zeroed device buffer with cuDF-safe padding for Arrow validity masks. +/// Allocate a zeroed bitmap with cuDF-safe padding. fn device_zeroed_byte_buffer( byte_len: usize, ctx: &mut CudaExecutionCtx, @@ -916,40 +916,44 @@ fn device_zeroed_byte_buffer( "zero-length validity buffers should be omitted" ); let allocation_len = byte_len.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); - let mut buffer = ctx.device_alloc::(allocation_len)?; - ctx.stream() - .memset_zeros(&mut buffer) - .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer: {err}"))?; - // The memset above zeroed the whole allocation, including cuDF tail padding. - Ok( - BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new_with_zeroed_tail(buffer, 0)?)) - .slice(0..byte_len), - ) + let buffer = ctx + .stream() + .alloc_zeros::(allocation_len) + .map_err(|err| vortex_err!("Failed to allocate zeroed Arrow validity buffer: {err}"))?; + // The whole allocation is zeroed, including cuDF tail padding. + Ok(BufferHandle::new_device( + CudaDeviceBuffer::new_with_zeroed_tail(buffer, 0)?.slice(0..byte_len), + )) } -/// Exports a matching-offset bitmap by reusing it or copying it into zero-padded storage. -fn export_arrow_validity_bitmap( +/// Export a bitmap with cuDF-safe storage, repacking it when the bit offsets differ. +fn export_arrow_bitmap( bitmap: &BufferHandle, input_offset: usize, len: usize, arrow_offset: usize, ctx: &mut CudaExecutionCtx, -) -> VortexResult> { +) -> VortexResult { if input_offset != arrow_offset { - return Ok(None); + return repack_arrow_bitmap(bitmap, input_offset, len, arrow_offset, ctx); } - let output_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + let output_bytes = arrow_bitmap_byte_len(len, arrow_offset)?; let allocation_bytes = output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); - if bitmap.has_zeroed_tail_padding(output_bytes, allocation_bytes)? { - return Ok(Some(bitmap.slice(0..output_bytes))); + // cuDF reads both validity and bool values as 4-byte-aligned mask words. + if bitmap + .cuda_device_ptr()? + .is_multiple_of(size_of::() as u64) + && bitmap.has_zeroed_tail_padding(output_bytes, allocation_bytes)? + { + return Ok(bitmap.slice(0..output_bytes)); } - copy_arrow_validity_buffer(bitmap, output_bytes, ctx).map(Some) + copy_arrow_bitmap(bitmap, output_bytes, ctx) } -/// Copies a validity bitmap into a new cuDF-padded buffer without shifting bits. -fn copy_arrow_validity_buffer( +/// Copy a bitmap into cuDF-padded storage without shifting bits. +fn copy_arrow_bitmap( input_buffer: &BufferHandle, output_bytes: usize, ctx: &mut CudaExecutionCtx, @@ -966,23 +970,17 @@ fn copy_arrow_validity_buffer( let allocation_bytes = output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); let mut output = ctx.device_alloc::(allocation_bytes)?; - ctx.stream() - .memset_zeros(&mut output) - .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer padding: {err}"))?; let input_view = input_buffer.cuda_view::()?.slice(0..output_bytes); - let mut output_view = output.slice_mut(0..output_bytes); ctx.stream() - .memcpy_dtod(&input_view, &mut output_view) + .memcpy_dtod(&input_view, &mut output) .map_err(|err| vortex_err!("Failed to copy Arrow validity buffer: {err}"))?; + // The copy initializes the logical bytes; only the tail needs zeroing. + zero_padding(ctx.stream(), &mut output, output_bytes)?; - Ok( - BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new_with_zeroed_tail( - output, - output_bytes, - )?)) - .slice(0..output_bytes), - ) + Ok(BufferHandle::new_device( + CudaDeviceBuffer::new_with_zeroed_tail(output, output_bytes)?.slice(0..output_bytes), + )) } pub fn count_arrow_validity_nulls( @@ -995,7 +993,7 @@ pub fn count_arrow_validity_nulls( return Ok(0); } - let expected_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + let expected_bytes = arrow_bitmap_byte_len(len, arrow_offset)?; vortex_ensure!( bitmap.len() >= expected_bytes, "Arrow validity bitmap has {} bytes, expected at least {expected_bytes}", @@ -1012,13 +1010,16 @@ pub fn count_arrow_validity_nulls( let arrow_offset = u64::try_from(arrow_offset)?; let kernel = ctx.load_function_with_suffixes("arrow_validity", &["count_valid"])?; - const COUNT_THREADS_PER_BLOCK: u32 = 256; - const MAX_COUNT_BLOCKS: u32 = 4096; - let num_blocks = u32::try_from(expected_bytes.div_ceil(COUNT_THREADS_PER_BLOCK as usize))? - .clamp(1, MAX_COUNT_BLOCKS); + #[allow( + clippy::cast_possible_truncation, + reason = "capped at MAX_BITMAP_BLOCKS" + )] + let num_blocks = expected_bytes + .div_ceil(BITMAP_THREADS_PER_BLOCK as usize) + .clamp(1, MAX_BITMAP_BLOCKS as usize) as u32; let config = LaunchConfig { grid_dim: (num_blocks, 1, 1), - block_dim: (COUNT_THREADS_PER_BLOCK, 1, 1), + block_dim: (BITMAP_THREADS_PER_BLOCK, 1, 1), shared_mem_bytes: 0, }; ctx.launch_kernel_config(&kernel, config, expected_bytes, |args| { @@ -1039,26 +1040,21 @@ pub fn count_arrow_validity_nulls( Ok(i64::try_from(len - valid_count)?) } -/// Repack a validity bitmap into Arrow layout without copying bitmap bits back to the CPU. -/// -/// Vortex bitmaps may start at any bit offset. Arrow exposes only a byte-addressed validity buffer -/// plus an array offset, so sliced compact exports need a GPU rewrite when either side has a -/// bit-level offset. The output handle keeps Arrow's logical byte length, while the backing -/// allocation is zero-padded to cuDF's mask allocation size for consumers that read full masks. -pub fn repack_arrow_validity_buffer( +/// Repack validity or bool values on the GPU to match Arrow's bit offset. +/// The output retains its logical byte length with zeroed, cuDF-sized tail padding. +pub fn repack_arrow_bitmap( input_buffer: &BufferHandle, input_offset: usize, len: usize, arrow_offset: usize, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let output_bytes = validity_bitmap_byte_len(len, arrow_offset)?; + let output_bytes = arrow_bitmap_byte_len(len, arrow_offset)?; vortex_ensure!( output_bytes > 0, "zero-length validity buffers should be omitted" ); - // The CUDA kernel writes the bitmap as u64 words, so round the logical byte length up to the - // number of words that cover the exported Arrow bytes. + // The nonzero byte count guarantees at least one u64 output word. let output_words = output_bytes.div_ceil(size_of::()); // `device_alloc::` takes a word count, while the padding policy is expressed in bytes. // Round up so the padded byte allocation is fully represented by whole u64 words. @@ -1066,49 +1062,49 @@ pub fn repack_arrow_validity_buffer( .next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) .div_ceil(size_of::()); - // The kernel loads the input bitmap as 64-bit words. - if !input_buffer - .cuda_device_ptr()? - .is_multiple_of(size_of::() as u64) - { - vortex_bail!("Arrow validity repack requires an 8-byte aligned device buffer"); - } + let expected_input_bytes = arrow_bitmap_byte_len(len, input_offset)?; + vortex_ensure!( + input_buffer.len() >= expected_input_bytes, + "Arrow validity bitmap has {} bytes, expected at least {expected_input_bytes}", + input_buffer.len() + ); - let mut output = ctx.device_alloc::(allocation_words.max(1))?; - // The repack kernel writes only the logical bitmap words. Zero the whole backing allocation so - // cuDF's padded mask reads see invalid rows, not uninitialized CUDA memory. - ctx.stream() - .memset_zeros(&mut output) - .map_err(|err| vortex_err!("Failed to zero Arrow validity buffer padding: {err}"))?; - - if output_words > 0 { - let input_view = input_buffer.cuda_view::()?; - let len = u64::try_from(len)?; - let input_offset = u64::try_from(input_offset)?; - let arrow_offset = u64::try_from(arrow_offset)?; - let input_bytes = u64::try_from(input_buffer.len())?; - - let kernel = ctx.load_function_with_suffixes("arrow_validity", &["repack"])?; - const REPACK_THREADS_PER_BLOCK: u32 = 256; - let num_blocks = u32::try_from(output_words.div_ceil(REPACK_THREADS_PER_BLOCK as usize))?; - let config = LaunchConfig { - grid_dim: (num_blocks, 1, 1), - block_dim: (REPACK_THREADS_PER_BLOCK, 1, 1), - shared_mem_bytes: 0, - }; - ctx.launch_kernel_config(&kernel, config, output_words, |args| { - args.arg(&input_view) - .arg(&mut output) - .arg(&len) - .arg(&input_offset) - .arg(&arrow_offset) - .arg(&input_bytes); - })?; - } + let mut output = ctx.device_alloc::(allocation_words)?; + + let (input_view, prefix_bytes) = cuda_aligned_bitmap_view(input_buffer)?; + let len = u64::try_from(len)?; + let input_offset = u64::try_from(input_offset)? + (prefix_bytes * 8) as u64; + let arrow_offset = u64::try_from(arrow_offset)?; + let input_bytes = u64::try_from(input_view.len())?; + + let kernel = ctx.load_function_with_suffixes("arrow_validity", &["repack"])?; + // The kernel's grid-stride loop covers words beyond the capped grid. + #[allow( + clippy::cast_possible_truncation, + reason = "capped at MAX_BITMAP_BLOCKS" + )] + let num_blocks = output_words + .div_ceil(BITMAP_THREADS_PER_BLOCK as usize) + .min(MAX_BITMAP_BLOCKS as usize) as u32; + let config = LaunchConfig { + grid_dim: (num_blocks, 1, 1), + block_dim: (BITMAP_THREADS_PER_BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + ctx.launch_kernel_config(&kernel, config, output_words, |args| { + args.arg(&input_view) + .arg(&mut output) + .arg(&len) + .arg(&input_offset) + .arg(&arrow_offset) + .arg(&input_bytes); + })?; - // The memset above zeroed all allocation bytes after the logical output. - let output_device = CudaDeviceBuffer::new_with_zeroed_tail(output, output_bytes)?; - Ok(BufferHandle::new_device(Arc::new(output_device)).slice(0..output_bytes)) + // The kernel masks unused bits in its last word; only unwritten words need zeroing. + zero_padding(ctx.stream(), &mut output, output_words)?; + Ok(BufferHandle::new_device( + CudaDeviceBuffer::new_with_zeroed_tail(output, output_bytes)?.slice(0..output_bytes), + )) } /// Export a Vortex list-view as an Arrow Device array with `List` layout. @@ -1502,9 +1498,11 @@ mod tests { use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; + use vortex::error::vortex_err; use vortex::extension::datetime::TimeUnit; use crate::CudaBufferExt; + use crate::CudaDeviceBuffer; use crate::CudaDispatchMode; use crate::CudaExecutionCtx; use crate::DictionaryExport; @@ -1514,8 +1512,11 @@ mod tests { use crate::arrow::DeviceArrayExt; use crate::arrow::PrivateData; use crate::arrow::arrow_schema_for_array; + use crate::arrow::canonical::BITMAP_THREADS_PER_BLOCK; + use crate::arrow::canonical::MAX_BITMAP_BLOCKS; + use crate::arrow::canonical::count_arrow_validity_nulls; use crate::arrow::canonical::export_arrow_validity_buffer; - use crate::arrow::canonical::repack_arrow_validity_buffer; + use crate::arrow::canonical::repack_arrow_bitmap; use crate::arrow::dictionary_tests::upload; use crate::arrow::tests::private_data_buffer_bytes; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; @@ -1524,6 +1525,15 @@ mod tests { use crate::session::CudaSession; use crate::session::VarBinExportLayout; + fn device_bool( + bits: BitBuffer, + validity: Validity, + ctx: &CudaExecutionCtx, + ) -> VortexResult { + let (offset, len, bytes) = bits.into_inner(); + let buffer = ctx.stream().copy_to_device_sync(bytes.as_ref())?; + BoolArray::try_new_from_handle(buffer, offset, len, validity) + } unsafe fn release_exported_array(array: *mut ArrowArray) { unsafe { if let Some(release) = (*array).release { @@ -1777,6 +1787,46 @@ mod tests { Ok(Buffer::::from_byte_buffer(private_data_buffer_bytes(array, buffer_idx)?).to_vec()) } + // Normal uploads add 64-byte padding that can hide out-of-bounds reads and missing export + // padding. Use exact-sized allocations so these tests exercise those cases. + fn upload_unpadded(bytes: &ByteBuffer, ctx: &CudaExecutionCtx) -> VortexResult { + let mut allocation = ctx.device_alloc::(bytes.len())?; + ctx.stream() + .memcpy_htod(bytes.as_ref(), &mut allocation) + .map_err(|err| vortex_err!("Failed to upload unpadded test buffer: {err}"))?; + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + allocation, + )))) + } + + /// Assert cuDF word alignment and zeroed tail padding, returning the logical bytes on the host. + /// With `exact_allocation`, also require the backing allocation to match the padded size. + fn assert_bitmap_padding( + buffer: &BufferHandle, + logical_bytes: usize, + exact_allocation: bool, + ) -> VortexResult { + assert!(buffer.is_on_device()); + assert_eq!(buffer.len(), logical_bytes); + let pointer = buffer.cuda_device_ptr()?; + assert!(pointer.is_multiple_of(size_of::() as u64)); + let padded_bytes = logical_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); + let backing = cuda_backing_allocation(buffer)?; + let start = usize::try_from(pointer - backing.cuda_device_ptr()?)?; + // Sliced exports may retain a larger allocation; repacked outputs must be exact-sized. + if exact_allocation { + assert_eq!(backing.len(), padded_bytes); + } + assert!(start + padded_bytes <= backing.len()); + let bytes = backing.try_to_host_sync()?; + assert!( + bytes[start + logical_bytes..start + padded_bytes] + .iter() + .all(|byte| *byte == 0) + ); + Ok(bytes.slice_unaligned(start..start + logical_bytes)) + } + // Assert Arrow Binary export uses the standard null bitmap, i32 offsets, and values layout. fn assert_binary_layout( array: &ArrowArray, @@ -2283,25 +2333,117 @@ mod tests { Ok(()) } + #[rstest] + #[case::host(3, false)] + #[case::device(3, true)] + #[case::empty(0, false)] #[crate::test] - async fn test_export_bool() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); + async fn test_export_bool(#[case] len: usize, #[case] on_device: bool) -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; - let array = BoolArray::from_iter([true, false, true]).into_array(); + let bits = BitBuffer::from_iter([true, false, true].into_iter().take(len)); + let array = if on_device { + device_bool(bits, Validity::NonNullable, &ctx)? + } else { + BoolArray::from(bits) + } + .into_array(); + let input_values = on_device.then(|| array.buffer_handles()[0].clone()); let mut device_array = array.export_device_array(&mut ctx).await?; - assert_eq!(device_array.array.length, 3); + assert_eq!(device_array.array.length, i64::try_from(len)?); + assert_eq!(device_array.array.offset, 0); assert_eq!(device_array.array.null_count, 0); assert_eq!(device_array.array.n_buffers, 2); assert_eq!(device_array.array.n_children, 0); assert!(device_array.array.release.is_some()); assert_eq!(device_array.device_type, ARROW_DEVICE_CUDA); + if let Some(input_values) = input_values { + // SAFETY: The live export owns PrivateData until release below. + let private = unsafe { &*device_array.array.private_data.cast::() }; + let values = private.buffers[1] + .as_ref() + .ok_or_else(|| vortex_err!("expected exported bool values"))?; + assert_eq!(values.cuda_device_ptr()?, input_values.cuda_device_ptr()?); + } + unsafe { release_exported_array(&raw mut device_array.array) }; Ok(()) } + #[rstest] + #[case::byte_aligned(8, 33, 0)] + #[case::bit_offset(13, 65, 0)] + #[case::last_byte(9, 1, 0)] + #[case::aligned_unpadded(64, 33, 0)] + #[case::middle_slice(13, 65, 17)] + #[case::word_boundary(31, 34, 0)] + #[case::below_padding_block(13, 499, 0)] + // 520 uploaded rows sliced at 13: offset 5, 507 rows, exactly 64 validity bytes. + #[case::full_padding_block(13, 507, 0)] + #[case::above_padding_block(13, 515, 0)] + #[crate::test] + async fn test_export_byte_sliced_bool_values( + #[case] start: usize, + #[case] len: usize, + #[case] suffix_len: usize, + #[values(false, true)] nullable: bool, + ) -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let source_len = start + len + suffix_len; + let source_bits = BitBuffer::from_iter((0..source_len).map(|idx| idx % 5 < 2)); + let values = upload_unpadded(source_bits.inner(), &ctx)?; + let valid_bits = BitBuffer::from_iter((0..source_len).map(|idx| idx % 3 != 0)); + let validity = if nullable { + Validity::Array( + device_bool(valid_bits.clone(), Validity::NonNullable, &ctx)?.into_array(), + ) + } else { + Validity::NonNullable + }; + let array = BoolArray::try_new_from_handle(values, 0, source_len, validity)?; + let range = start..start + len; + let mut exported = array + .slice(range.clone())? + .export_device_array(&mut ctx) + .await?; + ctx.synchronize_stream()?; + + let offset = usize::try_from(exported.array.offset)?; + assert_eq!(offset, start % 8); + assert_eq!(exported.array.length, i64::try_from(len)?); + assert_eq!(exported.array.n_buffers, 2); + // SAFETY: The export owns live PrivateData until release below. + let private = unsafe { &*exported.array.private_data.cast::() }; + let logical_bytes = (offset + len).div_ceil(8); + let padded_bytes = logical_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); + for (index, expected) in [(1, source_bits), (0, valid_bits)] { + if index == 0 && !nullable { + assert!(private.buffers[0].is_none()); + continue; + } + let buffer = private.buffers[index] + .as_ref() + .ok_or_else(|| vortex_err!("expected exported bool buffer {index}"))?; + let bytes = assert_bitmap_padding(buffer, logical_bytes, false)?; + assert!(buffer.has_zeroed_tail_padding(logical_bytes, padded_bytes)?); + assert_eq!( + BitBuffer::new(bytes, offset + len).slice(offset..offset + len), + expected.slice(range.clone()) + ); + } + let null_count = if nullable { + range.filter(|idx| idx % 3 == 0).count() + } else { + 0 + }; + assert_eq!(exported.array.null_count, i64::try_from(null_count)?); + // SAFETY: This is the sole release of the successfully exported array. + unsafe { release_exported_array(&raw mut exported.array) }; + Ok(()) + } + #[crate::test] async fn test_export_varbinview_opt_in() -> VortexResult<()> { let mut ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBinView)?; @@ -3410,77 +3552,156 @@ mod tests { Ok(()) } - #[rstest::rstest] - #[case::input_ahead_of_arrow(5, 3, 9)] - #[case::arrow_ahead_of_input(3, 70, 9)] - #[case::equal_offsets(7, 7, 9)] - #[case::byte_aligned_input(0, 9, 9)] - #[case::word_aligned_offsets(64, 128, 130)] - #[case::multi_word(13, 0, 301)] #[crate::test] - async fn test_repack_arrow_validity_buffer_offsets( + async fn test_count_arrow_validity_nulls_capped_grid() -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let len = 10_001usize; + let arrow_offset = 13; + let bitmap_bits = arrow_offset + len; + // Dirty prefix and tail bits must not count as valid rows. + let source = BitBuffer::from_iter( + std::iter::repeat_n(true, arrow_offset) + .chain((0..len).map(|idx| idx % 3 != 0)) + .chain(std::iter::repeat_n( + true, + bitmap_bits.next_multiple_of(8) - bitmap_bits, + )), + ); + let input = upload_unpadded(source.inner(), &ctx)?; + assert!(input.len() > (BITMAP_THREADS_PER_BLOCK * MAX_BITMAP_BLOCKS) as usize); + + let null_count = count_arrow_validity_nulls(&input, len, arrow_offset, &mut ctx)?; + let expected_nulls = (0..len).filter(|idx| idx % 3 == 0).count(); + assert_eq!(null_count, i64::try_from(expected_nulls)?); + Ok(()) + } + + #[rstest] + #[case::input_ahead_of_arrow(5, 3, 9, 0)] + #[case::arrow_ahead_of_input(3, 70, 9, 1)] + #[case::equal_offsets(7, 7, 9, 7)] + #[case::byte_aligned_input(0, 9, 9, 1)] + #[case::word_aligned_offsets(64, 128, 130, 0)] + #[case::multi_word(13, 0, 301, 15)] + #[case::single_tail_bit(7, 0, 1, 1)] + #[case::exact_word(0, 0, 64, 0)] + #[case::word_boundary(7, 3, 65, 7)] + #[case::below_padding_block(7, 3, 501, 1)] + #[case::full_padding_block(7, 3, 509, 7)] + #[case::above_padding_block(7, 3, 517, 1)] + #[case::multi_block(13, 5, 20_001, 0)] + // The test cap of two blocks forces multiple grid-stride iterations. + #[case::capped_grid(13, 5, 100_001, 7)] + #[crate::test] + async fn test_repack_arrow_bitmap_offsets( #[case] input_offset: usize, #[case] arrow_offset: usize, #[case] len: usize, + #[case] byte_offset: usize, ) -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; - let logical_bits = (0..len).map(|idx| idx % 3 != 0).collect::>(); - // All-true filler before the slice would leak into the output if offsets were mishandled. + let logical_bits = (0..len).map(|idx| idx % 3 != 0); + // Dirty prefix and tail bits must not leak into the exported rows or padding. + let source_bits = byte_offset * 8 + input_offset + len; let source = BitBuffer::from_iter( - std::iter::repeat_n(true, input_offset).chain(logical_bits.iter().copied()), + std::iter::repeat_n(true, byte_offset * 8 + input_offset) + .chain(logical_bits.clone()) + .chain(std::iter::repeat_n( + true, + source_bits.next_multiple_of(8) - source_bits, + )), + ); + let input_bytes = source.inner(); + let input_buffer = + upload_unpadded(input_bytes, &ctx)?.slice(byte_offset..input_bytes.len()); + assert_eq!( + input_buffer.cuda_device_ptr()? % 8, + u64::try_from(byte_offset % 8)? ); - let sliced = source.slice(input_offset..input_offset + len); - // BitBuffer rebases whole bytes into the backing buffer, keeping the bit offset below 8. - let (input_offset, _, input_buffer) = sliced.into_inner(); - - let input_buffer = ctx - .ensure_on_device(BufferHandle::new_host(input_buffer)) - .await?; let output_bits = len + arrow_offset; - let output = - repack_arrow_validity_buffer(&input_buffer, input_offset, len, arrow_offset, &mut ctx)?; + let output = repack_arrow_bitmap(&input_buffer, input_offset, len, arrow_offset, &mut ctx)?; ctx.synchronize_stream()?; - let actual = BitBuffer::new(output.to_host_sync(), output_bits) - .iter() - .collect::>(); - let expected = std::iter::repeat_n(false, arrow_offset) - .chain(logical_bits) - .collect::>(); - assert_eq!(actual, expected); + let expected = + BitBuffer::from_iter(std::iter::repeat_n(false, arrow_offset).chain(logical_bits)); + assert_eq!( + assert_bitmap_padding(&output, output_bits.div_ceil(8), true)?, + expected.into_inner().2 + ); Ok(()) } + #[rstest] #[crate::test] - async fn test_repack_arrow_validity_buffer_zeroes_padding() -> VortexResult<()> { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); - - let len = 9; - let arrow_offset = 3; - let source = BitBuffer::from_iter(std::iter::repeat_n(true, len)); - let (input_offset, _, input_buffer) = source.into_inner(); - let input_buffer = ctx - .ensure_on_device(BufferHandle::new_host(input_buffer)) + async fn test_export_byte_sliced_device_validity( + #[values(false, true)] boolean_values: bool, + ) -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let len = 100; + let range = 13..78; + let valid_bits = BitBuffer::from_iter((0..len).map(|idx| idx % 3 != 0)); + let validity = Validity::Array( + device_bool(valid_bits.clone(), Validity::NonNullable, &ctx)?.into_array(), + ); + let array = if boolean_values { + // Values and validity have different bit offsets, requiring validity repacking + // even when the exported values retain their own Arrow offset. + let values = BitBuffer::from_iter((0..len + 3).map(|idx| idx % 2 == 0)); + device_bool(values.slice(3..len + 3), validity, &ctx)?.into_array() + } else { + let values = Buffer::from_iter(0..i32::try_from(len)?); + let values = ctx.stream().copy_to_device_sync(values.as_ref())?; + PrimitiveArray::from_buffer_handle(values, PType::I32, validity).into_array() + }; + let mut exported = array + .slice(range.clone())? + .export_device_array(&mut ctx) .await?; - let output_bytes = (len + arrow_offset).div_ceil(8); - - let output = - repack_arrow_validity_buffer(&input_buffer, input_offset, len, arrow_offset, &mut ctx)?; ctx.synchronize_stream()?; - assert_eq!(output.len(), output_bytes); - let backing = cuda_backing_allocation(&output)?; - let backing_bytes = backing.to_host_sync(); + assert_eq!(exported.array.length, i64::try_from(range.len())?); + assert_eq!(exported.array.offset, 0); + assert_eq!( + exported.array.null_count, + i64::try_from(range.clone().filter(|idx| idx % 3 == 0).count())? + ); assert_eq!( - backing_bytes.len(), - output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + BitBuffer::new(private_data_buffer_bytes(&exported.array, 0)?, range.len()), + valid_bits.slice(range.clone()) ); - assert!(backing_bytes[output_bytes..].iter().all(|byte| *byte == 0)); + let values = private_data_buffer_bytes(&exported.array, 1)?; + if boolean_values { + assert_eq!( + BitBuffer::new(values, range.len()), + BitBuffer::from_iter(range.clone().map(|idx| (idx + 3) % 2 == 0)) + ); + } else { + assert_eq!( + Buffer::::from_byte_buffer(values), + Buffer::from_iter(i32::try_from(range.start)?..i32::try_from(range.end)?) + ); + } + // SAFETY: This is the sole release of the successfully exported array. + unsafe { release_exported_array(&raw mut exported.array) }; + Ok(()) + } + #[rstest] + #[case::empty(0, 0)] + #[case::truncated(7, 2)] + #[case::overflow(usize::MAX, 1)] + #[crate::test] + async fn test_repack_arrow_bitmap_rejects_invalid_range( + #[case] input_offset: usize, + #[case] len: usize, + ) -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let input = ctx + .ensure_on_device(BufferHandle::new_host(ByteBuffer::from(vec![0xff]))) + .await?; + assert!(repack_arrow_bitmap(&input, input_offset, len, 0, &mut ctx).is_err()); Ok(()) } @@ -3503,66 +3724,53 @@ mod tests { assert_eq!(null_count, 1); let buffer = buffer.vortex_expect("nullable validity should export a null buffer"); let output_bytes = (len + arrow_offset).div_ceil(8); - assert_eq!(buffer.len(), output_bytes); - let actual = BitBuffer::new(buffer.to_host_sync(), len + arrow_offset) - .iter() - .collect::>(); - assert_eq!(actual, [true, false, true]); - - let backing = cuda_backing_allocation(&buffer)?; - let backing_bytes = backing.to_host_sync(); + let bytes = assert_bitmap_padding(&buffer, output_bytes, true)?; assert_eq!( - backing_bytes.len(), - output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + BitBuffer::new(bytes, len + arrow_offset), + BitBuffer::from_iter([true, false, true]) ); - assert!(backing_bytes[output_bytes..].iter().all(|byte| *byte == 0)); Ok(()) } + #[rstest] + #[case::unsliced(0)] + #[case::sliced(8)] #[crate::test] - async fn test_export_validity_buffer_reuses_matching_padded_device_bitmap() -> VortexResult<()> - { - let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("failed to create execution context"); - - let len = 3; - let source = BitBuffer::from_iter([true, false, true]); - let (input_offset, _, input_buffer) = source.into_inner(); - let input_buffer = ctx - .ensure_on_device(BufferHandle::new_host(input_buffer)) - .await?; - let input_ptr = input_buffer.cuda_device_ptr()?; - let validity = BoolArray::new_handle( - input_buffer.clone(), - input_offset, - len, - Validity::NonNullable, - ) - .into_array(); - + async fn test_export_validity_buffer_reuses_matching_padded_device_bitmap( + #[case] byte_offset: usize, + ) -> VortexResult<()> { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let len = 449; + let logical_bits = (0..len).map(|idx| idx % 3 != 0); + let source = BitBuffer::from_iter( + std::iter::repeat_n(true, byte_offset * 8).chain(logical_bits.clone()), + ); + let input = ctx.stream().copy_to_device_sync(source.inner().as_ref())?; + let input_ptr = input.cuda_device_ptr()? + byte_offset as u64; + // Slicing to the end preserves the zeroed tail, measured from the allocation base. + let input = input.slice(byte_offset..input.len()); + let validity = + BoolArray::try_new_from_handle(input, 0, len, Validity::NonNullable)?.into_array(); let (buffer, null_count) = - export_arrow_validity_buffer(Validity::Array(validity), len, input_offset, &mut ctx) - .await?; - ctx.synchronize_stream()?; + export_arrow_validity_buffer(Validity::Array(validity), len, 0, &mut ctx).await?; - assert_eq!(null_count, 1); + assert_eq!( + null_count, + i64::try_from(logical_bits.clone().filter(|bit| !bit).count())? + ); let buffer = buffer.vortex_expect("nullable validity should export a null buffer"); assert_eq!(buffer.cuda_device_ptr()?, input_ptr); - assert_eq!(buffer.len(), (len + input_offset).div_ceil(8)); - let actual = BitBuffer::new(buffer.to_host_sync(), len + input_offset) - .iter() - .collect::>(); - let expected = std::iter::repeat_n(false, input_offset) - .chain([true, false, true]) - .collect::>(); - assert_eq!(actual, expected); - + let bytes = assert_bitmap_padding(&buffer, len.div_ceil(8), false)?; + assert_eq!( + BitBuffer::new(bytes, len), + BitBuffer::from_iter(logical_bits) + ); Ok(()) } #[crate::test] - async fn test_export_validity_buffer_repacks_matching_offset_without_tail_padding() + async fn test_export_validity_buffer_copies_matching_offset_without_tail_padding() -> VortexResult<()> { let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -3591,19 +3799,12 @@ mod tests { let buffer = buffer.vortex_expect("nullable validity should export a null buffer"); assert_ne!(buffer.cuda_device_ptr()?, input_ptr); let output_bytes = (len + input_offset).div_ceil(8); - assert_eq!(buffer.len(), output_bytes); - let actual = BitBuffer::new(buffer.to_host_sync(), len + input_offset) - .iter() - .collect::>(); - let expected = std::iter::repeat_n(false, input_offset) - .chain([true, false, true]) - .collect::>(); - assert_eq!(actual, expected); - - let backing = cuda_backing_allocation(&buffer)?; + let bytes = assert_bitmap_padding(&buffer, output_bytes, true)?; assert_eq!( - backing.len(), - output_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) + BitBuffer::new(bytes, len + input_offset), + BitBuffer::from_iter( + std::iter::repeat_n(false, input_offset).chain([true, false, true]) + ) ); Ok(()) @@ -3641,14 +3842,8 @@ mod tests { assert_eq!(null_count, i64::try_from(len)?); let buffer = buffer.vortex_expect("all-false validity should export a null buffer"); - let bytes = buffer.to_host_sync(); - assert_eq!(bytes.len(), (len + arrow_offset).div_ceil(8)); + let bytes = assert_bitmap_padding(&buffer, (len + arrow_offset).div_ceil(8), true)?; assert!(bytes.iter().all(|byte| *byte == 0)); - let backing = cuda_backing_allocation(&buffer)?; - assert_eq!( - backing.len(), - bytes.len().next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING) - ); Ok(()) } diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index b7e584c32f7..6bf5d96fc53 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -87,7 +87,7 @@ use arrow_c_abi::ArrowSchema; #[doc(hidden)] pub mod test_harness { pub use crate::arrow::canonical::count_arrow_validity_nulls; - pub use crate::arrow::canonical::repack_arrow_validity_buffer; + pub use crate::arrow::canonical::repack_arrow_bitmap; } /// CUDA device memory. diff --git a/vortex-cuda/src/device_buffer.rs b/vortex-cuda/src/device_buffer.rs index 602fa16ba94..e5ffddac43e 100644 --- a/vortex-cuda/src/device_buffer.rs +++ b/vortex-cuda/src/device_buffer.rs @@ -184,6 +184,26 @@ impl CudaDeviceBuffer { } } +/// Include the slice's leading bytes up to an 8-byte boundary, without extending its tail. +/// CUDA allocations are aligned, so the prefix remains within the backing allocation. +pub(crate) fn cuda_aligned_bitmap_view( + handle: &BufferHandle, +) -> VortexResult<(CudaView<'_, u8>, usize)> { + let device_buffer = handle + .as_device_opt() + .ok_or_else(|| vortex_err!("Buffer is not on device"))?; + let cuda_buf = device_buffer + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("expected CudaDeviceBuffer, was {device_buffer:?}"))?; + let prefix_bytes = cuda_buf.offset % size_of::(); + let view = cuda_buf + .allocation + .as_bytes_view() + .slice(cuda_buf.offset - prefix_bytes..cuda_buf.offset + cuda_buf.len); + Ok((view, prefix_bytes)) +} + #[cfg(test)] pub(crate) fn cuda_backing_allocation(handle: &BufferHandle) -> VortexResult { let device_buffer = handle diff --git a/vortex-cuda/src/kernel/arrays/dict.rs b/vortex-cuda/src/kernel/arrays/dict.rs index dea04a0cd0e..8034ae23589 100644 --- a/vortex-cuda/src/kernel/arrays/dict.rs +++ b/vortex-cuda/src/kernel/arrays/dict.rs @@ -138,7 +138,7 @@ async fn execute_dict_bool_typed( // Each CUDA thread owns complete output bytes, avoiding races between threads writing // different bits in the same byte. The kernel handles the final partial byte explicitly. let output_bytes = codes_len.div_ceil(8); - let mut output_slice = ctx.device_alloc::(output_bytes)?; + let mut output_slice = ctx.stream().device_alloc_bitmap(output_bytes)?; let values_view = values_device.cuda_view::()?; let codes_view = codes_device.cuda_view::()?; @@ -155,9 +155,9 @@ async fn execute_dict_bool_typed( .arg(&mut output_slice); })?; - let output_device = CudaDeviceBuffer::new(output_slice); + let output_device = CudaDeviceBuffer::new_with_zeroed_tail(output_slice, output_bytes)?; Ok(Canonical::Bool(BoolArray::new_handle( - BufferHandle::new_device(Arc::new(output_device)), + BufferHandle::new_device(Arc::new(output_device)).slice(0..output_bytes), 0, codes_len, output_validity, @@ -389,6 +389,9 @@ mod tests { use super::*; use crate::CanonicalCudaExt; + use crate::arrow::DeviceArrayExt; + use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; + use crate::device_buffer::cuda_backing_allocation; use crate::session::CudaSession; /// Copy a CUDA primitive array result to host memory. @@ -406,8 +409,7 @@ mod tests { let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); - // Slicing leaves the dictionary values at a non-zero bit offset. Thirteen codes also - // exercise a final partial output byte. + // Slicing leaves a non-zero input bit offset; thirteen codes exercise a partial output byte. let values = BoolArray::from_iter([ false, true, false, true, false, true, true, false, true, false, ]) @@ -417,6 +419,7 @@ mod tests { Buffer::from(vec![0u8, 1, 2, 3, 4, 3, 2, 1, 0, 4, 1, 3, 0]), NonNullable, ); + let len = codes.len(); let expected = DictArray::try_new(codes.clone().into_array(), values.clone())?.into_array(); let codes_handle = cuda_ctx @@ -426,14 +429,36 @@ mod tests { PrimitiveArray::from_buffer_handle(codes_handle, codes.ptype(), codes.validity()?); let dict = DictArray::try_new(device_codes.into_array(), values)?.into_array(); - let actual = DictExecutor - .execute(dict, &mut cuda_ctx) - .await? - .into_host() - .await? - .into_bool(); - - assert_arrays_eq!(actual.into_array(), expected, &mut ctx); + let actual = DictExecutor.execute(dict, &mut cuda_ctx).await?; + let bits = actual.clone().into_bool().into_data().into_parts(len).bits; + let logical_bytes = len.div_ceil(8); + let padded_bytes = logical_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); + assert_eq!(bits.len(), logical_bytes); + + assert!(bits.has_zeroed_tail_padding(logical_bytes, padded_bytes)?); + let backing = cuda_backing_allocation(&bits)?; + assert_eq!(backing.len(), padded_bytes); + let bytes = backing.try_to_host()?.await?; + assert!(bytes[logical_bytes..].iter().all(|byte| *byte == 0)); + + let bits_ptr = bits.cuda_device_ptr()?; + let mut exported = actual + .clone() + .into_array() + .export_device_array(&mut cuda_ctx) + .await?; + // SAFETY: The live bool export has a host-resident buffer table with values at index 1. + let exported_bits_ptr = unsafe { *exported.array.buffers.add(1) } as u64; + let release = exported + .array + .release + .vortex_expect("missing Arrow release callback"); + // SAFETY: This is the sole release of the successfully exported array. + unsafe { release(&raw mut exported.array) }; + assert_eq!(exported_bits_ptr, bits_ptr); + + let actual = actual.into_host().await?.into_array(); + assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index da6973ff11c..f932cfb9287 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -164,7 +164,7 @@ async fn decode_runend_typed(output_bytes)?; + let mut validity_out = ctx.stream().device_alloc_bitmap(output_bytes)?; let validity_offset_u64 = validity_meta.offset() as u64; let ends_ptype = E::PTYPE.to_string(); @@ -180,8 +180,10 @@ async fn decode_runend_typed(ends: Vec, values: Vec, ctx: &mut ExecutionCtx) -> RunEndArray @@ -390,12 +395,38 @@ mod tests { // host-resident array, hiding a GPU failure. let gpu_result = RunEndExecutor .execute(runend_array.clone().into_array(), &mut cuda_ctx) - .await - .vortex_expect("GPU decompression failed") - .into_host() - .await? - .into_array(); - + .await?; + let Validity::Array(validity) = gpu_result.clone().into_primitive().validity()? else { + vortex_bail!("expected expanded validity bitmap"); + }; + let bits = &validity.buffer_handles()[0]; + let logical_bytes = runend_array.len().div_ceil(8); + let padded_bytes = logical_bytes.next_multiple_of(CUDF_VALIDITY_BUFFER_PADDING); + assert_eq!(bits.len(), logical_bytes); + + assert!(bits.has_zeroed_tail_padding(logical_bytes, padded_bytes)?); + let backing = cuda_backing_allocation(bits)?; + assert_eq!(backing.len(), padded_bytes); + let bytes = backing.try_to_host()?.await?; + assert!(bytes[logical_bytes..].iter().all(|byte| *byte == 0)); + + let bits_ptr = bits.cuda_device_ptr()?; + let mut exported = gpu_result + .clone() + .into_array() + .export_device_array(&mut cuda_ctx) + .await?; + // SAFETY: The live primitive export has a host-resident buffer table with validity at index 0. + let exported_bits_ptr = unsafe { *exported.array.buffers } as u64; + let release = exported + .array + .release + .vortex_expect("missing Arrow release callback"); + // SAFETY: This is the sole release of the successfully exported array. + unsafe { release(&raw mut exported.array) }; + assert_eq!(exported_bits_ptr, bits_ptr); + + let gpu_result = gpu_result.into_host().await?.into_array(); assert_arrays_eq!(runend_array, gpu_result, &mut ctx); Ok(()) diff --git a/vortex-cuda/src/stream.rs b/vortex-cuda/src/stream.rs index d54f8d0940b..ec07e2a7329 100644 --- a/vortex-cuda/src/stream.rs +++ b/vortex-cuda/src/stream.rs @@ -61,6 +61,16 @@ impl VortexCudaStream { } } + /// Allocates bitmap bytes with cuDF-sized padding, zeroing only the tail. + /// + /// The producer must initialize `byte_count` bytes on this stream and leave the padding untouched. + pub(crate) fn device_alloc_bitmap(&self, byte_count: usize) -> VortexResult> { + let allocation_len = padded_device_allocation_len::(byte_count)?; + let mut buffer = self.device_alloc::(allocation_len)?; + zero_padding(self, &mut buffer, byte_count)?; + Ok(buffer) + } + /// Copies host data to the device. /// /// Allocates device memory, schedules an async copy, and returns a future @@ -157,20 +167,21 @@ fn padded_device_allocation_len(byte_count: usize) -> VortexResult { Ok(min_allocation_bytes.div_ceil(element_size)) } -/// Zeroes the allocation tail after the copied values. +/// Zeroes the allocation tail after `initialized_count` elements. /// -/// Returned handles are sliced to the copied byte count; the trailing padding -/// exists so padded mask reads stay within the backing allocation. -fn zero_padding( +/// Copies or kernels must initialize the preceding elements on the same stream. +/// cuDF reads bitmaps in whole mask words, which can extend past the final logical byte. +/// Allocations include tail padding for those reads; this zeroes that padding. +pub(crate) fn zero_padding( stream: &VortexCudaStream, cuda_slice: &mut CudaSlice, - copied_len: usize, + initialized_count: usize, ) -> VortexResult<()> { - if copied_len >= cuda_slice.len() { + if initialized_count >= cuda_slice.len() { return Ok(()); } - let mut padding = cuda_slice.slice_mut(copied_len..); + let mut padding = cuda_slice.slice_mut(initialized_count..); stream .memset_zeros(&mut padding) .map_err(|e| vortex_err!("Failed to zero device buffer padding: {}", e))