Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions vortex-cuda/kernels/src/dynamic_dispatch.cu
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,17 @@ __device__ inline void dynamic_source_op(const T *__restrict input,
return;
}

case SourceOp::SEQUENCE: {
// Generate a linear sequence: value[i] = base + i * multiplier.
// Used for SequenceArray (e.g. monotonic run-end endpoints).
const T base = static_cast<T>(source_op.params.sequence.base);
const T mul = static_cast<T>(source_op.params.sequence.multiplier);
for (uint32_t i = threadIdx.x; i < chunk_len; i += blockDim.x) {
smem_output[i] = base + static_cast<T>(chunk_start + i) * mul;
}
break;
}

default:
__builtin_unreachable();
}
Expand Down
8 changes: 7 additions & 1 deletion vortex-cuda/kernels/src/dynamic_dispatch.h
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,16 @@ union SourceParams {
uint64_t num_runs;
uint64_t offset;
} runend;

/// Generate a linear sequence: `value[i] = base + i * multiplier`.
struct SequenceParams {
int64_t base;
int64_t multiplier;
} sequence;
};

struct SourceOp {
enum SourceOpCode { BITUNPACK, LOAD, RUNEND } op_code;
enum SourceOpCode { BITUNPACK, LOAD, RUNEND, SEQUENCE } op_code;
union SourceParams params;
};

Expand Down
67 changes: 67 additions & 0 deletions vortex-cuda/src/dynamic_dispatch/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,17 @@ impl SourceOp {
},
}
}

/// Generate a linear sequence: `value[i] = base + i * multiplier`.
/// Used for SequenceArray (e.g. monotonic run-end endpoints).
pub fn sequence(base: i64, multiplier: i64) -> Self {
Self {
op_code: SourceOp_SourceOpCode_SEQUENCE,
params: SourceParams {
sequence: SourceParams_SequenceParams { base, multiplier },
},
}
}
}

impl ScalarOp {
Expand Down Expand Up @@ -1004,4 +1015,60 @@ mod tests {

Ok(())
}

#[rstest]
#[case(0u32, 1u32, 100)]
#[case(5u32, 3u32, 2048)]
#[case(0u32, 1u32, 4096)]
#[case(100u32, 7u32, 5000)]
#[crate::test]
fn test_sequence_unsigned(
#[case] base: u32,
#[case] multiplier: u32,
#[case] len: usize,
) -> VortexResult<()> {
use vortex::dtype::Nullability;
use vortex::encodings::sequence::SequenceArray;

let expected: Vec<u32> = (0..len).map(|i| base + (i as u32) * multiplier).collect();

let seq = SequenceArray::try_new_typed(base, multiplier, Nullability::NonNullable, len)?;

let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?;
let (plan, _bufs) = build_plan(&seq.into_array(), &cuda_ctx)?;

let actual = run_dynamic_dispatch_plan(&cuda_ctx, expected.len(), &plan)?;
assert_eq!(actual, expected);

Ok(())
}

#[rstest]
#[case(0i32, 1i32, 100)]
#[case(-10i32, 3i32, 2048)]
#[case(100i32, -1i32, 100)]
#[case(-500i32, -7i32, 50)]
#[case(0i32, 1i32, 5000)]
#[crate::test]
fn test_sequence_signed(
#[case] base: i32,
#[case] multiplier: i32,
#[case] len: usize,
) -> VortexResult<()> {
use vortex::dtype::Nullability;
use vortex::encodings::sequence::SequenceArray;

let expected: Vec<i32> = (0..len).map(|i| base + (i as i32) * multiplier).collect();

let seq = SequenceArray::try_new_typed(base, multiplier, Nullability::NonNullable, len)?;

let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?;
let (plan, _bufs) = build_plan(&seq.into_array(), &cuda_ctx)?;

let actual_u32 = run_dynamic_dispatch_plan(&cuda_ctx, expected.len(), &plan)?;
let actual: Vec<i32> = actual_u32.into_iter().map(|v| v as i32).collect();
assert_eq!(actual, expected);

Ok(())
}
}
24 changes: 24 additions & 0 deletions vortex-cuda/src/dynamic_dispatch/plan_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ use vortex::encodings::fastlanes::FoR;
use vortex::encodings::fastlanes::FoRArray;
use vortex::encodings::runend::RunEnd;
use vortex::encodings::runend::RunEndArrayParts;
use vortex::encodings::sequence::Sequence;
use vortex::encodings::sequence::SequenceArrayParts;
use vortex::encodings::zigzag::ZigZag;
use vortex::error::VortexResult;
use vortex::error::vortex_bail;
Expand Down Expand Up @@ -82,6 +84,7 @@ struct Pipeline {
/// - `ALPArray` → recurse + `ALP` scalar op (f32 only, no patches)
/// - `DictArray` → input stage for values + recurse codes + `DICT` scalar op
/// - `RunEndArray` → input stages for ends/values + `RUNEND` source
/// - `SequenceArray` → `SEQUENCE` source (integer ptypes only)
/// - `SliceArray` → resolve via child's slice reduce/kernel
///
/// # Limitations
Expand Down Expand Up @@ -158,6 +161,8 @@ impl PlanBuilderState<'_> {
self.walk_primitive(array)
} else if id == Slice::ID {
self.walk_slice(array)
} else if id == Sequence::ID {
self.walk_sequence(array)
} else {
vortex_bail!(
"Encoding {:?} not supported by dynamic dispatch plan builder",
Expand Down Expand Up @@ -305,6 +310,25 @@ impl PlanBuilderState<'_> {
Ok(pipeline)
}

/// SequenceArray → SEQUENCE source op
///
/// Generates `value[i] = base + i * multiplier` on the GPU.
fn walk_sequence(&mut self, array: ArrayRef) -> VortexResult<Pipeline> {
let seq = array
.try_into::<Sequence>()
.map_err(|_| vortex_err!("Expected SequenceArray"))?;
let SequenceArrayParts {
base, multiplier, ..
} = seq.into_parts();

Ok(Pipeline {
source: SourceOp::sequence(base.cast()?, multiplier.cast()?),
scalar_ops: vec![],
// SEQUENCE does not have an input pointer.
input_ptr: 0,
})
}

/// RunEndArray → add input stages for ends and values, RUNEND source op.
fn walk_runend(&mut self, array: ArrayRef) -> VortexResult<Pipeline> {
let re = array
Expand Down
Loading