apply alp patches in parallel in cuda (#8793)

We used to walk through the patches after decoding ALP then apply them
one by one. Now we scatter them among workers and apply these in
parallel

---------

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-cuda/benches/dynamic_dispatch_cuda.rs

@@ -7,6 +7,7 @@
 
 mod bench_config;
 
+use std::f64::consts::PI;
 use std::marker::PhantomData;
 use std::mem::size_of;
 use std::os::raw::c_void;
@@ -39,6 +40,7 @@ use vortex::dtype::NativePType;
 use vortex::dtype::PType;
 use vortex::encodings::alp::ALP;
 use vortex::encodings::alp::ALPArrayExt;
+use vortex::encodings::alp::ALPArrayOwnedExt;
 use vortex::encodings::alp::ALPArraySlotsExt;
 use vortex::encodings::alp::ALPFloat;
 use vortex::encodings::alp::Exponents;
@@ -745,40 +747,50 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) {
     for (len, len_str) in BENCH_SIZES {
         group.throughput(Throughput::Bytes((len * size_of::<f64>()) as u64));
 
-        // Generate f64 values that ALP-encode without patches.
-        let floats: Vec<f64> = (0..*len)
-            .map(|i| <f64 as ALPFloat>::decode_single(10 + (i as i64 % 64), exponents))
-            .collect();
-        let float_prim = PrimitiveArray::new(Buffer::from(floats), NonNullable);
-
-        // Encode: ALP → FoR → BitPacked
-        let alp =
-            alp_encode(float_prim.as_view(), Some(exponents), &mut ctx).vortex_expect("alp_encode");
-        assert!(alp.patches().is_none());
-        let for_arr = FoRData::encode(
-            alp.encoded()
-                .clone()
-                .execute::<PrimitiveArray>(&mut ctx)
-                .vortex_expect("to primitive"),
-            &mut ctx,
-        )
-        .vortex_expect("for encode");
-        let bp = BitPackedData::encode(for_arr.encoded(), bit_width, &mut ctx)
-            .vortex_expect("bitpack encode");
-
-        let tree = ALP::new(
-            FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone())
-                .vortex_expect("for_new")
-                .into_array(),
-            exponents,
-            None,
-        );
-        let array = tree.into_array();
+        for (patch_interval, benchmark_name) in [
+            (None, "cuda/alp_for_bp_6bw_f64/dispatch_f64"),
+            (
+                Some(100),
+                "cuda/alp_for_bp_6bw_f64_1pct_patches/dispatch_f64",
+            ),
+        ] {
+            let floats: Vec<f64> = (0..*len)
+                .map(|i| {
+                    if patch_interval.is_some_and(|interval| i % interval == 0) {
+                        PI
+                    } else {
+                        <f64 as ALPFloat>::decode_single(10 + (i as i64 % 64), exponents)
+                    }
+                })
+                .collect();
+            let float_prim = PrimitiveArray::new(Buffer::from(floats), NonNullable);
+
+            // Encode: ALP → FoR → BitPacked, preserving ALP's exception patches.
+            let alp = alp_encode(float_prim.as_view(), Some(exponents), &mut ctx)
+                .vortex_expect("alp_encode");
+            assert_eq!(alp.patches().is_some(), patch_interval.is_some());
+            let (encoded, alp_exponents, patches) = alp.into_parts();
+            let for_arr = FoRData::encode(
+                encoded
+                    .execute::<PrimitiveArray>(&mut ctx)
+                    .vortex_expect("to primitive"),
+                &mut ctx,
+            )
+            .vortex_expect("for encode");
+            let bp = BitPackedData::encode(for_arr.encoded(), bit_width, &mut ctx)
+                .vortex_expect("bitpack encode");
+            assert!(bp.patches().is_none(), "expected only ALP patches");
+
+            let tree = ALP::new(
+                FoR::try_new(bp.into_array(), for_arr.reference_scalar().clone())
+                    .vortex_expect("for_new")
+                    .into_array(),
+                alp_exponents,
+                patches,
+            );
+            let array = tree.into_array();
 
-        group.bench_with_input(
-            BenchmarkId::new("cuda/alp_for_bp_6bw_f64/dispatch_f64", len_str),
-            len,
-            |b, &n| {
+            group.bench_with_input(BenchmarkId::new(benchmark_name, len_str), len, |b, &n| {
                 let mut cuda_ctx =
                     CudaSession::create_execution_ctx(&cuda_session()).vortex_expect("ctx");
 
@@ -791,8 +803,8 @@ fn bench_alp_for_bitpacked_f64(c: &mut Criterion) {
                     }
                     total_time
                 });
-            },
-        );
+            });
+        }
     }
 
     group.finish();

