fix[gpu]: handle sliced BP arrays in CUDA (#7912)

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

vortex-cuda/kernels/src/patches.h

@@ -14,7 +14,7 @@ typedef enum { CO_U8 = 0, CO_U16 = 1, CO_U32 = 2, CO_U64 = 3 } ChunkOffsetType;
 
 static const uint32_t PATCH_DERIVE_INDICES_BASE = UINT32_MAX;
 
-/// GPU-resident patches for fused exception patching during bit-unpacking.
+/// GPU-resident patches for fused CUDA exception patching.
 ///
 /// Patches are stored in sorted order within each chunk. The chunk_offsets
 /// array maps each chunk to the start of its range in the indices/values arrays.

vortex-cuda/src/dynamic_dispatch/mod.rs

@@ -2680,11 +2680,6 @@ mod tests {
         #[case] len: usize,
         #[case] slice_range: Option<Range<usize>>,
     ) -> VortexResult<()> {
-        // TODO(#7839): BitPacked SliceReduce returns None when patches are present,
-        // producing SliceArray instead of BitPacked. CUDA cannot handle this yet.
-        if true {
-            return Ok(());
-        }
         let bit_width: u8 = 4;
         let max_val = (1u32 << bit_width) - 1;
         let values: Vec<u32> = (0..len)
@@ -2766,12 +2761,6 @@ mod tests {
 
     #[crate::test]
     async fn test_for_bitpacked_with_patches_sliced() -> VortexResult<()> {
-        // TODO(#7839): BitPacked SliceReduce returns None when patches are present,
-        // producing SliceArray instead of BitPacked. CUDA cannot handle this yet.
-        if true {
-            return Ok(());
-        }
-
         let len = 5000;
         let bit_width: u8 = 6;
         let reference = 42u32;

vortex-cuda/src/dynamic_dispatch/plan_builder.rs

@@ -52,6 +52,7 @@ use super::tag_to_ptype;
 use crate::CudaBufferExt;
 use crate::CudaExecutionCtx;
 use crate::executor::CudaDispatchMode;
+use crate::kernel::bitpacked_slice_view;
 use crate::kernel::load_patches_to_gpu;
 
 /// A plan whose source buffers have been copied to the device, ready for kernel launch.
@@ -154,12 +155,13 @@ pub fn has_standalone_kernel(array: &ArrayRef) -> bool {
 
 /// Patch payload attached to the op that consumes it.
 ///
-/// `slice` is a logical output range to apply when materializing the patch descriptor on the GPU.
-/// This lets the planner avoid calling `Patches::slice` when patch metadata may already be device-resident.
+/// `range` is the logical output range to apply when materializing the patch descriptor on the GPU.
+/// This lets the planner avoid calling `Patches::slice` when patch metadata may already be
+/// device-resident.
 #[derive(Clone)]
 struct PlanPatches {
     patches: Patches,
-    slice: Option<Range<usize>>,
+    range: Option<Range<usize>>,
 }
 
 /// An unmaterialized stage: a source op, scalar ops, and optional source buffer reference.
@@ -191,6 +193,11 @@ impl Stage {
             source_ptype,
         }
     }
+
+    fn with_source_patches(mut self, source_patches: Option<PlanPatches>) -> Self {
+        self.source_patches = source_patches;
+        self
+    }
 }
 
 type SmemByteOffset = u32;
@@ -409,7 +416,7 @@ impl FusedPlan {
             // Upload source patches (e.g. BitPacked exceptions).
             if let Some(patches) = &stage.source_patches {
                 let (ptr, bufs) =
-                    load_patches_to_gpu(&patches.patches, patches.slice.clone(), ctx).await?;
+                    load_patches_to_gpu(&patches.patches, patches.range.clone(), ctx).await?;
                 source.params.bitunpack.patches_ptr = ptr;
                 device_buffers.extend(bufs);
             }
@@ -419,7 +426,7 @@ impl FusedPlan {
             for (mut op, patches) in stage.scalar_ops.clone() {
                 if let Some(patches) = &patches {
                     let (ptr, bufs) =
-                        load_patches_to_gpu(&patches.patches, patches.slice.clone(), ctx).await?;
+                        load_patches_to_gpu(&patches.patches, patches.range.clone(), ctx).await?;
                     op.params.alp.patches_ptr = ptr;
                     device_buffers.extend(bufs);
                 }
@@ -503,9 +510,8 @@ impl FusedPlan {
 
     /// SliceArray → resolve the slice via reduce/execute rules.
     ///
-    /// When the plan builder encounters a `SliceArray`, it resolves the slice
-    /// by invoking the child's `reduce_parent`. If that fails (e.g. ALP
-    /// doesn't implement it), we manually slice the child's sub-arrays.
+    /// When the plan builder encounters a `SliceArray`, it first asks the child to reduce the
+    /// slice. If reduction fails, the planner falls back to encoding-specific handling.
     fn walk_slice(
         &mut self,
         array: ArrayRef,
@@ -518,6 +524,30 @@ impl FusedPlan {
             return self.walk(reduced, pending_subtrees);
         }
 
+        // BitPacked with patches does not reduce through Slice. Slice the
+        // packed buffer here, and defer patch slicing to CUDA materialization.
+        if child.encoding_id() == BitPacked.id() {
+            let bp = child.as_::<BitPacked>();
+            let offset = slice_arr.data().slice_range().start;
+            let len = array.len();
+            let (packed, bitpacked_offset, patch_range) = bitpacked_slice_view(bp, offset, len)?;
+
+            let source_ptype = ptype_to_tag(PType::try_from(bp.dtype()).map_err(|_| {
+                vortex_err!("BitPacked must have primitive dtype, got {:?}", bp.dtype())
+            })?);
+            let buf_index = self.source_buffers.len();
+            self.source_buffers.push(Some(packed));
+            return Ok(Stage::new(
+                SourceOp::bitunpack(bp.bit_width(), bitpacked_offset),
+                Some(buf_index),
+                source_ptype,
+            )
+            .with_source_patches(bp.patches().map(|patches| PlanPatches {
+                patches,
+                range: Some(patch_range),
+            })));
+        }
+
         // ALP doesn't implement reduce_parent. Slice the encoded child here,
         // and defer patch slicing to CUDA materialization so device-resident
         // patch buffers stay on device.
@@ -530,7 +560,7 @@ impl FusedPlan {
                 sliced_encoded,
                 alp.patches().map(|patches| PlanPatches {
                     patches,
-                    slice: Some(offset..offset + len),
+                    range: Some(offset..offset + len),
                 }),
                 alp.exponents(),
                 pending_subtrees,
@@ -562,16 +592,15 @@ impl FusedPlan {
         })?);
         let buf_index = self.source_buffers.len();
         self.source_buffers.push(Some(bp.packed().clone()));
-        let mut stage = Stage::new(
+        Ok(Stage::new(
             SourceOp::bitunpack(bp.bit_width(), bp.offset()),
             Some(buf_index),
             source_ptype,
-        );
-        stage.source_patches = bp.patches().map(|patches| PlanPatches {
+        )
+        .with_source_patches(bp.patches().map(|patches| PlanPatches {
             patches,
-            slice: None,
-        });
-        Ok(stage)
+            range: None,
+        })))
     }
 
     fn walk_for(
@@ -629,7 +658,7 @@ impl FusedPlan {
             alp.encoded().clone(),
             alp.patches().map(|patches| PlanPatches {
                 patches,
-                slice: None,
+                range: None,
             }),
             alp.exponents(),
             pending_subtrees,

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

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::fmt::Debug;
+use std::ops::Range;
 
 use async_trait::async_trait;
 use cudarc::driver::CudaFunction;
@@ -10,14 +11,20 @@ use cudarc::driver::LaunchConfig;
 use cudarc::driver::PushKernelArg;
 use tracing::instrument;
 use vortex::array::ArrayRef;
+use vortex::array::ArrayVTable;
+use vortex::array::ArrayView;
 use vortex::array::Canonical;
 use vortex::array::arrays::PrimitiveArray;
+use vortex::array::arrays::Slice;
+use vortex::array::arrays::slice::SliceArrayExt;
 use vortex::array::buffer::BufferHandle;
 use vortex::array::buffer::DeviceBufferExt;
 use vortex::array::match_each_integer_ptype;
+use vortex::array::patches::PATCH_CHUNK_SIZE;
 use vortex::dtype::NativePType;
 use vortex::encodings::fastlanes::BitPacked;
 use vortex::encodings::fastlanes::BitPackedArray;
+use vortex::encodings::fastlanes::BitPackedArrayExt;
 use vortex::encodings::fastlanes::BitPackedDataParts;
 use vortex::encodings::fastlanes::unpack_iter::BitPacked as BitPackedUnpack;
 use vortex::error::VortexResult;
@@ -30,14 +37,71 @@ use crate::executor::CudaExecute;
 use crate::executor::CudaExecutionCtx;
 use crate::kernel::patches::build_gpu_patches;
 use crate::kernel::patches::types::load_device_patches;
+use crate::kernel::patches::types::slice_device_patches;
 
 /// CUDA decoder for bit-packed arrays.
 #[derive(Debug)]
 pub(crate) struct BitPackedExecutor;
 
+/// Build the packed buffer view for decoding `Slice(BitPacked)`.
+///
+/// Bit-unpack kernels decode full FastLanes chunks, so the packed buffer is
+/// widened to chunk boundaries and `offset` is converted into the in-chunk
+/// starting position. The returned logical range is passed to patch
+/// materialization so exception metadata is sliced consistently.
+pub(crate) fn bitpacked_slice_view(
+    bp: ArrayView<'_, BitPacked>,
+    offset: usize,
+    len: usize,
+) -> VortexResult<(BufferHandle, u16, Range<usize>)> {
+    let patch_range = offset..offset + len;
+    let offset_start = patch_range.start + bp.offset() as usize;
+    let offset_stop = offset_start + len;
+    let bitpacked_offset = offset_start % PATCH_CHUNK_SIZE;
+    let block_start = offset_start - bitpacked_offset;
+    let block_stop = offset_stop.div_ceil(PATCH_CHUNK_SIZE) * PATCH_CHUNK_SIZE;
+
+    let encoded_start = (block_start / 8) * bp.bit_width() as usize;
+    let encoded_stop = (block_stop / 8) * bp.bit_width() as usize;
+
+    Ok((
+        bp.packed().slice(encoded_start..encoded_stop),
+        u16::try_from(bitpacked_offset)?,
+        patch_range,
+    ))
+}
+
 impl BitPackedExecutor {
-    fn try_specialize(array: ArrayRef) -> Option<BitPackedArray> {
-        array.try_downcast::<BitPacked>().ok()
+    fn try_specialize(
+        array: ArrayRef,
+    ) -> VortexResult<Option<(BitPackedArray, Option<Range<usize>>)>> {
+        if let Ok(array) = array.clone().try_downcast::<BitPacked>() {
+            return Ok(Some((array, None)));
+        }
+
+        let Some(slice) = array.as_opt::<Slice>() else {
+            return Ok(None);
+        };
+        let child = slice.child();
+        if child.encoding_id() != BitPacked.id() {
+            return Ok(None);
+        }
+
+        let bp = child.as_::<BitPacked>();
+        let offset = slice.data().slice_range().start;
+        let len = array.len();
+        let (packed, bitpacked_offset, patch_range) = bitpacked_slice_view(bp, offset, len)?;
+        let sliced = BitPacked::try_new(
+            packed,
+            bp.ptype(bp.dtype()),
+            child.validity()?.slice(patch_range.clone())?,
+            bp.patches(),
+            bp.bit_width(),
+            len,
+            bitpacked_offset,
+        )?;
+
+        Ok(Some((sliced, Some(patch_range))))
     }
 }
 
@@ -49,11 +113,12 @@ impl CudaExecute for BitPackedExecutor {
         array: ArrayRef,
         ctx: &mut CudaExecutionCtx,
     ) -> VortexResult<Canonical> {
-        let array =
-            Self::try_specialize(array).ok_or_else(|| vortex_err!("Expected BitPackedArray"))?;
+        let (array, patch_range) =
+            Self::try_specialize(array)?.ok_or_else(|| vortex_err!("Expected BitPackedArray"))?;
+        let ptype = array.ptype(array.dtype());
 
-        match_each_integer_ptype!(array.ptype(array.dtype()), |A| {
-            decode_bitpacked::<A>(array, A::default(), ctx).await
+        match_each_integer_ptype!(ptype, |A| {
+            decode_bitpacked::<A>(array, A::default(), patch_range, ctx).await
         })
     }
 }
@@ -88,6 +153,7 @@ pub fn bitpacked_cuda_launch_config(output_width: usize, len: usize) -> VortexRe
 pub(crate) async fn decode_bitpacked<A>(
     array: BitPackedArray,
     reference: A,
+    patch_range: Option<Range<usize>>,
     ctx: &mut CudaExecutionCtx,
 ) -> VortexResult<Canonical>
 where
@@ -122,7 +188,11 @@ where
 
     // We hold this here to keep the device buffers alive.
     let device_patches = if let Some(patches) = patches {
-        Some(load_device_patches(&patches, ctx).await?)
+        let mut device_patches = load_device_patches(&patches, ctx).await?;
+        if let Some(range) = patch_range {
+            slice_device_patches(&patches, range, &mut device_patches);
+        }
+        Some(device_patches)
     } else {
         None
     };
@@ -551,17 +621,53 @@ mod tests {
         Ok(())
     }
 
+    #[rstest]
+    #[case::direct(None, 4096, None, 0, 4096 * 9 / 8)]
+    #[case::mid_chunk(Some(67..3969), 3902, Some(67..3969), 67, 4096 * 9 / 8)]
+    #[case::chunk_aligned(Some(1024..3072), 2048, Some(1024..3072), 0, 2048 * 9 / 8)]
+    #[case::tail_chunk(Some(3000..4096), 1096, Some(3000..4096), 952, 2048 * 9 / 8)]
+    #[crate::test]
+    fn test_bitunpack_try_specialize_slices(
+        #[case] range: Option<Range<usize>>,
+        #[case] expected_len: usize,
+        #[case] expected_patch_range: Option<Range<usize>>,
+        #[case] expected_offset: u16,
+        #[case] expected_packed_len: usize,
+    ) -> VortexResult<()> {
+        let values = PrimitiveArray::new(
+            (0u16..4096)
+                .map(|i| if i % 1000 == 0 { 600 } else { i % 512 })
+                .collect::<Buffer<_>>(),
+            NonNullable,
+        );
+        let bitpacked = BitPacked::encode(
+            &values.into_array(),
+            9,
+            &mut LEGACY_SESSION.create_execution_ctx(),
+        )?;
+        assert!(bitpacked.patches().is_some());
+        let array = if let Some(range) = range {
+            bitpacked.into_array().slice(range)?
+        } else {
+            bitpacked.into_array()
+        };
+
+        let (specialized, patch_range) =
+            BitPackedExecutor::try_specialize(array)?.vortex_expect("expected BitPacked input");
+
+        assert_eq!(specialized.len(), expected_len);
+        assert_eq!(specialized.offset(), expected_offset);
+        assert_eq!(specialized.packed().len(), expected_packed_len);
+        assert_eq!(patch_range, expected_patch_range);
+
+        Ok(())
+    }
+
     /// Test slicing a bitpacked array with patches where the slice boundary
     /// falls in the middle of a chunk's patch range, creating a non-zero
     /// offset_within_chunk.
     #[crate::test]
     fn test_cuda_bitunpack_sliced_patches_offset_within_chunk() -> VortexResult<()> {
-        // TODO(#7839): BitPacked SliceReduce returns None when patches are present,
-        // producing SliceArray instead of BitPacked. CUDA cannot handle this yet.
-        if true {
-            return Ok(());
-        }
-
         let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
             .vortex_expect("failed to create execution context");
 
@@ -604,12 +710,6 @@ mod tests {
     /// Test slicing a bitpacked array multiple times, accumulating offset_within_chunk.
     #[crate::test]
     fn test_cuda_bitunpack_double_sliced_patches() -> VortexResult<()> {
-        // TODO(#7839): BitPacked SliceReduce returns None when patches are present,
-        // producing SliceArray instead of BitPacked. CUDA cannot handle this yet.
-        if true {
-            return Ok(());
-        }
-
         let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
             .vortex_expect("failed to create execution context");
 
@@ -664,12 +764,6 @@ mod tests {
     /// Test slicing to skip an entire chunk's worth of patches.
     #[crate::test]
     fn test_cuda_bitunpack_sliced_skip_first_chunk_patches() -> VortexResult<()> {
-        // TODO(#7839): BitPacked SliceReduce returns None when patches are present,
-        // producing SliceArray instead of BitPacked. CUDA cannot handle this yet.
-        if true {
-            return Ok(());
-        }
-
         let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
             .vortex_expect("failed to create execution context");