feat[cuda]: SequenceArray (#6232)

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

Cargo.lock

@@ -59,13 +59,22 @@ pub struct SequenceMetadata {
     multiplier: Option<vortex_proto::scalar::ScalarValue>,
 }
 
+/// Components of [`SequenceArray`].
+pub struct SequenceArrayParts {
+    pub base: PValue,
+    pub multiplier: PValue,
+    pub len: usize,
+    pub ptype: PType,
+    pub nullability: Nullability,
+}
+
 #[derive(Clone, Debug)]
 /// An array representing the equation `A[i] = base + i * multiplier`.
 pub struct SequenceArray {
     base: PValue,
     multiplier: PValue,
     dtype: DType,
-    pub(crate) length: usize,
+    pub(crate) len: usize,
     stats_set: ArrayStats,
 }
 
@@ -124,7 +133,7 @@ impl SequenceArray {
             base,
             multiplier,
             dtype,
-            length,
+            len: length,
             // TODO(joe): add stats, on construct or on use?
             stats_set: Default::default(),
         }
@@ -164,7 +173,7 @@ impl SequenceArray {
     }
 
     pub(crate) fn index_value(&self, idx: usize) -> PValue {
-        assert!(idx < self.length, "index_value({idx}): index out of bounds");
+        assert!(idx < self.len, "index_value({idx}): index out of bounds");
 
         match_each_native_ptype!(self.ptype(), |P| {
             let base = self.base.cast::<P>();
@@ -177,9 +186,19 @@ impl SequenceArray {
 
     /// Returns the validated final value of a sequence array
     pub fn last(&self) -> PValue {
-        Self::try_last(self.base, self.multiplier, self.ptype(), self.length)
+        Self::try_last(self.base, self.multiplier, self.ptype(), self.len)
             .vortex_expect("validated array")
     }
+
+    pub fn into_parts(self) -> SequenceArrayParts {
+        SequenceArrayParts {
+            base: self.base,
+            multiplier: self.multiplier,
+            len: self.len,
+            ptype: self.dtype.as_ptype(),
+            nullability: self.dtype.nullability(),
+        }
+    }
 }
 
 impl VTable for SequenceVTable {
@@ -355,7 +374,7 @@ fn execute_iter<P: NativePType, I: Iterator<Item = usize>>(
 
 impl BaseArrayVTable<SequenceVTable> for SequenceVTable {
     fn len(array: &SequenceArray) -> usize {
-        array.length
+        array.len
     }
 
     fn dtype(array: &SequenceArray) -> &DType {
@@ -374,14 +393,14 @@ impl BaseArrayVTable<SequenceVTable> for SequenceVTable {
         array.base.hash(state);
         array.multiplier.hash(state);
         array.dtype.hash(state);
-        array.length.hash(state);
+        array.len.hash(state);
     }
 
     fn array_eq(array: &SequenceArray, other: &SequenceArray, _precision: Precision) -> bool {
         array.base == other.base
             && array.multiplier == other.multiplier
             && array.dtype == other.dtype
-            && array.length == other.length
+            && array.len == other.len
     }
 }
 

encodings/sequence/src/array.rs

@@ -89,7 +89,7 @@ impl ExecuteParentKernel<SequenceVTable> for SequenceCompareKernel {
             // Constant is null - result is all null for comparisons
             let nullability = array.dtype().nullability() | constant.dtype().nullability();
             let result_array =
-                ConstantArray::new(Scalar::null(DType::Bool(nullability)), array.length).to_array();
+                ConstantArray::new(Scalar::null(DType::Bool(nullability)), array.len).to_array();
             return Ok(Some(result_array.execute(ctx)?));
         };
 
@@ -125,22 +125,22 @@ fn compare_eq_neq(
 
     // Check if there exists an integer solution to const = base + idx * multiplier
     let Some(set_idx) =
-        find_intersection_scalar(array.base(), array.multiplier(), array.length, constant)
+        find_intersection_scalar(array.base(), array.multiplier(), array.len, constant)
     else {
         let result_array = ConstantArray::new(
             Scalar::new(DType::Bool(nullability), not_match_val.into()),
-            array.length,
+            array.len,
         )
         .to_array();
         return Ok(Some(result_array.execute(ctx)?));
     };
     let idx = set_idx as u64;
-    let len = array.length as u64;
+    let len = array.len as u64;
 
     if len == 1 && set_idx == 0 {
         let result_array = ConstantArray::new(
             Scalar::new(DType::Bool(nullability), match_val.into()),
-            array.length,
+            array.len,
         )
         .to_array();
         return Ok(Some(result_array.execute(ctx)?));
@@ -179,31 +179,31 @@ fn compare_ordering(
     let transition = find_transition_point(
         array.base(),
         array.multiplier(),
-        array.length,
+        array.len,
         constant,
         operator,
     );
 
     let result_array = match transition {
         Transition::AllTrue => ConstantArray::new(
             Scalar::new(DType::Bool(nullability), true.into()),
-            array.length,
+            array.len,
         )
         .to_array(),
         Transition::AllFalse => ConstantArray::new(
             Scalar::new(DType::Bool(nullability), false.into()),
-            array.length,
+            array.len,
         )
         .to_array(),
         Transition::FalseToTrue(idx) => {
             // [0..idx) is false, [idx..len) is true
-            let ends = buffer![idx as u64, array.length as u64].into_array();
+            let ends = buffer![idx as u64, array.len as u64].into_array();
             let values = BoolArray::new(bitbuffer![false, true], nullability.into()).into_array();
             RunEndArray::try_new(ends, values)?.into_array()
         }
         Transition::TrueToFalse(idx) => {
             // [0..idx) is true, [idx..len) is false
-            let ends = buffer![idx as u64, array.length as u64].into_array();
+            let ends = buffer![idx as u64, array.len as u64].into_array();
             let values = BoolArray::new(bitbuffer![true, false], nullability.into()).into_array();
             RunEndArray::try_new(ends, values)?.into_array()
         }

encodings/sequence/src/kernel.rs

@@ -9,6 +9,7 @@ mod kernel;
 /// Represents the equation A\[i\] = a * i + b.
 /// This can be used for compression, fast comparisons and also for row ids.
 pub use array::SequenceArray;
+pub use array::SequenceArrayParts;
 /// Represents the equation A\[i\] = a * i + b.
 /// This can be used for compression, fast comparisons and also for row ids.
 pub use array::SequenceVTable;

encodings/sequence/src/lib.rs

@@ -23,7 +23,7 @@ _test-harness = []
 [dependencies]
 arc-swap = { workspace = true }
 async-trait = { workspace = true }
-cudarc = { workspace = true }
+cudarc = { workspace = true, features = ["f16"] }
 fastlanes = { workspace = true }
 futures = { workspace = true, features = ["executor"] }
 kanal = { workspace = true }
@@ -41,6 +41,7 @@ vortex-fastlanes = { workspace = true }
 vortex-io = { workspace = true }
 vortex-mask = { workspace = true }
 vortex-nvcomp = { path = "nvcomp" }
+vortex-sequence = { workspace = true }
 vortex-session = { workspace = true }
 vortex-utils = { workspace = true }
 vortex-zigzag = { workspace = true }

vortex-cuda/Cargo.toml

@@ -15,7 +15,7 @@ __device__ void patches(
 ) {
     const uint64_t worker = blockIdx.x * blockDim.x + threadIdx.x;
     const uint64_t startElem = START_ELEM(worker, patchesLen);
-    const uint64_t stopElem = START_ELEM(worker, patchesLen);
+    const uint64_t stopElem = STOP_ELEM(worker, patchesLen);
 
     if (startElem >= patchesLen) {
         return;

vortex-cuda/kernels/src/patches.cu

@@ -0,0 +1,42 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#include "config.cuh"
+
+template<typename ValueT>
+__device__ void sequence(
+    ValueT *const output,
+    ValueT base,
+    ValueT multiplier,
+    uint64_t len
+) {
+    const uint64_t worker = blockIdx.x * blockDim.x + threadIdx.x;
+
+    const uint64_t startElem = START_ELEM(worker, len);
+    const uint64_t stopElem = STOP_ELEM(worker, len);
+
+    for (uint64_t idx = startElem; idx < stopElem; idx++) {
+        output[idx] = static_cast<ValueT>(idx) * multiplier + base;
+    }
+}
+
+#define GENERATE_KERNEL(ValueT, suffix) \
+extern "C" __global__ void sequence_##suffix( \
+    ValueT *const output, \
+    ValueT base, \
+    ValueT multiplier, \
+    uint64_t len \
+) { \
+    sequence(output, base, multiplier, len); \
+}
+
+GENERATE_KERNEL(uint8_t, u8);
+GENERATE_KERNEL(uint16_t, u16);
+GENERATE_KERNEL(uint32_t, u32);
+GENERATE_KERNEL(uint64_t, u64);
+GENERATE_KERNEL(int8_t, i8);
+GENERATE_KERNEL(int16_t, i16);
+GENERATE_KERNEL(int32_t, i32);
+GENERATE_KERNEL(int64_t, i64);
+GENERATE_KERNEL(float, f32);
+GENERATE_KERNEL(double, f64);

vortex-cuda/kernels/src/sequence.cu

@@ -5,13 +5,15 @@ mod alp;
 mod bitpacked;
 mod decimal_byte_parts;
 mod for_;
+mod sequence;
 mod zigzag;
 mod zstd;
 
 pub use alp::ALPExecutor;
 pub use bitpacked::BitPackedExecutor;
 pub use decimal_byte_parts::DecimalBytePartsExecutor;
 pub use for_::FoRExecutor;
+pub use sequence::SequenceExecutor;
 pub use zigzag::ZigZagExecutor;
 pub use zstd::ZstdExecutor;
 pub use zstd::ZstdKernelPrep;