vortex-cuda/kernels/src/dynamic_dispatch.cu

@@ -117,12 +117,8 @@ __shared__ uint64_t runend_cursors[BLOCK_SIZE];
 // ═══════════════════════════════════════════════════════════════════════════
 
 /// Apply one scalar operation to N values in registers.
-///
-/// `abs_pos` is the absolute output position of the first value to process.
-/// It is used by scalar operations that apply patches, e.g. ALP.
 template <typename T, uint32_t N>
-__device__ inline void
-scalar_op(T *values, const struct ScalarOp &op, char *__restrict smem, uint64_t abs_pos = 0) {
+__device__ inline void scalar_op(T *values, const struct ScalarOp &op, char *__restrict smem) {
     switch (op.op_code) {
     case ScalarOp::FOR: {
         const T ref = static_cast<T>(op.params.frame_of_ref.reference);
@@ -165,30 +161,9 @@ scalar_op(T *values, const struct ScalarOp &op, char *__restrict smem, uint64_t
                 values[i] = static_cast<T>(__double_as_longlong(r));
             }
         }
-        // Apply ALP patches: override positions whose float value couldn't
-        // be reconstructed through the ALP encode/decode cycle.
-        // Per-value cursor — with a slice offset, a tile's N values can
-        // straddle two FL chunks, so each value needs its own lookup.
-        if (op.params.alp.patches_ptr != 0) {
-            const auto &patches = *reinterpret_cast<const GPUPatches *>(op.params.alp.patches_ptr);
-            const uint32_t chunk_start = patches.offset / FL_CHUNK;
-#pragma unroll
-            for (uint32_t i = 0; i < N; ++i) {
-                uint64_t my_pos = (N > 1) ? abs_pos + i * blockDim.x + threadIdx.x : abs_pos;
-                uint64_t orig = my_pos + patches.offset;
-                uint32_t chunk = static_cast<uint32_t>(orig / FL_CHUNK) - chunk_start;
-                uint32_t within = static_cast<uint32_t>(orig % FL_CHUNK);
-                PatchesCursor<T> cursor(patches, chunk, 0, 1);
-                auto patch = cursor.next();
-                while (patch.index != FL_CHUNK) {
-                    if (patch.index == within) {
-                        values[i] = patch.value;
-                        break;
-                    }
-                    patch = cursor.next();
-                }
-            }
-        }
+        // ALP patches are scattered cooperatively after the stage has decoded.
+        // Looking up patches here would restart a chunk cursor for every value,
+        // turning sparse exception handling into O(values * patches).
         break;
     }
     case ScalarOp::DICT: {
@@ -225,6 +200,58 @@ scatter_patches_chunk(const GPUPatches &patches, T *__restrict out, uint32_t chu
     }
 }
 
+/// Scatter patches overlapping a logical output range.
+///
+/// `logical_start` and `range_len` are relative to the sliced array represented by
+/// `patches`. Patch indices are stored in the coordinate space of the original array,
+/// so add `patches.offset` while locating chunks and subtract it when addressing `out`.
+/// Every thread cooperates on each overlapping FastLanes chunk.
+template <typename T>
+__device__ inline void scatter_patches_range(const GPUPatches &patches,
+                                             T *__restrict out,
+                                             uint64_t logical_start,
+                                             uint32_t range_len) {
+    if (range_len == 0) {
+        return;
+    }
+
+    const uint64_t original_start = logical_start + patches.offset;
+    const uint64_t original_end = original_start + range_len;
+    const uint32_t patch_chunk_start = patches.offset / FL_CHUNK;
+    const uint32_t first_chunk = static_cast<uint32_t>(original_start / FL_CHUNK);
+    const uint32_t last_chunk = static_cast<uint32_t>((original_end - 1) / FL_CHUNK);
+
+    for (uint32_t original_chunk = first_chunk; original_chunk <= last_chunk; ++original_chunk) {
+        const uint32_t chunk = original_chunk - patch_chunk_start;
+        PatchesCursor<T> cursor(patches, chunk, threadIdx.x, blockDim.x);
+        auto patch = cursor.next();
+        while (patch.index != FL_CHUNK) {
+            const uint64_t original_pos = static_cast<uint64_t>(original_chunk) * FL_CHUNK + patch.index;
+            if (original_pos >= original_start && original_pos < original_end) {
+                out[original_pos - patches.offset] = patch.value;
+            }
+            patch = cursor.next();
+        }
+    }
+}
+
+/// Apply patch payloads attached to scalar operations after their stage has decoded.
+template <typename T>
+__device__ inline void
+scatter_scalar_patches(const Stage &stage, T *__restrict out, uint64_t logical_start, uint32_t range_len) {
+    for (uint8_t op_idx = 0; op_idx < stage.num_scalar_ops; ++op_idx) {
+        const auto &op = stage.scalar_ops[op_idx];
+        if (op.op_code == ScalarOp::ALP && op.params.alp.patches_ptr != 0) {
+            const auto &patches = *reinterpret_cast<const GPUPatches *>(op.params.alp.patches_ptr);
+            // All ordinary writes must complete before exception values overwrite them.
+            __syncthreads();
+            scatter_patches_range<T>(patches, out, logical_start, range_len);
+            // A later stage may consume this patched shared-memory output.
+            __syncthreads();
+        }
+    }
+}
+
 // ═══════════════════════════════════════════════════════════════════════════
 // Source ops
 // ═══════════════════════════════════════════════════════════════════════════
@@ -491,7 +518,7 @@ __device__ void execute_output_stage(T *__restrict output,
                                           smem);
 
             for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) {
-                scalar_op<T, VALUES_PER_TILE>(values, stage.scalar_ops[op], smem, tile_start);
+                scalar_op<T, VALUES_PER_TILE>(values, stage.scalar_ops[op], smem);
             }
 
 #pragma unroll
@@ -513,7 +540,7 @@ __device__ void execute_output_stage(T *__restrict output,
             source_op<T, 1>(&val, src, raw_input, ptype, smem_src, i, gpos, smem);
 
             for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) {
-                scalar_op<T, 1>(&val, stage.scalar_ops[op], smem, gpos);
+                scalar_op<T, 1>(&val, stage.scalar_ops[op], smem);
             }
             __stcs(&output[gpos], val);
         }
