From 8c180a7618ad343631ff37080c8c076bf7e02463 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 17 Aug 2026 21:03:39 +0100 Subject: [PATCH 1/2] Keep GPU patch indices on device Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/benches/dynamic_dispatch_cuda.rs | 4 +- vortex-cuda/kernels/src/patches.cuh | 56 ++++++++++-------- vortex-cuda/kernels/src/patches.h | 9 +-- vortex-cuda/src/kernel/encodings/bitpacked.rs | 59 +++++++++++++++++++ vortex-cuda/src/kernel/patches/mod.rs | 30 +++++----- vortex-cuda/src/kernel/patches/types.rs | 38 +++++------- 6 files changed, 128 insertions(+), 68 deletions(-) diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index 2d9db5ddbbd..d9115ba08d8 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/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, diff --git a/vortex-cuda/kernels/src/patches.cuh b/vortex-cuda/kernels/src/patches.cuh index 076bf11e72c..8acb068a89b 100644 --- a/vortex-cuda/kernels/src/patches.cuh +++ b/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(patches.chunk_offsets)[idx]; - case CO_U16: - return reinterpret_cast(patches.chunk_offsets)[idx]; - case CO_U32: - return reinterpret_cast(patches.chunk_offsets)[idx]; - case CO_U64: - return static_cast(reinterpret_cast(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(values)[idx]; + case UNSIGNED_U16: + return reinterpret_cast(values)[idx]; + case UNSIGNED_U32: + return reinterpret_cast(values)[idx]; + case UNSIGNED_U64: + return static_cast(reinterpret_cast(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(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(*indices - chunk_base); + uint16_t within_chunk = + static_cast(load_unsigned(indices, indices_type, index) - chunk_base); Patch 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; diff --git a/vortex-cuda/kernels/src/patches.h b/vortex-cuda/kernels/src/patches.h index 32dfa0de2cc..6f520adc3ee 100644 --- a/vortex-cuda/kernels/src/patches.h +++ b/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; diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 86b7a88b276..3ec161a6223 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -234,6 +234,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 +308,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 = (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)] diff --git a/vortex-cuda/src/kernel/patches/mod.rs b/vortex-cuda/src/kernel/patches/mod.rs index 7c651b7e507..727ef4b7611 100644 --- a/vortex-cuda/src/kernel/patches/mod.rs +++ b/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 { +/// Convert a [`PType`] to the corresponding [`UnsignedType`] for GPU patches. +pub(crate) fn ptype_to_unsigned_type(ptype: PType) -> VortexResult { 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, diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 3bfe2270b66..da0058a63fc 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/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 { 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,7 +72,7 @@ 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() @@ -76,23 +80,7 @@ pub(crate) async fn load_device_patches( .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 = Buffer::from_byte_buffer(indices_buf); - let mut dst: BufferMut = 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; From 67c91104ae8fa96e23a9315ee8df6174d065075c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 17 Aug 2026 21:04:03 +0100 Subject: [PATCH 2/2] Avoid synchronizing fused GPU patch decodes Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/src/kernel/encodings/bitpacked.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 3ec161a6223..da67326e893 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/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::(offset..(offset + len)));