feat[cuda]: device exports for varbinview (#6338)

Need to copy to VarBin before sending to cudf

---------

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

vortex-cuda/kernels/src/varbinview.cuh

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#pragma once
+
+#include <stdint.h>
+
+// Maximum length of inlined string.
+constexpr int32_t MAX_INLINED_SIZE = 12;
+
+// a byte buffer holding string data
+typedef uint8_t* Buffer;
+
+// an i32 offsets buffer
+typedef int32_t* Offsets;
+
+struct InlinedBinaryView {
+    int32_t size;
+    uint8_t data[12];
+};
+
+struct RefBinaryView {
+    int32_t size;
+    uint8_t prefix[4];
+    int32_t index;
+    int32_t offset;
+};
+
+// The BinaryView type is how we access values.
+union alignas(int64_t) BinaryView {
+    InlinedBinaryView inlined;
+    RefBinaryView ref;
+};

vortex-cuda/kernels/src/varbinview_compute_offsets.cu

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#include "config.cuh"
+#include "varbinview.cuh"
+
+// single-threaded, compute offsets
+extern "C" __global__ void varbinview_compute_offsets(
+    const BinaryView *views,
+    int64_t num_strings,
+    Offsets out_offsets,
+    int32_t *last_offset
+) {
+    const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    // force execution to be single-threaded to compute the prefix
+    // sum.
+    // TODO(aduffy): we could do this with a CUB kernel instead.
+    //  Check the profiles later to see where this shows up.
+    if (tid != 0) {
+        return;
+    }
+
+    int32_t offset = 0;
+    out_offsets[0] = 0;
+    for (int i = 0; i < num_strings; i++) {
+        offset += views[i].inlined.size;
+        out_offsets[i + 1] = offset;
+    }
+
+    *last_offset = offset;
+}

vortex-cuda/kernels/src/varbinview_copy_strings.cu

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#include "config.cuh"
+#include "varbinview.cuh"
+
+// Lookup a string from a binary view, copying it into
+// a destination buffer.
+__device__ void copy_string_to_dst(
+    BinaryView& view,
+    Buffer *buffers,
+    uint8_t *dst
+) {
+    int32_t size = view.inlined.size;
+    uint8_t *src;
+    if (size <= MAX_INLINED_SIZE) {
+        // TODO(aduffy): use uint64_t loads instead?
+        src = view.inlined.data;
+    } else {
+        auto ref = view.ref;
+        src = buffers[ref.index] + ref.offset;
+    }
+    memcpy(dst, src, size);
+}
+
+extern "C" __global__ void varbinview_copy_strings(
+    int64_t len,
+    BinaryView* views,
+    Buffer* buffers,
+    Buffer dst_buffer,
+    Offsets dst_offsets
+) {
+    const int64_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    // Each thread is responsible for copying a single string.
+    // Any excess threads do no work.
+    if (tid >= len) {
+        return;
+    }
+
+    auto view = views[tid];
+    int32_t offset = dst_offsets[tid];
+    uint8_t *dst = &dst_buffer[offset];
+
+    copy_string_to_dst(view, buffers, dst);
+}

vortex-cuda/src/arrow/canonical.rs

@@ -12,7 +12,7 @@ use vortex_array::arrays::PrimitiveArrayParts;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::StructArrayParts;
 use vortex_array::buffer::BufferHandle;
-use vortex_array::validity::Validity;
+use vortex_array::vtable::ValidityHelper;
 use vortex_dtype::DecimalType;
 use vortex_dtype::datetime::AnyTemporal;
 use vortex_error::VortexResult;
@@ -26,6 +26,10 @@ use crate::arrow::DeviceType;
 use crate::arrow::ExportDeviceArray;
 use crate::arrow::PrivateData;
 use crate::arrow::SyncEvent;
+use crate::arrow::check_validity_empty;
+use crate::arrow::ensure_device_resident;
+use crate::arrow::varbinview::BinaryParts;
+use crate::arrow::varbinview::copy_varbinview_to_varbin;
 use crate::executor::CudaArrayExt;
 
 /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by
@@ -43,11 +47,11 @@ impl ExportDeviceArray for CanonicalDeviceArrayExport {
     ) -> VortexResult<ArrowDeviceArray> {
         let cuda_array = array.execute_cuda(ctx).await?;
 
-        let (arrow_array, sync_event) = export_canonical(cuda_array, ctx).await?;
+        let (arrow_array, _) = export_canonical(cuda_array, ctx).await?;
 
         Ok(ArrowDeviceArray {
             array: arrow_array,
-            sync_event,
+            sync_event: None,
             device_id: ctx.stream().context().ordinal() as i64,
             device_type: DeviceType::Cuda,
             _reserved: Default::default(),
@@ -68,7 +72,7 @@ fn export_canonical(
                     buffer, validity, ..
                 } = primitive.into_parts();
 
-                check_validity_empty(validity)?;
+                check_validity_empty(&validity)?;
 
                 let buffer = ensure_device_resident(buffer, ctx).await?;
 
@@ -96,7 +100,7 @@ fn export_canonical(
                 } = decimal.into_parts();
 
                 // verify that there is no null buffer
-                check_validity_empty(validity)?;
+                check_validity_empty(&validity)?;
 
                 // TODO(aduffy): GPU kernel for upcasting.
                 vortex_ensure!(
@@ -120,12 +124,11 @@ fn export_canonical(
                     buffer, validity, ..
                 } = values.into_parts();
 
-                check_validity_empty(validity)?;
+                check_validity_empty(&validity)?;
 
                 let buffer = ensure_device_resident(buffer, ctx).await?;
                 export_fixed_size(buffer, len, 0, ctx)
             }
-
             Canonical::Bool(bool_array) => {
                 let BoolArrayParts {
                     bits,
@@ -135,10 +138,40 @@ fn export_canonical(
                     ..
                 } = bool_array.into_parts();
 
-                check_validity_empty(validity)?;
+                check_validity_empty(&validity)?;
 
                 export_fixed_size(bits, len, offset, ctx)
             }
+            Canonical::VarBinView(varbinview) => {
+                let len = varbinview.len();
+                check_validity_empty(varbinview.validity())?;
+
+                let BinaryParts { offsets, bytes } =
+                    copy_varbinview_to_varbin(varbinview, ctx).await?;
+
+                let offsets = ensure_device_resident(offsets, ctx).await?;
+                let bytes = ensure_device_resident(bytes, ctx).await?;
+
+                let buffers = vec![None, Some(offsets), Some(bytes)];
+                let mut private_data = PrivateData::new(buffers, vec![], ctx)?;
+                let sync_event = private_data.sync_event();
+                //
+                let arrow_array = ArrowArray {
+                    length: len as i64,
+                    null_count: 0,
+                    offset: 0,
+                    // 1 (optional) buffer for nulls, one buffer for the data
+                    n_buffers: 2,
+                    buffers: private_data.buffer_ptrs.as_mut_ptr(),
+                    n_children: 0,
+                    children: std::ptr::null_mut(),
+                    release: Some(release_array),
+                    dictionary: std::ptr::null_mut(),
+                    private_data: Box::into_raw(private_data).cast(),
+                };
+
+                Ok((arrow_array, sync_event))
+            }
             // TODO(aduffy): implement VarBinView. cudf doesn't support it, so we need to
             //  execute a kernel to translate from VarBinView -> VarBin.
             c => todo!("support for exporting {} arrays", c.dtype()),
@@ -155,7 +188,7 @@ async fn export_struct(
         validity, fields, ..
     } = array.into_parts();
 
-    check_validity_empty(validity)?;
+    check_validity_empty(&validity)?;
 
     // We need the children to be held across await points.
     let mut children = Vec::with_capacity(fields.len());
@@ -220,26 +253,6 @@ fn export_fixed_size(
     Ok((arrow_array, sync_event))
 }
 
-/// Check that the validity buffer is empty and does not need to be copied over the device boundary.
-fn check_validity_empty(validity: Validity) -> VortexResult<()> {
-    if let Validity::AllInvalid | Validity::Array(_) = validity {
-        vortex_bail!("Exporting array with non-trivial validity not supported yet")
-    }
-
-    Ok(())
-}
-
-async fn ensure_device_resident(
-    buffer_handle: BufferHandle,
-    ctx: &mut CudaExecutionCtx,
-) -> VortexResult<BufferHandle> {
-    if buffer_handle.is_on_device() {
-        Ok(buffer_handle)
-    } else {
-        ctx.move_to_device(buffer_handle)?.await
-    }
-}
-
 // export some nested data
 
 unsafe extern "C" fn release_array(array: *mut ArrowArray) {

vortex-cuda/src/arrow/mod.rs

@@ -9,6 +9,7 @@
 //! More documentation at <https://arrow.apache.org/docs/format/CDeviceDataInterface.html>
 
 mod canonical;
+mod varbinview;
 
 use std::ffi::c_void;
 use std::fmt::Debug;
@@ -24,7 +25,9 @@ use cudarc::runtime::sys::cudaEvent_t;
 use vortex_array::Array;
 use vortex_array::ArrayRef;
 use vortex_array::buffer::BufferHandle;
+use vortex_array::validity::Validity;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
 
 use crate::CudaBufferExt;
@@ -213,3 +216,23 @@ pub trait ExportDeviceArray: Debug + Send + Sync + 'static {
         ctx: &mut CudaExecutionCtx,
     ) -> VortexResult<ArrowDeviceArray>;
 }
+
+/// Check that the validity buffer is empty and does not need to be copied over the device boundary.
+pub(crate) fn check_validity_empty(validity: &Validity) -> VortexResult<()> {
+    if let Validity::AllInvalid | Validity::Array(_) = validity {
+        vortex_bail!("Exporting array with non-trivial validity not supported yet")
+    }
+
+    Ok(())
+}
+
+pub(crate) async fn ensure_device_resident(
+    buffer_handle: BufferHandle,
+    ctx: &mut CudaExecutionCtx,
+) -> VortexResult<BufferHandle> {
+    if buffer_handle.is_on_device() {
+        Ok(buffer_handle)
+    } else {
+        ctx.move_to_device(buffer_handle)?.await
+    }
+}