@@ -525,6 +552,8 @@ __device__ void execute_output_stage(T *__restrict output,
         }
         elem_idx += chunk_len;
     }
+
+    scatter_scalar_patches<T>(stage, output, block_start, block_len);
 }
 
 // ═══════════════════════════════════════════════════════════════════════════
@@ -559,7 +588,7 @@ __device__ void execute_input_stage(const Stage &stage, char *__restrict smem) {
             for (uint32_t elem_idx = threadIdx.x; elem_idx < stage.len; elem_idx += blockDim.x) {
                 T val = smem_out[elem_idx];
                 for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) {
-                    scalar_op<T, 1>(&val, stage.scalar_ops[op], smem, elem_idx);
+                    scalar_op<T, 1>(&val, stage.scalar_ops[op], smem);
                 }
                 smem_out[elem_idx] = val;
             }
@@ -583,14 +612,16 @@ __device__ void execute_input_stage(const Stage &stage, char *__restrict smem) {
             T val;
             source_op<T, 1>(&val, src, raw_input, stage.source_ptype, nullptr, 0, elem_idx, smem);
             for (uint8_t op = 0; op < stage.num_scalar_ops; ++op) {
-                scalar_op<T, 1>(&val, stage.scalar_ops[op], smem, elem_idx);
+                scalar_op<T, 1>(&val, stage.scalar_ops[op], smem);
             }
             smem_out[elem_idx] = val;
         }
         // Write barrier: smem region is fully populated for subsequent
         // stages to read.
         __syncthreads();
     }
+
+    scatter_scalar_patches<T>(stage, smem_out, 0, stage.len);
 }
 
 // ═══════════════════════════════════════════════════════════════════════════