GPU: keep patch indices on device

vortex-cuda/benches/dynamic_dispatch_cuda.rs

@@ -410,7 +410,8 @@ mod standalone {
     struct NullGpuPatches {
         chunk_offsets: *mut c_void,
         chunk_offset_type: u32,
-        indices: *mut u32,
+        indices_type: u32,
+        indices: *mut c_void,
         values: *mut c_void,
         offset: u32,
         offset_within_chunk: u32,
@@ -425,6 +426,7 @@ mod standalone {
         const NULL: Self = Self {
             chunk_offsets: ptr::null_mut(),
             chunk_offset_type: 2,
+            indices_type: 2,
             indices: ptr::null_mut(),
             values: ptr::null_mut(),
             offset: 0,

vortex-cuda/kernels/src/patches.cuh

@@ -6,21 +6,25 @@
 #include "fastlanes_common.cuh"
 #include "patches.h"
 
-/// Load a chunk offset value, dispatching on the runtime type.
-__device__ inline uint32_t load_chunk_offset(const GPUPatches &patches, uint32_t idx) {
-    switch (patches.chunk_offset_type) {
-    case CO_U8:
-        return reinterpret_cast<const uint8_t *>(patches.chunk_offsets)[idx];
-    case CO_U16:
-        return reinterpret_cast<const uint16_t *>(patches.chunk_offsets)[idx];
-    case CO_U32:
-        return reinterpret_cast<const uint32_t *>(patches.chunk_offsets)[idx];
-    case CO_U64:
-        return static_cast<uint32_t>(reinterpret_cast<const uint64_t *>(patches.chunk_offsets)[idx]);
+/// Load an unsigned integer value, dispatching on the runtime type.
+__device__ inline uint32_t load_unsigned(const void *values, UnsignedType type, uint32_t idx) {
+    switch (type) {
+    case UNSIGNED_U8:
+        return reinterpret_cast<const uint8_t *>(values)[idx];
+    case UNSIGNED_U16:
+        return reinterpret_cast<const uint16_t *>(values)[idx];
+    case UNSIGNED_U32:
+        return reinterpret_cast<const uint32_t *>(values)[idx];
+    case UNSIGNED_U64:
+        return static_cast<uint32_t>(reinterpret_cast<const uint64_t *>(values)[idx]);
     }
     return 0;
 }
 
+__device__ inline uint32_t load_chunk_offset(const GPUPatches &patches, uint32_t idx) {
+    return load_unsigned(patches.chunk_offsets, patches.chunk_offset_type, idx);
+}
+
 /// A single patch: a within-chunk index and its replacement value.
 /// A sentinel patch has index == FL_CHUNK, which can never match a valid
 /// within-chunk position (0–FL_CHUNK-1).
@@ -49,17 +53,14 @@ public:
     /// Construct a cursor for this thread's portion of patches in the chunk.
     __device__
     PatchesCursor(const GPUPatches &patches, uint32_t chunk, uint32_t thread_idx, uint32_t n_threads) {
-        if (patches.chunk_offsets == nullptr) {
-            indices = nullptr;
-            values = nullptr;
-            remaining = 0;
-            return;
-        }
+        indices = nullptr;
+        indices_type = UNSIGNED_U32;
+        index = 0;
+        values = nullptr;
+        remaining = 0;
+        chunk_base = 0;
 
-        if (chunk >= patches.n_chunks) {
-            indices = nullptr;
-            values = nullptr;
-            remaining = 0;
+        if (patches.chunk_offsets == nullptr || chunk >= patches.n_chunks) {
             return;
         }
 
@@ -94,7 +95,9 @@ public:
 
         uint32_t start = patches_start_idx + my_start;
         remaining = my_end - my_start;
-        indices = patches.indices + start;
+        indices = patches.indices;
+        indices_type = patches.indices_type;
+        index = start;
         values = reinterpret_cast<const T *>(patches.values) + start;
 
         // The iterator returns indices relative to the start of the chunk.
@@ -110,16 +113,19 @@ public:
         if (remaining == 0) {
             return {FL_CHUNK, T {}};
         }
-        uint16_t within_chunk = static_cast<uint16_t>(*indices - chunk_base);
+        uint16_t within_chunk =
+            static_cast<uint16_t>(load_unsigned(indices, indices_type, index) - chunk_base);
         Patch<T> patch = {within_chunk, *values};
-        indices++;
+        index++;
         values++;
         remaining--;
         return patch;
     }
 
 private:
-    const uint32_t *indices;
+    const void *indices;
+    UnsignedType indices_type;
+    uint32_t index;
     const T *values;
     uint32_t remaining;
     uint32_t chunk_base;

vortex-cuda/kernels/src/patches.h

@@ -9,8 +9,8 @@
 extern "C" {
 #endif
 
-/// Type tag for chunk_offsets pointer.
-typedef enum { CO_U8 = 0, CO_U16 = 1, CO_U32 = 2, CO_U64 = 3 } ChunkOffsetType;
+/// Type tag for an unsigned integer pointer.
+typedef enum { UNSIGNED_U8 = 0, UNSIGNED_U16 = 1, UNSIGNED_U32 = 2, UNSIGNED_U64 = 3 } UnsignedType;
 
 static const uint32_t PATCH_DERIVE_INDICES_BASE = UINT32_MAX;
 
@@ -24,8 +24,9 @@ static const uint32_t PATCH_DERIVE_INDICES_BASE = UINT32_MAX;
 /// A NULL chunk_offsets pointer indicates no patches are present.
 typedef struct {
     void *chunk_offsets;
-    ChunkOffsetType chunk_offset_type;
-    uint32_t *indices;
+    UnsignedType chunk_offset_type;
+    UnsignedType indices_type;
+    void *indices;
     void *values;
     uint32_t offset;
     uint32_t offset_within_chunk;

vortex-cuda/src/kernel/encodings/bitpacked.rs

@@ -206,11 +206,6 @@ where
             .arg(&patches_arg);
     })?;
 
-    // Patch-free decodes need no host synchronization.
-    if device_patches.is_some() {
-        ctx.synchronize_stream()?;
-    }
-
     let output_buf = CudaDeviceBuffer::new(output_slice);
     let output_handle =
         BufferHandle::new_device(output_buf.slice_typed::<A>(offset..(offset + len)));
@@ -234,6 +229,7 @@ mod tests {
     use vortex::array::validity::Validity::NonNullable;
     use vortex::buffer::Buffer;
     use vortex::buffer::buffer;
+    use vortex::dtype::PType;
     use vortex::encodings::fastlanes::BitPackedArrayExt;
     use vortex::error::VortexExpect;
     use vortex_array::VortexSessionExecute;
@@ -307,6 +303,64 @@ mod tests {
         Ok(())
     }
 
+    #[rstest]
+    #[case::u8(PrimitiveArray::from_iter([0u8, 100, 200]).into_array())]
+    #[case::u16(PrimitiveArray::from_iter([0u16, 100, 200]).into_array())]
+    #[case::u32(PrimitiveArray::from_iter([0u32, 100, 200]).into_array())]
+    #[case::u64(PrimitiveArray::from_iter([0u64, 100, 200]).into_array())]
+    #[crate::test]
+    fn test_cuda_bitunpack_native_patch_index_widths(
+        #[case] indices: ArrayRef,
+    ) -> 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 mut values: Vec<u16> = (0..1024).map(|i| i % 16).collect();
+        values[0] = 500;
+        values[100] = 600;
+        values[200] = 700;
+        let expected = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array();
+
+        let encoded = BitPacked::encode(&expected, 4, &mut ctx)?;
+        let BitPackedDataParts {
+            offset,
+            bit_width,
+            len,
+            packed,
+            patches,
+            validity,
+        } = BitPacked::into_parts(encoded);
+        let original_patches = patches.vortex_expect("expected patches");
+        let native_patches = vortex_array::patches::Patches::new(
+            len,
+            0,
+            indices,
+            original_patches.values().clone(),
+            original_patches.chunk_offsets().clone(),
+        )?;
+        let encoded = BitPacked::try_new(
+            packed,
+            PType::U16,
+            validity,
+            Some(native_patches),
+            bit_width,
+            len,
+            offset,
+        )?;
+
+        let gpu_result = block_on(async {
+            BitPackedExecutor
+                .execute(encoded.into_array(), &mut cuda_ctx)
+                .await
+                .vortex_expect("GPU decompression failed")
+                .into_host()
+                .await
+                .map(|array| array.into_array())
+        })?;
+        assert_arrays_eq!(expected, gpu_result, &mut ctx);
+        Ok(())
+    }
+
     #[rstest]
     #[case::bw_1(1)]
     #[case::bw_2(2)]

vortex-cuda/src/kernel/patches/mod.rs

@@ -27,13 +27,13 @@ use crate::CudaBufferExt;
 use crate::CudaDeviceBuffer;
 use crate::CudaExecutionCtx;
 use crate::executor::CudaArrayExt;
-use crate::kernel::patches::gpu::ChunkOffsetType;
-use crate::kernel::patches::gpu::ChunkOffsetType_CO_U8;
-use crate::kernel::patches::gpu::ChunkOffsetType_CO_U16;
-use crate::kernel::patches::gpu::ChunkOffsetType_CO_U32;
-use crate::kernel::patches::gpu::ChunkOffsetType_CO_U64;
 use crate::kernel::patches::gpu::GPUPatches;
 use crate::kernel::patches::gpu::PATCH_DERIVE_INDICES_BASE;
+use crate::kernel::patches::gpu::UnsignedType;
+use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U8;
+use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U16;
+use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U32;
+use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U64;
 use crate::kernel::patches::types::DevicePatches;
 
 // Safe because `GPUPatches` contains only raw pointers, POD integers, and an enum.
@@ -44,7 +44,8 @@ impl GPUPatches {
     /// `chunk_offsets` pointer is the signal `PatchesCursor` checks for.
     pub(crate) const NULL_PATCHES: Self = Self {
         chunk_offsets: std::ptr::null_mut(),
-        chunk_offset_type: ChunkOffsetType_CO_U32,
+        chunk_offset_type: UnsignedType_UNSIGNED_U32,
+        indices_type: UnsignedType_UNSIGNED_U32,
         indices: std::ptr::null_mut(),
         values: std::ptr::null_mut(),
         offset: 0,
@@ -55,14 +56,14 @@ impl GPUPatches {
     };
 }
 
-/// Convert a [`PType`] to the corresponding [`ChunkOffsetType`] for GPU patches.
-pub(crate) fn ptype_to_chunk_offset_type(ptype: PType) -> VortexResult<ChunkOffsetType> {
+/// Convert a [`PType`] to the corresponding [`UnsignedType`] for GPU patches.
+pub(crate) fn ptype_to_unsigned_type(ptype: PType) -> VortexResult<UnsignedType> {
     match ptype {
-        PType::U8 => Ok(ChunkOffsetType_CO_U8),
-        PType::U16 => Ok(ChunkOffsetType_CO_U16),
-        PType::U32 => Ok(ChunkOffsetType_CO_U32),
-        PType::U64 => Ok(ChunkOffsetType_CO_U64),
-        _ => vortex_bail!("Invalid PType for chunk_offsets: {:?}", ptype),
+        PType::U8 => Ok(UnsignedType_UNSIGNED_U8),
+        PType::U16 => Ok(UnsignedType_UNSIGNED_U16),
+        PType::U32 => Ok(UnsignedType_UNSIGNED_U32),
+        PType::U64 => Ok(UnsignedType_UNSIGNED_U64),
+        _ => vortex_bail!("Invalid unsigned PType: {:?}", ptype),
     }
 }
 
@@ -77,7 +78,8 @@ pub(crate) fn build_gpu_patches(
     match device_patches {
         Some(p) => Ok(GPUPatches {
             chunk_offsets: p.chunk_offsets.cuda_device_ptr()? as _,
-            chunk_offset_type: ptype_to_chunk_offset_type(p.chunk_offset_ptype)?,
+            chunk_offset_type: ptype_to_unsigned_type(p.chunk_offset_ptype)?,
+            indices_type: ptype_to_unsigned_type(p.indices_ptype)?,
             indices: p.indices.cuda_device_ptr()? as _,
             values: p.values.cuda_device_ptr()? as _,
             offset: p.offset as u32,

vortex-cuda/src/kernel/patches/types.rs

@@ -6,31 +6,29 @@
 use std::mem::size_of;
 use std::ops::Range;
 
-use num_traits::ToPrimitive;
 use vortex::array::buffer::BufferHandle;
 use vortex::buffer::Alignment;
-use vortex::buffer::Buffer;
-use vortex::buffer::BufferMut;
 use vortex::buffer::ByteBufferMut;
 use vortex::dtype::PType;
-use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_array::patches::PATCH_CHUNK_SIZE;
 use vortex_array::patches::Patches;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
 
 use crate::CudaBufferExt;
 use crate::CudaExecutionCtx;
 use crate::executor::CudaArrayExt;
 use crate::kernel::patches::gpu::GPUPatches;
 use crate::kernel::patches::gpu::PATCH_DERIVE_INDICES_BASE;
-use crate::kernel::patches::ptype_to_chunk_offset_type;
+use crate::kernel::patches::ptype_to_unsigned_type;
 
 /// A set of device-resident patches.
 pub struct DevicePatches {
     pub(crate) chunk_offsets: BufferHandle,
     pub(crate) chunk_offset_ptype: PType,
     pub(crate) indices: BufferHandle,
+    pub(crate) indices_ptype: PType,
     pub(crate) values: BufferHandle,
     pub(crate) offset: usize,
     pub(crate) offset_within_chunk: usize,
@@ -55,6 +53,12 @@ pub(crate) async fn load_device_patches(
     ctx: &mut CudaExecutionCtx,
 ) -> VortexResult<DevicePatches> {
     let offset = patches.offset();
+    vortex_ensure!(
+        offset
+            .checked_add(patches.array_len())
+            .is_some_and(|end| end <= u32::MAX as usize),
+        "CUDA patches require offset + array length to fit in u32"
+    );
     let offset_within_chunk = patches.offset_within_chunk().unwrap_or_default();
     // Get or compute chunk_offsets
     let Some(co) = patches.chunk_offsets() else {
@@ -68,31 +72,15 @@ pub(crate) async fn load_device_patches(
         (co_canonical.buffer_handle().clone(), ptype, len)
     };
 
-    // Load indices - must be converted to u32 for GPU use
+    // Load indices at their native width.
     let indices = patches
         .indices()
         .clone()
         .execute_cuda(ctx)
         .await?
         .into_primitive();
     let indices_ptype = indices.ptype();
-    #[expect(clippy::expect_used)]
-    let indices = if indices_ptype == PType::U32 {
-        indices.buffer_handle().clone()
-    } else {
-        // Convert indices to u32
-        let indices_buf = indices.buffer_handle().to_host().await;
-        let indices_u32 = match_each_unsigned_integer_ptype!(indices_ptype, |I| {
-            let src: Buffer<I> = Buffer::from_byte_buffer(indices_buf);
-            let mut dst: BufferMut<u32> = BufferMut::with_capacity(src.len());
-            for &idx in src.as_slice() {
-                // Indices are limited to u32 range for GPU
-                dst.push(idx.to_u32().expect("index should fit in u32"));
-            }
-            dst.freeze()
-        });
-        BufferHandle::new_host(indices_u32.into_byte_buffer())
-    };
+    let indices = indices.buffer_handle().clone();
 
     // Load values
     let values = patches
@@ -113,6 +101,7 @@ pub(crate) async fn load_device_patches(
         chunk_offsets,
         chunk_offset_ptype,
         indices,
+        indices_ptype,
         values,
         offset,
         offset_within_chunk,
@@ -134,7 +123,8 @@ fn build_gpu_patches(
     // chunk_offset_type and indices) which would be UB when serialized.
     let mut gpu_patches: GPUPatches = unsafe { std::mem::zeroed() };
     gpu_patches.chunk_offsets = dp.chunk_offsets.cuda_device_ptr()? as _;
-    gpu_patches.chunk_offset_type = ptype_to_chunk_offset_type(dp.chunk_offset_ptype)?;
+    gpu_patches.chunk_offset_type = ptype_to_unsigned_type(dp.chunk_offset_ptype)?;
+    gpu_patches.indices_type = ptype_to_unsigned_type(dp.indices_ptype)?;
     gpu_patches.indices = dp.indices.cuda_device_ptr()? as _;
     gpu_patches.values = dp.values.cuda_device_ptr()? as _;
     gpu_patches.offset = dp.offset as u32;