diff --git a/Cargo.lock b/Cargo.lock index d9f80fc40b7..2bd30332122 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9770,6 +9770,7 @@ dependencies = [ "vortex-error", "vortex-fsst", "vortex-nvcomp", + "vortex-onpair", ] [[package]] diff --git a/encodings/onpair/src/array.rs b/encodings/onpair/src/array.rs index 168155cc471..1ed2a64f04d 100644 --- a/encodings/onpair/src/array.rs +++ b/encodings/onpair/src/array.rs @@ -238,7 +238,7 @@ fn build_dictionary( /// dictionary is memoized in [`OnPairData`]. Once cached, subsequent calls — /// including on arrays derived by slice / filter / cast, which share the cell — /// pay neither cost again. -pub(crate) fn dict_view<'a>( +pub fn dict_view<'a>( array: ArrayView<'a, OnPair>, ctx: &mut ExecutionCtx, ) -> VortexResult> { diff --git a/encodings/onpair/src/decode.rs b/encodings/onpair/src/decode.rs index 96add100e36..7aac8c3963f 100644 --- a/encodings/onpair/src/decode.rs +++ b/encodings/onpair/src/decode.rs @@ -32,10 +32,6 @@ pub(crate) fn collect_widened( .into_buffer::()) } -/// Read one `codes_offsets` boundary by point lookup. This decodes at most a -/// single chunk of the child — never the whole per-row offsets array — so the -/// callers that only need a row window (`scalar_at`, the canonical decode's -/// start/end bounds) don't pay to materialise every boundary. pub(crate) fn code_boundary_at( codes_offsets: &ArrayRef, index: usize, diff --git a/encodings/onpair/src/lib.rs b/encodings/onpair/src/lib.rs index 8b91154fe39..f9d8be5ff0d 100644 --- a/encodings/onpair/src/lib.rs +++ b/encodings/onpair/src/lib.rs @@ -22,9 +22,12 @@ mod tests; pub use array::*; pub use compress::*; +pub use onpair::CompactDictionaryView; pub use onpair::Config; pub use onpair::DEFAULT_CONFIG; +pub use onpair::DictionaryView; pub use onpair::Error as OnPairError; +pub use onpair::MAX_TOKEN_SIZE; pub use onpair::MaxDictBits; pub use onpair::Threshold; use vortex_array::session::ArraySessionExt; diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index a0c57478ce1..0364a0c455f 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -43,6 +43,7 @@ vortex-cub = { path = "cub" } vortex-cuda-macros = { workspace = true } vortex-error = { workspace = true, features = ["object_store"] } vortex-nvcomp = { path = "nvcomp" } +vortex-onpair = { workspace = true } [target.'cfg(target_os = "linux")'.dependencies] rustix = { workspace = true } @@ -108,6 +109,10 @@ harness = false name = "fsst_cuda" harness = false +[[bench]] +name = "onpair_cuda" +harness = false + [[bench]] name = "list_view_cuda" harness = false diff --git a/vortex-cuda/benches/arrow_binary_cuda.rs b/vortex-cuda/benches/arrow_binary_cuda.rs index 6668abd3866..91d625abf6f 100644 --- a/vortex-cuda/benches/arrow_binary_cuda.rs +++ b/vortex-cuda/benches/arrow_binary_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::cast_possible_truncation)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; diff --git a/vortex-cuda/benches/bench_config/mod.rs b/vortex-cuda/benches/bench_config/mod.rs index 75905fee19d..909018a434f 100644 --- a/vortex-cuda/benches/bench_config/mod.rs +++ b/vortex-cuda/benches/bench_config/mod.rs @@ -9,6 +9,10 @@ use criterion::Criterion; /// /// 100M elements keeps every kernel above ~500 µs, well above the /// ~15 µs CUDA driver noise floor that caused 15-45% swings at 10M. +// Each bench binary includes this module textually, and the string-heavy +// benches (fsst, onpair, arrow_binary, list_view) define their own smaller +// sizes instead of this const — allow it to go unused in those binaries. +#[allow(dead_code)] pub const BENCH_SIZES: &[(usize, &str)] = &[(100_000_000, "100M")]; /// Returns a [`Criterion`] configuration tuned for CUDA benchmarks. diff --git a/vortex-cuda/benches/fsst_cuda.rs b/vortex-cuda/benches/fsst_cuda.rs index f15e5ddc242..c5e660979bb 100644 --- a/vortex-cuda/benches/fsst_cuda.rs +++ b/vortex-cuda/benches/fsst_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::unwrap_used)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; diff --git a/vortex-cuda/benches/list_view_cuda.rs b/vortex-cuda/benches/list_view_cuda.rs index 4ce365210d6..9fb84f7ff49 100644 --- a/vortex-cuda/benches/list_view_cuda.rs +++ b/vortex-cuda/benches/list_view_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::cast_possible_truncation)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; diff --git a/vortex-cuda/benches/onpair_cuda.rs b/vortex-cuda/benches/onpair_cuda.rs new file mode 100644 index 00000000000..0aa716f0c0a --- /dev/null +++ b/vortex-cuda/benches/onpair_cuda.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA benchmarks for OnPair decompression. + +#![expect(clippy::unwrap_used)] + +mod bench_config; +mod timed_launch_strategy; + +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use criterion::BenchmarkId; +use criterion::Criterion; +use criterion::Throughput; +use futures::executor::block_on; +use vortex::array::ArrayRef; +use vortex::array::IntoArray; +use vortex::array::arrays::VarBinArray; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +use vortex::error::VortexExpect; +use vortex_cuda::CudaDispatchMode; +use vortex_cuda::CudaSession; +use vortex_cuda::executor::CudaArrayExt; +use vortex_cuda_macros::cuda_available; +use vortex_cuda_macros::cuda_not_available; +use vortex_onpair::DEFAULT_CONFIG; +use vortex_onpair::onpair_compress; + +use crate::timed_launch_strategy::TimedLaunchStrategy; + +// Bench-local size instead of the workspace 100M default: each input is a +// URL-shaped string, much heavier per-element than the fixed-width primitives +// other kernels benchmark. +const BENCH_SIZES: &[(usize, &str)] = &[(10_000_000, "10M")]; + +struct OnPairBenchFixture { + array: ArrayRef, + uncompressed_size: u64, +} + +fn make_fixture(n: usize) -> OnPairBenchFixture { + let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec = (0..n) + .map(|i| format!("https://www.example.com/path/{i}/segment?q={}", i % 97)) + .collect(); + let uncompressed_size = strings.iter().map(|s| s.len() as u64).sum(); + let varbin = VarBinArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let array = onpair_compress(&varbin, DEFAULT_CONFIG, setup_ctx.execution_ctx()) + .vortex_expect("OnPair compression failed"); + + OnPairBenchFixture { + array, + uncompressed_size, + } +} + +fn benchmark_onpair_cuda_decompress(c: &mut Criterion) { + let mut group = c.benchmark_group("cuda"); + + for &(n, len_str) in BENCH_SIZES { + let fixture = make_fixture(n); + + group.throughput(Throughput::Bytes(fixture.uncompressed_size)); + group.bench_with_input( + BenchmarkId::new("cuda/onpair/decompress_to_varbinview", len_str), + &fixture.array, + |b, onpair_array| { + b.iter_custom(|iters| { + let timed = TimedLaunchStrategy::default(); + let timer = timed.timer(); + + let mut cuda_ctx = + CudaSession::create_execution_ctx(&vortex_cuda::cuda_session()) + .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) + .with_launch_strategy(Arc::new(timed)); + + for _ in 0..iters { + block_on(onpair_array.clone().execute_cuda(&mut cuda_ctx)).unwrap(); + } + Duration::from_nanos(timer.load(Ordering::Relaxed)) + }); + }, + ); + } + + group.finish(); +} + +criterion::criterion_group! { + name = benches; + config = bench_config::cuda_bench_config(); + targets = benchmark_onpair_cuda_decompress +} + +#[cuda_available] +criterion::criterion_main!(benches); + +#[cuda_not_available] +fn main() {} diff --git a/vortex-cuda/cub/build.rs b/vortex-cuda/cub/build.rs index 5888d9234dd..d051de0b6b9 100644 --- a/vortex-cuda/cub/build.rs +++ b/vortex-cuda/cub/build.rs @@ -102,6 +102,7 @@ fn generate_rust_bindings(kernels_dir: &Path, out_dir: &Path) { .allowlist_function("filter_bytemask_.*") .allowlist_function("filter_bitmask_.*") .allowlist_function("scan_exclusive_sum_.*") + .allowlist_function("onpair_.*") // Allow CUDA types .allowlist_type("cudaError_t") // Blocklist cudaStream_t and define it manually as an opaque pointer diff --git a/vortex-cuda/cub/kernels/filter.h b/vortex-cuda/cub/kernels/filter.h index c49dc62faed..0fae1b639a8 100644 --- a/vortex-cuda/cub/kernels/filter.h +++ b/vortex-cuda/cub/kernels/filter.h @@ -99,6 +99,26 @@ cudaError_t scan_exclusive_sum_i64(void *d_temp, int64_t num_items, cudaStream_t stream); +// Fused OnPair per-batch offsets regeneration: a single sweep kernel reduces +// each 128-token batch's decoded size and exclusive-scans the sizes in-kernel +// via decoupled look-back, writing `num_batches + 1` offsets whose last +// element is the total decoded byte count. `code_width` selects the code +// stream's element size in bytes (1 or 2). A code outside the dictionary +// raises `*status` to 1 and contributes zero bytes. +cudaError_t onpair_batch_offsets_temp_size(size_t *temp_bytes, int64_t num_batches); + +cudaError_t onpair_batch_offsets(void *d_temp, + size_t temp_bytes, + const void *codes, + uint32_t code_width, + const uint8_t *lens, + uint32_t dict_size, + uint64_t total_tokens, + uint64_t *chunk_offsets, + uint32_t *status, + int64_t num_batches, + cudaStream_t stream); + #ifdef __cplusplus } #endif diff --git a/vortex-cuda/cub/kernels/onpair.cu b/vortex-cuda/cub/kernels/onpair.cu new file mode 100644 index 00000000000..65336ec0886 --- /dev/null +++ b/vortex-cuda/cub/kernels/onpair.cu @@ -0,0 +1,218 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +// OnPair per-batch output-offsets regeneration, fused into a single sweep +// kernel: every warp reduces one 128-token batch's decoded size from the codes +// and the length LUT, the first warp scans the block's batch sums, and CUB's +// decoupled look-back (`ScanTileState` + `TilePrefixCallbackOp` — the same +// machinery `cub::DeviceScan` is built on) resolves the running prefix of all +// preceding tiles in-kernel. The per-batch sizes live only in registers and +// shared memory; only the offsets are written to global memory. +// +// `onpair_batch_offsets` is the sole entry point. The tile-state init launch +// it performs first is CUB's own requirement — `DeviceScan` runs the identical +// init internally before its sweep. + +#include +#include +#include +#include + +namespace { + +// One warp reduces one 128-token batch (4 tokens per lane); a 512-thread block +// covers a tile of 16 batches. The decode kernel shares this geometry. +constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; +constexpr uint32_t ONPAIR_BLOCK_THREADS = 512; +constexpr uint32_t ONPAIR_WARPS_PER_BLOCK = ONPAIR_BLOCK_THREADS / 32; + +using OnPairTileState = cub::ScanTileState; + +// Version-stable u64 addition functor (avoids deprecated cub thread operators). +struct SumU64 { + __host__ __device__ inline uint64_t operator()(uint64_t a, uint64_t b) const { + return a + b; + } +}; + +using OnPairPrefixOp = cub::TilePrefixCallbackOp; + +__global__ void onpair_batch_offsets_init(OnPairTileState tile_state, int num_tiles) { + tile_state.InitializeStatus(num_tiles); +} + +// The fused sweep, instantiated for the two code widths OnPair stores (u16 +// natively, u8 when the compressor narrowed the codes). A code outside the +// dictionary raises `status` to 1 and contributes zero bytes: the host must +// check the flag before trusting the offsets and before launching the decode +// kernel, whose dictionary gathers are unchecked. +template +__global__ +__launch_bounds__(ONPAIR_BLOCK_THREADS) void onpair_batch_offsets_sweep(const CodeT *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, + uint64_t total_tokens, + uint64_t *__restrict chunk_offsets, + uint32_t *__restrict status, + OnPairTileState tile_state, + int64_t num_batches) { + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp = threadIdx.x >> 5; + const int tile_idx = static_cast(blockIdx.x); + const int64_t batch = (int64_t)tile_idx * ONPAIR_WARPS_PER_BLOCK + (int64_t)warp; + + // Warp-parallel reduction of this batch's decoded size. Code reads are + // lane-consecutive (coalesced); the length LUT is small and + // cache-resident. Batches past the end (last tile) contribute zero. + uint32_t batch_bytes = 0; + if (batch < num_batches) { + const uint64_t base = (uint64_t)batch * (uint64_t)ONPAIR_TOKENS_PER_BATCH; +#pragma unroll + for (uint8_t token = 0; token < 4; ++token) { + const uint64_t i = base + (uint64_t)lane + (uint64_t)(token * 32); + if (i < total_tokens) { + const uint32_t code = (uint32_t)codes[i]; + if (code < dict_size) { + batch_bytes += (uint32_t)lens[code]; + } else { + atomicMax(status, 1u); + } + } + } + } +#pragma unroll + for (uint8_t offset = 16; offset > 0; offset >>= 1) { + batch_bytes += __shfl_down_sync(0xffffffffu, batch_bytes, offset); + } + + __shared__ uint64_t warp_sums[ONPAIR_WARPS_PER_BLOCK]; + __shared__ uint64_t warp_excl[ONPAIR_WARPS_PER_BLOCK]; + __shared__ typename OnPairPrefixOp::TempStorage prefix_storage; + if (lane == 0) { + warp_sums[warp] = (uint64_t)batch_bytes; + } + __syncthreads(); + + // The first warp scans the tile's batch sums and resolves the running + // prefix of all preceding tiles via decoupled look-back. + if (warp == 0) { + const uint64_t lane_sum = (lane < ONPAIR_WARPS_PER_BLOCK) ? warp_sums[lane] : 0; + uint64_t inclusive = lane_sum; +#pragma unroll + for (uint8_t offset = 1; offset < 32; offset <<= 1) { + const uint64_t shifted = __shfl_up_sync(0xffffffffu, inclusive, offset); + if (lane >= offset) { + inclusive += shifted; + } + } + const uint64_t aggregate = __shfl_sync(0xffffffffu, inclusive, 31); + + uint64_t prefix = 0; + if (tile_idx == 0) { + if (lane == 0) { + tile_state.SetInclusive(0, aggregate); + } + } else { + // Collective over the first warp, as BlockScan would invoke it. + // The callback's return value is only defined in lane 0 (its + // look-back window reduction is a WarpReduce); broadcast it, the + // same way BlockScan shares the prefix before applying it. + OnPairPrefixOp prefix_op(tile_state, prefix_storage, SumU64(), tile_idx); + const uint64_t lane0_prefix = prefix_op(aggregate); + prefix = __shfl_sync(0xffffffffu, lane0_prefix, 0); + } + if (lane < ONPAIR_WARPS_PER_BLOCK) { + warp_excl[lane] = prefix + inclusive - lane_sum; + } + } + __syncthreads(); + + if (lane == 0 && batch < num_batches) { + chunk_offsets[batch] = warp_excl[warp]; + // The trailing slot holds the total decoded byte count. + if (batch == num_batches - 1) { + chunk_offsets[num_batches] = warp_excl[warp] + warp_sums[warp]; + } + } +} + +int onpair_num_tiles(int64_t num_batches) { + return static_cast((num_batches + ONPAIR_WARPS_PER_BLOCK - 1) / ONPAIR_WARPS_PER_BLOCK); +} + +} // namespace + +// Query the look-back tile-state storage for `num_batches` batches. +extern "C" cudaError_t onpair_batch_offsets_temp_size(size_t *temp_bytes, int64_t num_batches) { + if (num_batches < 0 || num_batches / ONPAIR_WARPS_PER_BLOCK >= INT_MAX) { + return cudaErrorInvalidValue; + } + return OnPairTileState::AllocationSize(onpair_num_tiles(num_batches), *temp_bytes); +} + +// Regenerate the OnPair decode kernel's per-batch output offsets on `stream` +// in one fused sweep, writing `chunk_offsets[0..num_batches]` where +// `chunk_offsets[b]` is the decoded byte count preceding batch `b` and +// `chunk_offsets[num_batches]` is the total. `code_width` selects the code +// stream's element size in bytes (1 or 2), so narrowed codes never need a +// widening pass. A code outside the dictionary raises `*status` to 1 and +// contributes zero bytes; the caller must check the flag before trusting the +// offsets. +extern "C" cudaError_t onpair_batch_offsets(void *d_temp, + size_t temp_bytes, + const void *codes, + uint32_t code_width, + const uint8_t *lens, + uint32_t dict_size, + uint64_t total_tokens, + uint64_t *chunk_offsets, + uint32_t *status, + int64_t num_batches, + cudaStream_t stream) { + if (num_batches <= 0 || num_batches / ONPAIR_WARPS_PER_BLOCK >= INT_MAX) { + return cudaErrorInvalidValue; + } + const int num_tiles = onpair_num_tiles(num_batches); + + OnPairTileState tile_state; + cudaError_t err = tile_state.Init(num_tiles, d_temp, temp_bytes); + if (err != cudaSuccess) { + return err; + } + + constexpr int INIT_THREADS = 128; + const int init_blocks = (num_tiles + INIT_THREADS - 1) / INIT_THREADS; + onpair_batch_offsets_init<<>>(tile_state, num_tiles); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return err; + } + + switch (code_width) { + case 1: + onpair_batch_offsets_sweep + <<>>(static_cast(codes), + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + tile_state, + num_batches); + break; + case 2: + onpair_batch_offsets_sweep + <<>>(static_cast(codes), + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + tile_state, + num_batches); + break; + default: + return cudaErrorInvalidValue; + } + return cudaGetLastError(); +} diff --git a/vortex-cuda/cub/src/lib.rs b/vortex-cuda/cub/src/lib.rs index f60a6075d29..7221f902e11 100644 --- a/vortex-cuda/cub/src/lib.rs +++ b/vortex-cuda/cub/src/lib.rs @@ -25,6 +25,7 @@ pub mod sys; mod error; pub mod filter; +pub mod onpair; pub mod scan; pub use error::CubError; diff --git a/vortex-cuda/cub/src/onpair.rs b/vortex-cuda/cub/src/onpair.rs new file mode 100644 index 00000000000..614b69c42b2 --- /dev/null +++ b/vortex-cuda/cub/src/onpair.rs @@ -0,0 +1,71 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Rust wrappers around the fused OnPair per-batch offsets regeneration. + +use std::ffi::c_void; + +use crate::cub_library; +use crate::error::CubError; +use crate::error::check_cuda_error; +pub use crate::sys::cudaStream_t; + +/// Get temporary storage size (the look-back tile state) for +/// [`batch_offsets`]. +pub fn batch_offsets_temp_size(num_batches: i64) -> Result { + let lib = cub_library()?; + let mut temp_bytes: usize = 0; + let err = unsafe { (lib.onpair_batch_offsets_temp_size)(&raw mut temp_bytes, num_batches) }; + check_cuda_error(err, "onpair_batch_offsets_temp_size")?; + Ok(temp_bytes) +} + +/// Regenerate the OnPair decode kernel's per-batch output offsets in one fused +/// sweep: each warp reduces one 128-token batch's decoded size and the +/// exclusive scan over the sizes runs in-kernel via decoupled look-back. +/// Writes `num_batches + 1` offsets; the last is the total decoded byte count. +/// `code_width` selects the code stream's element size in bytes (1 or 2). A +/// code outside the dictionary raises `*status` to 1 and contributes zero +/// bytes. +/// +/// # Safety +/// +/// All device pointers must be valid and properly sized: +/// - `d_temp` must have at least `temp_bytes` bytes allocated. +/// - `codes` must have at least `total_tokens` elements of `code_width` bytes +/// each, with `total_tokens <= num_batches * 128`. +/// - `lens` must have at least `dict_size` bytes. +/// - `chunk_offsets` must have at least `num_batches + 1` `u64` values. +/// - `status` must point to a valid device `u32`. +#[allow(clippy::too_many_arguments)] +pub unsafe fn batch_offsets( + d_temp: *mut c_void, + temp_bytes: usize, + codes: *const c_void, + code_width: u32, + lens: *const u8, + dict_size: u32, + total_tokens: u64, + chunk_offsets: *mut u64, + status: *mut u32, + num_batches: i64, + stream: cudaStream_t, +) -> Result<(), CubError> { + let lib = cub_library()?; + let err = unsafe { + (lib.onpair_batch_offsets)( + d_temp, + temp_bytes, + codes, + code_width, + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + num_batches, + stream, + ) + }; + check_cuda_error(err, "onpair_batch_offsets") +} diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu new file mode 100644 index 00000000000..e15cfec7391 --- /dev/null +++ b/vortex-cuda/kernels/src/onpair.cu @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "config.cuh" + +#include + +// Support kernels for OnPair GPU decompression. The per-batch output-offsets +// regeneration lives in the CUB shim (`cub/kernels/onpair.cu`), where the +// reduction and exclusive scan run as one fused sweep; this module holds the +// window-bounds and view-construction kernels launched around the decode. + +// Tokens per decode batch. Must match `ONPAIR_TOKENS_PER_BATCH` in +// `cub/kernels/onpair.cu` and `TOKENS_PER_BATCH` in the Rust launch code. +constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; + +// The token window of this array's rows, resolved entirely on device from the +// (possibly slice-narrowed, possibly device-resident) `codes_offsets` child: +// the offsets are nondecreasing, so the window's min and max are its first +// and last elements. Writes `bounds[0] = codes_offsets[0]` and +// `bounds[1] = codes_offsets[last]`; a signed negative offset sign-extends +// huge and fails the host's post-readback range validation. +#define GENERATE_TOKEN_BOUNDS_KERNEL(suffix, OffsetT) \ + extern "C" __global__ void onpair_token_bounds_##suffix(const OffsetT *__restrict codes_offsets, \ + uint64_t last, \ + uint64_t *__restrict bounds) { \ + if (threadIdx.x == 0 && blockIdx.x == 0) { \ + bounds[0] = (uint64_t)codes_offsets[0]; \ + bounds[1] = (uint64_t)codes_offsets[last]; \ + } \ + } + +GENERATE_TOKEN_BOUNDS_KERNEL(i8, int8_t) +GENERATE_TOKEN_BOUNDS_KERNEL(i16, int16_t) +GENERATE_TOKEN_BOUNDS_KERNEL(i32, int32_t) +GENERATE_TOKEN_BOUNDS_KERNEL(i64, int64_t) +GENERATE_TOKEN_BOUNDS_KERNEL(u8, uint8_t) +GENERATE_TOKEN_BOUNDS_KERNEL(u16, uint16_t) +GENERATE_TOKEN_BOUNDS_KERNEL(u32, uint32_t) +GENERATE_TOKEN_BOUNDS_KERNEL(u64, uint64_t) + +// Byte positions of the visible code window's bounds in the full decoded +// stream. A sliced array keeps its whole `codes` child, so the decode runs +// over the full stream; each boundary's byte position is the whole-batch +// prefix from `chunk_offsets` plus a warp reduction over the boundary batch's +// head `[batch_start, boundary)`. Launched after the offsets sweep and the +// token-bounds resolution with one 32-thread block per boundary; the token +// boundaries are read from `bounds[0..2)` (clamped to `total_tokens` so a +// corrupt offset cannot read out of bounds — the host rejects it after +// readback) and the byte positions are written to `bounds[2 + blockIdx.x]`. +// A code outside the dictionary contributes zero bytes (the sweep already +// raised the status flag the host checks before trusting `bounds`). `CodeT` +// is the code stream's element type (u16 natively, u8 when the compressor +// narrowed the codes). +template +__device__ inline void onpair_window_offsets_body(const CodeT *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, + const uint64_t *__restrict chunk_offsets, + uint64_t total_tokens, + uint64_t *__restrict bounds) { + const uint64_t requested = bounds[blockIdx.x]; + const uint64_t boundary = requested < total_tokens ? requested : total_tokens; + const uint64_t batch = boundary / ONPAIR_TOKENS_PER_BATCH; + const uint64_t batch_base = batch * ONPAIR_TOKENS_PER_BATCH; + const uint32_t lane = threadIdx.x & 31u; + + uint32_t partial = 0; +#pragma unroll + for (uint8_t token = 0; token < 4; ++token) { + const uint64_t i = batch_base + lane + (uint64_t)(token * 32u); + if (i < boundary) { + const uint32_t code = (uint32_t)codes[i]; + if (code < dict_size) { + partial += (uint32_t)lens[code]; + } + } + } +#pragma unroll + for (uint8_t offset = 16; offset > 0; offset >>= 1) { + partial += __shfl_down_sync(0xffffffffu, partial, offset); + } + if (lane == 0) { + bounds[2 + blockIdx.x] = chunk_offsets[batch] + (uint64_t)partial; + } +} + +extern "C" __global__ void onpair_window_offsets_u8(const uint8_t *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, + const uint64_t *__restrict chunk_offsets, + uint64_t total_tokens, + uint64_t *__restrict bounds) { + onpair_window_offsets_body(codes, lens, dict_size, chunk_offsets, total_tokens, bounds); +} + +extern "C" __global__ void onpair_window_offsets_u16(const uint16_t *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, + const uint64_t *__restrict chunk_offsets, + uint64_t total_tokens, + uint64_t *__restrict bounds) { + onpair_window_offsets_body(codes, lens, dict_size, chunk_offsets, total_tokens, bounds); +} + +// Arrow/Vortex variable-length view records are 16 bytes. Values up to 12 bytes +// are stored inline after the u32 length. Longer values store their first four +// bytes, backing-buffer index, and byte offset. +constexpr uint8_t MAX_INLINED_SIZE = 12; + +// Build one BinaryView over the flat decoded byte stream. Row `rid`'s bytes +// are `output_bytes[row_offsets[rid]..row_offsets[rid + 1])`; the offsets are +// the Arrow i32 offsets built on device from the lengths child. The Rust +// caller only launches this when the window fits a single backing buffer +// (`MAX_BUFFER_LEN`, i32::MAX), so every offset is non-negative and fits the +// view's u32 fields. +__device__ inline void onpair_write_view(const int32_t *__restrict row_offsets, + const uint8_t *__restrict output_bytes, + uint4 *__restrict views, + uint64_t rid) { + const uint32_t start = (uint32_t)row_offsets[rid]; + const uint32_t len = (uint32_t)row_offsets[rid + 1] - start; + if (len <= MAX_INLINED_SIZE) { + uint32_t words[3] = {0, 0, 0}; +#pragma unroll + for (uint8_t i = 0; i < MAX_INLINED_SIZE; i++) { + if (i < len) { + words[i >> 2] |= (uint32_t)output_bytes[start + i] << (8u * (i & 3u)); + } + } + views[rid] = make_uint4(len, words[0], words[1], words[2]); + return; + } + + const uint32_t prefix = (uint32_t)output_bytes[start] | ((uint32_t)output_bytes[start + 1] << 8u) | + ((uint32_t)output_bytes[start + 2] << 16u) | + ((uint32_t)output_bytes[start + 3] << 24u); + views[rid] = make_uint4(len, prefix, 0, start); +} + +extern "C" __global__ void onpair_build_views(const int32_t *__restrict row_offsets, + const uint8_t *__restrict output_bytes, + uint4 *__restrict views, + uint64_t num_rows) { + const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; + const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; + const uint64_t block_end = + (block_start + elements_per_block < num_rows) ? (block_start + elements_per_block) : num_rows; + for (uint64_t rid = block_start + threadIdx.x; rid < block_end; rid += blockDim.x) { + onpair_write_view(row_offsets, output_bytes, views, rid); + } +} diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu new file mode 100644 index 00000000000..4143c75c4e9 --- /dev/null +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -0,0 +1,240 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include +#include +#include +#include + +// OnPair decompress — 4 tokens/thread, split-read dictionary. +// +// One warp decodes one 128-token batch (4 tokens per lane) in four phases: +// +// 1. Load — each lane fetches its 4 codes, their first-8 dictionary bytes +// (`dict_s8`), and their lengths. +// 2. Scan — a warp prefix-scan of the lengths positions every token within +// the batch and yields the batch's total decoded size. +// 3. Stage — token bytes are gathered into a per-warp shared staging buffer; +// only the rare `len > 8` tokens touch the full 16-byte-row +// `dict_padded`. +// 4. Drain — the staged bytes stream to global output as an aligned `uint4` +// body between a byte head and tail. +// +// Baseline `onpair_shmem_4tpt` is L1/TEX-cache-request bound on the per-token +// 16-byte `uint4` gather into the 64 KB padded dict, where the dict L1 hit rate +// is only ~31% (the 64 KB dict thrashes against the streaming codes/output). +// +// Most tokens are short (mean dict len ~6). This variant reads the common case +// from the **32 KB** `dict_s8` array (first 8 bytes/entry, `uint2`) and only +// touches the 64 KB `dict_padded` for the rare `len > 8` tokens. Halving the +// hot dict working set aims to raise the dict L1 hit rate, cutting L2 sectors +// and L1/TEX-request pressure. As a bonus, holding `uint2 lo[4]` (32 B) instead +// of `uint4 t[4]` (64 B) lowers register pressure. + +#ifndef WARPS_PER_BLOCK_MAX +#define WARPS_PER_BLOCK_MAX 16u +#endif +// The launch bounds are an exception, not a pattern: this kernel was tuned on +// GH200 where NCU showed it register-capped at 64 (50% occupancy); bounding to +// 2 blocks/SM was benchmarked as the best trade against spills. Re-evaluate on +// new architectures rather than copying this to other kernels. +#ifndef ONPAIR_LAUNCH_BOUNDS +#define ONPAIR_LAUNCH_BOUNDS __launch_bounds__(512, 2) +#endif +#define WARP_BUF_BYTES 2080u + +__device__ inline uint32_t onpair_warp_inclusive_scan_u32(uint32_t x, uint32_t lane) { + constexpr unsigned mask = 0xffffffffu; +#pragma unroll + for (uint8_t offset = 1; offset < 32; offset <<= 1) { + uint32_t y = __shfl_up_sync(mask, x, offset); + if (lane >= offset) { + x += y; + } + } + return x; +} + +// One lane's slice of a 128-token batch: 4 tokens, strided a warp apart. +struct OnPairTokens { + // First 8 dictionary bytes of each token (the common-case read). + uint2 lo[4]; + // Dictionary codes, kept for the rare `len > 8` high-byte gather. + uint32_t code[4]; + // Decoded token lengths. + uint32_t len[4]; +}; + +// Phase 1 — load this lane's 4 (code, dict_s8 bytes, length) triples. Tokens +// past the end of the stream load as empty. `CodeT` is the code stream's +// element type (u16 natively, u8 when the compressor narrowed the codes). +template +__device__ inline OnPairTokens onpair_load_tokens(const CodeT *__restrict codes, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict lens, + uint64_t base_i, + uint64_t total_tokens) { + OnPairTokens t; +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint64_t i = base_i + (uint64_t)(k * 32); + if (i < total_tokens) { + const uint32_t code = (uint32_t)codes[i]; + t.code[k] = code; + t.lo[k] = *reinterpret_cast(dict_s8 + (size_t)code * 8u); + t.len[k] = (uint32_t)lens[code]; + } else { + t.code[k] = 0u; + t.lo[k] = make_uint2(0u, 0u); + t.len[k] = 0u; + } + } + return t; +} + +// Phase 2 — position every token within the batch: `excl[k]` is the exclusive +// prefix (the token's staging offset) via 4 chained warp scans of the lengths. +// Returns the batch's total decoded byte count. +__device__ inline uint32_t onpair_scan_offsets(const uint32_t (&len)[4], uint32_t lane, uint32_t (&excl)[4]) { + constexpr unsigned mask = 0xffffffffu; + uint32_t acc_base = 0u; +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint32_t incl = onpair_warp_inclusive_scan_u32(len[k], lane); + excl[k] = acc_base + (incl - len[k]); + acc_base += __shfl_sync(mask, incl, 31); + } + return acc_base; +} + +// Phase 3 — gather each token's bytes into the warp's shared staging buffer at +// its scanned offset. The common case writes only the 8 `dict_s8` bytes +// already in registers; tokens longer than 8 bytes take the rare path through +// the full padded dictionary. +__device__ inline void onpair_stage_tokens(const OnPairTokens &t, + const uint32_t (&excl)[4], + const uint8_t *__restrict dict_padded, + uint8_t *__restrict s_buf) { +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint32_t len = t.len[k]; + if (len == 0u) { + continue; + } + const uint32_t base = excl[k]; + const uint8_t *lob = reinterpret_cast(&t.lo[k]); + const uint32_t nlo = len < 8u ? len : 8u; +#pragma unroll + for (uint8_t j = 0; j < 8; ++j) { + if (j < nlo) { + s_buf[base + j] = lob[j]; + } + } + if (len > 8u) { + // Rare path: high bytes from the full padded dict. + const uint2 hi = *reinterpret_cast(dict_padded + (size_t)t.code[k] * 16u + 8u); + const uint8_t *hib = reinterpret_cast(&hi); +#pragma unroll + for (uint8_t j = 0; j < 8; ++j) { + if (8u + j < len) { + s_buf[base + 8 + j] = hib[j]; + } + } + } + } +} + +// Phase 4 — copy the staged batch to global output: a byte head up to the +// first 16-aligned output address, an aligned `uint4` body with streaming +// stores, and a byte tail. The caller offset `s_buf` so shared and global +// 16-byte alignment phases match. +__device__ inline void onpair_drain(const uint8_t *__restrict s_buf, + uint8_t *__restrict output_bytes, + uint64_t out_start, + uint32_t head_pre, + uint32_t warp_total, + uint32_t lane) { + const uint32_t head = head_pre < warp_total ? head_pre : warp_total; + if (lane < head) { + output_bytes[out_start + (uint64_t)lane] = s_buf[lane]; + } + if (head >= warp_total) { + return; + } + + const uint32_t body_chunks = (warp_total - head) >> 4; + for (uint8_t k = (uint8_t)lane; k < body_chunks; k += 32u) { + const uint32_t off = head + k * 16u; + const uint4 v = *reinterpret_cast(s_buf + off); + __stcs(reinterpret_cast(output_bytes + out_start + off), v); + } + + const uint32_t tail_start = head + (body_chunks << 4); + if (lane < warp_total - tail_start) { + output_bytes[out_start + (uint64_t)tail_start + (uint64_t)lane] = s_buf[tail_start + lane]; + } +} + +template +__device__ inline void onpair_decode_body(const CodeT *__restrict codes, + const uint64_t *__restrict chunk_offsets, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict dict_padded, + const uint8_t *__restrict lens, + uint8_t *__restrict output_bytes, + uint64_t total_tokens) { + const uint32_t lane = threadIdx.x & 31; + const uint32_t warp_id = threadIdx.x >> 5; + const uint64_t chunk = (uint64_t)blockIdx.x * (uint64_t)(blockDim.x >> 5) + (uint64_t)warp_id; + if (chunk * 128u >= total_tokens) { + return; + } + + __shared__ __align__(16) uint8_t s_buf_all[WARPS_PER_BLOCK_MAX * WARP_BUF_BYTES]; + uint8_t *s_buf_base = &s_buf_all[warp_id * WARP_BUF_BYTES]; + + const OnPairTokens t = + onpair_load_tokens(codes, dict_s8, lens, chunk * 128u + (uint64_t)lane, total_tokens); + + uint32_t excl[4]; + const uint32_t warp_total = onpair_scan_offsets(t.len, lane, excl); + + // Offset the staging buffer by (out_start % 16) so the drain's global + // 16-byte stores land aligned when copied from 16-aligned shared reads. + const uint64_t out_start = chunk_offsets[chunk]; + const uint32_t head_pre = (16u - (uint32_t)(out_start & 15u)) & 15u; + uint8_t *s_buf = s_buf_base + ((16u - head_pre) & 15u); + + onpair_stage_tokens(t, excl, dict_padded, s_buf); + __syncwarp(); + + onpair_drain(s_buf, output_bytes, out_start, head_pre, warp_total, lane); +} + +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void +onpair_shmem_4tpt_split8read_u8(const uint8_t *__restrict codes, + const uint64_t *__restrict chunk_offsets, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict dict_padded, + const uint8_t *__restrict lens, + uint8_t *__restrict output_bytes, + uint64_t total_tokens) { + onpair_decode_body(codes, chunk_offsets, dict_s8, dict_padded, lens, output_bytes, total_tokens); +} + +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void +onpair_shmem_4tpt_split8read_u16(const uint16_t *__restrict codes, + const uint64_t *__restrict chunk_offsets, + const uint8_t *__restrict dict_s8, + const uint8_t *__restrict dict_padded, + const uint8_t *__restrict lens, + uint8_t *__restrict output_bytes, + uint64_t total_tokens) { + onpair_decode_body(codes, + chunk_offsets, + dict_s8, + dict_padded, + lens, + output_bytes, + total_tokens); +} diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 48c5042716b..f6180d162e2 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -55,12 +55,12 @@ use vortex::dtype::Nullability; use vortex::dtype::PType; use vortex::dtype::i256; use vortex::encodings::fsst::FSST; -use vortex::encodings::fsst::FSSTArray; use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; +use vortex_onpair::OnPair; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -81,8 +81,9 @@ use crate::cub::exclusive_sum_i32; use crate::device_buffer::CUDF_VALIDITY_BUFFER_PADDING; use crate::executor::CudaArrayExt; use crate::executor::execute_validity_cuda; -use crate::kernel::FSSTVarBin; +use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; +use crate::kernel::decode_onpair_varbin; /// An implementation of `ExportDeviceArray` that exports Vortex arrays to `ArrowDeviceArray` by /// first decoding the array on the GPU and then converting the canonical type to the nearest @@ -229,11 +230,23 @@ fn export_array( // `CudaDispatchMode` only governs `execute_cuda`'s fused-vs-standalone planning. let array = match array.try_downcast::() { Ok(fsst) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { - return export_fsst_varbin(fsst, ctx).await; + let decoded = decode_fsst_varbin(fsst, ctx).await?; + return export_decoded_varbin(decoded, ctx).await; } Ok(fsst) => fsst.into_array(), Err(array) => array, }; + // OnPair takes the same offset-based export shortcut as FSST. + let array = match array.try_downcast::() { + Ok(onpair) + if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => + { + let decoded = decode_onpair_varbin(onpair, ctx).await?; + return export_decoded_varbin(decoded, ctx).await; + } + Ok(onpair) => onpair.into_array(), + Err(array) => array, + }; let cuda_array = array.execute_cuda(ctx).await?; export_canonical(cuda_array, ctx).await @@ -581,20 +594,22 @@ async fn export_varbin( export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) } -async fn export_fsst_varbin( - fsst: FSSTArray, +/// Export an offset-based decompression result (FSST or OnPair) with the +/// standard Arrow `Utf8`/`Binary` layout. +async fn export_decoded_varbin( + decoded: DecodedVarBin, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { - let FSSTVarBin { + let DecodedVarBin { dtype, len, offsets, values, validity, - } = decode_fsst_varbin(fsst, ctx).await?; + } = decoded; vortex_ensure!( matches!(dtype, DType::Utf8(_) | DType::Binary(_)), - "FSST produced invalid variable-length dtype {dtype}" + "offset-based decode produced invalid variable-length dtype {dtype}" ); let (validity_buffer, null_count) = export_arrow_validity_buffer(validity, len, 0, ctx).await?; export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 4b6009e7cea..0612a13fc07 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -8,13 +8,70 @@ use std::ffi::c_void; use cudarc::driver::CudaSlice; use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; +use vortex::array::buffer::BufferHandle; use vortex::error::VortexResult; use vortex::error::vortex_err; +use vortex_cub::onpair; use vortex_cub::scan; use vortex_cub::scan::cudaStream_t; +use crate::CudaBufferExt; use crate::CudaExecutionCtx; +/// Regenerate the OnPair decode kernel's per-batch output offsets in one +/// fused sweep (see `cub/kernels/onpair.cu`): the per-batch decoded-size +/// reduction and the exclusive scan over the sizes run in a single kernel via +/// decoupled look-back. `code_width` selects the code stream's element size in +/// bytes (1 or 2). Returns `num_batches + 1` offsets; the last is the total +/// decoded byte count. A code outside the dictionary raises `status` to 1; +/// the caller must check the flag before trusting the offsets. +#[allow(clippy::too_many_arguments)] +pub(crate) fn onpair_batch_offsets( + codes: &BufferHandle, + code_width: u32, + lens: &BufferHandle, + dict_size: u32, + num_tokens: usize, + num_batches: usize, + status: &mut CudaSlice, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let num_batches_i64 = i64::try_from(num_batches)?; + let temp_bytes = onpair::batch_offsets_temp_size(num_batches_i64) + .map_err(|err| vortex_err!("CUB onpair_batch_offsets_temp_size failed: {err}"))?; + + let mut temp = ctx.device_alloc::(temp_bytes.max(1))?; + let mut chunk_offsets = ctx.device_alloc::(num_batches + 1)?; + let codes_ptr = codes.cuda_device_ptr()?; + let lens_ptr = lens.cuda_device_ptr()?; + let total_tokens = u64::try_from(num_tokens)?; + let stream = ctx.stream(); + let stream_ptr = stream.cu_stream() as cudaStream_t; + let (status_ptr, record_status) = status.device_ptr_mut(stream); + let (offsets_ptr, record_offsets) = chunk_offsets.device_ptr_mut(stream); + let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); + + ctx.launch_external(num_tokens, || unsafe { + onpair::batch_offsets( + temp_ptr as *mut c_void, + temp_bytes, + codes_ptr as *const c_void, + code_width, + lens_ptr as *const u8, + dict_size, + total_tokens, + offsets_ptr as *mut u64, + status_ptr as *mut u32, + num_batches_i64, + stream_ptr, + ) + .map_err(|err| vortex_err!("CUB onpair_batch_offsets failed: {err}")) + })?; + drop((record_status, record_offsets, record_temp)); + + Ok(chunk_offsets) +} + pub(crate) fn exclusive_sum_i32( input: &CudaSlice, len: usize, @@ -47,3 +104,138 @@ pub(crate) fn exclusive_sum_i32( Ok(output) } + +#[cfg(test)] +mod tests { + use vortex::error::VortexExpect; + + use super::*; + use crate::session::CudaSession; + + /// Upload synthetic codes and lengths, regenerate the chunk offsets, and + /// read them back together with the status flag. The code width follows + /// the element type. + async fn batch_offsets_roundtrip( + codes: Vec, + lens: Vec, + ) -> VortexResult<(Vec, u32)> + where + C: cudarc::driver::DeviceRepr + + cudarc::driver::ValidAsZeroBits + + std::fmt::Debug + + Send + + Sync + + 'static, + { + let mut ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let num_tokens = codes.len(); + let num_batches = num_tokens.div_ceil(128); + let dict_size = u32::try_from(lens.len())?; + let code_width = u32::try_from(size_of::())?; + + let codes_dev = ctx.copy_to_device(codes)?.await?; + let lens_dev = ctx.copy_to_device(lens)?.await?; + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|e| vortex_err!("Failed to zero status flag: {e}"))?; + + let offsets = onpair_batch_offsets( + &codes_dev, + code_width, + &lens_dev, + dict_size, + num_tokens, + num_batches, + &mut status, + &mut ctx, + )?; + + let offsets = ctx + .stream() + .clone_dtoh(&offsets) + .map_err(|e| vortex_err!("Failed to copy offsets to host: {e}"))?; + let status = ctx + .stream() + .clone_dtoh(&status) + .map_err(|e| vortex_err!("Failed to copy status to host: {e}"))?; + Ok((offsets, status[0])) + } + + /// The exclusive prefix at 128-token boundaries, plus the trailing total. + fn host_reference(codes: &[u16], lens: &[u8]) -> Vec { + let mut expected = Vec::with_capacity(codes.len().div_ceil(128) + 1); + expected.push(0u64); + let mut acc = 0u64; + for (i, &code) in codes.iter().enumerate() { + acc += u64::from(lens[code as usize]); + if (i + 1) % 128 == 0 { + expected.push(acc); + } + } + if !codes.len().is_multiple_of(128) { + expected.push(acc); + } + expected + } + + /// A single partial batch: one tile, no look-back. + #[crate::test] + async fn test_onpair_batch_offsets_single_batch() -> VortexResult<()> { + let lens: Vec = (1..=16).collect(); + let codes: Vec = (0..100u16).map(|i| i % 16).collect(); + let expected = host_reference(&codes, &lens); + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 0); + assert_eq!(offsets, expected); + Ok(()) + } + + /// Many look-back tiles with a ragged tail batch; the offsets must match + /// a host prefix sum sampled at 128-token boundaries. Regression test for + /// the look-back prefix being defined only in lane 0. + #[crate::test] + async fn test_onpair_batch_offsets_multi_tile() -> VortexResult<()> { + let lens: Vec = (1..=16u8).cycle().take(300).collect(); + let codes: Vec = (0..2000u32 * 128 - 57) + .map(|i| u16::try_from(i * 31 % 300).vortex_expect("bounded by dictionary size")) + .collect(); + let expected = host_reference(&codes, &lens); + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 0); + assert_eq!(offsets, expected); + Ok(()) + } + + /// u8 codes dispatch the narrow sweep instantiation. + #[crate::test] + async fn test_onpair_batch_offsets_u8_codes() -> VortexResult<()> { + let lens: Vec = (1..=16).collect(); + let codes: Vec = (0..300u32) + .map(|i| u8::try_from(i * 7 % 16).vortex_expect("bounded by dictionary size")) + .collect(); + let widened: Vec = codes.iter().map(|&c| u16::from(c)).collect(); + let expected = host_reference(&widened, &lens); + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 0); + assert_eq!(offsets, expected); + Ok(()) + } + + /// A code outside the dictionary raises the status flag and contributes + /// zero bytes. + #[crate::test] + async fn test_onpair_batch_offsets_flags_out_of_range_code() -> VortexResult<()> { + let lens = vec![2u8; 4]; + let mut codes = vec![1u16; 200]; + codes[130] = 9; + + let (offsets, status) = batch_offsets_roundtrip(codes, lens).await?; + assert_eq!(status, 1); + assert_eq!(offsets, vec![0, 256, 256 + 71 * 2]); + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 1e18dc5fec9..767934d94a0 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -48,8 +48,10 @@ use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; use crate::executor::execute_validity_cuda; -/// Device-resident offset-based result of FSST decompression. -pub(crate) struct FSSTVarBin { +/// Device-resident offset-based (Arrow `Utf8`/`Binary`) decompression result: +/// i32 offsets plus a contiguous values heap. Produced by the FSST and OnPair +/// varbin decoders for the offset-based Arrow export path. +pub(crate) struct DecodedVarBin { pub(crate) dtype: DType, pub(crate) len: usize, pub(crate) offsets: BufferHandle, @@ -163,7 +165,7 @@ impl CudaExecute for FSSTExecutor { pub(crate) async fn decode_fsst_varbin( fsst: FSSTArray, ctx: &mut CudaExecutionCtx, -) -> VortexResult { +) -> VortexResult { let dtype = fsst.dtype().clone(); let validity = fsst.codes().validity()?; let len = fsst.len(); @@ -188,7 +190,7 @@ pub(crate) async fn decode_fsst_varbin( if total_size == 0 { let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); let values = BufferHandle::new_device(allocation.slice(0..0)); - return Ok(FSSTVarBin { + return Ok(DecodedVarBin { dtype, len, offsets: output_offsets, @@ -208,7 +210,7 @@ async fn decode_fsst_varbin_typed( output_offsets: BufferHandle, total_size: usize, ctx: &mut CudaExecutionCtx, -) -> VortexResult +) -> VortexResult where U: NativePType + DeviceRepr + Send + Sync + 'static, { @@ -264,7 +266,7 @@ where .arg(&len_u64); })?; - Ok(FSSTVarBin { + Ok(DecodedVarBin { dtype, len, offsets: output_offsets, diff --git a/vortex-cuda/src/kernel/encodings/mod.rs b/vortex-cuda/src/kernel/encodings/mod.rs index 6ad37e7ead1..8d33433a2db 100644 --- a/vortex-cuda/src/kernel/encodings/mod.rs +++ b/vortex-cuda/src/kernel/encodings/mod.rs @@ -7,6 +7,7 @@ mod date_time_parts; mod decimal_byte_parts; mod for_; mod fsst; +mod onpair; mod runend; mod sequence; mod zigzag; @@ -20,9 +21,11 @@ pub(crate) use bitpacked::bitpacked_slice_view; pub(crate) use date_time_parts::DateTimePartsExecutor; pub(crate) use decimal_byte_parts::DecimalBytePartsExecutor; pub(crate) use for_::FoRExecutor; +pub(crate) use fsst::DecodedVarBin; pub(crate) use fsst::FSSTExecutor; -pub(crate) use fsst::FSSTVarBin; pub(crate) use fsst::decode_fsst_varbin; +pub(crate) use onpair::OnPairExecutor; +pub(crate) use onpair::decode_onpair_varbin; pub(crate) use runend::RunEndExecutor; pub(crate) use sequence::SequenceExecutor; pub(crate) use zigzag::ZigZagExecutor; diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs new file mode 100644 index 00000000000..38cc51a3afb --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -0,0 +1,1002 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA executor for OnPair decompression. +//! +//! Decoding runs on the GPU over the full token stream — a sliced array keeps +//! its whole `codes` child, and buffers never round-trip between host and +//! device to cut a window out of it: +//! +//! 1. `onpair_batch_offsets` (in the CUB shim) regenerates the per-batch +//! output offsets (`chunk_offsets`) the decode kernel positions its writes +//! with: one fused sweep reduces every 128-token batch's decoded size and +//! exclusive-scans the sizes in-kernel via decoupled look-back. +//! 2. `onpair_token_bounds` reads this array's token window from the +//! device-resident `codes_offsets` child — the offsets are nondecreasing, +//! so the window's min and max are its first and last elements — and +//! `onpair_window_offsets` resolves the window's byte positions inside the +//! decoded stream: the whole-batch prefix from `chunk_offsets` plus a +//! partial-batch reduction over each boundary batch's head. No boundary is +//! read on host before the single gating readback. +//! 3. `onpair_shmem_4tpt_split8read` gathers each token's bytes from the +//! split dictionary layout and scatters them to the output byte stream; +//! the window is then exposed as a zero-copy device slice of the heap. +//! +//! Every kernel that reads the codes is instantiated for the two widths +//! OnPair stores (u16 natively, u8 when the compressor narrowed the codes), +//! so the code stream is decompressed on device and never widened. +//! +//! The heap size and window bounds come from the GPU; each output path builds +//! its row offsets from the lengths child with [`i32_offsets_from_lengths`] +//! on device and validates them against the window size before they index the +//! decoded bytes. The result is exposed either as a canonical `VarBinView` +//! (views built on-device by `onpair_build_views`, or on host for windows +//! that exceed a single backing buffer — the only path that materialises the +//! lengths) or as Arrow-compatible i32 offsets plus values via +//! [`decode_onpair_varbin`] — mirroring the FSST varbin path. + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DeviceRepr; +use cudarc::driver::LaunchConfig; +use cudarc::driver::PushKernelArg; +use num_traits::AsPrimitive; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; +use vortex::array::arrays::varbinview::build_views::build_views; +use vortex::array::buffer::BufferHandle; +use vortex::array::buffer::DeviceBuffer; +use vortex::array::match_each_integer_ptype; +use vortex::array::validity::Validity; +use vortex::dtype::DType; +use vortex::dtype::NativePType; +use vortex::dtype::PType; +use vortex::error::VortexExpect; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; +use vortex::error::vortex_err; +use vortex_array::ArrayView; +use vortex_onpair::DictionaryView; +use vortex_onpair::MAX_TOKEN_SIZE; +use vortex_onpair::OnPair; +use vortex_onpair::OnPairArray; +use vortex_onpair::OnPairArrayExt; +use vortex_onpair::OnPairArraySlotsExt; +use vortex_onpair::dict_view; + +use crate::CanonicalCudaExt; +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::arrow::I32Offsets; +use crate::arrow::i32_offsets_from_lengths; +use crate::cub::onpair_batch_offsets; +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::kernel::encodings::DecodedVarBin; + +// The kernels fix the dictionary row stride at 16 bytes (two `uint2` reads). +const _: () = assert!(MAX_TOKEN_SIZE == 16); + +/// Tokens per decode batch: one decode-kernel warp emits 128 tokens (4 per +/// thread). Must match `ONPAIR_TOKENS_PER_BATCH` in `cub/kernels/onpair.cu`. +const TOKENS_PER_BATCH: usize = 128; +/// Threads per block for the warp-per-batch kernels (16 warps). +const BLOCK_THREADS: u32 = 512; +const WARPS_PER_BLOCK: usize = (BLOCK_THREADS / 32) as usize; + +/// Launch config for the warp-per-batch kernels: one warp per 128-token batch. +fn batch_launch_config(num_batches: usize) -> VortexResult { + let grid_dim = u32::try_from(num_batches.div_ceil(WARPS_PER_BLOCK))?; + Ok(LaunchConfig { + grid_dim: (grid_dim, 1, 1), + block_dim: (BLOCK_THREADS, 1, 1), + shared_mem_bytes: 0, + }) +} + +/// CUDA decoder for OnPair. +#[derive(Debug)] +pub(crate) struct OnPairExecutor; + +#[async_trait] +impl CudaExecute for OnPairExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let onpair = array + .as_typed::() + .ok_or_else(|| vortex_err!("Expected OnPairArray"))?; + decode_onpair(onpair, ctx).await + } +} + +/// Checked host sum of the per-row decoded lengths. A negative length +/// sign-extends and surfaces as overflow here or as a mismatch against the +/// GPU-computed window size. +fn sum_lengths(lengths: &PrimitiveArray) -> VortexResult { + match_each_integer_ptype!(lengths.ptype(), |P| { + let mut acc = 0u64; + for &length in lengths.as_slice::

() { + acc = acc + .checked_add(AsPrimitive::::as_(length)) + .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; + } + Ok(acc) + }) +} + +/// The row offsets index the decoded heap, so the lengths must account for +/// exactly the bytes the codes decoded to. +fn ensure_lengths_match(row_total: u64, total_size: usize) -> VortexResult<()> { + vortex_ensure!( + row_total == total_size as u64, + "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" + ); + Ok(()) +} + +/// All-empty output: `num_rows` inline empty views and no backing buffers. +async fn empty_views( + num_rows: usize, + dtype: DType, + validity: Validity, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let views = ctx.copy_to_device(vec![0i128; num_rows])?.await?; + Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })) +} + +/// The device-staged compressed token stream: the full code stream at its +/// native width, the split dictionary layout, and the regenerated per-batch +/// output offsets. +struct StagedCodes { + codes: BufferHandle, + dict_s8: BufferHandle, + dict_padded: BufferHandle, + lens: BufferHandle, + dict_size: u32, + /// Exclusive per-batch output offsets, `num_batches + 1` entries; the last + /// is the total decoded byte count of the full code stream. + chunk_offsets: CudaSlice, + num_batches: usize, + num_tokens: usize, + launch_config: LaunchConfig, +} + +/// The shared result of the OnPair GPU decode pipeline. +struct OnPairDecoded { + /// This array's rows' decoded bytes: a zero-copy device slice of the full + /// decoded heap, bounded by the on-device window-offsets resolution. + bytes: BufferHandle, + /// Byte size of the window, computed on device. + total_size: usize, + /// Per-row lengths, resident wherever `execute_cuda` produced them. The + /// varbin path and the canonical fast path build their row offsets from + /// them on device; only the host rollover path materialises them. + lengths: PrimitiveArray, +} + +/// Stage this array's device-decompressed codes and dictionary and regenerate +/// the decode kernel's per-batch output offsets from them in one fused sweep +/// (see [`onpair_batch_offsets`]). The caller has validated that the codes +/// are u8 or u16; the sweep reads them at their native width. +async fn stage_codes( + onpair: ArrayView<'_, OnPair>, + codes: PrimitiveArray, + status: &mut CudaSlice, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let num_tokens = codes.len(); + let code_width = u32::try_from(codes.ptype().byte_width())?; + let PrimitiveDataParts { + buffer: codes_buffer, + .. + } = codes.into_data_parts(); + + // Stage the dictionary in the decode kernel's split layout: fixed 16-byte + // rows (`dict_padded`, the rare `len > 8` read), the first 8 bytes of every + // row (`dict_s8`, the common-case read), and the per-code lengths. + let dict = dict_view(onpair, ctx.execution_ctx())?; + let dict_size = dict.num_tokens(); + let dict_size_u32 = u32::try_from(dict_size)?; + let mut dict_padded = vec![0u8; dict_size * MAX_TOKEN_SIZE]; + let mut dict_s8 = vec![0u8; dict_size * 8]; + let mut lens = vec![0u8; dict_size]; + for code in 0..dict_size { + let token = + dict.token(u16::try_from(code).vortex_expect("dictionary has at most 2^16 tokens")); + let len = token.len(); + lens[code] = u8::try_from(len).vortex_expect("token length is at most MAX_TOKEN_SIZE"); + dict_padded[code * MAX_TOKEN_SIZE..code * MAX_TOKEN_SIZE + len].copy_from_slice(token); + let head = len.min(8); + dict_s8[code * 8..code * 8 + head].copy_from_slice(&token[..head]); + } + + let (codes_dev, s8_dev, padded_dev, lens_dev) = futures::try_join!( + ctx.ensure_on_device(codes_buffer), + ctx.copy_to_device(dict_s8)?, + ctx.copy_to_device(dict_padded)?, + ctx.copy_to_device(lens)?, + )?; + + let num_batches = num_tokens.div_ceil(TOKENS_PER_BATCH); + let launch_config = batch_launch_config(num_batches)?; + let chunk_offsets = onpair_batch_offsets( + &codes_dev, + code_width, + &lens_dev, + dict_size_u32, + num_tokens, + num_batches, + status, + ctx, + )?; + + Ok(StagedCodes { + codes: codes_dev, + dict_s8: s8_dev, + dict_padded: padded_dev, + lens: lens_dev, + dict_size: dict_size_u32, + chunk_offsets, + num_batches, + num_tokens, + launch_config, + }) +} + +/// Run the OnPair decode pipeline over the full token stream: stage the codes +/// and dictionary, regenerate the per-batch output offsets on the device, +/// validate the compressed stream, and decode the flat byte stream. A sliced +/// array keeps its whole `codes` child, so the decode runs unwindowed and this +/// array's rows are exposed as a zero-copy device slice of the decoded heap, +/// bounded by the on-device `onpair_window_offsets` resolution — the codes +/// never round-trip through the host. Returns `Ok(None)` when there is +/// nothing to decode: the array is empty, every row is null, or the code +/// window is empty. +async fn decode_onpair_bytes( + onpair: ArrayView<'_, OnPair>, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let num_rows = onpair.len(); + + if num_rows == 0 { + return Ok(None); + } + + // Every row null (cheap metadata check): nothing to decode. A sliced + // all-null window usually carries a validity child instead of the + // `AllInvalid` marker and is caught when the token window resolves empty. + if onpair.array_validity().definitely_all_null() { + return Ok(None); + } + + let lengths = onpair + .uncompressed_lengths() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + + // No codes at all (e.g. every row empty): the child's length is host + // metadata, so this early-out costs no device read. + if onpair.codes().is_empty() { + ensure_zero_lengths(lengths).await?; + return Ok(None); + } + + // Decompress the per-row code boundaries on device; the token window is + // resolved from them by a kernel, never by host scalar reads. + let codes_offsets = onpair + .codes_offsets() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + + // Decompress the codes child on device. The kernels are instantiated for + // the two widths OnPair stores — u16 natively, u8 when the compressor + // narrowed the codes — so no widening pass is needed. + let codes = onpair + .codes() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + match codes.ptype() { + PType::U8 => decode_window::(onpair, codes, codes_offsets, lengths, ctx).await, + PType::U16 => decode_window::(onpair, codes, codes_offsets, lengths, ctx).await, + other => vortex_bail!("OnPair codes must decompress to u8 or u16, got {other}"), + } +} + +/// Cold path: the window has no codes, so the rows must decode to zero bytes. +async fn ensure_zero_lengths(lengths: PrimitiveArray) -> VortexResult<()> { + let lengths = Canonical::Primitive(lengths) + .into_host() + .await? + .into_primitive(); + let total = sum_lengths(&lengths)?; + vortex_ensure!( + total == 0, + "OnPair records {total} decoded bytes but has no codes" + ); + Ok(()) +} + +/// Stage the codes at their native width `C`, resolve the window bounds on +/// device, validate the stream, and decode the full token stream. Returns +/// `Ok(None)` when the token window turns out to be empty. +async fn decode_window( + onpair: ArrayView<'_, OnPair>, + codes: PrimitiveArray, + codes_offsets: PrimitiveArray, + lengths: PrimitiveArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> +where + C: NativePType + DeviceRepr + Send + Sync + 'static, +{ + // Corruption flag raised by the batch-sizes kernel for a code outside the + // dictionary; checked before the unchecked decode kernel is allowed to run. + let mut status = ctx.device_alloc::(1)?; + ctx.stream() + .memset_zeros(&mut status) + .map_err(|e| vortex_err!("Failed to zero OnPair status flag: {e}"))?; + + let staged = stage_codes(onpair, codes, &mut status, ctx).await?; + + let ptype = C::PTYPE.to_string(); + let num_tokens_u64 = u64::try_from(staged.num_tokens)?; + let codes_view = staged.codes.cuda_view::()?; + let s8_view = staged.dict_s8.cuda_view::()?; + let padded_view = staged.dict_padded.cuda_view::()?; + let lens_view = staged.lens.cuda_view::()?; + + // The window bounds scratch: the token-bounds kernel writes the token + // window into slots 0..2 and the window-offsets kernel resolves the byte + // window into slots 2..4. + let mut bounds = ctx.device_alloc::(4)?; + + // Token bounds of this array's rows in the full code stream, read on + // device from the (possibly slice-narrowed) `codes_offsets` child: the + // offsets are nondecreasing, so the window's min and max are its first + // and last elements. + let offsets_ptype = codes_offsets.ptype(); + let last_boundary = u64::try_from(codes_offsets.len().saturating_sub(1))?; + let PrimitiveDataParts { + buffer: offsets_buffer, + .. + } = codes_offsets.into_data_parts(); + let offsets_dev = ctx.ensure_on_device(offsets_buffer).await?; + let bounds_fn = + ctx.load_function_with_suffixes("onpair", &["token_bounds", &offsets_ptype.to_string()])?; + match_each_integer_ptype!(offsets_ptype, |O| { + let offsets_view = offsets_dev.cuda_view::()?; + ctx.launch_kernel_config( + &bounds_fn, + LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }, + 1, + |args| { + args.arg(&offsets_view).arg(&last_boundary).arg(&mut bounds); + }, + )?; + }); + + // Byte positions of the code window inside the decoded stream, computed + // on device: the whole-batch prefix from `chunk_offsets` plus a + // partial-batch reduction over each boundary batch's head. + let window_fn = ctx.load_function_with_suffixes("onpair", &["window_offsets", &ptype])?; + ctx.launch_kernel_config( + &window_fn, + LaunchConfig { + grid_dim: (2, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }, + 2, + |args| { + args.arg(&codes_view) + .arg(&lens_view) + .arg(&staged.dict_size) + .arg(&staged.chunk_offsets) + .arg(&num_tokens_u64) + .arg(&mut bounds); + }, + )?; + + // One synchronizing readback gates the decode kernel — whose dictionary + // gathers and output scatters are unchecked — and yields the GPU-computed + // heap size, token window, and byte window. The lengths child is + // validated against the window size by whichever output path materialises + // row offsets. + let status = ctx + .stream() + .clone_dtoh(&status) + .map_err(|e| vortex_err!("Failed to copy OnPair status flag to host: {e}"))?; + if status.first().copied().unwrap_or(1) != 0 { + vortex_bail!("OnPair code out of dictionary range"); + } + let chunk_total = ctx + .stream() + .clone_dtoh( + &staged + .chunk_offsets + .slice(staged.num_batches..staged.num_batches + 1), + ) + .map_err(|e| vortex_err!("Failed to copy OnPair decoded size to host: {e}"))? + .first() + .copied() + .ok_or_else(|| vortex_err!("OnPair batch offset scan returned no total"))?; + let heap_size = usize::try_from(chunk_total)?; + let bounds = ctx + .stream() + .clone_dtoh(&bounds) + .map_err(|e| vortex_err!("Failed to copy OnPair window bounds to host: {e}"))?; + let [token_start, token_end, byte_start, byte_end] = bounds[..] else { + vortex_bail!("OnPair window resolution returned no bounds"); + }; + vortex_ensure!( + token_start <= token_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + token_end <= num_tokens_u64, + "OnPair codes_offsets end {token_end} exceeds codes len {num_tokens_u64}" + ); + if token_start == token_end { + // No codes in the window (e.g. a slice covering only null rows). + ensure_zero_lengths(lengths).await?; + return Ok(None); + } + let byte_start = usize::try_from(byte_start)?; + let byte_end = usize::try_from(byte_end)?; + vortex_ensure!( + byte_start <= byte_end && byte_end <= heap_size, + "OnPair window bounds [{byte_start}, {byte_end}) exceed decoded heap size {heap_size}" + ); + let total_size = byte_end - byte_start; + // A conformant dictionary has no zero-length tokens, so a non-empty code + // window decodes to at least one byte. + vortex_ensure!(total_size > 0, "OnPair has codes but decodes to zero bytes"); + + // Decode the full stream. The kernel's drain gates 16-byte stores on + // `out_start % 16` relative to the buffer base, so the base must be + // 16-aligned. + let mut bytes = ctx.device_alloc::(heap_size)?; + let (bytes_base_ptr, _) = bytes.device_ptr(ctx.stream()); + assert_eq!( + bytes_base_ptr % 16, + 0, + "output base not 16-aligned: {bytes_base_ptr:#x}", + ); + + let decode_fn = ctx.load_function_with_suffixes("onpair_shmem_4tpt_split8read", &[&ptype])?; + ctx.launch_kernel_config( + &decode_fn, + staged.launch_config, + staged.num_tokens, + |args| { + args.arg(&codes_view) + .arg(&staged.chunk_offsets) + .arg(&s8_view) + .arg(&padded_view) + .arg(&lens_view) + .arg(&mut bytes) + .arg(&num_tokens_u64); + }, + )?; + + // This array's rows as a zero-copy device slice of the decoded heap. + let heap = CudaDeviceBuffer::new(bytes); + Ok(Some(OnPairDecoded { + bytes: BufferHandle::new_device(heap.slice(byte_start..byte_end)), + total_size, + lengths, + })) +} + +async fn decode_onpair( + onpair: ArrayView<'_, OnPair>, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = onpair.dtype().clone(); + let validity = onpair.array_validity(); + let num_rows = onpair.len(); + + if onpair.is_empty() { + return Ok(Canonical::empty(&dtype)); + } + + let Some(decoded) = decode_onpair_bytes(onpair, ctx).await? else { + return empty_views(num_rows, dtype, validity, ctx).await; + }; + let OnPairDecoded { + bytes, + total_size, + lengths, + } = decoded; + + // Fast path: the decoded window fits a single BinaryView backing buffer + // (`MAX_BUFFER_LEN`, i32::MAX), so the per-row offsets fit Arrow's i32 + // range and build on device from the device-resident lengths — nothing + // touches the host. + if total_size <= MAX_BUFFER_LEN { + let I32Offsets { + buffer: row_offsets, + total, + } = i32_offsets_from_lengths(lengths, ctx).await?; + ensure_lengths_match(u64::try_from(total)?, total_size)?; + let row_offsets_view = row_offsets.cuda_view::()?; + let bytes_view = bytes.cuda_view::()?; + let mut device_views = ctx.device_alloc::(num_rows)?; + let num_rows_u64 = u64::try_from(num_rows)?; + let build_views_fn = ctx.load_function_with_suffixes("onpair", &["build_views"])?; + ctx.launch_kernel(&build_views_fn, num_rows, |args| { + args.arg(&row_offsets_view) + .arg(&bytes_view) + .arg(&mut device_views) + .arg(&num_rows_u64); + })?; + + let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([bytes]), dtype, validity) + })); + } + + // BinaryView offsets are u32. Windows that need multiple backing buffers + // roll the decoded bytes over on host, mirroring the CPU canonical path; + // only here do the lengths leave the device. The host views index the + // copied window, so validate the lengths first. + let lengths = Canonical::Primitive(lengths) + .into_host() + .await? + .into_primitive(); + ensure_lengths_match(sum_lengths(&lengths)?, total_size)?; + let host_bytes = bytes.try_to_host()?.await?; + + let (buffers, views) = match_each_integer_ptype!(lengths.ptype(), |P| { + build_views( + 0, + MAX_BUFFER_LEN, + host_bytes.into_mut(), + lengths.as_slice::

