Execute Widen arrays on CUDA

vortex-cuda/kernels/src/decimal_cast.cu

@@ -6,17 +6,16 @@
 #include <stdint.h>
 #include <type_traits>
 
-// 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.
+// The Rust side rejects narrowing casts, because detecting overflow on the device would require
+// synchronizing an overflow flag back to the host.
 
-// Low 64-bit conversion for Decimal32/64 outputs.
+// Low 64-bit conversion for 16-, 32-, and 64-bit outputs.
 template <typename Input>
 __device__ __forceinline__ int64_t decimal_to_i64(Input value) {
     if constexpr (std::is_same_v<Input, int128_t>) {
@@ -42,10 +41,12 @@ __device__ __forceinline__ int128_t decimal_to_i128(Input value) {
     }
 }
 
-// Convert one value to the Arrow schema's physical width.
+// Convert one value to the output physical width.
 template <typename Output, typename Input>
 __device__ __forceinline__ Output decimal_cast_value(Input value) {
-    if constexpr (std::is_same_v<Output, int32_t>) {
+    if constexpr (std::is_same_v<Output, int16_t>) {
+        return static_cast<int16_t>(decimal_to_i64(value));
+    } else if constexpr (std::is_same_v<Output, int32_t>) {
         return static_cast<int32_t>(decimal_to_i64(value));
     } else if constexpr (std::is_same_v<Output, int64_t>) {
         return decimal_to_i64(value);
@@ -63,7 +64,7 @@ __device__ __forceinline__ Output decimal_cast_value(Input value) {
     }
 }
 
-// Cast a contiguous values buffer to the Arrow schema's physical width.
+// Cast a contiguous values buffer to the output physical width.
 template <typename Input, typename Output>
 __device__ void
 decimal_cast_device(const Input *__restrict input, Output *__restrict output, uint64_t array_len) {
@@ -76,8 +77,13 @@ decimal_cast_device(const Input *__restrict input, Output *__restrict output, ui
     }
 }
 
-// Generate Decimal32/64/128/256 cast kernels for one input storage type.
+// Generate the i16, i32, i64, i128, and i256 cast kernels for one input storage type.
 #define GENERATE_DECIMAL_CAST_KERNELS(input_suffix, InputType)                                               \
+    extern "C" __global__ void decimal_cast_##input_suffix##_i16(const InputType *__restrict input,          \
+                                                                 int16_t *__restrict output,                 \
+                                                                 uint64_t array_len) {                       \
+        decimal_cast_device(input, output, array_len);                                                       \
+    }                                                                                                        \
     extern "C" __global__ void decimal_cast_##input_suffix##_i32(const InputType *__restrict input,          \
                                                                  int32_t *__restrict output,                 \
                                                                  uint64_t array_len) {                       \

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 <type_traits>
+
+template <typename InputT, typename OutputT>
+struct WidenOp {
+    static_assert(sizeof(InputT) < sizeof(OutputT));
+    static_assert(std::is_signed_v<InputT> == std::is_signed_v<OutputT>);
+
+    __device__ OutputT operator()(InputT value) const {
+        return static_cast<OutputT>(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<InputT, OutputT> {});                                \
+    }
+
+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)

vortex-cuda/src/arrow/canonical.rs

@@ -44,15 +44,12 @@ 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;
 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;
@@ -61,7 +58,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;
@@ -91,6 +87,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;
@@ -474,73 +471,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::<S>(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<S>(
-    values: BufferHandle,
-    target_type: DecimalType,
-    len: usize,
-    ctx: &mut CudaExecutionCtx,
-) -> VortexResult<BufferHandle>
-where
-    S: NativeDecimalType + DeviceRepr,
-{
-    match target_type {
-        DecimalType::I32 => decimal_cast::<S, i32>(values, len, ctx).await,
-        DecimalType::I64 => decimal_cast::<S, i64>(values, len, ctx).await,
-        DecimalType::I128 => decimal_cast::<S, i128>(values, len, ctx).await,
-        DecimalType::I256 => decimal_cast::<S, i256>(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<S, D>(
-    values: BufferHandle,
-    len: usize,
-    ctx: &mut CudaExecutionCtx,
-) -> VortexResult<BufferHandle>
-where
-    S: NativeDecimalType + DeviceRepr,
-    D: NativeDecimalType + DeviceRepr,
-{
-    if len == 0 {
-        return ctx
-            .ensure_on_device(BufferHandle::new_host(
-                Buffer::<D>::empty().into_byte_buffer(),
-            ))
-            .await;
-    }
-
-    let mut output_buffer = ctx.device_alloc::<D>(len)?;
-
-    let values_view = values.cuda_view::<S>()?;
-    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.

vortex-cuda/src/dynamic_dispatch/plan_builder.rs

@@ -130,6 +130,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_::<FoR>().references().is_global())
         || id == ZigZag.id()
         || id == Primitive.id()

vortex-cuda/src/kernel/decimal_cast.rs

@@ -0,0 +1,122 @@
+// 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<BufferHandle> {
+    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::<S>(values, target_type, len, ctx).await
+    })
+}
+
+/// Dispatch [`decimal_cast`] from `S` on the runtime `target_type`.
+async fn decimal_cast_to<S>(
+    values: BufferHandle,
+    target_type: DecimalType,
+    len: usize,
+    ctx: &mut CudaExecutionCtx,
+) -> VortexResult<BufferHandle>
+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::<S, i16>(values, len, ctx).await,
+        DecimalType::I32 => decimal_cast::<S, i32>(values, len, ctx).await,
+        DecimalType::I64 => decimal_cast::<S, i64>(values, len, ctx).await,
+        DecimalType::I128 => decimal_cast::<S, i128>(values, len, ctx).await,
+        DecimalType::I256 => decimal_cast::<S, i256>(values, len, ctx).await,
+    }
+}
+
+/// Launch the `decimal_cast_<S>_<D>` kernel, which casts device values of type `S` to type `D`.
+///
+/// The cast must widen. A narrowing cast truncates each value with no overflow check.
+async fn decimal_cast<S, D>(
+    values: BufferHandle,
+    len: usize,
+    ctx: &mut CudaExecutionCtx,
+) -> VortexResult<BufferHandle>
+where
+    S: NativeDecimalType + DeviceRepr,
+    D: NativeDecimalType + DeviceRepr,
+{
+    if len == 0 {
+        return ctx
+            .ensure_on_device(BufferHandle::new_host(
+                Buffer::<D>::empty().into_byte_buffer(),
+            ))
+            .await;
+    }
+
+    let mut output_buffer = ctx.device_alloc::<D>(len)?;
+
+    let values_view = values.cuda_view::<S>()?;
+    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,
+    )?)))
+}

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;