diff --git a/vortex-cuda/kernels/src/decimal_cast.cu b/vortex-cuda/kernels/src/decimal_cast.cu index 45c8c75e2a0..88618176a94 100644 --- a/vortex-cuda/kernels/src/decimal_cast.cu +++ b/vortex-cuda/kernels/src/decimal_cast.cu @@ -6,64 +6,43 @@ #include #include -// Arrow decimal schemas fix the physical values buffer width: -// - Decimal32: 4 bytes per value. -// - Decimal64: 8 bytes per value. -// - Decimal128: 16 bytes per value. -// - Decimal256: 32 bytes per value. +// Widen decimal values buffers to a wider physical width. Two callers use these kernels: +// - The Arrow export casts to the width that the Arrow schema fixes: 4 bytes for Decimal32, +// 8 bytes for Decimal64, 16 bytes for Decimal128, and 32 bytes for Decimal256. +// - The Widen executor casts a narrow decimal child to the canonical width of its decimal +// dtype, from 2 to 32 bytes. // -// Vortex storage width can differ, so export casts to the schema-implied width. -// Rust-side export rejects narrowing casts because detecting overflow on-device -// would require synchronizing an overflow flag back to the host. +// Only widening pairs are instantiated. A narrowing cast can overflow, and detecting that on the +// device would require synchronizing an overflow flag back to the host, so the Rust side rejects +// narrowing before it loads a kernel. -// Low 64-bit conversion for Decimal32/64 outputs. -template -__device__ __forceinline__ int64_t decimal_to_i64(Input value) { - if constexpr (std::is_same_v) { - return value.lo; - } else if constexpr (std::is_same_v) { - return value.parts[0]; - } else { - return static_cast(value); - } -} - -// 128-bit conversion for Decimal128 outputs. -template -__device__ __forceinline__ int128_t decimal_to_i128(Input value) { - if constexpr (std::is_same_v) { - return value; - } else if constexpr (std::is_same_v) { - return int128_t {value.parts[0], value.parts[1]}; - } else { - const int64_t lo = static_cast(value); - const int64_t hi = value < 0 ? -1 : 0; - return int128_t {lo, hi}; - } -} - -// Convert one value to the Arrow schema's physical width. +// Sign-extend one value to the wider output width in a single step, with no intermediate width. +// Native outputs use one `static_cast`. The 128-bit and 256-bit outputs keep the input's words +// and fill the upper words with its sign. template __device__ __forceinline__ Output decimal_cast_value(Input value) { - if constexpr (std::is_same_v) { - return static_cast(decimal_to_i64(value)); - } else if constexpr (std::is_same_v) { - return decimal_to_i64(value); - } else if constexpr (std::is_same_v) { - return decimal_to_i128(value); + static_assert(sizeof(Input) < sizeof(Output), "decimal casts on the GPU only widen"); + + if constexpr (std::is_same_v) { + static_assert(std::is_integral_v && std::is_signed_v); + const int64_t lo = static_cast(value); + return int128_t {lo, lo < 0 ? -1 : 0}; + } else if constexpr (std::is_same_v && std::is_same_v) { + const int64_t sign = value.hi < 0 ? -1 : 0; + return int256_t {{value.lo, value.hi, sign, sign}}; + } else if constexpr (std::is_same_v) { + static_assert(std::is_integral_v && std::is_signed_v); + const int64_t lo = static_cast(value); + const int64_t sign = lo < 0 ? -1 : 0; + return int256_t {{lo, sign, sign, sign}}; } else { - static_assert(std::is_same_v); - if constexpr (std::is_same_v) { - return value; - } else { - const int128_t value128 = decimal_to_i128(value); - const int64_t sign = value128.hi < 0 ? -1 : 0; - return int256_t {{value128.lo, value128.hi, sign, sign}}; - } + static_assert(std::is_integral_v && std::is_signed_v); + static_assert(std::is_integral_v && std::is_signed_v); + return static_cast(value); } } -// Cast a contiguous values buffer to the Arrow schema's physical width. +// Cast a contiguous values buffer to the output physical width. template __device__ void decimal_cast_device(const Input *__restrict input, Output *__restrict output, uint64_t array_len) { @@ -76,28 +55,26 @@ decimal_cast_device(const Input *__restrict input, Output *__restrict output, ui } } -// Generate Decimal32/64/128/256 cast kernels for one input storage type. -#define GENERATE_DECIMAL_CAST_KERNELS(input_suffix, InputType) \ - extern "C" __global__ void decimal_cast_##input_suffix##_i32(const InputType *__restrict input, \ - int32_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i64(const InputType *__restrict input, \ - int64_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i128(const InputType *__restrict input, \ - int128_t *__restrict output, \ - uint64_t array_len) { \ - decimal_cast_device(input, output, array_len); \ - } \ - extern "C" __global__ void decimal_cast_##input_suffix##_i256(const InputType *__restrict input, \ - int256_t *__restrict output, \ - uint64_t array_len) { \ +#define GENERATE_DECIMAL_CAST(input_suffix, InputType, output_suffix, OutputType) \ + extern "C" __global__ void decimal_cast_##input_suffix##_##output_suffix( \ + const InputType *__restrict input, \ + OutputType *__restrict output, \ + uint64_t array_len) { \ decimal_cast_device(input, output, array_len); \ } -FOR_EACH_SIGNED_INT(GENERATE_DECIMAL_CAST_KERNELS) -FOR_EACH_LARGE_DECIMAL(GENERATE_DECIMAL_CAST_KERNELS) +GENERATE_DECIMAL_CAST(i8, int8_t, i16, int16_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i32, int32_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i8, int8_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i32, int32_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i16, int16_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i64, int64_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i32, int32_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i64, int64_t, i128, int128_t) +GENERATE_DECIMAL_CAST(i64, int64_t, i256, int256_t) +GENERATE_DECIMAL_CAST(i128, int128_t, i256, int256_t) diff --git a/vortex-cuda/kernels/src/widen.cu b/vortex-cuda/kernels/src/widen.cu new file mode 100644 index 00000000000..9654b9db3a0 --- /dev/null +++ b/vortex-cuda/kernels/src/widen.cu @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "scalar_kernel.cuh" +#include "types.cuh" +#include + +template +struct WidenOp { + static_assert(sizeof(InputT) < sizeof(OutputT)); + static_assert(std::is_signed_v == std::is_signed_v); + + __device__ OutputT operator()(InputT value) const { + return static_cast(value); + } +}; + +#define GENERATE_WIDEN(input_suffix, InputT, output_suffix, OutputT) \ + extern "C" __global__ void widen_##input_suffix##_##output_suffix(const InputT *__restrict input, \ + OutputT *__restrict output, \ + uint64_t array_len) { \ + scalar_kernel(input, output, array_len, WidenOp {}); \ + } + +GENERATE_WIDEN(i8, int8_t, i16, int16_t) +GENERATE_WIDEN(i8, int8_t, i32, int32_t) +GENERATE_WIDEN(i8, int8_t, i64, int64_t) +GENERATE_WIDEN(i16, int16_t, i32, int32_t) +GENERATE_WIDEN(i16, int16_t, i64, int64_t) +GENERATE_WIDEN(i32, int32_t, i64, int64_t) +GENERATE_WIDEN(u8, uint8_t, u16, uint16_t) +GENERATE_WIDEN(u8, uint8_t, u32, uint32_t) +GENERATE_WIDEN(u8, uint8_t, u64, uint64_t) +GENERATE_WIDEN(u16, uint16_t, u32, uint32_t) +GENERATE_WIDEN(u16, uint16_t, u64, uint64_t) +GENERATE_WIDEN(u32, uint32_t, u64, uint64_t) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index b3c87d1abbf..67dfe0cd9ae 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -44,16 +44,13 @@ 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::optimizer::ArrayOptimizer; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::DecimalType; -use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; -use vortex::dtype::i256; use vortex::encodings::fsst::FSST; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -62,7 +59,6 @@ use vortex::error::vortex_ensure_eq; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; use vortex_cuda_driver::CudaSlice; -use vortex_cuda_driver::DeviceRepr; use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::PushKernelArg; use vortex_onpair::OnPair; @@ -92,6 +88,7 @@ use crate::executor::execute_validity_cuda; use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; use crate::kernel::decode_onpair_varbin; +use crate::kernel::widen_decimal_values; use crate::stream::zero_padding; const BITMAP_THREADS_PER_BLOCK: u32 = 256; @@ -475,73 +472,11 @@ async fn export_decimal_values( "cannot export decimal values from {values_type} storage to Arrow {target_type}: narrowing would require a device-to-host overflow check", ); } - let values = ctx.ensure_on_device(values).await?; if values_type == target_type { - return Ok(values); + return ctx.ensure_on_device(values).await; } - match_each_decimal_value_type!(values_type, |S| { - export_decimal_values_from::(values, target_type, len, ctx).await - }) -} - -/// Dispatch from a concrete storage type `S` to the Arrow-required output width. -async fn export_decimal_values_from( - values: BufferHandle, - target_type: DecimalType, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, -{ - match target_type { - DecimalType::I32 => decimal_cast::(values, len, ctx).await, - DecimalType::I64 => decimal_cast::(values, len, ctx).await, - DecimalType::I128 => decimal_cast::(values, len, ctx).await, - DecimalType::I256 => decimal_cast::(values, len, ctx).await, - target_type => { - vortex_bail!("cannot export DecimalArray as Arrow decimal value type {target_type}") - } - } -} - -/// Launches the CUDA kernel that casts from Vortex storage type `S` to Arrow output type `D`. -/// -/// The caller must ensure this is not a narrowing cast. -async fn decimal_cast( - values: BufferHandle, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, - D: NativeDecimalType + DeviceRepr, -{ - if len == 0 { - return ctx - .ensure_on_device(BufferHandle::new_host( - Buffer::::empty().into_byte_buffer(), - )) - .await; - } - - let mut output_buffer = ctx.device_alloc::(len)?; - - let values_view = values.cuda_view::()?; - let len_u64 = len as u64; - let cuda_function = ctx.load_function_with_suffixes( - "decimal_cast", - &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], - )?; - - ctx.launch_kernel(&cuda_function, len, |args| { - args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); - })?; - - Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( - output_buffer, - )?))) + widen_decimal_values(values, values_type, target_type, len, ctx).await } /// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index d0405f7ee41..9498f1c8419 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -133,6 +133,8 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool { _ => false, }; } + // TODO(connor): Add Widen fusion after fixing sign extension of signed narrow BitPacked and + // FoR children in dynamic_dispatch.cu. Those paths currently zero-extend to the output width. (id == FoR.id() && array.as_::().references().is_global()) || id == ZigZag.id() || id == Primitive.id() diff --git a/vortex-cuda/src/kernel/decimal_cast.rs b/vortex-cuda/src/kernel/decimal_cast.rs new file mode 100644 index 00000000000..8f1fb0d9f87 --- /dev/null +++ b/vortex-cuda/src/kernel/decimal_cast.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Widening casts of decimal values buffers on the device. +//! +//! The Arrow export widens decimal values to the width that the Arrow schema fixes, and the +//! Widen executor widens a narrow decimal child to the canonical width of its decimal dtype. Both +//! launch the `decimal_cast` kernels through this module. +//! +//! These casts only widen. A narrowing cast can overflow, and to detect that on the device, the +//! kernel would need an overflow flag that synchronizes with the host. So this module rejects a +//! narrowing cast, and each caller decides how to handle values that are wider than its target. + +use std::sync::Arc; + +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_decimal_value_type; +use vortex::buffer::Buffer; +use vortex::dtype::DecimalType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::i256; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_ensure_eq; +use vortex::error::vortex_panic; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; + +/// Widen `len` decimal values from `values_type` to the wider `target_type` on the device. +/// +/// The values can be on the host or on the device, and the result is on the device. Each value is +/// sign-extended, so the result holds the same values as the input. +/// +/// # Errors +/// +/// Returns an error if `target_type` is not wider than `values_type`, if the buffer does not +/// contain `len` complete values, or if the copy to the device or the kernel launch fails. +pub(crate) async fn widen_decimal_values( + values: BufferHandle, + values_type: DecimalType, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!( + values_type < target_type, + "a decimal cast on the GPU must widen to a type wider than {values_type}, got \ + {target_type}" + ); + vortex_ensure_eq!(values.len(), len * values_type.byte_width()); + + let values = ctx.ensure_on_device(values).await?; + + match_each_decimal_value_type!(values_type, |S| { + decimal_cast_to::(values, target_type, len, ctx).await + }) +} + +/// Dispatch [`decimal_cast`] from `S` on the runtime `target_type`. +async fn decimal_cast_to( + values: BufferHandle, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, +{ + match target_type { + DecimalType::I8 => vortex_panic!( + "widen_decimal_values ensures a target wider than the values type, and no decimal \ + type is narrower than i8, got {target_type}" + ), + DecimalType::I16 => decimal_cast::(values, len, ctx).await, + DecimalType::I32 => decimal_cast::(values, len, ctx).await, + DecimalType::I64 => decimal_cast::(values, len, ctx).await, + DecimalType::I128 => decimal_cast::(values, len, ctx).await, + DecimalType::I256 => decimal_cast::(values, len, ctx).await, + } +} + +/// Launch the `decimal_cast__` kernel, which sign-extends device values of type `S` to the +/// wider type `D`. +/// +/// `D` **must** be wider than `S`. The `decimal_cast.cu` kernels only instantiate widening pairs, so +/// loading any other pair fails. [`widen_decimal_values`] checks this before it dispatches. +async fn decimal_cast( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, + D: NativeDecimalType + DeviceRepr, +{ + if len == 0 { + return ctx + .ensure_on_device(BufferHandle::new_host( + Buffer::::empty().into_byte_buffer(), + )) + .await; + } + + let mut output_buffer = ctx.device_alloc::(len)?; + + let values_view = values.cuda_view::()?; + let len_u64 = len as u64; + let cuda_function = ctx.load_function_with_suffixes( + "decimal_cast", + &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], + )?; + + ctx.launch_kernel(&cuda_function, len, |args| { + args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output_buffer, + )?))) +} diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index fb292bc445c..fbc2defb696 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -1,6 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! CUDA executors for encoded arrays. +//! +//! Executors decode children through the CUDA execution context and return canonical arrays. + mod alp; mod bitpacked; mod date_time_parts; @@ -11,6 +15,7 @@ mod fsst; mod onpair; mod runend; mod sequence; +mod widen; mod zigzag; mod zstd; mod zstd_buffers; @@ -29,6 +34,7 @@ pub(crate) use onpair::OnPairExecutor; pub(crate) use onpair::decode_onpair_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; +pub(crate) use widen::WidenExecutor; pub(crate) use zigzag::ZigZagExecutor; pub(crate) use zstd::ZstdExecutor; pub use zstd::ZstdKernelPrep; diff --git a/vortex-cuda/src/kernel/encodings/widen/mod.rs b/vortex-cuda/src/kernel/encodings/widen/mod.rs new file mode 100644 index 00000000000..be324b868b7 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/widen/mod.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Restores the logical width of integer and decimal [`Widen`] arrays on the device. +//! +//! Children execute at their own width before widening, so compressed signed values retain their +//! sign. Execution keeps the child's validity and allocates a wider values buffer, unless decimal +//! storage already has the output width. + +use std::sync::Arc; + +use async_trait::async_trait; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::Widen; +use vortex::array::arrays::widen::WidenArraySlotsExt; +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_integer_ptype; +use vortex::dtype::DType; +use vortex::dtype::NativePType; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::kernel::widen_decimal_values; + +/// CUDA execution for integer and decimal arrays stored at a narrower width. +#[derive(Debug)] +pub(crate) struct WidenExecutor; + +#[async_trait] +impl CudaExecute for WidenExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let array = array.try_downcast::().map_err(|array| { + vortex_err!( + "WidenExecutor requires a Widen array, got {}", + array.encoding_id() + ) + })?; + let child = array.values().clone().execute_cuda(ctx).await?; + + match (array.dtype(), child) { + (DType::Primitive(output_ptype, _), Canonical::Primitive(child)) => { + let len = child.len(); + let parts = child.into_data_parts(); + let values = ctx.ensure_on_device(parts.buffer).await?; + let output = match_each_integer_ptype!(parts.ptype, |S| { + match_each_integer_ptype!(*output_ptype, |D| { + widen_integers::(values, len, ctx)? + }) + }); + + Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + output, + *output_ptype, + parts.validity, + ))) + } + (DType::Decimal(decimal_dtype, _), Canonical::Decimal(child)) => { + let len = child.len(); + let parts = child.into_data_parts(); + let target_type = decimal_dtype.canonical_values_type(); + vortex_ensure!( + parts.values_type <= target_type, + "CUDA Widen requires decimal storage no wider than {target_type}, got {}", + parts.values_type + ); + + let values = if parts.values_type == target_type { + ctx.ensure_on_device(parts.values).await? + } else { + widen_decimal_values(parts.values, parts.values_type, target_type, len, ctx) + .await? + }; + + Ok(Canonical::Decimal(DecimalArray::try_new_handle( + values, + target_type, + *decimal_dtype, + parts.validity, + )?)) + } + (dtype, child) => vortex_bail!( + "CUDA Widen requires an integer or decimal child matching {dtype}, got {}", + child.dtype() + ), + } + } +} + +/// Widens `len` device values of type `S` into a new buffer of type `D`. +/// +/// `D` must be strictly wider than `S` with the same signedness, as the [`Widen`] invariant +/// guarantees. The `widen.cu` kernels only instantiate those pairs. The other monomorphized pairs +/// cannot run because a valid `Widen` never dispatches to them. +fn widen_integers( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativePType + DeviceRepr, + D: NativePType + DeviceRepr, +{ + let input = values.cuda_view::()?; + let mut output = ctx.device_alloc::(len)?; + let len_u64 = len as u64; + let function = ctx.load_function("widen", &[S::PTYPE, D::PTYPE])?; + ctx.launch_kernel(&function, len, |args| { + args.arg(&input).arg(&mut output).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output, + )?))) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-cuda/src/kernel/encodings/widen/tests.rs b/vortex-cuda/src/kernel/encodings/widen/tests.rs new file mode 100644 index 00000000000..9b0781eeb73 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/widen/tests.rs @@ -0,0 +1,246 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Device execution and CPU agreement for Widen arrays. + +use num_traits::Bounded; +use rstest::rstest; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::IntoArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::WidenArray; +use vortex::array::arrays::decimal::DecimalArrayExt; +use vortex::array::assert_arrays_eq; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::dtype::DType; +use vortex::dtype::DecimalDType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::NativePType; +use vortex::dtype::Nullability; +use vortex::dtype::PType; +use vortex::encodings::fastlanes::FoR; +use vortex::error::VortexResult; + +use crate::CanonicalCudaExt; +use crate::CudaDispatchMode; +use crate::CudaExecutionCtx; +use crate::CudaSession; +use crate::cuda_session; +use crate::executor::CudaArrayExt; + +/// Copies fixture values to the device while preserving their dtype and validity. +async fn upload(child: ArrayRef, ctx: &mut CudaExecutionCtx) -> VortexResult { + match child.execute::(ctx.execution_ctx())? { + Canonical::Primitive(child) => { + let parts = child.into_data_parts(); + Ok(PrimitiveArray::from_buffer_handle( + ctx.ensure_on_device(parts.buffer).await?, + parts.ptype, + parts.validity, + ) + .into_array()) + } + Canonical::Decimal(child) => { + let parts = child.into_data_parts(); + Ok(DecimalArray::try_new_handle( + ctx.ensure_on_device(parts.values).await?, + parts.values_type, + parts.decimal_dtype, + parts.validity, + )? + .into_array()) + } + _ => unreachable!("Widen test fixtures are primitive or decimal"), + } +} + +async fn check_widen( + host_child: ArrayRef, + device_child: ArrayRef, + dtype: DType, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<()> { + let host = WidenArray::try_new(host_child, dtype.clone())?.into_array(); + let expected = host.execute::(ctx.execution_ctx())?.into_array(); + let device = WidenArray::try_new(device_child, dtype.clone())?.into_array(); + assert!(!device.is_host()); + + let result = device.execute_cuda(ctx).await?; + assert_eq!(result.dtype(), &dtype); + assert_eq!(result.len(), expected.len()); + match &result { + Canonical::Primitive(array) => assert!(!array.buffer_handle().is_on_host()), + Canonical::Decimal(array) => { + assert!(!array.buffer_handle().is_on_host()); + assert_eq!( + array.values_type(), + dtype + .as_decimal_opt() + .expect("the result is decimal") + .canonical_values_type() + ); + } + _ => unreachable!("Widen returns primitive or decimal arrays"), + } + + let actual = result.into_host(ctx.execution_ctx()).await?.into_array(); + assert_arrays_eq!(actual, expected, ctx.execution_ctx()); + + Ok(()) +} + +fn validity(nullable: bool, len: usize) -> Validity { + if nullable { + Validity::from_iter((0..len).map(|i| i % 3 != 0)) + } else { + Validity::NonNullable + } +} + +#[rstest] +#[case::i8_i16(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I16)] +#[case::i8_i32(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I32)] +#[case::i8_i64(Buffer::from_iter([i8::MIN, i8::MAX, -1, 0, 1]), PType::I64)] +#[case::i16_i32(Buffer::from_iter([i16::MIN, i16::MAX, -1, 0, 1]), PType::I32)] +#[case::i16_i64(Buffer::from_iter([i16::MIN, i16::MAX, -1, 0, 1]), PType::I64)] +#[case::i32_i64(Buffer::from_iter([i32::MIN, i32::MAX, -1, 0, 1]), PType::I64)] +#[case::u8_u16(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U16)] +#[case::u8_u32(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U32)] +#[case::u8_u64(Buffer::from_iter([u8::MIN, u8::MAX, 0, 1, 1 << 7]), PType::U64)] +#[case::u16_u32(Buffer::from_iter([u16::MIN, u16::MAX, 0, 1, 1 << 15]), PType::U32)] +#[case::u16_u64(Buffer::from_iter([u16::MIN, u16::MAX, 0, 1, 1 << 15]), PType::U64)] +#[case::u32_u64(Buffer::from_iter([u32::MIN, u32::MAX, 0, 1, 1 << 31]), PType::U64)] +#[crate::test] +async fn test_integer_pairs( + #[case] values: Buffer, + #[case] output: PType, + #[values(false, true)] nullable: bool, + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let validity = validity(nullable, values.len()); + let child = PrimitiveArray::new(values, validity).into_array(); + let dtype = DType::Primitive(output, child.dtype().nullability()); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, dtype, &mut ctx).await +} + +#[rstest] +#[case::i8_p4([-99i8, 99], 2, 4)] +#[case::i8_p9([-99i8, 99], 2, 9)] +#[case::i8_p18([-99i8, 99], 2, 18)] +#[case::i8_p38([-99i8, 99], 2, 38)] +#[case::i8_p76([-99i8, 99], 2, 76)] +#[case::i16_p9([-9_999i16, 9_999], 4, 9)] +#[case::i16_p18([-9_999i16, 9_999], 4, 18)] +#[case::i16_p38([-9_999i16, 9_999], 4, 38)] +#[case::i16_p76([-9_999i16, 9_999], 4, 76)] +#[case::i32_p18([-999_999_999i32, 999_999_999], 9, 18)] +#[case::i32_p38([-999_999_999i32, 999_999_999], 9, 38)] +#[case::i32_p76([-999_999_999i32, 999_999_999], 9, 76)] +#[case::i64_p38([-999_999_999_999_999_999i64, 999_999_999_999_999_999], 18, 38)] +#[case::i64_p76([-999_999_999_999_999_999i64, 999_999_999_999_999_999], 18, 76)] +#[case::i128_p76([-(1i128 << 100), (1i128 << 100)], 38, 76)] +#[crate::test] +async fn test_decimal_pairs( + #[case] bounds: [T; 2], + #[case] child_precision: u8, + #[case] precision: u8, + #[values(false, true)] nullable: bool, + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let mut values = [ + T::default(), + bounds[0], + T::default(), + T::default(), + bounds[1], + ]; + if nullable { + values[0] = T::min_value(); + values[3] = T::max_value(); + } + + let child = DecimalArray::new( + Buffer::from_iter(values), + DecimalDType::new(child_precision, 1), + validity(nullable, 5), + ) + .into_array(); + let dtype = DType::Decimal(DecimalDType::new(precision, 1), child.dtype().nullability()); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, dtype, &mut ctx).await +} + +#[rstest] +#[crate::test] +async fn test_compressed_child( + #[values(CudaDispatchMode::Auto, CudaDispatchMode::StandaloneOnly)] mode: CudaDispatchMode, +) -> VortexResult<()> { + let child = PrimitiveArray::from_iter([0i8, 1, 50, 99]).into_array(); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?.with_dispatch_mode(mode); + let device_child = upload(child.clone(), &mut ctx).await?; + let host = FoR::try_new(child, (-100i8).into())?.into_array(); + let device = FoR::try_new(device_child, (-100i8).into())?.into_array(); + + check_widen(host, device, PType::I64.into(), &mut ctx).await +} + +#[rstest] +#[case::one(1)] +#[case::before_vector(15)] +#[case::vector(16)] +#[case::after_vector(17)] +#[case::before_block(2047)] +#[case::block(2048)] +#[case::after_block(2049)] +#[crate::test] +async fn test_kernel_boundaries(#[case] len: usize) -> VortexResult<()> { + let child = PrimitiveArray::from_iter((0..len).map(|i| if i % 2 == 0 { -99i8 } else { 99 })) + .into_array(); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?; + let device_child = upload(child.clone(), &mut ctx).await?; + + check_widen(child, device_child, PType::I64.into(), &mut ctx).await +} + +#[rstest] +#[case::equal_width(false)] +#[case::narrowing(true)] +#[crate::test] +async fn test_legacy_decimal_storage(#[case] narrowing: bool) -> VortexResult<()> { + let dtype = DecimalDType::new(9, 1); + let child = if narrowing { + DecimalArray::from_iter([-99i128, 0, 99], dtype) + } else { + DecimalArray::from_iter([-99i64, 0, 99], dtype) + } + .into_array(); + let logical = DType::Decimal(DecimalDType::new(18, 1), Nullability::NonNullable); + let mut ctx = CudaSession::create_execution_ctx(&cuda_session())?; + let device_child = upload(child.clone(), &mut ctx).await?; + + if narrowing { + let array = WidenArray::try_new(device_child, logical)?.into_array(); + let error = array + .execute_cuda(&mut ctx) + .await + .expect_err("i128 storage requires narrowing to i64"); + assert!( + error + .to_string() + .contains("CUDA Widen requires decimal storage no wider than i64, got i128"), + "{error}" + ); + } else { + check_widen(child, device_child, logical, &mut ctx).await?; + } + + Ok(()) +} diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index f1eb24b1fad..94d15ffa6d0 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -19,6 +19,10 @@ use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::sys::CUevent_flags; mod arrays; + +mod decimal_cast; +pub(crate) use decimal_cast::widen_decimal_values; + // Fatbins let the CUDA driver select among the configured real and virtual architectures. mod embedded_kernels { include!(concat!(env!("OUT_DIR"), "/embedded_kernels.rs")); diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index bdf9af1421b..0e60e61b2c1 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -54,6 +54,7 @@ use kernel::OnPairExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; +use kernel::WidenExecutor; use kernel::ZigZagExecutor; use kernel::ZstdBuffersExecutor; use kernel::ZstdExecutor; @@ -80,6 +81,7 @@ use vortex::array::arrays::Filter; use vortex::array::arrays::List; use vortex::array::arrays::Masked; use vortex::array::arrays::Slice; +use vortex::array::arrays::Widen; use vortex::encodings::alp::ALP; use vortex::encodings::datetime_parts::DateTimeParts; use vortex::encodings::decimal_byte_parts::DecimalByteParts; @@ -124,6 +126,7 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(OnPair.id(), &OnPairExecutor); session.register_kernel(RunEnd.id(), &RunEndExecutor); session.register_kernel(Sequence.id(), &SequenceExecutor); + session.register_kernel(Widen.id(), &WidenExecutor); session.register_kernel(ZigZag.id(), &ZigZagExecutor); session.register_kernel(Zstd.id(), &ZstdExecutor); session.register_kernel(ZstdBuffers.id(), &ZstdBuffersExecutor);