(), + ) + }); + + Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_unchecked(views, Arc::from(buffers), dtype, validity) + })) +} + +/// Decode OnPair directly into Arrow-compatible i32 offsets and contiguous +/// values on device, mirroring the FSST varbin path. +pub(crate) async fn decode_onpair_varbin( + onpair: OnPairArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = onpair.dtype().clone(); + let validity = onpair.array_validity(); + let len = onpair.len(); + + let Some(decoded) = decode_onpair_bytes(onpair.as_view(), ctx).await? else { + // Zero decoded bytes: all-zero offsets and an empty values heap. + let offsets = ctx.copy_to_device(vec![0i32; len + 1])?.await?; + let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); + let values = BufferHandle::new_device(allocation.slice(0..0)); + return Ok(DecodedVarBin { + dtype, + len, + offsets, + values, + validity, + }); + }; + + let OnPairDecoded { + bytes, + total_size, + lengths, + } = decoded; + + // Build the Arrow i32 offsets from the lengths on device; this also + // rejects windows beyond Arrow's i32 offset range. + let I32Offsets { + buffer: offsets, + total, + } = i32_offsets_from_lengths(lengths, ctx).await?; + ensure_lengths_match(u64::try_from(total)?, total_size)?; + + Ok(DecodedVarBin { + dtype, + len, + offsets, + values: bytes, + validity, + }) +} + +#[cfg(test)] +mod tests { + use arrow_schema::DataType; + use arrow_schema::Field; + use rstest::rstest; + use vortex::array::IntoArray; + use vortex::array::arrays::VarBinArray; + use vortex::array::assert_arrays_eq; + use vortex::buffer::Buffer; + use vortex::dtype::Nullability; + use vortex::error::VortexExpect; + use vortex_array::VortexSessionExecute; + use vortex_onpair::DEFAULT_CONFIG; + use vortex_onpair::onpair_compress; + + use super::*; + use crate::CanonicalCudaExt; + use crate::arrow::DeviceArrayExt; + use crate::arrow::release_device_array; + use crate::arrow::release_schema; + use crate::session::CudaSession; + use crate::session::VarBinExportLayout; + + fn cuda_ctx_with_varbin_layout(layout: VarBinExportLayout) -> VortexResult { + let session = vortex::array::array_session() + .with_some(CudaSession::try_default()?.with_varbin_export_layout(layout)); + CudaSession::create_execution_ctx(&session) + } + + fn assert_device_resident(canonical: &Canonical) { + let varbinview = canonical.as_varbinview(); + assert!(varbinview.views_handle().is_on_device()); + assert!( + varbinview + .data_buffers() + .iter() + .all(BufferHandle::is_on_device) + ); + } + + fn compress_onpair( + strings: Vec>, + dtype: DType, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let varbin = VarBinArray::from_iter(strings, dtype).into_array(); + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, ctx.execution_ctx())?; + vortex_ensure!( + onpair.as_opt::().is_some(), + "expected OnPair array, got {}", + onpair.encoding_id() + ); + Ok(onpair) + } + + #[rstest] + #[case::binary_non_null( + vec![Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex onpair test string"[..])], + DType::Binary(Nullability::NonNullable), + )] + #[case::utf8_non_null( + vec![Some(&b"the quick brown fox"[..]), + Some(&b"jumps over the lazy dog"[..]), + Some(&b"hello world"[..]), + Some(&b"vortex onpair test string"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_inline_boundary( + vec![Some(&b""[..]), + Some(&b"123456789012"[..]), + Some(&b"1234567890123"[..]), + Some(&b"this is another outlined value"[..])], + DType::Utf8(Nullability::NonNullable), + )] + #[case::utf8_partial_nulls( + vec![Some(&b"alpha"[..]), None, Some(&b"gamma"[..]), None, Some(&b"epsilon"[..])], + DType::Utf8(Nullability::Nullable), + )] + #[case::binary_all_empty( + vec![Some(&b""[..]), Some(&b""[..]), Some(&b""[..])], + DType::Binary(Nullability::NonNullable), + )] + #[crate::test] + async fn test_cuda_onpair_decompression_roundtrip( + #[case] strings: Vec>, + #[case] dtype: DType, + ) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let onpair = compress_onpair(strings, dtype.clone(), &mut cuda_ctx)?; + + let gpu_result = OnPairExecutor + .execute(onpair.clone(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed"); + assert_eq!(gpu_result.dtype(), &dtype); + assert_device_resident(&gpu_result); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(onpair, host_result, &mut ctx); + Ok(()) + } + + /// A slice keeps the whole `codes` child and narrows only `codes_offsets`, + /// so this exercises the nonzero `code_start` window. + #[crate::test] + async fn test_cuda_onpair_decompression_sliced() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let values = vec![ + Some(&b"before the window"[..]), + None, + Some(&b"the quick brown fox"[..]), + None, + Some(&b"after the window"[..]), + ]; + let onpair = compress_onpair(values, DType::Utf8(Nullability::Nullable), &mut cuda_ctx)?; + let sliced = onpair.slice(1..4)?; + + let gpu_result = OnPairExecutor + .execute(sliced.clone(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + /// A slice covering only null rows decodes zero bytes and takes the + /// empty-views path. + #[crate::test] + async fn test_cuda_onpair_decompression_null_slice() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let values = vec![Some(&b"alpha"[..]), None, None, Some(&b"omega"[..])]; + let onpair = compress_onpair(values, DType::Utf8(Nullability::Nullable), &mut cuda_ctx)?; + let sliced = onpair.slice(1..3)?; + + let gpu_result = OnPairExecutor + .execute(sliced.clone(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + /// Exercises many 128-token batches and the multi-block decode grid. + #[crate::test] + async fn test_cuda_onpair_decompression_roundtrip_large() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec = (0..100_000) + .map(|i| format!("https://www.example.com/path/{i}/segment?q={}", i % 97)) + .collect(); + let varbin = VarBinArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, cuda_ctx.execution_ctx())?; + + let gpu_result = OnPairExecutor + .execute(onpair.clone(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed"); + assert_device_resident(&gpu_result); + + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(onpair, host_result, &mut ctx); + Ok(()) + } + + /// A slice deep into a large array: both code-window boundaries land + /// mid-batch in non-zero batches, exercising the on-device window-bounds + /// resolution (whole-batch prefix plus partial-batch reduction) and the + /// zero-copy window slice of the full decoded heap. + #[crate::test] + async fn test_cuda_onpair_decompression_sliced_large() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec = (0..40_000) + .map(|i| format!("https://www.example.com/path/{i}/segment?q={}", i % 97)) + .collect(); + let varbin = VarBinArray::from_iter( + strings.iter().map(|s| Some(s.as_str())), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, cuda_ctx.execution_ctx())?; + let sliced = onpair.slice(19_997..20_101)?; + + let gpu_result = OnPairExecutor + .execute(sliced.clone(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(sliced, host_result, &mut ctx); + Ok(()) + } + + /// Codes narrowed to u8 dispatch the u8 kernel instantiations end to end. + /// A trained dictionary always holds the 256 single-byte tokens sorted + /// among its merges, so real merge codes never fit u8; the u8-addressable + /// case is the minimal alphabet-only dictionary, where token id `b` is + /// exactly the byte `b` and every row is coded byte per byte. + #[crate::test] + async fn test_cuda_onpair_decompression_u8_codes() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let strings: Vec<&[u8]> = [ + &b"tokenized token stream"[..], + b"tokenized", + b"token stream", + b"stream of tokens", + ] + .into_iter() + .cycle() + .take(800) + .collect(); + + // The alphabet-only compact dictionary: the 256 single-byte tokens + // (sorted by construction) plus the trailing read padding. + let mut dict_bytes: Vec = (0..=u8::MAX).collect(); + dict_bytes.resize(255 + MAX_TOKEN_SIZE, 0); + let dict_offsets: Vec = (0..=256).collect(); + + let codes: Vec = strings.concat(); + let mut codes_offsets = vec![0u32]; + let mut lengths = Vec::with_capacity(strings.len()); + let mut acc = 0u32; + for s in &strings { + let len = u32::try_from(s.len())?; + lengths.push(len); + acc += len; + codes_offsets.push(acc); + } + + let onpair = OnPair::try_new( + DType::Utf8(Nullability::NonNullable), + BufferHandle::new_host(Buffer::from(dict_bytes).into_byte_buffer()), + Buffer::from(dict_offsets).into_array(), + Buffer::from(codes).into_array(), + Buffer::from(codes_offsets).into_array(), + Buffer::from(lengths).into_array(), + Validity::NonNullable, + )?; + + let expected = VarBinArray::from_iter( + strings.iter().map(|s| Some(*s)), + DType::Utf8(Nullability::NonNullable), + ) + .into_array(); + + let gpu_result = OnPairExecutor + .execute(onpair.into_array(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(expected, host_result, &mut ctx); + Ok(()) + } + + #[crate::test] + async fn test_cuda_onpair_direct_varbin_output() -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(VarBinExportLayout::VarBin)?; + let values: [&[u8]; 3] = [ + b"", + b"short", + b"this value is stored directly in the values buffer", + ]; + let onpair = compress_onpair( + values.iter().map(|v| Some(*v)).collect(), + DType::Utf8(Nullability::NonNullable), + &mut cuda_ctx, + )? + .try_downcast::() + .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; + + let output = decode_onpair_varbin(onpair, &mut cuda_ctx).await?; + assert_eq!(output.dtype, DType::Utf8(Nullability::NonNullable)); + assert_eq!(output.len, values.len()); + assert!(output.offsets.is_on_device()); + assert!(output.values.is_on_device()); + + let offsets = Buffer::::from_byte_buffer(output.offsets.try_to_host()?.await?); + assert_eq!( + offsets.as_slice(), + &[0, 0, 5, i32::try_from(5 + values[2].len())?,] + ); + assert_eq!( + output.values.try_to_host()?.await?.as_ref(), + values.concat() + ); + Ok(()) + } + + #[rstest] + #[case::binary( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Binary, + 3 + )] + #[case::utf8( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBin, + DataType::Utf8, + 3 + )] + #[case::binary_view( + DType::Binary(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::BinaryView, + 4 + )] + #[case::utf8_view( + DType::Utf8(Nullability::NonNullable), + VarBinExportLayout::VarBinView, + DataType::Utf8View, + 4 + )] + #[crate::test] + async fn test_cuda_onpair_arrow_export_uses_dtype_layout( + #[case] dtype: DType, + #[case] layout: VarBinExportLayout, + #[case] expected_data_type: DataType, + #[case] expected_n_buffers: i64, + ) -> VortexResult<()> { + let mut cuda_ctx = cuda_ctx_with_varbin_layout(layout)?; + let values = vec![ + Some(&b"short"[..]), + Some(&b"this value is stored out of line"[..]), + ]; + let onpair = compress_onpair(values, dtype, &mut cuda_ctx)?; + + let mut exported = onpair + .export_device_array_with_schema(&mut cuda_ctx) + .await?; + assert_eq!( + Field::try_from(&exported.schema)?, + Field::new("", expected_data_type, false) + ); + assert_eq!(exported.array.array.n_buffers, expected_n_buffers); + + release_device_array(&mut exported.array); + release_schema(&mut exported.schema); + Ok(()) + } +} diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 3c712d20fb8..d2b3ea3d22a 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -48,6 +48,7 @@ use kernel::FSSTExecutor; use kernel::FilterExecutor; use kernel::FoRExecutor; pub use kernel::LaunchStrategy; +use kernel::OnPairExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; @@ -90,6 +91,7 @@ use vortex::encodings::zstd::ZstdBuffers; #[cfg(test)] use vortex_cuda_macros::test; pub use vortex_nvcomp as nvcomp; +use vortex_onpair::OnPair; use crate::kernel::SequenceExecutor; use crate::kernel::SliceExecutor; @@ -118,6 +120,7 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(Shared.id(), &SharedExecutor); session.register_kernel(FoR.id(), &FoRExecutor); session.register_kernel(FSST.id(), &FSSTExecutor); + session.register_kernel(OnPair.id(), &OnPairExecutor); session.register_kernel(RunEnd.id(), &RunEndExecutor); session.register_kernel(Sequence.id(), &SequenceExecutor); session.register_kernel(ZigZag.id(), &ZigZagExecutor);