feat: dyn dispatch plan

Author: 0ax1

vortex-cuda/benches/dynamic_dispatch_cuda.rs

@@ -26,10 +26,9 @@ use vortex_cuda::CudaExecutionCtx;
 use vortex_cuda::CudaSession;
 use vortex_cuda::bitpacked_cuda_kernel;
 use vortex_cuda::bitpacked_cuda_launch_config;
-use vortex_cuda::dynamic_dispatch_op::DynamicOp;
-use vortex_cuda::dynamic_dispatch_op::DynamicOpCode_ALP;
-use vortex_cuda::dynamic_dispatch_op::DynamicOpCode_BITUNPACK;
-use vortex_cuda::dynamic_dispatch_op::DynamicOpCode_FOR;
+use vortex_cuda::dynamic_dispatch::DynamicDispatchPlan;
+use vortex_cuda::dynamic_dispatch::ScalarOp;
+use vortex_cuda::dynamic_dispatch::SourceOp;
 use vortex_cuda_macros::cuda_available;
 use vortex_cuda_macros::cuda_not_available;
 use vortex_dtype::PType;
@@ -54,10 +53,6 @@ const ALP_E: f32 = 1.0;
 // Helpers
 // ---------------------------------------------------------------------------
 
-fn pack_alp_f32_param(f: f32, e: f32) -> u64 {
-    (e.to_bits() as u64) << 32 | f.to_bits() as u64
-}
-
 /// Helper: launch a single FoR kernel on a device buffer (in-place).
 fn launch_for_kernel(
     cuda_ctx: &mut CudaExecutionCtx,
@@ -107,12 +102,11 @@ fn run_dynamic_dispatch_timed(
     input_ptr: u64,
     output_ptr: u64,
     array_len: usize,
-    device_ops: &Arc<cudarc::driver::CudaSlice<DynamicOp>>,
-    num_ops: u8,
+    device_plan: &Arc<cudarc::driver::CudaSlice<DynamicDispatchPlan>>,
 ) -> VortexResult<Duration> {
     let cuda_function = cuda_ctx.load_function("dynamic_dispatch", &["u32"])?;
     let array_len_u64 = array_len as u64;
-    let ops_ptr = device_ops.device_ptr(cuda_ctx.stream()).0;
+    let plan_ptr = device_plan.device_ptr(cuda_ctx.stream()).0;
 
     let stream = cuda_ctx.stream();
     let ctx = stream.context();
@@ -127,8 +121,7 @@ fn run_dynamic_dispatch_timed(
     launch_builder.arg(&input_ptr);
     launch_builder.arg(&output_ptr);
     launch_builder.arg(&array_len_u64);
-    launch_builder.arg(&ops_ptr);
-    launch_builder.arg(&num_ops);
+    launch_builder.arg(&plan_ptr);
 
     let num_blocks = array_len.div_ceil(2048) as u32;
     let config = LaunchConfig {
@@ -275,15 +268,14 @@ fn bench_bitunpack_for_separate(c: &mut Criterion) {
 }
 
 // ============================================================================
-// Benchmark: BitUnpack + FoR — single fused dynamic scalar_decode launch
+// Benchmark: BitUnpack + FoR — single fused dynamic dispatch launch
 // ============================================================================
 
 /// Run a fused dynamic_dispatch launch on a bitpacked array, returning GPU time.
 fn run_dynamic_dispatch_bitpacked_timed(
     cuda_ctx: &mut CudaExecutionCtx,
     bitpacked_array: &BitPackedArray,
-    device_ops: &Arc<cudarc::driver::CudaSlice<DynamicOp>>,
-    num_ops: u8,
+    device_plan: &Arc<cudarc::driver::CudaSlice<DynamicDispatchPlan>>,
 ) -> VortexResult<Duration> {
     let packed = bitpacked_array.packed().clone();
     let len = bitpacked_array.len();
@@ -314,24 +306,17 @@ fn run_dynamic_dispatch_bitpacked_timed(
         .synchronize()
         .map_err(|e| vortex_err!("failed to synchronize stream: {:?}", e))?;
 
-    run_dynamic_dispatch_timed(cuda_ctx, input_ptr, output_ptr, len, device_ops, num_ops)
+    run_dynamic_dispatch_timed(cuda_ctx, input_ptr, output_ptr, len, device_plan)
 }
 
 fn bench_bitunpack_for_dynamic_dispatch(c: &mut Criterion) {
     let mut group = c.benchmark_group("bitunpack_for");
     group.sample_size(10);
 
-    // ops = [BITUNPACK(bit_width=BIT_WIDTH), FOR(REFERENCE_VALUE)]
-    let ops = vec![
-        DynamicOp {
-            op: DynamicOpCode_BITUNPACK,
-            param: BIT_WIDTH as u64,
-        },
-        DynamicOp {
-            op: DynamicOpCode_FOR,
-            param: REFERENCE_VALUE as u64,
-        },
-    ];
+    let plan = DynamicDispatchPlan::new(
+        SourceOp::bitunpack(BIT_WIDTH),
+        &[ScalarOp::frame_of_ref(REFERENCE_VALUE as u64)],
+    );
 
     for (len, len_str) in BENCH_ARGS {
         group.throughput(Throughput::Bytes((len * size_of::<u32>()) as u64));
@@ -350,11 +335,11 @@ fn bench_bitunpack_for_dynamic_dispatch(c: &mut Criterion) {
                     .load_function("dynamic_dispatch", &["u32"])
                     .vortex_expect("failed to preload dynamic_dispatch kernel");
 
-                let device_ops = Arc::new(
+                let device_plan = Arc::new(
                     cuda_ctx
                         .stream()
-                        .clone_htod(ops.as_slice())
-                        .expect("failed to copy ops to device"),
+                        .clone_htod(std::slice::from_ref(&plan))
+                        .expect("failed to copy plan to device"),
                 );
 
                 b.iter_custom(|iters| {
@@ -364,8 +349,7 @@ fn bench_bitunpack_for_dynamic_dispatch(c: &mut Criterion) {
                         let kernel_time = run_dynamic_dispatch_bitpacked_timed(
                             &mut cuda_ctx,
                             array,
-                            &device_ops,
-                            ops.len() as u8,
+                            &device_plan,
                         )
                         .vortex_expect("bitunpack+for dynamic_dispatch failed");
                         total_time += kernel_time;
@@ -388,21 +372,13 @@ fn bench_bitunpack_for_alp_dynamic_dispatch(c: &mut Criterion) {
     let mut group = c.benchmark_group("bitunpack_for_alp");
     group.sample_size(10);
 
-    // ops = [BITUNPACK(bit_width), FOR(reference), ALP(f, e)]
-    let ops = vec![
-        DynamicOp {
-            op: DynamicOpCode_BITUNPACK,
-            param: BIT_WIDTH as u64,
-        },
-        DynamicOp {
-            op: DynamicOpCode_FOR,
-            param: REFERENCE_VALUE as u64,
-        },
-        DynamicOp {
-            op: DynamicOpCode_ALP,
-            param: pack_alp_f32_param(ALP_F, ALP_E),
-        },
-    ];
+    let plan = DynamicDispatchPlan::new(
+        SourceOp::bitunpack(BIT_WIDTH),
+        &[
+            ScalarOp::frame_of_ref(REFERENCE_VALUE as u64),
+            ScalarOp::alp(ALP_F, ALP_E),
+        ],
+    );
 
     for (len, len_str) in BENCH_ARGS {
         group.throughput(Throughput::Bytes((len * size_of::<u32>()) as u64));
@@ -421,11 +397,11 @@ fn bench_bitunpack_for_alp_dynamic_dispatch(c: &mut Criterion) {
                     .load_function("dynamic_dispatch", &["u32"])
                     .vortex_expect("failed to preload dynamic_dispatch kernel");
 
-                let device_ops = Arc::new(
+                let device_plan = Arc::new(
                     cuda_ctx
                         .stream()
-                        .clone_htod(ops.as_slice())
-                        .expect("failed to copy ops to device"),
+                        .clone_htod(std::slice::from_ref(&plan))
+                        .expect("failed to copy plan to device"),
                 );
 
                 b.iter_custom(|iters| {
@@ -435,8 +411,7 @@ fn bench_bitunpack_for_alp_dynamic_dispatch(c: &mut Criterion) {
                         let kernel_time = run_dynamic_dispatch_bitpacked_timed(
                             &mut cuda_ctx,
                             array,
-                            &device_ops,
-                            ops.len() as u8,
+                            &device_plan,
                         )
                         .vortex_expect("bitunpack+for+alp dynamic_dispatch failed");
                         total_time += kernel_time;

vortex-cuda/build.rs

@@ -182,24 +182,32 @@ fn nvcc_compile_ptx(
 
 /// Generate bindings for the dynamic dispatch shared header.
 ///
-/// `DynamicOp` and `DynamicOpCode` are shared between CUDA kernels
+/// `DynamicDispatchPlan` and related types are shared between CUDA kernels
 /// and Rust host code.
 fn generate_dynamic_dispatch_bindings(kernels_src: &Path, out_dir: &Path) {
     let header = kernels_src.join("dynamic_dispatch.h");
     println!("cargo:rerun-if-changed={}", header.display());
 
     let bindings = bindgen::Builder::default()
         .header(header.to_string_lossy())
-        .allowlist_type("DynamicOp")
-        .allowlist_type("DynamicOpCode")
+        .allowlist_type("DynamicDispatchPlan")
+        .allowlist_type("SourceOp")
+        .allowlist_type("SourceOpCode")
+        .allowlist_type("SourceParams")
+        .allowlist_type("BitunpackParams")
+        .allowlist_type("ScalarOp")
+        .allowlist_type("ScalarOpCode")
+        .allowlist_type("ScalarParams")
+        .allowlist_type("FoRParams")
+        .allowlist_type("AlpParams")
         .derive_copy(true)
         .derive_debug(true)
         .generate()
         .expect("Failed to generate dynamic_dispatch bindings");
 
     bindings
-        .write_to_file(out_dir.join("dynamic_dispatch_op.rs"))
-        .expect("Failed to write dynamic_dispatch_op.rs");
+        .write_to_file(out_dir.join("dynamic_dispatch.rs"))
+        .expect("Failed to write dynamic_dispatch.rs");
 }
 
 /// Check if CUDA is available based on nvcc.

vortex-cuda/kernels/src/dynamic_dispatch.cu

@@ -2,8 +2,8 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 // Dynamic dispatch kernel: decodes an array by applying a sequence of operations
-// in a single kernel launch. The first op may optionally be a "source" op, e.g. bitunpack.
-// Subsequent transform ops are applied element-wise in registers.
+// in a single kernel launch. The source op fills shared memory (e.g. bitunpack),
+// then scalar ops are applied element-wise in registers (e.g. FoR, zigzag, ALP).
 
 #include <assert.h>
 #include <cuda.h>
@@ -17,26 +17,19 @@
 #include "dynamic_dispatch.h"
 #include "types.cuh"
 
-constexpr uint8_t MAX_DECODE_OPS = 8;
 constexpr uint32_t FL_CHUNK_SIZE = 1024;
 
-__device__ __forceinline__ bool is_source_op(enum DynamicOpCode op) {
-    return op == BITUNPACK;
-}
-
 template <typename T>
-__device__ __forceinline__ T apply_scalar_op(T value, const DynamicOp &op) {
-    switch (op.op) {
-    case FOR: {
-        return value + static_cast<T>(op.param);
+__device__ __forceinline__ T apply_scalar_op(T value, const struct ScalarOp &op) {
+    switch (op.op_code) {
+    case ScalarOp::FOR: {
+        return value + static_cast<T>(op.params.frame_of_ref.reference);
     }
-    case ZIGZAG: {
+    case ScalarOp::ZIGZAG: {
         return (value >> 1) ^ static_cast<T>(-(value & 1));
     }
-    case ALP: {
-        float f_val = __uint_as_float(static_cast<uint32_t>(op.param));
-        float e_val = __uint_as_float(static_cast<uint32_t>(op.param >> 32));
-        float result = static_cast<float>(static_cast<int32_t>(value)) * f_val * e_val;
+    case ScalarOp::ALP: {
+        float result = static_cast<float>(static_cast<int32_t>(value)) * op.params.alp.f * op.params.alp.e;
         return static_cast<T>(__float_as_uint(result));
     }
     default: __builtin_unreachable();
@@ -67,13 +60,13 @@ BITUNPACK_LANE(64, uint64_t, int64_t)
 template <typename T>
 __device__ __forceinline__ void source_fill_op(const T *__restrict input, T *__restrict smem,
                                                uint64_t chunk_start, uint32_t chunk_len,
-                                               const DynamicOp &source_op) {
+                                               const struct SourceOp &source_op) {
     constexpr uint32_t T_BITS = sizeof(T) * 8;
     constexpr uint32_t FL_LANES = FL_CHUNK_SIZE / T_BITS;
 
-    switch (source_op.op) {
-    case BITUNPACK: {
-        const uint32_t bit_width = static_cast<uint32_t>(source_op.param);
+    switch (source_op.op_code) {
+    case SourceOp::BITUNPACK: {
+        const uint32_t bit_width = source_op.params.bitunpack.bit_width;
         const uint32_t packed_words_per_chunk = FL_LANES * bit_width;
         const uint64_t chunk_idx = chunk_start / FL_CHUNK_SIZE;
         const T *packed_chunk = input + chunk_idx * packed_words_per_chunk;
@@ -82,44 +75,45 @@ __device__ __forceinline__ void source_fill_op(const T *__restrict input, T *__r
         }
         break;
     }
-    default:
-        for (uint32_t elem_idx = threadIdx.x; elem_idx < chunk_len; elem_idx += blockDim.x) {
-            smem[elem_idx] = input[chunk_start + elem_idx];
-        }
-        break;
+    default: __builtin_unreachable();
     }
 }
 
 template <typename T>
 __device__ void dynamic_dispatch_impl(const T *__restrict input, T *__restrict output, uint64_t array_len,
-                                      const DynamicOp *__restrict ops, uint8_t num_ops) {
-    assert(num_ops <= MAX_DECODE_OPS);
-
+                                      const struct DynamicDispatchPlan *__restrict plan) {
     constexpr uint32_t ELEMENTS_PER_BLOCK = 2048;
     constexpr uint32_t VALUES_PER_LOOP = 32 / sizeof(T);
 
-    __shared__ DynamicOp smem_ops[MAX_DECODE_OPS];
+    __shared__ struct SourceOp smem_source;
+    __shared__ uint8_t smem_num_scalar_ops;
+    __shared__ struct ScalarOp smem_scalar_ops[MAX_SCALAR_OPS];
     __shared__ T smem_values[FL_CHUNK_SIZE];
 
-    // Cache ops in shared memory.
-    if (threadIdx.x < num_ops) {
-        smem_ops[threadIdx.x] = ops[threadIdx.x];
+    // Cache the plan in shared memory.
+    if (threadIdx.x == 0) {
+        smem_source = plan->source;
+        smem_num_scalar_ops = plan->num_scalar_ops;
+    }
+    if (threadIdx.x < plan->num_scalar_ops) {
+        smem_scalar_ops[threadIdx.x] = plan->scalar_ops[threadIdx.x];
     }
     __syncthreads();
 
+    assert(smem_num_scalar_ops <= MAX_SCALAR_OPS);
+
     const uint64_t block_start = static_cast<uint64_t>(blockIdx.x) * ELEMENTS_PER_BLOCK;
     const uint64_t block_end = min(block_start + ELEMENTS_PER_BLOCK, array_len);
 
     for (uint64_t chunk_start = block_start; chunk_start < block_end; chunk_start += FL_CHUNK_SIZE) {
         const uint32_t chunk_len =
             static_cast<uint32_t>(min(static_cast<uint64_t>(FL_CHUNK_SIZE), block_end - chunk_start));
 
-        source_fill_op<T>(input, smem_values, chunk_start, chunk_len, smem_ops[0]);
+        source_fill_op<T>(input, smem_values, chunk_start, chunk_len, smem_source);
         __syncthreads();
 
         const uint32_t tile_size = blockDim.x * VALUES_PER_LOOP;
         const uint32_t num_full_tiles = chunk_len / tile_size;
-        const uint8_t scalar_op_start_idx = is_source_op(smem_ops[0].op);
 
         for (uint32_t tile = 0; tile < num_full_tiles; ++tile) {
             const uint32_t tile_base = tile * tile_size;
@@ -134,12 +128,12 @@ __device__ void dynamic_dispatch_impl(const T *__restrict input, T *__restrict o
                 values[idx] = smem_values[tile_base + idx * blockDim.x + threadIdx.x];
             }
 
-            for (uint8_t op_idx = scalar_op_start_idx; op_idx < num_ops; ++op_idx) {
-                const DynamicOp &decode_op = smem_ops[op_idx];
+            for (uint8_t op_idx = 0; op_idx < smem_num_scalar_ops; ++op_idx) {
+                const struct ScalarOp &scalar_op = smem_scalar_ops[op_idx];
 
                 #pragma unroll
                 for (uint32_t idx = 0; idx < VALUES_PER_LOOP; ++idx) {
-                    values[idx] = apply_scalar_op(values[idx], decode_op);
+                    values[idx] = apply_scalar_op(values[idx], scalar_op);
                 }
             }
 
@@ -153,8 +147,8 @@ __device__ void dynamic_dispatch_impl(const T *__restrict input, T *__restrict o
         const uint32_t rem_start = num_full_tiles * tile_size;
         for (uint32_t elem_idx = rem_start + threadIdx.x; elem_idx < chunk_len; elem_idx += blockDim.x) {
             T val = smem_values[elem_idx];
-            for (uint8_t op_idx = scalar_op_start_idx; op_idx < num_ops; ++op_idx) {
-                val = apply_scalar_op(val, smem_ops[op_idx]);
+            for (uint8_t op_idx = 0; op_idx < smem_num_scalar_ops; ++op_idx) {
+                val = apply_scalar_op(val, smem_scalar_ops[op_idx]);
             }
             output[chunk_start + elem_idx] = val;
         }
@@ -166,8 +160,8 @@ __device__ void dynamic_dispatch_impl(const T *__restrict input, T *__restrict o
 #define GENERATE_DYNAMIC_DISPATCH_KERNEL(suffix, Type)                                                       \
     extern "C" __global__ void dynamic_dispatch_##suffix(const Type *__restrict input,                       \
                                                          Type *__restrict output, uint64_t array_len,        \
-                                                         const DynamicOp *__restrict ops, uint8_t num_ops) {  \
-        dynamic_dispatch_impl<Type>(input, output, array_len, ops, num_ops);                                 \
+                                                         const struct DynamicDispatchPlan *__restrict plan) { \
+        dynamic_dispatch_impl<Type>(input, output, array_len, plan);                                         \
     }
 
 FOR_EACH_INTEGER(GENERATE_DYNAMIC_DISPATCH_KERNEL)