diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index cd97755e636..606a7596ed1 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -32,6 +32,13 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: "true" + # The Public BI datasets are converted from CSV with the DuckDB CLI. + - name: Install DuckDB + run: | + wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.5/duckdb_cli-linux-amd64.zip | funzip > duckdb + chmod +x duckdb + echo "$PWD" >> "$GITHUB_PATH" + - uses: ./.github/actions/system-info - name: Display NVIDIA GPU details run: | diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index bf2d3efc1db..8fe0909e3e1 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -15,8 +15,9 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format cargo run -p compress-bench --profile release_debug ``` -GPU decompression is opt-in and runs only the existing benchmark names allow-listed in -`src/main.rs`: +GPU decompression is opt-in and runs the full compress suite, including the +pcodec-hosted `airquality` and `rplace` datasets. Use `--datasets` to narrow it to a +subset: ```bash cargo run -p compress-bench --profile release_debug \ diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1b1603e52c8..8d89a9ea726 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -178,14 +178,6 @@ async fn run_compress( // ), ]; - // Add an existing benchmark name here only after its CUDA-compatible compression and - // decompression kernels have been verified end to end. - #[expect( - clippy::useless_vec, - reason = "this is an intentionally incremental allow-list of benchmark names" - )] - let gpu_decompress_benchmarks = vec!["TPC-H l_comment canonical"]; - let datasets: Vec<&dyn Dataset> = [ &TaxiData as &dyn Dataset, PBI_DATASETS.get(Arade), @@ -207,11 +199,11 @@ async fn run_compress( .into_iter() .chain(structlistofints.iter().map(|d| d as &dyn Dataset)) .filter(|d| { - if gpu_decompress && !gpu_decompress_benchmarks.contains(&d.name()) { - return false; - } if let Some(filter) = datasets_filter.as_ref() { filter.is_match(d.name()) + } else if gpu_decompress { + // The GPU suite runs every dataset, including the pcodec-hosted ones. + true } else { // These download data from pcodec's public bucket, presumably creating egress charges // for pcodec. As such, we do not run in CI. diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 53cfc1d2be4..43947893cb0 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -175,6 +175,7 @@ impl BtrBlocksCompressorBuilder { float::ALPRDScheme.id(), float::FloatRLEScheme.id(), float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), string::StringDictScheme.id(), binary::BinaryDictScheme.id(), ]; @@ -270,6 +271,22 @@ mod tests { ); } + /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. + #[test] + fn cuda_compatible_excludes_every_sparse_scheme() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + for excluded in [ + integer::SparseScheme.id(), + float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), + ] { + assert!( + !builder.schemes.iter().any(|s| s.id() == excluded), + "{excluded} should be excluded" + ); + } + } + #[test] fn cuda_compatible_uses_fsst_for_strings() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-cuda/kernels/src/date_time_parts.cu b/vortex-cuda/kernels/src/date_time_parts.cu index ccb3e614991..a556ee5c57e 100644 --- a/vortex-cuda/kernels/src/date_time_parts.cu +++ b/vortex-cuda/kernels/src/date_time_parts.cu @@ -43,22 +43,34 @@ __device__ void date_time_parts(const DaysT *__restrict days, } #define EXPAND_DAYS(X) \ + X(u8, uint8_t) \ + X(u16, uint16_t) \ + X(u32, uint32_t) \ + X(u64, uint64_t) \ X(i8, int8_t) \ X(i16, int16_t) \ X(i32, int32_t) \ X(i64, int64_t) #define EXPAND_SUBSECONDS(d, DT, s, ST) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u8, uint8_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u16, uint16_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u32, uint32_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u64, uint64_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i8, int8_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i16, int16_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i32, int32_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i64, int64_t) #define EXPAND_SECONDS(d, DT) \ + EXPAND_SUBSECONDS(d, DT, u8, uint8_t) \ + EXPAND_SUBSECONDS(d, DT, u16, uint16_t) \ + EXPAND_SUBSECONDS(d, DT, u32, uint32_t) \ + EXPAND_SUBSECONDS(d, DT, u64, uint64_t) \ EXPAND_SUBSECONDS(d, DT, i8, int8_t) \ EXPAND_SUBSECONDS(d, DT, i16, int16_t) \ EXPAND_SUBSECONDS(d, DT, i32, int32_t) \ EXPAND_SUBSECONDS(d, DT, i64, int64_t) -// Generate all 64 kernels (4³) +// Generate all 512 kernels (8³: every signed and unsigned integer width per component) EXPAND_DAYS(EXPAND_SECONDS) diff --git a/vortex-cuda/kernels/src/list.cu b/vortex-cuda/kernels/src/list.cu new file mode 100644 index 00000000000..3e87e267299 --- /dev/null +++ b/vortex-cuda/kernels/src/list.cu @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "config.cuh" +#include "types.cuh" + +// Converts Arrow-style `List` offsets into `ListView` offset/size pairs. +// +// `List` stores `list_len + 1` monotonically increasing offsets; a `ListView` stores one offset +// and one size per list. Both outputs are written by the same thread so the two views of a list +// are always produced together. +template +__device__ void list_views(const OffsetT *const __restrict offsets, + OffsetT *const __restrict out_offsets, + OffsetT *const __restrict out_sizes, + uint64_t list_len) { + const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; + const uint64_t block_end = min(block_start + elements_per_block, list_len); + + for (uint64_t idx = block_start + threadIdx.x; idx < block_end; idx += blockDim.x) { + const OffsetT start = offsets[idx]; + out_offsets[idx] = start; + out_sizes[idx] = static_cast(offsets[idx + 1] - start); + } +} + +#define GENERATE_LIST_VIEWS_KERNEL(offset_suffix, OffsetT) \ + extern "C" __global__ void list_views_##offset_suffix(const OffsetT *const __restrict offsets, \ + OffsetT *const __restrict out_offsets, \ + OffsetT *const __restrict out_sizes, \ + uint64_t list_len) { \ + list_views(offsets, out_offsets, out_sizes, list_len); \ + } + +FOR_EACH_INTEGER(GENERATE_LIST_VIEWS_KERNEL) diff --git a/vortex-cuda/kernels/src/runend.cu b/vortex-cuda/kernels/src/runend.cu index a3f1d245dbe..7472e7dafba 100644 --- a/vortex-cuda/kernels/src/runend.cu +++ b/vortex-cuda/kernels/src/runend.cu @@ -127,6 +127,45 @@ __device__ void runend_decode_kernel(const EndsT *const __restrict ends, } } +// Expands run-end encoded validity bits into a packed output bitmap. +// +// Mirrors `runend_decode_kernel`, but each thread owns one complete output byte so that +// threads never race on bits within the same byte. Runs are located with a global binary +// search per element rather than the shared-memory cache: validity expansion runs once per +// array and is not the decode hot path. +template +__device__ void runend_bool_kernel(const EndsT *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values, + uint64_t values_bit_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output) { + const uint64_t output_bytes = (output_len + 7) / 8; + const uint32_t elements_per_block = blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * elements_per_block; + const uint64_t block_end = min(block_start + elements_per_block, output_bytes); + + for (uint64_t byte_idx = block_start + threadIdx.x; byte_idx < block_end; byte_idx += blockDim.x) { + const uint64_t row_start = byte_idx * 8; + uint8_t packed = 0; +#pragma unroll + for (uint32_t bit = 0; bit < 8; ++bit) { + const uint64_t row = row_start + bit; + if (row < output_len) { + uint64_t run_idx = upper_bound(ends, num_runs, row + offset); + if (run_idx >= num_runs) { + run_idx = num_runs - 1; + } + const uint64_t value_idx = values_bit_offset + run_idx; + const uint8_t value = (values[value_idx / 8] >> (value_idx % 8)) & 1; + packed |= static_cast(value << bit); + } + } + output[byte_idx] = packed; + } +} + #define GENERATE_RUNEND_KERNEL(value_suffix, ValueType, ends_suffix, EndsType) \ extern "C" __global__ void runend_##value_suffix##_##ends_suffix( \ const EndsType *const __restrict ends, \ @@ -155,3 +194,21 @@ GENERATE_RUNEND_KERNELS_FOR_VALUE(i64, int64_t) GENERATE_RUNEND_KERNELS_FOR_VALUE(f16, __half) GENERATE_RUNEND_KERNELS_FOR_VALUE(f32, float) GENERATE_RUNEND_KERNELS_FOR_VALUE(f64, double) + +#define GENERATE_RUNEND_BOOL_KERNEL(ends_suffix, EndsType) \ + extern "C" __global__ void runend_bool_##ends_suffix(const EndsType *const __restrict ends, \ + uint64_t num_runs, \ + const uint8_t *const __restrict values, \ + uint64_t values_bit_offset, \ + uint64_t offset, \ + uint64_t output_len, \ + uint8_t *const __restrict output) { \ + runend_bool_kernel(ends, num_runs, values, values_bit_offset, offset, output_len, output); \ + } + +// Validity bitmaps use a different physical layout and launch unit, but dispatch over the +// same run-end index types. +GENERATE_RUNEND_BOOL_KERNEL(u8, uint8_t) +GENERATE_RUNEND_BOOL_KERNEL(u16, uint16_t) +GENERATE_RUNEND_BOOL_KERNEL(u32, uint32_t) +GENERATE_RUNEND_BOOL_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 9f1fce7e68e..12689912e39 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -5,29 +5,67 @@ use std::sync::Arc; use async_trait::async_trait; use futures::future::try_join_all; +use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; +use vortex::array::arrays::Bool; use vortex::array::arrays::BoolArray; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::ExtensionArray; +use vortex::array::arrays::ListViewArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::VarBinViewArray; use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::decimal::DecimalDataParts; use vortex::array::arrays::extension::ExtensionArrayExt; +use vortex::array::arrays::listview::ListViewDataParts; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::arrays::struct_::StructDataParts; use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::legacy_session; +use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; use vortex::error::VortexResult; +/// Copy a canonical child array to the host. +async fn child_into_host(child: ArrayRef) -> VortexResult { + #[allow(clippy::disallowed_methods)] + Ok(child + .execute::(&mut legacy_session().create_execution_ctx())? + .into_host() + .await? + .into_array()) +} + +/// Copy a `Validity::Array` bitmap back to the host. +/// +/// Canonical arrays keep their validity as a separate child array, so moving only the data +/// buffer to the host leaves a device-resident bitmap behind that CPU compute cannot read. +async fn validity_into_host(validity: Validity) -> VortexResult { + let Validity::Array(array) = validity else { + return Ok(validity); + }; + + let Ok(bools) = array.clone().try_downcast::() else { + return Ok(Validity::Array(array)); + }; + + let len = bools.len(); + let inner_validity = bools.validity()?; + let BoolDataParts { bits, meta } = bools.into_data().into_parts(len); + let bits = BitBuffer::new_with_offset(bits.try_into_host()?.await?, meta.len(), meta.offset()); + + Ok(Validity::Array( + BoolArray::new(bits, inner_validity).into_array(), + )) +} + /// Move all canonical data from to_host from device. #[async_trait] pub trait CanonicalCudaExt { @@ -50,6 +88,7 @@ impl CanonicalCudaExt for Canonical { validity, .. } = struct_array.into_data_parts(); + let validity = validity_into_host(validity).await?; let mut host_fields = vec![]; for field in fields.iter() { @@ -75,7 +114,7 @@ impl CanonicalCudaExt for Canonical { // NOTE: update to copy to host when adding buffer handle. // Also update other method to copy validity to host. let len = bool.len(); - let validity = bool.validity()?; + let validity = validity_into_host(bool.validity()?).await?; let BoolDataParts { bits, meta } = bool.into_data().into_parts(len); let bits = BitBuffer::new_with_offset( @@ -92,6 +131,7 @@ impl CanonicalCudaExt for Canonical { validity, .. } = prim.into_data_parts(); + let validity = validity_into_host(validity).await?; Ok(Canonical::Primitive(PrimitiveArray::from_byte_buffer( buffer.try_into_host()?.await?, ptype, @@ -106,6 +146,7 @@ impl CanonicalCudaExt for Canonical { validity, .. } = decimal.into_data_parts(); + let validity = validity_into_host(validity).await?; Ok(Canonical::Decimal(unsafe { DecimalArray::new_unchecked_handle( BufferHandle::new_host(values.try_into_host()?.await?), @@ -122,6 +163,7 @@ impl CanonicalCudaExt for Canonical { validity, dtype, } = varbinview.into_data_parts(); + let validity = validity_into_host(validity).await?; // Copy all device views to host let host_views = views.try_into_host()?.await?; @@ -140,6 +182,23 @@ impl CanonicalCudaExt for Canonical { VarBinViewArray::new_unchecked(host_views, host_buffers, dtype, validity) })) } + Canonical::List(list) => { + let ListViewDataParts { + elements, + offsets, + sizes, + validity, + .. + } = list.into_data_parts(); + let validity = validity_into_host(validity).await?; + + Ok(Canonical::List(ListViewArray::try_new( + child_into_host(elements).await?, + child_into_host(offsets).await?, + child_into_host(sizes).await?, + validity, + )?)) + } Canonical::Extension(ext) => { // Copy the storage array to host and rewrap in ExtensionArray. let host_storage = ext diff --git a/vortex-cuda/src/kernel/arrays/list.rs b/vortex-cuda/src/kernel/arrays/list.rs new file mode 100644 index 00000000000..d6aab480e4e --- /dev/null +++ b/vortex-cuda/src/kernel/arrays/list.rs @@ -0,0 +1,214 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use async_trait::async_trait; +use cudarc::driver::DeviceRepr; +use cudarc::driver::PushKernelArg; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::arrays::List; +use vortex::array::arrays::ListViewArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::list::ListArrayExt; +use vortex::array::arrays::list::ListArraySlotsExt; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_integer_ptype; +use vortex::array::validity::Validity; +use vortex::dtype::NativePType; +use vortex::dtype::Nullability; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; + +/// CUDA executor for `ListArray`. +/// +/// `List` stores `len + 1` Arrow-style offsets; its canonical form, `ListView`, stores one +/// offset and one size per list. Decode the elements on the GPU and derive the view pair from +/// the offsets with a single kernel. +#[derive(Debug)] +pub(crate) struct ListExecutor; + +#[async_trait] +impl CudaExecute for ListExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let list = array + .try_downcast::() + .map_err(|_| vortex_err!("Expected ListArray"))?; + + let list_len = list.len(); + let validity = execute_validity_cuda(list.list_validity(), list_len, ctx).await?; + let elements = list + .elements() + .clone() + .execute_cuda(ctx) + .await? + .into_array(); + + if list_len == 0 { + let empty = PrimitiveArray::empty::(Nullability::NonNullable); + return Ok(Canonical::List(ListViewArray::try_new( + elements, + empty.clone().into_array(), + empty.into_array(), + validity, + )?)); + } + + let offsets = list + .offsets() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + vortex_ensure!( + offsets.len() == list_len + 1, + "ListArray must have {} offsets, got {}", + list_len + 1, + offsets.len() + ); + + let offsets_ptype = offsets.ptype(); + match_each_integer_ptype!(offsets_ptype, |O| { + list_views_typed::(offsets, elements, validity, list_len, ctx).await + }) + } +} + +async fn list_views_typed( + offsets: PrimitiveArray, + elements: ArrayRef, + validity: Validity, + list_len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let PrimitiveDataParts { + buffer: offsets_buffer, + .. + } = offsets.into_data_parts(); + + let offsets_device = ctx.ensure_on_device(offsets_buffer).await?; + let offsets_view = offsets_device.cuda_view::()?; + + let mut view_offsets = ctx.device_alloc::(list_len)?; + let mut view_sizes = ctx.device_alloc::(list_len)?; + let list_len_u64 = list_len as u64; + + let offsets_ptype = O::PTYPE.to_string(); + let cuda_function = ctx.load_function_with_suffixes("list", &["views", &offsets_ptype])?; + ctx.launch_kernel(&cuda_function, list_len, |args| { + args.arg(&offsets_view) + .arg(&mut view_offsets) + .arg(&mut view_sizes) + .arg(&list_len_u64); + })?; + + let view_offsets = PrimitiveArray::from_buffer_handle( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(view_offsets))), + O::PTYPE, + Validity::NonNullable, + ); + let view_sizes = PrimitiveArray::from_buffer_handle( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(view_sizes))), + O::PTYPE, + Validity::NonNullable, + ); + + Ok(Canonical::List(ListViewArray::try_new( + elements, + view_offsets.into_array(), + view_sizes.into_array(), + validity, + )?)) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::ListArray; + use vortex::array::assert_arrays_eq; + use vortex::buffer::Buffer; + use vortex::buffer::buffer; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + + use super::*; + use crate::CanonicalCudaExt; + use crate::session::CudaSession; + + #[rstest] + #[case::single_run(vec![0i32, 2, 5, 9])] + #[case::empty_lists(vec![0i32, 0, 3, 3])] + #[crate::test] + async fn test_cuda_list_decompression(#[case] offsets: Vec) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let element_count = *offsets.last().vortex_expect("offsets are non-empty"); + let elements = PrimitiveArray::new( + (0..element_count).collect::>(), + Validity::NonNullable, + ) + .into_array(); + let offsets_array = + PrimitiveArray::new(Buffer::from(offsets), Validity::NonNullable).into_array(); + let list = ListArray::try_new(elements, offsets_array, Validity::NonNullable)?; + + let gpu_result = ListExecutor + .execute(list.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(list, gpu_result, &mut ctx); + + Ok(()) + } + + #[crate::test] + async fn test_cuda_list_with_nulls() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let elements = + PrimitiveArray::new(buffer![10i32, 20, 30, 40], Validity::NonNullable).into_array(); + let offsets = + PrimitiveArray::new(buffer![0i32, 2, 2, 4], Validity::NonNullable).into_array(); + let validity = + Validity::Array(BoolArray::from_iter([true, false, true].into_iter()).into_array()); + let list = ListArray::try_new(elements, offsets, validity)?; + + let gpu_result = ListExecutor + .execute(list.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(list, gpu_result, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/arrays/masked.rs b/vortex-cuda/src/kernel/arrays/masked.rs new file mode 100644 index 00000000000..b739ed2bf60 --- /dev/null +++ b/vortex-cuda/src/kernel/arrays/masked.rs @@ -0,0 +1,99 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use async_trait::async_trait; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::Masked; +use vortex::array::arrays::masked::MaskedArrayExt; +use vortex::array::arrays::masked::MaskedArraySlotsExt; +use vortex::array::arrays::masked::mask_validity_canonical; +use vortex::array::validity::Validity; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_err; + +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; + +/// CUDA executor for MaskedArray. +/// +/// A `MaskedArray` is a child array that carries no nulls of its own, plus the validity +/// bitmap that supplies them. Decode the child on the GPU, decode the mask on the GPU, and +/// attach the mask to the result. +#[derive(Debug)] +pub(crate) struct MaskedExecutor; + +#[async_trait] +impl CudaExecute for MaskedExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let masked = array + .try_downcast::() + .map_err(|_| vortex_err!("Expected MaskedArray"))?; + + let len = masked.len(); + let validity = masked.masked_validity(); + + // `MaskedArray` guarantees its child holds no nulls, so the mask alone determines the + // output validity. Combining two device-resident bitmaps would need a CPU compute pass. + if matches!(masked.child().validity()?, Validity::Array(_)) { + vortex_bail!( + "MaskedArray child carries a per-element validity bitmap, which cannot be combined with the mask on the GPU" + ); + } + + let child = masked.child().clone().execute_cuda(ctx).await?; + + let validity = execute_validity_cuda(validity, len, ctx).await?; + mask_validity_canonical(child, validity, ctx.execution_ctx()) + } +} + +#[cfg(test)] +mod tests { + use vortex::array::IntoArray; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::MaskedArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::assert_arrays_eq; + use vortex::buffer::buffer; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + + use super::*; + use crate::CanonicalCudaExt; + use crate::session::CudaSession; + + #[crate::test] + async fn test_cuda_masked_applies_validity() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let child = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let validity = Validity::Array( + BoolArray::from_iter([true, false, true, true].into_iter()).into_array(), + ); + let masked = MaskedArray::try_new(child, validity)?; + + let gpu_result = MaskedExecutor + .execute(masked.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(masked, gpu_result, &mut ctx); + + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/arrays/mod.rs b/vortex-cuda/src/kernel/arrays/mod.rs index ab81934bb27..56d4c483480 100644 --- a/vortex-cuda/src/kernel/arrays/mod.rs +++ b/vortex-cuda/src/kernel/arrays/mod.rs @@ -3,8 +3,12 @@ mod constant; mod dict; +mod list; +mod masked; mod shared; pub(crate) use constant::ConstantNumericExecutor; pub(crate) use dict::DictExecutor; +pub(crate) use list::ListExecutor; +pub(crate) use masked::MaskedExecutor; pub(crate) use shared::SharedExecutor; diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index bff691e262b..57f487c5896 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -15,7 +15,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::TemporalArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::buffer::BufferHandle; -use vortex::array::match_each_signed_integer_ptype; +use vortex::array::match_each_integer_ptype; use vortex::array::validity::Validity; use vortex::dtype::DType; use vortex::dtype::NativePType; @@ -110,9 +110,9 @@ impl CudaExecute for DateTimePartsExecutor { let seconds_ptype = seconds_prim.ptype(); let subseconds_ptype = subseconds_prim.ptype(); - match_each_signed_integer_ptype!(days_ptype, |DaysT| { - match_each_signed_integer_ptype!(seconds_ptype, |SecondsT| { - match_each_signed_integer_ptype!(subseconds_ptype, |SubsecondsT| { + match_each_integer_ptype!(days_ptype, |DaysT| { + match_each_integer_ptype!(seconds_ptype, |SecondsT| { + match_each_integer_ptype!(subseconds_ptype, |SubsecondsT| { decode_datetimeparts_typed::( days_prim, seconds_prim, @@ -297,6 +297,46 @@ mod tests { Ok(()) } + /// Compression picks the narrowest ptype per component, so unsigned components are + /// common in real data (the taxi benchmark produces `u16` seconds). + #[crate::test] + async fn test_cuda_datetimeparts_unsigned_components() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let len = 3; + let days_arr = PrimitiveArray::new(buffer![1u32, 2, 3], Validity::NonNullable).into_array(); + let seconds_arr = + PrimitiveArray::new(buffer![3600u16, 0, 60], Validity::NonNullable).into_array(); + let subseconds_arr = + PrimitiveArray::new(buffer![250u8, 0, 99], Validity::NonNullable).into_array(); + + let temporal = TemporalArray::new_timestamp( + PrimitiveArray::new(buffer![0i64; len], Validity::NonNullable).into_array(), + TimeUnit::Milliseconds, + None, + ); + let dtp_array = DateTimeParts::try_new( + temporal.dtype().clone(), + days_arr, + seconds_arr, + subseconds_arr, + )?; + + let gpu_result = DateTimePartsExecutor + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); + + Ok(()) + } + #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 36ceb8c7b8b..da6973ff11c 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -10,8 +10,10 @@ use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::ConstantArray; use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::match_each_native_ptype; @@ -149,10 +151,41 @@ async fn decode_runend_typed { unreachable!("AllInvalid should be handled by RunEndExecutor::execute") } - Validity::Array(_) => { - vortex_bail!( - "RunEnd GPU decoding does not yet support per-element validity in values; falling back to CPU" - ); + Validity::Array(validity) => { + // Expand the per-run validity bitmap through the same run mapping as the values. + let validity_bools = validity.execute_cuda(ctx).await?.into_bool(); + let validity_len = validity_bools.len(); + let BoolDataParts { + bits: validity_bits, + meta: validity_meta, + } = validity_bools.into_data().into_parts(validity_len); + let validity_device = ctx.ensure_on_device(validity_bits).await?; + let validity_view = validity_device.cuda_view::()?; + + // Each thread owns a whole output byte, so threads never race on bits in one byte. + let output_bytes = output_len.div_ceil(8); + let mut validity_out = ctx.device_alloc::(output_bytes)?; + let validity_offset_u64 = validity_meta.offset() as u64; + + let ends_ptype = E::PTYPE.to_string(); + let validity_function = + ctx.load_function_with_suffixes("runend", &["bool", &ends_ptype])?; + ctx.launch_kernel(&validity_function, output_bytes, |args| { + args.arg(&ends_view) + .arg(&num_runs_u64) + .arg(&validity_view) + .arg(&validity_offset_u64) + .arg(&offset_u64) + .arg(&output_len_u64) + .arg(&mut validity_out); + })?; + + let validity_buffer = + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(validity_out))); + Validity::Array( + BoolArray::new_handle(validity_buffer, 0, output_len, Validity::NonNullable) + .into_array(), + ) } }; @@ -182,7 +215,6 @@ mod tests { use super::*; use crate::CanonicalCudaExt; - use crate::executor::CudaArrayExt; use crate::session::CudaSession; fn make_runend_array(ends: Vec, values: Vec, ctx: &mut ExecutionCtx) -> RunEndArray @@ -303,7 +335,7 @@ mod tests { } #[crate::test] - async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { + async fn test_cuda_runend_nullable_values() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -318,13 +350,48 @@ mod tests { PrimitiveArray::new(Buffer::from(vec![10i32, 0, 30]), validity).into_array(); let runend_array = RunEnd::new(ends_array, values_array, cuda_ctx.execution_ctx()); - // execute_cuda should fall back to CPU and still produce the correct result. - let gpu_result = runend_array - .clone() - .into_array() - .execute_cuda(&mut cuda_ctx) + // The GPU expands the per-run validity bitmap through the run mapping. + // Call the executor directly: `execute_cuda` would silently fall back to CPU for a + // 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(); + + assert_arrays_eq!(runend_array, gpu_result, &mut ctx); + + Ok(()) + } + + /// Validity expansion packs bits a byte at a time, so exercise a run layout whose runs + /// straddle output byte boundaries. + #[crate::test] + async fn test_cuda_runend_nullable_values_across_byte_boundaries() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let num_runs = 301; + let ends: Vec = (1..=num_runs).map(|run| run * 3).collect(); + let values: Vec = (0..num_runs as i32).collect(); + let validity = Validity::Array( + BoolArray::from_iter((0..num_runs).map(|run| run % 3 != 0)).into_array(), + ); + + let ends_array = + PrimitiveArray::new(Buffer::from(ends), Validity::NonNullable).into_array(); + let values_array = PrimitiveArray::new(Buffer::from(values), validity).into_array(); + let runend_array = RunEnd::new(ends_array, values_array, cuda_ctx.execution_ctx()); + + // Call the executor directly: `execute_cuda` would silently fall back to CPU for a + // host-resident array, hiding a GPU failure. + let gpu_result = RunEndExecutor + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await - .vortex_expect("GPU/CPU fallback should succeed") + .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index 36735024c7f..19da5a1f6dd 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -31,6 +31,8 @@ mod slice; pub(crate) use arrays::ConstantNumericExecutor; pub(crate) use arrays::DictExecutor; +pub(crate) use arrays::ListExecutor; +pub(crate) use arrays::MaskedExecutor; pub(crate) use arrays::SharedExecutor; pub use encodings::ZstdKernelPrep; pub use encodings::zstd_kernel_prepare; diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 3c712d20fb8..62f8c95f185 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -48,6 +48,8 @@ use kernel::FSSTExecutor; use kernel::FilterExecutor; use kernel::FoRExecutor; pub use kernel::LaunchStrategy; +use kernel::ListExecutor; +use kernel::MaskedExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; @@ -73,6 +75,8 @@ use vortex::array::ArrayVTable; use vortex::array::arrays::Constant; use vortex::array::arrays::Dict; use vortex::array::arrays::Filter; +use vortex::array::arrays::List; +use vortex::array::arrays::Masked; use vortex::array::arrays::Shared; use vortex::array::arrays::Slice; use vortex::encodings::alp::ALP; @@ -115,6 +119,8 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(DateTimeParts.id(), &DateTimePartsExecutor); session.register_kernel(DecimalByteParts.id(), &DecimalBytePartsExecutor); session.register_kernel(Dict.id(), &DictExecutor); + session.register_kernel(List.id(), &ListExecutor); + session.register_kernel(Masked.id(), &MaskedExecutor); session.register_kernel(Shared.id(), &SharedExecutor); session.register_kernel(FoR.id(), &FoRExecutor); session.register_kernel(FSST.id(), &FSSTExecutor);