From 042d01c5d2d35a5b4abc501e8a96171b8f580df0 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 15:32:34 +0100 Subject: [PATCH 01/11] Add CUDA decompression kernels for OnPair Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + encodings/onpair/src/array.rs | 2 +- encodings/onpair/src/lib.rs | 3 + vortex-cuda/Cargo.toml | 5 + vortex-cuda/benches/onpair_cuda.rs | 111 +++ vortex-cuda/kernels/src/onpair.cu | 162 ++++ .../src/onpair_shmem_4tpt_split8read.cu | 145 ++++ vortex-cuda/src/arrow/canonical.rs | 36 +- vortex-cuda/src/cub.rs | 38 + vortex-cuda/src/kernel/encodings/fsst.rs | 14 +- vortex-cuda/src/kernel/encodings/mod.rs | 5 +- vortex-cuda/src/kernel/encodings/onpair.rs | 796 ++++++++++++++++++ vortex-cuda/src/lib.rs | 3 + 13 files changed, 1311 insertions(+), 10 deletions(-) create mode 100644 vortex-cuda/benches/onpair_cuda.rs create mode 100644 vortex-cuda/kernels/src/onpair.cu create mode 100644 vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu create mode 100644 vortex-cuda/src/kernel/encodings/onpair.rs 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/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/onpair_cuda.rs b/vortex-cuda/benches/onpair_cuda.rs new file mode 100644 index 00000000000..4a43b2e307b --- /dev/null +++ b/vortex-cuda/benches/onpair_cuda.rs @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA benchmarks for OnPair decompression. + +#![expect(clippy::unwrap_used)] + +#[allow(dead_code)] +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_DICT12_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_DICT12_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/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu new file mode 100644 index 00000000000..03e260e984c --- /dev/null +++ b/vortex-cuda/kernels/src/onpair.cu @@ -0,0 +1,162 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include "config.cuh" + +#include + +// Support kernels for OnPair GPU decompression. +// +// The decode kernel (`onpair_shmem_4tpt_split8read.cu`) consumes per-batch output +// offsets: `chunk_offsets[b]` is the count of decoded bytes preceding the b-th +// 128-token batch. Vortex does not store those offsets; they are regenerated on +// the GPU at decode time. `onpair_batch_sizes` reduces each batch's decoded size +// from the codes and the per-token length LUT, and a CUB exclusive scan over the +// result yields `chunk_offsets`. This scans only the compressed codes — it +// touches neither the dictionary bytes nor the output. + +// Tokens per decode batch: one warp of the decode kernel emits 128 tokens +// (4 tokens/thread). Must match the decode kernel's layout. +constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; + +// One warp per 128-token batch: sums `lens[codes[t]]` over the batch's (up to) +// 128 tokens and writes the total to `batch_sizes[b]`. Code reads are +// lane-consecutive (coalesced); the length LUT is small and cache-resident. +// +// A code outside the dictionary raises `status` to 1 and contributes zero +// bytes: the host must check the flag before trusting `batch_sizes` and before +// launching the decode kernel, whose dictionary gathers are unchecked. +extern "C" __global__ void onpair_batch_sizes(const uint16_t *__restrict codes, + const uint8_t *__restrict lens, + uint32_t dict_size, uint64_t total_tokens, + uint64_t *__restrict batch_sizes, + uint32_t *__restrict status) { + const int lane = threadIdx.x & 31; + const uint32_t warp = threadIdx.x >> 5; + const uint64_t b = (uint64_t)blockIdx.x * (uint64_t)(blockDim.x >> 5) + (uint64_t)warp; + const uint64_t base = b * (uint64_t)ONPAIR_TOKENS_PER_BATCH; + if (base >= total_tokens) { + return; + } + uint32_t s = 0; +#pragma unroll + for (int k = 0; k < 4; ++k) { + const uint64_t i = base + (uint64_t)lane + (uint64_t)(k * 32); + if (i < total_tokens) { + const uint32_t code = (uint32_t)codes[i]; + if (code < dict_size) { + s += (uint32_t)lens[code]; + } else { + atomicMax(status, 1u); + } + } + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + s += __shfl_down_sync(0xffffffffu, s, offset); + } + if (lane == 0) { + batch_sizes[b] = (uint64_t)s; + } +} + +// Widen the per-row decoded lengths to the u64 scan input `row_sizes`. A CUB +// exclusive scan over the result (with one extra zeroed slot) yields the u64 +// per-row output offsets and, in the last slot, the total decoded byte count. +// A negative length raises `status` to 2 and contributes zero bytes; the host +// must check the flag before trusting the offsets. +template +__device__ inline void onpair_row_sizes_impl(const T *__restrict lengths, + uint64_t *__restrict row_sizes, + uint32_t *__restrict status, 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 i = block_start + threadIdx.x; i < block_end; i += blockDim.x) { + T len = lengths[i]; + if constexpr (static_cast(-1) < static_cast(0)) { + if (len < static_cast(0)) { + atomicMax(status, 2u); + len = static_cast(0); + } + } + row_sizes[i] = (uint64_t)len; + } +} + +#define GENERATE_ONPAIR_ROW_SIZES_KERNEL(suffix, Type) \ + extern "C" __global__ void onpair_row_sizes_##suffix( \ + const Type *__restrict lengths, uint64_t *__restrict row_sizes, \ + uint32_t *__restrict status, uint64_t num_rows) { \ + onpair_row_sizes_impl(lengths, row_sizes, status, num_rows); \ + } + +GENERATE_ONPAIR_ROW_SIZES_KERNEL(u8, uint8_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(u16, uint16_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(u32, uint32_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(u64, uint64_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(i8, int8_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(i16, int16_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(i32, int32_t) +GENERATE_ONPAIR_ROW_SIZES_KERNEL(i64, int64_t) + +// Narrow the u64 row offsets to the i32 Arrow `Utf8`/`Binary` offsets buffer. +// The host only launches this after checking the total decoded size fits i32, +// and offsets are nondecreasing, so every value fits. +extern "C" __global__ void onpair_offsets_to_i32(const uint64_t *__restrict row_offsets, + int32_t *__restrict arrow_offsets, + uint64_t num_offsets) { + 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_offsets) + ? (block_start + elements_per_block) + : num_offsets; + for (uint64_t i = block_start + threadIdx.x; i < block_end; i += blockDim.x) { + arrow_offsets[i] = (int32_t)row_offsets[i]; + } +} + +// 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 uint32_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 Rust caller only +// launches this when every offset fits the view's u32 fields and the decoded +// heap is exposed as backing buffer zero. +__device__ inline void onpair_write_view(const uint64_t *__restrict row_offsets, + const uint8_t *__restrict output_bytes, + uint4 *__restrict views, uint64_t rid) { + const uint64_t start = 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 (uint32_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, (uint32_t)start); +} + +extern "C" __global__ void onpair_build_views(const uint64_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..6d81ba2e898 --- /dev/null +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include +#include +#include +#include + +// OnPair decompress — 4 tokens/thread, split-read dictionary. +// +// 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. +// +// Identical scan/drain to `onpair_shmem_4tpt`; only the token-byte source +// changes. + +#ifndef WARPS_PER_BLOCK_MAX +#define WARPS_PER_BLOCK_MAX 16u +#endif +#ifndef ONPAIR_LAUNCH_BOUNDS +#define ONPAIR_LAUNCH_BOUNDS __launch_bounds__(512, 2) +#endif +#define WARP_BUF_BYTES 2080u + +__device__ inline uint32_t warp_inclusive_scan_u32_s8r(uint32_t x, int lane) { + constexpr unsigned mask = 0xffffffffu; +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + uint32_t y = __shfl_up_sync(mask, x, offset); + if (lane >= offset) { + x += y; + } + } + return x; +} + +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( + 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) { + constexpr unsigned mask = 0xffffffffu; + const int 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 uint64_t base_i = chunk * 128u + (uint64_t)lane; + uint2 lo[4]; + uint32_t c[4]; + uint32_t l[4]; +#pragma unroll + for (int 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]; + c[k] = code; + lo[k] = *reinterpret_cast(dict_s8 + (size_t)code * 8u); + l[k] = (uint32_t)lens[code]; + } else { + c[k] = 0u; + lo[k] = make_uint2(0u, 0u); + l[k] = 0u; + } + } + + uint32_t excl[4]; + uint32_t acc_base = 0u; +#pragma unroll + for (int k = 0; k < 4; ++k) { + const uint32_t incl = warp_inclusive_scan_u32_s8r(l[k], lane); + excl[k] = acc_base + (incl - l[k]); + acc_base += __shfl_sync(mask, incl, 31); + } + const uint32_t warp_total = acc_base; + + 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); + +#pragma unroll + for (int k = 0; k < 4; ++k) { + const uint32_t len = l[k]; + if (len == 0u) { + continue; + } + const uint32_t base = excl[k]; + const uint8_t *lob = reinterpret_cast(&lo[k]); + const uint32_t nlo = len < 8u ? len : 8u; +#pragma unroll + for (int j = 0; j < 8; ++j) { + if (j < (int)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)c[k] * 16u + 8u); + const uint8_t *hib = reinterpret_cast(&hi); +#pragma unroll + for (int j = 0; j < 8; ++j) { + if (8 + j < (int)len) { + s_buf[base + 8 + j] = hib[j]; + } + } + } + } + __syncwarp(); + + const uint32_t head = head_pre < warp_total ? head_pre : warp_total; + if ((uint32_t)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 (uint32_t k = 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 ((uint32_t)lane < warp_total - tail_start) { + output_bytes[out_start + (uint64_t)tail_start + (uint64_t)lane] = + s_buf[tail_start + lane]; + } +} diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 48c5042716b..f6ddf6ea78f 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -61,6 +61,8 @@ 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 vortex_onpair::OnPairArray; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -81,8 +83,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 @@ -234,6 +237,16 @@ fn export_array( 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 => + { + return export_onpair_varbin(onpair, ctx).await; + } + Ok(onpair) => onpair.into_array(), + Err(array) => array, + }; let cuda_array = array.execute_cuda(ctx).await?; export_canonical(cuda_array, ctx).await @@ -585,7 +598,7 @@ async fn export_fsst_varbin( fsst: FSSTArray, ctx: &mut CudaExecutionCtx, ) -> VortexResult<(ArrowArray, SyncEvent)> { - let FSSTVarBin { + let DecodedVarBin { dtype, len, offsets, @@ -600,6 +613,25 @@ async fn export_fsst_varbin( export_varbin_buffers(len, validity_buffer, null_count, offsets, values, ctx) } +async fn export_onpair_varbin( + onpair: OnPairArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult<(ArrowArray, SyncEvent)> { + let DecodedVarBin { + dtype, + len, + offsets, + values, + validity, + } = decode_onpair_varbin(onpair, ctx).await?; + vortex_ensure!( + matches!(dtype, DType::Utf8(_) | DType::Binary(_)), + "OnPair 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) +} + fn export_varbin_buffers( len: usize, validity_buffer: Option, diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 4b6009e7cea..2b2ee76cfac 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -15,6 +15,44 @@ use vortex_cub::scan::cudaStream_t; use crate::CudaExecutionCtx; +/// CUB `DeviceScan::ExclusiveSum` over device-resident `u64` values. +/// +/// Runs through the `i64` CUB instantiation: callers pass non-negative counts +/// whose prefix sums stay below `i64::MAX`, where two's complement `i64` and +/// `u64` addition produce identical bit patterns. +pub(crate) fn exclusive_sum_u64( + input: &CudaSlice, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let len_i64 = i64::try_from(len)?; + let temp_bytes = scan::exclusive_sum_i64_temp_size(len_i64) + .map_err(|err| vortex_err!("CUB scan_exclusive_sum_i64_temp_size failed: {err}"))?; + + let mut temp = ctx.device_alloc::(temp_bytes.max(1))?; + let mut output = ctx.device_alloc::(len)?; + let stream = ctx.stream(); + let stream_ptr = stream.cu_stream() as cudaStream_t; + let (input_ptr, record_input) = input.device_ptr(stream); + let (output_ptr, record_output) = output.device_ptr_mut(stream); + let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); + + ctx.launch_external(len, || unsafe { + scan::exclusive_sum_i64( + temp_ptr as *mut c_void, + temp_bytes, + input_ptr as *const i64, + output_ptr as *mut i64, + len_i64, + stream_ptr, + ) + .map_err(|err| vortex_err!("CUB scan_exclusive_sum_i64 failed: {err}")) + })?; + drop((record_input, record_output, record_temp)); + + Ok(output) +} + pub(crate) fn exclusive_sum_i32( input: &CudaSlice, len: usize, 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..60eb13c44f1 --- /dev/null +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -0,0 +1,796 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA executor for OnPair decompression. +//! +//! Decoding runs entirely on the GPU over the flat token stream: +//! +//! 1. `onpair_row_sizes` widens the per-row decoded lengths and a CUB +//! exclusive scan turns them into per-row output offsets (the last element +//! is the total decoded byte count). +//! 2. `onpair_batch_sizes` reduces the decoded byte size of every 128-token +//! batch from the codes and the per-token length LUT, and a second CUB +//! exclusive scan regenerates the per-batch output offsets +//! (`chunk_offsets`) the decode kernel positions its writes with. +//! 3. `onpair_shmem_4tpt_split8read` gathers each token's bytes from the +//! split dictionary layout and scatters them to the output byte stream. +//! +//! The result is exposed either as a canonical `VarBinView` (views built +//! on-device by `onpair_build_views`, or on host for heaps that exceed a +//! single backing buffer) 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::LaunchConfig; +use cudarc::driver::PushKernelArg; +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::builtins::ArrayBuiltins; +use vortex::array::match_each_integer_ptype; +use vortex::array::validity::Validity; +use vortex::buffer::Alignment; +use vortex::dtype::DType; +use vortex::dtype::Nullability; +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_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::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::cub::exclusive_sum_u64; +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 `kernels/src/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; + +/// `status` value raised by `onpair_batch_sizes` for a code outside the +/// dictionary. +const STATUS_CODE_OUT_OF_RANGE: u32 = 1; +/// `status` value raised by `onpair_row_sizes` for a negative decoded length. +const STATUS_NEGATIVE_LENGTH: u32 = 2; + +/// 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 + .try_downcast::() + .map_err(|_| vortex_err!("Expected OnPairArray"))?; + decode_onpair(onpair, ctx).await + } +} + +/// Read one `codes_offsets` boundary by point lookup, so a sliced array never +/// materialises the whole per-row offsets child just to bound its codes. +fn code_boundary( + codes_offsets: &ArrayRef, + index: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + codes_offsets + .execute_scalar(index, ctx.execution_ctx())? + .as_primitive() + .as_::() + .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) +} + +/// 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 u16 code window, the split +/// dictionary layout, and the regenerated per-batch output offsets. +struct StagedCodes { + codes: BufferHandle, + dict_s8: BufferHandle, + dict_padded: BufferHandle, + lens: BufferHandle, + /// Exclusive per-batch output offsets, `num_batches + 1` entries; the last + /// is the total decoded byte count of the code window. + chunk_offsets: CudaSlice, + num_batches: usize, + num_tokens: usize, + launch_config: LaunchConfig, +} + +/// The shared result of the OnPair GPU decode pipeline. +struct OnPairDecoded { + /// Exclusive per-row byte offsets over `bytes`, `num_rows + 1` entries on + /// device; the last is `total_size`. + row_offsets: CudaSlice, + /// The flat decoded byte stream. + bytes: CudaSlice, + /// Total decoded byte count. + total_size: usize, + /// Host-resident per-row lengths, for the host rollover path. + lengths: PrimitiveArray, +} + +/// Stage this array's code window and dictionary on the device and regenerate +/// the decode kernel's per-batch output offsets from them (steps 1–2 of the +/// pipeline: `onpair_batch_sizes` + CUB exclusive scan). +async fn stage_codes( + onpair: &OnPairArray, + code_start: usize, + code_end: usize, + status: &CudaSlice, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + // Widen this array's code window to the decode kernel's u16 ABI. + let codes = onpair + .codes() + .slice(code_start..code_end)? + .cast(DType::Primitive(PType::U16, Nullability::NonNullable))? + .execute::(ctx.execution_ctx())? + .into_buffer::(); + let num_tokens = codes.len(); + let num_tokens_u64 = u64::try_from(num_tokens)?; + + // 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.as_view(), 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.copy_to_device(codes)?, + 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)?; + + // Per-batch decoded sizes. One extra zeroed slot makes the exclusive + // scan's last element the total decoded byte count. + let mut batch_sizes = ctx.device_alloc::(num_batches + 1)?; + ctx.stream() + .memset_zeros(&mut batch_sizes) + .map_err(|e| vortex_err!("Failed to zero OnPair batch sizes: {e}"))?; + + let codes_view = codes_dev.cuda_view::()?; + let lens_view = lens_dev.cuda_view::()?; + let batch_sizes_fn = ctx.load_function_with_suffixes("onpair", &["batch_sizes"])?; + ctx.launch_kernel_config(&batch_sizes_fn, launch_config, num_tokens, |args| { + args.arg(&codes_view) + .arg(&lens_view) + .arg(&dict_size_u32) + .arg(&num_tokens_u64) + .arg(&batch_sizes) + .arg(status); + })?; + + let chunk_offsets = exclusive_sum_u64(&batch_sizes, num_batches + 1, ctx)?; + + Ok(StagedCodes { + codes: codes_dev, + dict_s8: s8_dev, + dict_padded: padded_dev, + lens: lens_dev, + chunk_offsets, + num_batches, + num_tokens, + launch_config, + }) +} + +/// Run the shared OnPair GPU decode pipeline: regenerate the row and batch +/// output offsets on the device, validate the compressed stream, and decode +/// the flat byte stream. Returns `Ok(None)` when the array decodes to zero +/// bytes. +async fn decode_onpair_bytes( + onpair: &OnPairArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let num_rows = onpair.len(); + + // Decompress the lengths child; its native-width buffer feeds the GPU + // row-offsets scan, and the host values serve the rare rollover path. + let lengths = onpair + .uncompressed_lengths() + .clone() + .execute::(ctx.execution_ctx())?; + let lengths_ptype = lengths.ptype(); + let PrimitiveDataParts { + buffer: lengths_buffer, + .. + } = lengths.clone().into_data_parts(); + let lengths_dev = ctx.ensure_on_device(lengths_buffer).await?; + + // Shared corruption flag for the size kernels; 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}"))?; + + // Row offsets on GPU: widen the lengths to u64 sizes (with one extra + // zeroed slot), then exclusive-scan; the last element is the total + // decoded byte count. + let mut row_sizes = ctx.device_alloc::(num_rows + 1)?; + ctx.stream() + .memset_zeros(&mut row_sizes) + .map_err(|e| vortex_err!("Failed to zero OnPair row sizes: {e}"))?; + let num_rows_u64 = u64::try_from(num_rows)?; + // The kernel is selected by the lengths ptype suffix and receives only the + // base pointer, so a byte-typed view of the native-width buffer suffices. + let lengths_view = lengths_dev.cuda_view::()?; + let row_sizes_fn = + ctx.load_function_with_suffixes("onpair", &["row_sizes", &lengths_ptype.to_string()])?; + ctx.launch_kernel(&row_sizes_fn, num_rows, |args| { + args.arg(&lengths_view) + .arg(&row_sizes) + .arg(&status) + .arg(&num_rows_u64); + })?; + let row_offsets = exclusive_sum_u64(&row_sizes, num_rows + 1, ctx)?; + + // `codes_offsets` may be a sliced view of the original; its first and last + // boundaries bound the contiguous run of `codes` belonging to this array's + // rows (`slice` keeps the full `codes` child and only narrows the offsets). + let code_start = code_boundary(onpair.codes_offsets(), 0, ctx)?; + let code_end = code_boundary(onpair.codes_offsets(), num_rows, ctx)?; + vortex_ensure!( + code_start <= code_end, + "OnPair codes_offsets must be nondecreasing" + ); + vortex_ensure!( + code_end <= onpair.codes().len(), + "OnPair codes_offsets end {} exceeds codes len {}", + code_end, + onpair.codes().len() + ); + + let mut staged = None; + if code_start < code_end { + staged = Some(stage_codes(onpair, code_start, code_end, &status, ctx).await?); + } + + // One synchronizing readback validates the compressed stream before the + // decode kernel — whose dictionary gathers and output scatters are + // unchecked — is allowed to run: no negative length, every code indexed + // the dictionary, and the codes decode to exactly the byte count the + // lengths record. + let status = ctx + .stream() + .clone_dtoh(&status) + .map_err(|e| vortex_err!("Failed to copy OnPair status flag to host: {e}"))?; + match status.first().copied().unwrap_or(STATUS_CODE_OUT_OF_RANGE) { + 0 => {} + STATUS_NEGATIVE_LENGTH => { + vortex_bail!("OnPair uncompressed length cannot be negative") + } + _ => vortex_bail!("OnPair code out of dictionary range"), + } + let row_total = ctx + .stream() + .clone_dtoh(&row_offsets.slice(num_rows..num_rows + 1)) + .map_err(|e| vortex_err!("Failed to copy OnPair decoded size to host: {e}"))? + .first() + .copied() + .ok_or_else(|| vortex_err!("OnPair row offset scan returned no total"))?; + let chunk_total = match &staged { + Some(staged) => 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"))?, + None => 0, + }; + vortex_ensure!( + row_total == chunk_total, + "OnPair codes decode to {chunk_total} bytes but uncompressed_lengths records {row_total}" + ); + if row_total == 0 { + return Ok(None); + } + let total_size = usize::try_from(row_total)?; + let Some(staged) = staged else { + vortex_bail!("OnPair records {total_size} decoded bytes but has no codes"); + }; + + // Decode. 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 bytes = ctx.device_alloc::(total_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 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::()?; + let decode_fn = ctx.load_function_with_suffixes("onpair_shmem_4tpt_split8read", &[])?; + 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(&bytes) + .arg(&num_tokens_u64); + }, + )?; + + Ok(Some(OnPairDecoded { + row_offsets, + bytes, + total_size, + lengths, + })) +} + +async fn decode_onpair(onpair: OnPairArray, 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)); + } + + if validity.definitely_all_null() { + return empty_views(num_rows, dtype, validity, ctx).await; + } + + let Some(decoded) = decode_onpair_bytes(&onpair, ctx).await? else { + return empty_views(num_rows, dtype, validity, ctx).await; + }; + let OnPairDecoded { + row_offsets, + bytes, + total_size, + lengths, + } = decoded; + + // Fast path: the decoded heap fits a single BinaryView backing buffer, so + // the per-row views build on-device straight from the row offsets. + if total_size <= MAX_BUFFER_LEN { + let 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) + .arg(&bytes) + .arg(&device_views) + .arg(&num_rows_u64); + })?; + + let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); + let bytes = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(bytes))); + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([bytes]), dtype, validity) + })); + } + + // BinaryView offsets are u32. Heaps that need multiple backing buffers + // roll the decoded bytes over on host, mirroring the CPU canonical path. + let host_bytes = CudaDeviceBuffer::new(bytes) + .copy_to_host(Alignment::new(1))? + .await?; + let host_bytes = host_bytes.slice(0..total_size); + + 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 decoded = if onpair.is_empty() || validity.definitely_all_null() { + None + } else { + decode_onpair_bytes(&onpair, ctx).await? + }; + + let Some(decoded) = decoded 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, + }); + }; + + vortex_ensure!( + i32::try_from(decoded.total_size).is_ok(), + "OnPair decoded size exceeds Arrow i32 offset range" + ); + + // Narrow the device row offsets to Arrow's i32 offsets; every value fits + // because the total was just checked and offsets are nondecreasing. + let arrow_offsets = ctx.device_alloc::(len + 1)?; + let num_offsets_u64 = u64::try_from(len + 1)?; + let offsets_fn = ctx.load_function_with_suffixes("onpair", &["offsets_to_i32"])?; + ctx.launch_kernel(&offsets_fn, len + 1, |args| { + args.arg(&decoded.row_offsets) + .arg(&arrow_offsets) + .arg(&num_offsets_u64); + })?; + + Ok(DecodedVarBin { + dtype, + len, + offsets: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(arrow_offsets))), + values: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(decoded.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::error::VortexExpect; + use vortex_array::VortexSessionExecute; + use vortex_onpair::DEFAULT_DICT12_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_DICT12_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_DICT12_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(()) + } + + #[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); From f24ca01ca940e9cabb97eb10e66d1ab9fc97bef1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 15:49:14 +0100 Subject: [PATCH 02/11] fuse Signed-off-by: Robert Kruszewski --- vortex-cuda/cub/kernels/filter.cu | 77 ++++++++++ vortex-cuda/cub/kernels/filter.h | 30 ++++ vortex-cuda/cub/src/scan.rs | 165 +++++++++++++++++++++ vortex-cuda/kernels/src/onpair.cu | 41 ----- vortex-cuda/src/cub.rs | 66 +++++++++ vortex-cuda/src/kernel/encodings/onpair.rs | 37 ++--- 6 files changed, 351 insertions(+), 65 deletions(-) diff --git a/vortex-cuda/cub/kernels/filter.cu b/vortex-cuda/cub/kernels/filter.cu index 73726b3a4a4..88db17279c8 100644 --- a/vortex-cuda/cub/kernels/filter.cu +++ b/vortex-cuda/cub/kernels/filter.cu @@ -218,3 +218,80 @@ static cudaError_t scan_exclusive_sum_temp_size_impl(size_t *temp_bytes, int64_t DEFINE_SCAN_EXCLUSIVE_SUM(i32, int32_t) DEFINE_SCAN_EXCLUSIVE_SUM(i64, int64_t) + +// Widening length functor for TransformInputIterator: index `i < num_rows` +// reads `lengths[i]` widened to u64; the final index contributes zero, so an +// exclusive scan over `num_rows + 1` items ends with the total. A negative +// length (signed types only) raises `*status` to 2 and contributes zero. +template +struct WideningLength { + const LenT *lengths; + int64_t num_rows; + uint32_t *status; + + __host__ __device__ inline uint64_t operator()(int64_t idx) const { + if (idx >= num_rows) { + return 0; + } + const LenT len = lengths[idx]; + if constexpr (static_cast(-1) < static_cast(0)) { + if (len < static_cast(0)) { +#ifdef __CUDA_ARCH__ + atomicMax(status, 2u); +#endif + return 0; + } + } + return static_cast(len); + } +}; + +template +using WideningLengthIterator = + thrust::transform_iterator, thrust::counting_iterator>; + +// Query CUB temporary storage for the fused widen + exclusive-sum scan. +template +static cudaError_t scan_exclusive_sum_lengths_temp_size_impl(size_t *temp_bytes, + int64_t num_offsets) { + WideningLengthIterator lengths_it(thrust::counting_iterator(0), + WideningLength{nullptr, 0, nullptr}); + size_t bytes = 0; + cudaError_t err = cub::DeviceScan::ExclusiveSum(nullptr, + bytes, + lengths_it, + static_cast(nullptr), + num_offsets); + *temp_bytes = bytes; + return err; +} + +// Fused widen + exclusive-sum: one CUB dispatch scans `num_offsets` values +// where value `i` is `lengths[i]` widened to u64 (zero for the final slot), +// writing u64 offsets whose last element is the total — no materialized +// widened input. +#define DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(suffix, Type) \ + extern "C" cudaError_t scan_exclusive_sum_lengths_##suffix##_temp_size(size_t *temp_bytes, \ + int64_t num_offsets) { \ + return scan_exclusive_sum_lengths_temp_size_impl(temp_bytes, num_offsets); \ + } \ + extern "C" cudaError_t scan_exclusive_sum_lengths_##suffix(void *d_temp, \ + size_t temp_bytes, \ + const Type *lengths, \ + uint64_t *d_out, \ + uint32_t *status, \ + int64_t num_offsets, \ + cudaStream_t stream) { \ + WideningLengthIterator lengths_it(thrust::counting_iterator(0), \ + WideningLength{lengths, num_offsets - 1, status}); \ + return cub::DeviceScan::ExclusiveSum(d_temp, temp_bytes, lengths_it, d_out, num_offsets, stream); \ + } + +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u8, uint8_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i8, int8_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u16, uint16_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i16, int16_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u32, uint32_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i32, int32_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u64, uint64_t) +DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i64, int64_t) diff --git a/vortex-cuda/cub/kernels/filter.h b/vortex-cuda/cub/kernels/filter.h index c49dc62faed..ab1f9d2429c 100644 --- a/vortex-cuda/cub/kernels/filter.h +++ b/vortex-cuda/cub/kernels/filter.h @@ -99,6 +99,36 @@ cudaError_t scan_exclusive_sum_i64(void *d_temp, int64_t num_items, cudaStream_t stream); +// Fused widen + exclusive-sum over per-row lengths: scans `num_offsets` +// (= num_rows + 1) values where value `i` is `lengths[i]` widened to u64 and +// the final slot contributes zero, so `d_out[num_offsets - 1]` is the total. +// A negative length (signed types only) raises `*status` to 2 and contributes +// zero bytes. +#define SCAN_LENGTHS_TYPE_TABLE(X) \ + X(u8, uint8_t) \ + X(i8, int8_t) \ + X(u16, uint16_t) \ + X(i16, int16_t) \ + X(u32, uint32_t) \ + X(i32, int32_t) \ + X(u64, uint64_t) \ + X(i64, int64_t) + +#define DECLARE_SCAN_EXCLUSIVE_SUM_LENGTHS(suffix, c_type) \ + cudaError_t scan_exclusive_sum_lengths_##suffix##_temp_size(size_t *temp_bytes, \ + int64_t num_offsets); \ + cudaError_t scan_exclusive_sum_lengths_##suffix(void *d_temp, \ + size_t temp_bytes, \ + const c_type *lengths, \ + uint64_t *d_out, \ + uint32_t *status, \ + int64_t num_offsets, \ + cudaStream_t stream); + +SCAN_LENGTHS_TYPE_TABLE(DECLARE_SCAN_EXCLUSIVE_SUM_LENGTHS) + +#undef DECLARE_SCAN_EXCLUSIVE_SUM_LENGTHS + #ifdef __cplusplus } #endif diff --git a/vortex-cuda/cub/src/scan.rs b/vortex-cuda/cub/src/scan.rs index 56b4e64b857..e17a201a7ac 100644 --- a/vortex-cuda/cub/src/scan.rs +++ b/vortex-cuda/cub/src/scan.rs @@ -10,6 +10,171 @@ use crate::error::CubError; use crate::error::check_cuda_error; pub use crate::sys::cudaStream_t; +/// Element type of the `lengths` input to [`exclusive_sum_lengths`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LengthType { + /// `u8` lengths. + U8, + /// `i8` lengths. + I8, + /// `u16` lengths. + U16, + /// `i16` lengths. + I16, + /// `u32` lengths. + U32, + /// `i32` lengths. + I32, + /// `u64` lengths. + U64, + /// `i64` lengths. + I64, +} + +/// Get temporary storage size for the fused widen + exclusive-sum scan +/// ([`exclusive_sum_lengths`]). +pub fn exclusive_sum_lengths_temp_size( + ty: LengthType, + num_offsets: i64, +) -> Result { + let lib = cub_library()?; + let mut temp_bytes: usize = 0; + let err = unsafe { + match ty { + LengthType::U8 => { + (lib.scan_exclusive_sum_lengths_u8_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::I8 => { + (lib.scan_exclusive_sum_lengths_i8_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::U16 => { + (lib.scan_exclusive_sum_lengths_u16_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::I16 => { + (lib.scan_exclusive_sum_lengths_i16_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::U32 => { + (lib.scan_exclusive_sum_lengths_u32_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::I32 => { + (lib.scan_exclusive_sum_lengths_i32_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::U64 => { + (lib.scan_exclusive_sum_lengths_u64_temp_size)(&raw mut temp_bytes, num_offsets) + } + LengthType::I64 => { + (lib.scan_exclusive_sum_lengths_i64_temp_size)(&raw mut temp_bytes, num_offsets) + } + } + }; + check_cuda_error(err, "scan_exclusive_sum_lengths_temp_size")?; + Ok(temp_bytes) +} + +/// Execute the fused widen + CUB `DeviceScan::ExclusiveSum` over per-row +/// lengths: scans `num_offsets` (= `num_rows + 1`) values where value `i` is +/// `lengths[i]` widened to u64 and the final slot contributes zero, so +/// `d_out[num_offsets - 1]` is the total. A negative length (signed types +/// only) raises `*status` to 2 and contributes zero bytes. +/// +/// # Safety +/// +/// All device pointers must be valid and properly sized: +/// - `d_temp` must have at least `temp_bytes` bytes allocated. +/// - `lengths` must have at least `num_offsets - 1` elements of type `ty`. +/// - `d_out` must have at least `num_offsets` `u64` values. +/// - `status` must point to a valid device `u32`. +#[allow(clippy::too_many_arguments)] +pub unsafe fn exclusive_sum_lengths( + ty: LengthType, + d_temp: *mut c_void, + temp_bytes: usize, + lengths: *const c_void, + d_out: *mut u64, + status: *mut u32, + num_offsets: i64, + stream: cudaStream_t, +) -> Result<(), CubError> { + let lib = cub_library()?; + let err = unsafe { + match ty { + LengthType::U8 => (lib.scan_exclusive_sum_lengths_u8)( + d_temp, + temp_bytes, + lengths as *const u8, + d_out, + status, + num_offsets, + stream, + ), + LengthType::I8 => (lib.scan_exclusive_sum_lengths_i8)( + d_temp, + temp_bytes, + lengths as *const i8, + d_out, + status, + num_offsets, + stream, + ), + LengthType::U16 => (lib.scan_exclusive_sum_lengths_u16)( + d_temp, + temp_bytes, + lengths as *const u16, + d_out, + status, + num_offsets, + stream, + ), + LengthType::I16 => (lib.scan_exclusive_sum_lengths_i16)( + d_temp, + temp_bytes, + lengths as *const i16, + d_out, + status, + num_offsets, + stream, + ), + LengthType::U32 => (lib.scan_exclusive_sum_lengths_u32)( + d_temp, + temp_bytes, + lengths as *const u32, + d_out, + status, + num_offsets, + stream, + ), + LengthType::I32 => (lib.scan_exclusive_sum_lengths_i32)( + d_temp, + temp_bytes, + lengths as *const i32, + d_out, + status, + num_offsets, + stream, + ), + LengthType::U64 => (lib.scan_exclusive_sum_lengths_u64)( + d_temp, + temp_bytes, + lengths as *const u64, + d_out, + status, + num_offsets, + stream, + ), + LengthType::I64 => (lib.scan_exclusive_sum_lengths_i64)( + d_temp, + temp_bytes, + lengths as *const i64, + d_out, + status, + num_offsets, + stream, + ), + } + }; + check_cuda_error(err, "scan_exclusive_sum_lengths") +} + /// Get temporary storage size for CUB `DeviceScan::ExclusiveSum`. pub fn exclusive_sum_i32_temp_size(num_items: i64) -> Result { let lib = cub_library()?; diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu index 03e260e984c..a45fc669149 100644 --- a/vortex-cuda/kernels/src/onpair.cu +++ b/vortex-cuda/kernels/src/onpair.cu @@ -60,47 +60,6 @@ extern "C" __global__ void onpair_batch_sizes(const uint16_t *__restrict codes, } } -// Widen the per-row decoded lengths to the u64 scan input `row_sizes`. A CUB -// exclusive scan over the result (with one extra zeroed slot) yields the u64 -// per-row output offsets and, in the last slot, the total decoded byte count. -// A negative length raises `status` to 2 and contributes zero bytes; the host -// must check the flag before trusting the offsets. -template -__device__ inline void onpair_row_sizes_impl(const T *__restrict lengths, - uint64_t *__restrict row_sizes, - uint32_t *__restrict status, 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 i = block_start + threadIdx.x; i < block_end; i += blockDim.x) { - T len = lengths[i]; - if constexpr (static_cast(-1) < static_cast(0)) { - if (len < static_cast(0)) { - atomicMax(status, 2u); - len = static_cast(0); - } - } - row_sizes[i] = (uint64_t)len; - } -} - -#define GENERATE_ONPAIR_ROW_SIZES_KERNEL(suffix, Type) \ - extern "C" __global__ void onpair_row_sizes_##suffix( \ - const Type *__restrict lengths, uint64_t *__restrict row_sizes, \ - uint32_t *__restrict status, uint64_t num_rows) { \ - onpair_row_sizes_impl(lengths, row_sizes, status, num_rows); \ - } - -GENERATE_ONPAIR_ROW_SIZES_KERNEL(u8, uint8_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(u16, uint16_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(u32, uint32_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(u64, uint64_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(i8, int8_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(i16, int16_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(i32, int32_t) -GENERATE_ONPAIR_ROW_SIZES_KERNEL(i64, int64_t) - // Narrow the u64 row offsets to the i32 Arrow `Utf8`/`Binary` offsets buffer. // The host only launches this after checking the total decoded size fits i32, // and offsets are nondecreasing, so every value fits. diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 2b2ee76cfac..5cca4c2e5f1 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -8,13 +8,79 @@ use std::ffi::c_void; use cudarc::driver::CudaSlice; use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; +use vortex::array::buffer::BufferHandle; +use vortex::dtype::PType; use vortex::error::VortexResult; +use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex_cub::scan; +use vortex_cub::scan::LengthType; use vortex_cub::scan::cudaStream_t; +use crate::CudaBufferExt; use crate::CudaExecutionCtx; +fn length_type(ptype: PType) -> VortexResult { + Ok(match ptype { + PType::U8 => LengthType::U8, + PType::I8 => LengthType::I8, + PType::U16 => LengthType::U16, + PType::I16 => LengthType::I16, + PType::U32 => LengthType::U32, + PType::I32 => LengthType::I32, + PType::U64 => LengthType::U64, + PType::I64 => LengthType::I64, + other => vortex_bail!("unsupported lengths ptype {other} for fused exclusive sum"), + }) +} + +/// Fused widen + CUB `DeviceScan::ExclusiveSum` over device-resident per-row +/// lengths of any integer width. +/// +/// One CUB dispatch scans the lengths through a widening transform iterator +/// and produces `num_rows + 1` u64 offsets whose last element is the total — +/// no separate widen kernel and no materialized u64 input. A negative length +/// raises `status` to 2 and contributes zero bytes; the caller must check the +/// flag before trusting the offsets. +pub(crate) fn exclusive_sum_lengths_u64( + lengths: &BufferHandle, + ptype: PType, + num_rows: usize, + status: &mut CudaSlice, + ctx: &mut CudaExecutionCtx, +) -> VortexResult> { + let ty = length_type(ptype)?; + let num_offsets = i64::try_from(num_rows + 1)?; + let temp_bytes = scan::exclusive_sum_lengths_temp_size(ty, num_offsets) + .map_err(|err| vortex_err!("CUB scan_exclusive_sum_lengths_temp_size failed: {err}"))?; + + let mut temp = ctx.device_alloc::(temp_bytes.max(1))?; + let mut output = ctx.device_alloc::(num_rows + 1)?; + let lengths_ptr = lengths.cuda_device_ptr()?; + let stream = ctx.stream(); + let stream_ptr = stream.cu_stream() as cudaStream_t; + let (status_ptr, record_status) = status.device_ptr_mut(stream); + let (output_ptr, record_output) = output.device_ptr_mut(stream); + let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); + + ctx.launch_external(num_rows + 1, || unsafe { + scan::exclusive_sum_lengths( + ty, + temp_ptr as *mut c_void, + temp_bytes, + lengths_ptr as *const c_void, + output_ptr as *mut u64, + status_ptr as *mut u32, + num_offsets, + stream_ptr, + ) + .map_err(|err| vortex_err!("CUB scan_exclusive_sum_lengths failed: {err}")) + })?; + drop((record_status, record_output, record_temp)); + + Ok(output) +} + /// CUB `DeviceScan::ExclusiveSum` over device-resident `u64` values. /// /// Runs through the `i64` CUB instantiation: callers pass non-negative counts diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index 60eb13c44f1..12321c4110e 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -5,9 +5,10 @@ //! //! Decoding runs entirely on the GPU over the flat token stream: //! -//! 1. `onpair_row_sizes` widens the per-row decoded lengths and a CUB -//! exclusive scan turns them into per-row output offsets (the last element -//! is the total decoded byte count). +//! 1. A fused CUB exclusive scan (widening transform iterator over the +//! native-width lengths) turns the per-row decoded lengths into per-row +//! output offsets in one dispatch (the last element is the total decoded +//! byte count). //! 2. `onpair_batch_sizes` reduces the decoded byte size of every 128-token //! batch from the codes and the per-token length LUT, and a second CUB //! exclusive scan regenerates the per-batch output offsets @@ -60,6 +61,7 @@ use vortex_onpair::dict_view; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; +use crate::cub::exclusive_sum_lengths_u64; use crate::cub::exclusive_sum_u64; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; @@ -78,7 +80,8 @@ const WARPS_PER_BLOCK: usize = (BLOCK_THREADS / 32) as usize; /// `status` value raised by `onpair_batch_sizes` for a code outside the /// dictionary. const STATUS_CODE_OUT_OF_RANGE: u32 = 1; -/// `status` value raised by `onpair_row_sizes` for a negative decoded length. +/// `status` value raised by the fused lengths scan for a negative decoded +/// length. const STATUS_NEGATIVE_LENGTH: u32 = 2; /// Launch config for the warp-per-batch kernels: one warp per 128-token batch. @@ -277,26 +280,12 @@ async fn decode_onpair_bytes( .memset_zeros(&mut status) .map_err(|e| vortex_err!("Failed to zero OnPair status flag: {e}"))?; - // Row offsets on GPU: widen the lengths to u64 sizes (with one extra - // zeroed slot), then exclusive-scan; the last element is the total - // decoded byte count. - let mut row_sizes = ctx.device_alloc::(num_rows + 1)?; - ctx.stream() - .memset_zeros(&mut row_sizes) - .map_err(|e| vortex_err!("Failed to zero OnPair row sizes: {e}"))?; - let num_rows_u64 = u64::try_from(num_rows)?; - // The kernel is selected by the lengths ptype suffix and receives only the - // base pointer, so a byte-typed view of the native-width buffer suffices. - let lengths_view = lengths_dev.cuda_view::()?; - let row_sizes_fn = - ctx.load_function_with_suffixes("onpair", &["row_sizes", &lengths_ptype.to_string()])?; - ctx.launch_kernel(&row_sizes_fn, num_rows, |args| { - args.arg(&lengths_view) - .arg(&row_sizes) - .arg(&status) - .arg(&num_rows_u64); - })?; - let row_offsets = exclusive_sum_u64(&row_sizes, num_rows + 1, ctx)?; + // Row offsets in one fused GPU dispatch: a CUB exclusive scan over a + // widening transform iterator of the native-width lengths. The last of + // the `num_rows + 1` outputs is the total decoded byte count; a negative + // length raises `status` and contributes zero bytes. + let row_offsets = + exclusive_sum_lengths_u64(&lengths_dev, lengths_ptype, num_rows, &mut status, ctx)?; // `codes_offsets` may be a sliced view of the original; its first and last // boundaries bound the contiguous run of `codes` belonging to this array's From d85260220476ea51e4e5ef2db94bdfd317ee66d8 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 16:14:52 +0100 Subject: [PATCH 03/11] fixes Signed-off-by: Robert Kruszewski --- vortex-cuda/benches/arrow_binary_cuda.rs | 1 - vortex-cuda/benches/bench_config/mod.rs | 4 ++++ vortex-cuda/benches/fsst_cuda.rs | 1 - vortex-cuda/benches/list_view_cuda.rs | 1 - vortex-cuda/benches/onpair_cuda.rs | 1 - 5 files changed, 4 insertions(+), 4 deletions(-) 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 index 4a43b2e307b..51e9d079351 100644 --- a/vortex-cuda/benches/onpair_cuda.rs +++ b/vortex-cuda/benches/onpair_cuda.rs @@ -5,7 +5,6 @@ #![expect(clippy::unwrap_used)] -#[allow(dead_code)] mod bench_config; mod timed_launch_strategy; From 2207cf690a0b1e4ef030f0f0a129b71b3e77eb30 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 17:45:57 +0100 Subject: [PATCH 04/11] less Signed-off-by: Robert Kruszewski --- vortex-cuda/cub/kernels/filter.cu | 77 ------- vortex-cuda/cub/kernels/filter.h | 30 --- vortex-cuda/cub/src/scan.rs | 165 -------------- vortex-cuda/kernels/src/onpair.cu | 16 -- .../src/onpair_shmem_4tpt_split8read.cu | 143 ++++++++---- vortex-cuda/src/cub.rs | 66 ------ vortex-cuda/src/kernel/encodings/onpair.rs | 210 +++++++++--------- 7 files changed, 199 insertions(+), 508 deletions(-) diff --git a/vortex-cuda/cub/kernels/filter.cu b/vortex-cuda/cub/kernels/filter.cu index 88db17279c8..73726b3a4a4 100644 --- a/vortex-cuda/cub/kernels/filter.cu +++ b/vortex-cuda/cub/kernels/filter.cu @@ -218,80 +218,3 @@ static cudaError_t scan_exclusive_sum_temp_size_impl(size_t *temp_bytes, int64_t DEFINE_SCAN_EXCLUSIVE_SUM(i32, int32_t) DEFINE_SCAN_EXCLUSIVE_SUM(i64, int64_t) - -// Widening length functor for TransformInputIterator: index `i < num_rows` -// reads `lengths[i]` widened to u64; the final index contributes zero, so an -// exclusive scan over `num_rows + 1` items ends with the total. A negative -// length (signed types only) raises `*status` to 2 and contributes zero. -template -struct WideningLength { - const LenT *lengths; - int64_t num_rows; - uint32_t *status; - - __host__ __device__ inline uint64_t operator()(int64_t idx) const { - if (idx >= num_rows) { - return 0; - } - const LenT len = lengths[idx]; - if constexpr (static_cast(-1) < static_cast(0)) { - if (len < static_cast(0)) { -#ifdef __CUDA_ARCH__ - atomicMax(status, 2u); -#endif - return 0; - } - } - return static_cast(len); - } -}; - -template -using WideningLengthIterator = - thrust::transform_iterator, thrust::counting_iterator>; - -// Query CUB temporary storage for the fused widen + exclusive-sum scan. -template -static cudaError_t scan_exclusive_sum_lengths_temp_size_impl(size_t *temp_bytes, - int64_t num_offsets) { - WideningLengthIterator lengths_it(thrust::counting_iterator(0), - WideningLength{nullptr, 0, nullptr}); - size_t bytes = 0; - cudaError_t err = cub::DeviceScan::ExclusiveSum(nullptr, - bytes, - lengths_it, - static_cast(nullptr), - num_offsets); - *temp_bytes = bytes; - return err; -} - -// Fused widen + exclusive-sum: one CUB dispatch scans `num_offsets` values -// where value `i` is `lengths[i]` widened to u64 (zero for the final slot), -// writing u64 offsets whose last element is the total — no materialized -// widened input. -#define DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(suffix, Type) \ - extern "C" cudaError_t scan_exclusive_sum_lengths_##suffix##_temp_size(size_t *temp_bytes, \ - int64_t num_offsets) { \ - return scan_exclusive_sum_lengths_temp_size_impl(temp_bytes, num_offsets); \ - } \ - extern "C" cudaError_t scan_exclusive_sum_lengths_##suffix(void *d_temp, \ - size_t temp_bytes, \ - const Type *lengths, \ - uint64_t *d_out, \ - uint32_t *status, \ - int64_t num_offsets, \ - cudaStream_t stream) { \ - WideningLengthIterator lengths_it(thrust::counting_iterator(0), \ - WideningLength{lengths, num_offsets - 1, status}); \ - return cub::DeviceScan::ExclusiveSum(d_temp, temp_bytes, lengths_it, d_out, num_offsets, stream); \ - } - -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u8, uint8_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i8, int8_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u16, uint16_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i16, int16_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u32, uint32_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i32, int32_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(u64, uint64_t) -DEFINE_SCAN_EXCLUSIVE_SUM_LENGTHS(i64, int64_t) diff --git a/vortex-cuda/cub/kernels/filter.h b/vortex-cuda/cub/kernels/filter.h index ab1f9d2429c..c49dc62faed 100644 --- a/vortex-cuda/cub/kernels/filter.h +++ b/vortex-cuda/cub/kernels/filter.h @@ -99,36 +99,6 @@ cudaError_t scan_exclusive_sum_i64(void *d_temp, int64_t num_items, cudaStream_t stream); -// Fused widen + exclusive-sum over per-row lengths: scans `num_offsets` -// (= num_rows + 1) values where value `i` is `lengths[i]` widened to u64 and -// the final slot contributes zero, so `d_out[num_offsets - 1]` is the total. -// A negative length (signed types only) raises `*status` to 2 and contributes -// zero bytes. -#define SCAN_LENGTHS_TYPE_TABLE(X) \ - X(u8, uint8_t) \ - X(i8, int8_t) \ - X(u16, uint16_t) \ - X(i16, int16_t) \ - X(u32, uint32_t) \ - X(i32, int32_t) \ - X(u64, uint64_t) \ - X(i64, int64_t) - -#define DECLARE_SCAN_EXCLUSIVE_SUM_LENGTHS(suffix, c_type) \ - cudaError_t scan_exclusive_sum_lengths_##suffix##_temp_size(size_t *temp_bytes, \ - int64_t num_offsets); \ - cudaError_t scan_exclusive_sum_lengths_##suffix(void *d_temp, \ - size_t temp_bytes, \ - const c_type *lengths, \ - uint64_t *d_out, \ - uint32_t *status, \ - int64_t num_offsets, \ - cudaStream_t stream); - -SCAN_LENGTHS_TYPE_TABLE(DECLARE_SCAN_EXCLUSIVE_SUM_LENGTHS) - -#undef DECLARE_SCAN_EXCLUSIVE_SUM_LENGTHS - #ifdef __cplusplus } #endif diff --git a/vortex-cuda/cub/src/scan.rs b/vortex-cuda/cub/src/scan.rs index e17a201a7ac..56b4e64b857 100644 --- a/vortex-cuda/cub/src/scan.rs +++ b/vortex-cuda/cub/src/scan.rs @@ -10,171 +10,6 @@ use crate::error::CubError; use crate::error::check_cuda_error; pub use crate::sys::cudaStream_t; -/// Element type of the `lengths` input to [`exclusive_sum_lengths`]. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum LengthType { - /// `u8` lengths. - U8, - /// `i8` lengths. - I8, - /// `u16` lengths. - U16, - /// `i16` lengths. - I16, - /// `u32` lengths. - U32, - /// `i32` lengths. - I32, - /// `u64` lengths. - U64, - /// `i64` lengths. - I64, -} - -/// Get temporary storage size for the fused widen + exclusive-sum scan -/// ([`exclusive_sum_lengths`]). -pub fn exclusive_sum_lengths_temp_size( - ty: LengthType, - num_offsets: i64, -) -> Result { - let lib = cub_library()?; - let mut temp_bytes: usize = 0; - let err = unsafe { - match ty { - LengthType::U8 => { - (lib.scan_exclusive_sum_lengths_u8_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::I8 => { - (lib.scan_exclusive_sum_lengths_i8_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::U16 => { - (lib.scan_exclusive_sum_lengths_u16_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::I16 => { - (lib.scan_exclusive_sum_lengths_i16_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::U32 => { - (lib.scan_exclusive_sum_lengths_u32_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::I32 => { - (lib.scan_exclusive_sum_lengths_i32_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::U64 => { - (lib.scan_exclusive_sum_lengths_u64_temp_size)(&raw mut temp_bytes, num_offsets) - } - LengthType::I64 => { - (lib.scan_exclusive_sum_lengths_i64_temp_size)(&raw mut temp_bytes, num_offsets) - } - } - }; - check_cuda_error(err, "scan_exclusive_sum_lengths_temp_size")?; - Ok(temp_bytes) -} - -/// Execute the fused widen + CUB `DeviceScan::ExclusiveSum` over per-row -/// lengths: scans `num_offsets` (= `num_rows + 1`) values where value `i` is -/// `lengths[i]` widened to u64 and the final slot contributes zero, so -/// `d_out[num_offsets - 1]` is the total. A negative length (signed types -/// only) raises `*status` to 2 and contributes zero bytes. -/// -/// # Safety -/// -/// All device pointers must be valid and properly sized: -/// - `d_temp` must have at least `temp_bytes` bytes allocated. -/// - `lengths` must have at least `num_offsets - 1` elements of type `ty`. -/// - `d_out` must have at least `num_offsets` `u64` values. -/// - `status` must point to a valid device `u32`. -#[allow(clippy::too_many_arguments)] -pub unsafe fn exclusive_sum_lengths( - ty: LengthType, - d_temp: *mut c_void, - temp_bytes: usize, - lengths: *const c_void, - d_out: *mut u64, - status: *mut u32, - num_offsets: i64, - stream: cudaStream_t, -) -> Result<(), CubError> { - let lib = cub_library()?; - let err = unsafe { - match ty { - LengthType::U8 => (lib.scan_exclusive_sum_lengths_u8)( - d_temp, - temp_bytes, - lengths as *const u8, - d_out, - status, - num_offsets, - stream, - ), - LengthType::I8 => (lib.scan_exclusive_sum_lengths_i8)( - d_temp, - temp_bytes, - lengths as *const i8, - d_out, - status, - num_offsets, - stream, - ), - LengthType::U16 => (lib.scan_exclusive_sum_lengths_u16)( - d_temp, - temp_bytes, - lengths as *const u16, - d_out, - status, - num_offsets, - stream, - ), - LengthType::I16 => (lib.scan_exclusive_sum_lengths_i16)( - d_temp, - temp_bytes, - lengths as *const i16, - d_out, - status, - num_offsets, - stream, - ), - LengthType::U32 => (lib.scan_exclusive_sum_lengths_u32)( - d_temp, - temp_bytes, - lengths as *const u32, - d_out, - status, - num_offsets, - stream, - ), - LengthType::I32 => (lib.scan_exclusive_sum_lengths_i32)( - d_temp, - temp_bytes, - lengths as *const i32, - d_out, - status, - num_offsets, - stream, - ), - LengthType::U64 => (lib.scan_exclusive_sum_lengths_u64)( - d_temp, - temp_bytes, - lengths as *const u64, - d_out, - status, - num_offsets, - stream, - ), - LengthType::I64 => (lib.scan_exclusive_sum_lengths_i64)( - d_temp, - temp_bytes, - lengths as *const i64, - d_out, - status, - num_offsets, - stream, - ), - } - }; - check_cuda_error(err, "scan_exclusive_sum_lengths") -} - /// Get temporary storage size for CUB `DeviceScan::ExclusiveSum`. pub fn exclusive_sum_i32_temp_size(num_items: i64) -> Result { let lib = cub_library()?; diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu index a45fc669149..c8e8b073afc 100644 --- a/vortex-cuda/kernels/src/onpair.cu +++ b/vortex-cuda/kernels/src/onpair.cu @@ -60,22 +60,6 @@ extern "C" __global__ void onpair_batch_sizes(const uint16_t *__restrict codes, } } -// Narrow the u64 row offsets to the i32 Arrow `Utf8`/`Binary` offsets buffer. -// The host only launches this after checking the total decoded size fits i32, -// and offsets are nondecreasing, so every value fits. -extern "C" __global__ void onpair_offsets_to_i32(const uint64_t *__restrict row_offsets, - int32_t *__restrict arrow_offsets, - uint64_t num_offsets) { - 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_offsets) - ? (block_start + elements_per_block) - : num_offsets; - for (uint64_t i = block_start + threadIdx.x; i < block_end; i += blockDim.x) { - arrow_offsets[i] = (int32_t)row_offsets[i]; - } -} - // 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. diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu index 6d81ba2e898..c53c87ff170 100644 --- a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -8,6 +8,18 @@ // 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). @@ -18,9 +30,6 @@ // 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. -// -// Identical scan/drain to `onpair_shmem_4tpt`; only the token-byte source -// changes. #ifndef WARPS_PER_BLOCK_MAX #define WARPS_PER_BLOCK_MAX 16u @@ -30,7 +39,7 @@ #endif #define WARP_BUF_BYTES 2080u -__device__ inline uint32_t warp_inclusive_scan_u32_s8r(uint32_t x, int lane) { +__device__ inline uint32_t onpair_warp_inclusive_scan_u32(uint32_t x, int lane) { constexpr unsigned mask = 0xffffffffu; #pragma unroll for (int offset = 1; offset < 32; offset <<= 1) { @@ -42,64 +51,71 @@ __device__ inline uint32_t warp_inclusive_scan_u32_s8r(uint32_t x, int lane) { return x; } -extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( - 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) { - constexpr unsigned mask = 0xffffffffu; - const int 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 uint64_t base_i = chunk * 128u + (uint64_t)lane; +// 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]; - uint32_t c[4]; - uint32_t l[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. +__device__ inline OnPairTokens onpair_load_tokens(const uint16_t *__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 (int 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]; - c[k] = code; - lo[k] = *reinterpret_cast(dict_s8 + (size_t)code * 8u); - l[k] = (uint32_t)lens[code]; + t.code[k] = code; + t.lo[k] = *reinterpret_cast(dict_s8 + (size_t)code * 8u); + t.len[k] = (uint32_t)lens[code]; } else { - c[k] = 0u; - lo[k] = make_uint2(0u, 0u); - l[k] = 0u; + t.code[k] = 0u; + t.lo[k] = make_uint2(0u, 0u); + t.len[k] = 0u; } } + return t; +} - uint32_t excl[4]; +// 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], int lane, + uint32_t (&excl)[4]) { + constexpr unsigned mask = 0xffffffffu; uint32_t acc_base = 0u; #pragma unroll for (int k = 0; k < 4; ++k) { - const uint32_t incl = warp_inclusive_scan_u32_s8r(l[k], lane); - excl[k] = acc_base + (incl - l[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); } - const uint32_t warp_total = acc_base; - - 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); + 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 (int k = 0; k < 4; ++k) { - const uint32_t len = l[k]; + const uint32_t len = t.len[k]; if (len == 0u) { continue; } const uint32_t base = excl[k]; - const uint8_t *lob = reinterpret_cast(&lo[k]); + const uint8_t *lob = reinterpret_cast(&t.lo[k]); const uint32_t nlo = len < 8u ? len : 8u; #pragma unroll for (int j = 0; j < 8; ++j) { @@ -110,7 +126,7 @@ extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( if (len > 8u) { // Rare path: high bytes from the full padded dict. const uint2 hi = - *reinterpret_cast(dict_padded + (size_t)c[k] * 16u + 8u); + *reinterpret_cast(dict_padded + (size_t)t.code[k] * 16u + 8u); const uint8_t *hib = reinterpret_cast(&hi); #pragma unroll for (int j = 0; j < 8; ++j) { @@ -120,8 +136,15 @@ extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( } } } - __syncwarp(); +} +// 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, int lane) { const uint32_t head = head_pre < warp_total ? head_pre : warp_total; if ((uint32_t)lane < head) { output_bytes[out_start + (uint64_t)lane] = s_buf[lane]; @@ -131,7 +154,7 @@ extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( } const uint32_t body_chunks = (warp_total - head) >> 4; - for (uint32_t k = lane; k < body_chunks; k += 32u) { + for (uint32_t k = (uint32_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); @@ -143,3 +166,37 @@ extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( s_buf[tail_start + lane]; } } + +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( + 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) { + const int 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); +} diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 5cca4c2e5f1..2b2ee76cfac 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -8,79 +8,13 @@ use std::ffi::c_void; use cudarc::driver::CudaSlice; use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; -use vortex::array::buffer::BufferHandle; -use vortex::dtype::PType; use vortex::error::VortexResult; -use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex_cub::scan; -use vortex_cub::scan::LengthType; use vortex_cub::scan::cudaStream_t; -use crate::CudaBufferExt; use crate::CudaExecutionCtx; -fn length_type(ptype: PType) -> VortexResult { - Ok(match ptype { - PType::U8 => LengthType::U8, - PType::I8 => LengthType::I8, - PType::U16 => LengthType::U16, - PType::I16 => LengthType::I16, - PType::U32 => LengthType::U32, - PType::I32 => LengthType::I32, - PType::U64 => LengthType::U64, - PType::I64 => LengthType::I64, - other => vortex_bail!("unsupported lengths ptype {other} for fused exclusive sum"), - }) -} - -/// Fused widen + CUB `DeviceScan::ExclusiveSum` over device-resident per-row -/// lengths of any integer width. -/// -/// One CUB dispatch scans the lengths through a widening transform iterator -/// and produces `num_rows + 1` u64 offsets whose last element is the total — -/// no separate widen kernel and no materialized u64 input. A negative length -/// raises `status` to 2 and contributes zero bytes; the caller must check the -/// flag before trusting the offsets. -pub(crate) fn exclusive_sum_lengths_u64( - lengths: &BufferHandle, - ptype: PType, - num_rows: usize, - status: &mut CudaSlice, - ctx: &mut CudaExecutionCtx, -) -> VortexResult> { - let ty = length_type(ptype)?; - let num_offsets = i64::try_from(num_rows + 1)?; - let temp_bytes = scan::exclusive_sum_lengths_temp_size(ty, num_offsets) - .map_err(|err| vortex_err!("CUB scan_exclusive_sum_lengths_temp_size failed: {err}"))?; - - let mut temp = ctx.device_alloc::(temp_bytes.max(1))?; - let mut output = ctx.device_alloc::(num_rows + 1)?; - let lengths_ptr = lengths.cuda_device_ptr()?; - let stream = ctx.stream(); - let stream_ptr = stream.cu_stream() as cudaStream_t; - let (status_ptr, record_status) = status.device_ptr_mut(stream); - let (output_ptr, record_output) = output.device_ptr_mut(stream); - let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); - - ctx.launch_external(num_rows + 1, || unsafe { - scan::exclusive_sum_lengths( - ty, - temp_ptr as *mut c_void, - temp_bytes, - lengths_ptr as *const c_void, - output_ptr as *mut u64, - status_ptr as *mut u32, - num_offsets, - stream_ptr, - ) - .map_err(|err| vortex_err!("CUB scan_exclusive_sum_lengths failed: {err}")) - })?; - drop((record_status, record_output, record_temp)); - - Ok(output) -} - /// CUB `DeviceScan::ExclusiveSum` over device-resident `u64` values. /// /// Runs through the `i64` CUB instantiation: callers pass non-negative counts diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index 12321c4110e..2ac05d593c2 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -3,23 +3,23 @@ //! CUDA executor for OnPair decompression. //! -//! Decoding runs entirely on the GPU over the flat token stream: +//! Decoding runs on the GPU over the flat token stream: //! -//! 1. A fused CUB exclusive scan (widening transform iterator over the -//! native-width lengths) turns the per-row decoded lengths into per-row -//! output offsets in one dispatch (the last element is the total decoded -//! byte count). -//! 2. `onpair_batch_sizes` reduces the decoded byte size of every 128-token -//! batch from the codes and the per-token length LUT, and a second CUB -//! exclusive scan regenerates the per-batch output offsets -//! (`chunk_offsets`) the decode kernel positions its writes with. -//! 3. `onpair_shmem_4tpt_split8read` gathers each token's bytes from the +//! 1. `onpair_batch_sizes` reduces the decoded byte size of every 128-token +//! batch from the codes and the per-token length LUT, and a CUB exclusive +//! scan regenerates the per-batch output offsets (`chunk_offsets`) the +//! decode kernel positions its writes with. +//! 2. `onpair_shmem_4tpt_split8read` gathers each token's bytes from the //! split dictionary layout and scatters them to the output byte stream. //! -//! The result is exposed either as a canonical `VarBinView` (views built -//! on-device by `onpair_build_views`, or on host for heaps that exceed a -//! single backing buffer) or as Arrow-compatible i32 offsets plus values via -//! [`decode_onpair_varbin`], mirroring the FSST varbin path. +//! The per-row lengths are only summed on the host (sizing the output and +//! cross-checking the code stream); per-row offsets are built solely by the +//! output path that needs them. The result is exposed either as a canonical +//! `VarBinView` (views built on-device by `onpair_build_views` from +//! host-prefix-summed offsets, or on host for heaps that exceed a single +//! backing buffer) or as Arrow-compatible i32 offsets plus values via +//! [`decode_onpair_varbin`], which builds the offsets on device with +//! [`i32_offsets_from_lengths`] — mirroring the FSST varbin path. use std::fmt::Debug; use std::sync::Arc; @@ -34,7 +34,6 @@ 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; @@ -61,7 +60,8 @@ use vortex_onpair::dict_view; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; -use crate::cub::exclusive_sum_lengths_u64; +use crate::arrow::I32Offsets; +use crate::arrow::i32_offsets_from_lengths; use crate::cub::exclusive_sum_u64; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; @@ -77,13 +77,6 @@ const TOKENS_PER_BATCH: usize = 128; const BLOCK_THREADS: u32 = 512; const WARPS_PER_BLOCK: usize = (BLOCK_THREADS / 32) as usize; -/// `status` value raised by `onpair_batch_sizes` for a code outside the -/// dictionary. -const STATUS_CODE_OUT_OF_RANGE: u32 = 1; -/// `status` value raised by the fused lengths scan for a negative decoded -/// length. -const STATUS_NEGATIVE_LENGTH: u32 = 2; - /// 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))?; @@ -157,14 +150,14 @@ struct StagedCodes { /// The shared result of the OnPair GPU decode pipeline. struct OnPairDecoded { - /// Exclusive per-row byte offsets over `bytes`, `num_rows + 1` entries on - /// device; the last is `total_size`. - row_offsets: CudaSlice, /// The flat decoded byte stream. bytes: CudaSlice, /// Total decoded byte count. total_size: usize, - /// Host-resident per-row lengths, for the host rollover path. + /// Host-resident per-row lengths. Each output path derives what it needs: + /// the views fast path prefix-sums them into row offsets, the varbin path + /// builds Arrow i32 offsets from them on device, and the rollover path + /// consumes them directly. lengths: PrimitiveArray, } @@ -250,42 +243,37 @@ async fn stage_codes( }) } -/// Run the shared OnPair GPU decode pipeline: regenerate the row and batch -/// output offsets on the device, validate the compressed stream, and decode -/// the flat byte stream. Returns `Ok(None)` when the array decodes to zero -/// bytes. +/// Run the OnPair decode pipeline: sum the per-row lengths on host, +/// regenerate the per-batch output offsets on the device, validate the +/// compressed stream, and decode the flat byte stream. Returns `Ok(None)` +/// when the array decodes to zero bytes. async fn decode_onpair_bytes( onpair: &OnPairArray, ctx: &mut CudaExecutionCtx, ) -> VortexResult> { let num_rows = onpair.len(); - // Decompress the lengths child; its native-width buffer feeds the GPU - // row-offsets scan, and the host values serve the rare rollover path. + // Sum the per-row decoded lengths on host — they are materialised there + // anyway. The total sizes the output allocation and cross-checks the code + // stream; per-row offsets are only built by the output paths that need + // them. let lengths = onpair .uncompressed_lengths() .clone() .execute::(ctx.execution_ctx())?; - let lengths_ptype = lengths.ptype(); - let PrimitiveDataParts { - buffer: lengths_buffer, - .. - } = lengths.clone().into_data_parts(); - let lengths_dev = ctx.ensure_on_device(lengths_buffer).await?; - - // Shared corruption flag for the size kernels; 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}"))?; - - // Row offsets in one fused GPU dispatch: a CUB exclusive scan over a - // widening transform iterator of the native-width lengths. The last of - // the `num_rows + 1` outputs is the total decoded byte count; a negative - // length raises `status` and contributes zero bytes. - let row_offsets = - exclusive_sum_lengths_u64(&lengths_dev, lengths_ptype, num_rows, &mut status, ctx)?; + let total_size: u64 = match_each_integer_ptype!(lengths.ptype(), |P| { + let mut acc = 0u64; + #[allow(clippy::unnecessary_cast)] + for &length in lengths.as_slice::

() { + let length = u64::try_from(length as i128) + .map_err(|_| vortex_err!("OnPair uncompressed length cannot be negative"))?; + acc = acc + .checked_add(length) + .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; + } + VortexResult::Ok(acc) + })?; + let total_size = usize::try_from(total_size)?; // `codes_offsets` may be a sliced view of the original; its first and last // boundaries bound the contiguous run of `codes` belonging to this array's @@ -303,59 +291,54 @@ async fn decode_onpair_bytes( onpair.codes().len() ); - let mut staged = None; - if code_start < code_end { - staged = Some(stage_codes(onpair, code_start, code_end, &status, ctx).await?); + if total_size == 0 { + // Every token decodes to at least one byte. + vortex_ensure!( + code_start == code_end, + "OnPair records zero decoded bytes but has codes" + ); + return Ok(None); } + vortex_ensure!( + code_start < code_end, + "OnPair records {total_size} decoded bytes but has no codes" + ); + + // 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, code_start, code_end, &status, ctx).await?; // One synchronizing readback validates the compressed stream before the // decode kernel — whose dictionary gathers and output scatters are - // unchecked — is allowed to run: no negative length, every code indexed - // the dictionary, and the codes decode to exactly the byte count the - // lengths record. + // unchecked — is allowed to run: every code indexed the dictionary, and + // the codes decode to exactly the byte count the lengths record. let status = ctx .stream() .clone_dtoh(&status) .map_err(|e| vortex_err!("Failed to copy OnPair status flag to host: {e}"))?; - match status.first().copied().unwrap_or(STATUS_CODE_OUT_OF_RANGE) { - 0 => {} - STATUS_NEGATIVE_LENGTH => { - vortex_bail!("OnPair uncompressed length cannot be negative") - } - _ => vortex_bail!("OnPair code out of dictionary range"), + if status.first().copied().unwrap_or(1) != 0 { + vortex_bail!("OnPair code out of dictionary range"); } - let row_total = ctx + let chunk_total = ctx .stream() - .clone_dtoh(&row_offsets.slice(num_rows..num_rows + 1)) + .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 row offset scan returned no total"))?; - let chunk_total = match &staged { - Some(staged) => 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"))?, - None => 0, - }; + .ok_or_else(|| vortex_err!("OnPair batch offset scan returned no total"))?; vortex_ensure!( - row_total == chunk_total, - "OnPair codes decode to {chunk_total} bytes but uncompressed_lengths records {row_total}" + chunk_total == total_size as u64, + "OnPair codes decode to {chunk_total} bytes but uncompressed_lengths records {total_size}" ); - if row_total == 0 { - return Ok(None); - } - let total_size = usize::try_from(row_total)?; - let Some(staged) = staged else { - vortex_bail!("OnPair records {total_size} decoded bytes but has no codes"); - }; // Decode. The kernel's drain gates 16-byte stores on `out_start % 16` // relative to the buffer base, so the base must be 16-aligned. @@ -389,7 +372,6 @@ async fn decode_onpair_bytes( )?; Ok(Some(OnPairDecoded { - row_offsets, bytes, total_size, lengths, @@ -413,20 +395,34 @@ async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> Vorte return empty_views(num_rows, dtype, validity, ctx).await; }; let OnPairDecoded { - row_offsets, bytes, total_size, lengths, } = decoded; // Fast path: the decoded heap fits a single BinaryView backing buffer, so - // the per-row views build on-device straight from the row offsets. + // the per-row views build on-device. Only this path needs the u64 row + // offsets: prefix-sum the lengths here (negatives were already rejected + // by the total sum) and stage them on device. if total_size <= MAX_BUFFER_LEN { + let row_offsets: Vec = match_each_integer_ptype!(lengths.ptype(), |P| { + let mut offsets = Vec::with_capacity(lengths.len() + 1); + let mut acc = 0u64; + offsets.push(0u64); + #[allow(clippy::unnecessary_cast)] + for &length in lengths.as_slice::

() { + acc += length as u64; + offsets.push(acc); + } + offsets + }); + let row_offsets_dev = ctx.copy_to_device(row_offsets)?.await?; + let row_offsets_view = row_offsets_dev.cuda_view::()?; let 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) + args.arg(&row_offsets_view) .arg(&bytes) .arg(&device_views) .arg(&num_rows_u64); @@ -490,26 +486,18 @@ pub(crate) async fn decode_onpair_varbin( }); }; - vortex_ensure!( - i32::try_from(decoded.total_size).is_ok(), - "OnPair decoded size exceeds Arrow i32 offset range" - ); - - // Narrow the device row offsets to Arrow's i32 offsets; every value fits - // because the total was just checked and offsets are nondecreasing. - let arrow_offsets = ctx.device_alloc::(len + 1)?; - let num_offsets_u64 = u64::try_from(len + 1)?; - let offsets_fn = ctx.load_function_with_suffixes("onpair", &["offsets_to_i32"])?; - ctx.launch_kernel(&offsets_fn, len + 1, |args| { - args.arg(&decoded.row_offsets) - .arg(&arrow_offsets) - .arg(&num_offsets_u64); - })?; + // Build the Arrow i32 offsets from the lengths on device; this also + // rejects heaps beyond Arrow's i32 offset range. + let I32Offsets { + buffer: offsets, + total, + } = i32_offsets_from_lengths(decoded.lengths.clone(), ctx).await?; + debug_assert_eq!(total, decoded.total_size); Ok(DecodedVarBin { dtype, len, - offsets: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(arrow_offsets))), + offsets, values: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(decoded.bytes))), validity, }) From 2aaa23106a56f0b9dc7618cb2384b14edeffe5e1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 22:39:00 +0100 Subject: [PATCH 05/11] formatted Signed-off-by: Robert Kruszewski --- vortex-cuda/kernels/src/onpair.cu | 15 +++++--- .../src/onpair_shmem_4tpt_split8read.cu | 38 ++++++++++--------- 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu index c8e8b073afc..08377bea86c 100644 --- a/vortex-cuda/kernels/src/onpair.cu +++ b/vortex-cuda/kernels/src/onpair.cu @@ -28,7 +28,8 @@ constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; // launching the decode kernel, whose dictionary gathers are unchecked. extern "C" __global__ void onpair_batch_sizes(const uint16_t *__restrict codes, const uint8_t *__restrict lens, - uint32_t dict_size, uint64_t total_tokens, + uint32_t dict_size, + uint64_t total_tokens, uint64_t *__restrict batch_sizes, uint32_t *__restrict status) { const int lane = threadIdx.x & 31; @@ -71,7 +72,8 @@ constexpr uint32_t MAX_INLINED_SIZE = 12; // heap is exposed as backing buffer zero. __device__ inline void onpair_write_view(const uint64_t *__restrict row_offsets, const uint8_t *__restrict output_bytes, - uint4 *__restrict views, uint64_t rid) { + uint4 *__restrict views, + uint64_t rid) { const uint64_t start = row_offsets[rid]; const uint32_t len = (uint32_t)(row_offsets[rid + 1] - start); if (len <= MAX_INLINED_SIZE) { @@ -86,15 +88,16 @@ __device__ inline void onpair_write_view(const uint64_t *__restrict row_offsets, 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); + 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, (uint32_t)start); } extern "C" __global__ void onpair_build_views(const uint64_t *__restrict row_offsets, const uint8_t *__restrict output_bytes, - uint4 *__restrict views, uint64_t num_rows) { + 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 = diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu index c53c87ff170..17db978923a 100644 --- a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -65,7 +65,8 @@ struct OnPairTokens { // past the end of the stream load as empty. __device__ inline OnPairTokens onpair_load_tokens(const uint16_t *__restrict codes, const uint8_t *__restrict dict_s8, - const uint8_t *__restrict lens, uint64_t base_i, + const uint8_t *__restrict lens, + uint64_t base_i, uint64_t total_tokens) { OnPairTokens t; #pragma unroll @@ -88,8 +89,7 @@ __device__ inline OnPairTokens onpair_load_tokens(const uint16_t *__restrict cod // 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], int lane, - uint32_t (&excl)[4]) { +__device__ inline uint32_t onpair_scan_offsets(const uint32_t (&len)[4], int lane, uint32_t (&excl)[4]) { constexpr unsigned mask = 0xffffffffu; uint32_t acc_base = 0u; #pragma unroll @@ -105,7 +105,8 @@ __device__ inline uint32_t onpair_scan_offsets(const uint32_t (&len)[4], int lan // 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], +__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 @@ -125,8 +126,7 @@ __device__ inline void onpair_stage_tokens(const OnPairTokens &t, const uint32_t } 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 uint2 hi = *reinterpret_cast(dict_padded + (size_t)t.code[k] * 16u + 8u); const uint8_t *hib = reinterpret_cast(&hi); #pragma unroll for (int j = 0; j < 8; ++j) { @@ -143,8 +143,11 @@ __device__ inline void onpair_stage_tokens(const OnPairTokens &t, const uint32_t // 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, int lane) { + uint8_t *__restrict output_bytes, + uint64_t out_start, + uint32_t head_pre, + uint32_t warp_total, + int lane) { const uint32_t head = head_pre < warp_total ? head_pre : warp_total; if ((uint32_t)lane < head) { output_bytes[out_start + (uint64_t)lane] = s_buf[lane]; @@ -162,20 +165,21 @@ __device__ inline void onpair_drain(const uint8_t *__restrict s_buf, const uint32_t tail_start = head + (body_chunks << 4); if ((uint32_t)lane < warp_total - tail_start) { - output_bytes[out_start + (uint64_t)tail_start + (uint64_t)lane] = - s_buf[tail_start + lane]; + output_bytes[out_start + (uint64_t)tail_start + (uint64_t)lane] = s_buf[tail_start + lane]; } } -extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void onpair_shmem_4tpt_split8read( - 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) { +extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void +onpair_shmem_4tpt_split8read(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) { const int 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; + const uint64_t chunk = (uint64_t)blockIdx.x * (uint64_t)(blockDim.x >> 5) + (uint64_t)warp_id; if (chunk * 128u >= total_tokens) { return; } From 43e97df965bb27573d80461483837393806b207d Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 12:18:26 +0100 Subject: [PATCH 06/11] less Signed-off-by: Robert Kruszewski --- vortex-cuda/src/arrow/canonical.rs | 37 ++++++++---------------------- 1 file changed, 10 insertions(+), 27 deletions(-) diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index f6ddf6ea78f..f6180d162e2 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -55,14 +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 vortex_onpair::OnPairArray; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; @@ -232,7 +230,8 @@ 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, @@ -242,7 +241,8 @@ fn export_array( Ok(onpair) if ctx.cuda_session().varbin_export_layout() == VarBinExportLayout::VarBin => { - return export_onpair_varbin(onpair, ctx).await; + let decoded = decode_onpair_varbin(onpair, ctx).await?; + return export_decoded_varbin(decoded, ctx).await; } Ok(onpair) => onpair.into_array(), Err(array) => array, @@ -594,8 +594,10 @@ 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 DecodedVarBin { @@ -604,29 +606,10 @@ async fn export_fsst_varbin( 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}" - ); - 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) -} - -async fn export_onpair_varbin( - onpair: OnPairArray, - ctx: &mut CudaExecutionCtx, -) -> VortexResult<(ArrowArray, SyncEvent)> { - let DecodedVarBin { - dtype, - len, - offsets, - values, - validity, - } = decode_onpair_varbin(onpair, ctx).await?; - vortex_ensure!( - matches!(dtype, DType::Utf8(_) | DType::Binary(_)), - "OnPair 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) From 01b2414a4a38aafd10e8feee6b1cac9c78841d2c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 14:50:41 +0100 Subject: [PATCH 07/11] less Signed-off-by: Robert Kruszewski --- vortex-cuda/cub/build.rs | 1 + vortex-cuda/cub/kernels/filter.h | 18 ++ vortex-cuda/cub/kernels/onpair.cu | 191 ++++++++++++++++++ vortex-cuda/cub/src/lib.rs | 1 + vortex-cuda/cub/src/onpair.rs | 68 +++++++ vortex-cuda/kernels/src/onpair.cu | 61 +----- .../src/onpair_shmem_4tpt_split8read.cu | 30 +-- vortex-cuda/src/cub.rs | 59 ++++-- vortex-cuda/src/kernel/encodings/onpair.rs | 74 +++---- 9 files changed, 365 insertions(+), 138 deletions(-) create mode 100644 vortex-cuda/cub/kernels/onpair.cu create mode 100644 vortex-cuda/cub/src/onpair.rs 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..384dfa7bcb7 100644 --- a/vortex-cuda/cub/kernels/filter.h +++ b/vortex-cuda/cub/kernels/filter.h @@ -99,6 +99,24 @@ 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. 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 uint16_t *codes, + 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..b95eefba6b5 --- /dev/null +++ b/vortex-cuda/cub/kernels/onpair.cu @@ -0,0 +1,191 @@ +// 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. 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. +__global__ +__launch_bounds__(ONPAIR_BLOCK_THREADS) void onpair_batch_offsets_sweep(const uint16_t *__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 s = 0; + if (batch < num_batches) { + const uint64_t base = (uint64_t)batch * (uint64_t)ONPAIR_TOKENS_PER_BATCH; +#pragma unroll + for (uint8_t k = 0; k < 4; ++k) { + const uint64_t i = base + (uint64_t)lane + (uint64_t)(k * 32); + if (i < total_tokens) { + const uint32_t code = (uint32_t)codes[i]; + if (code < dict_size) { + s += (uint32_t)lens[code]; + } else { + atomicMax(status, 1u); + } + } + } + } +#pragma unroll + for (uint8_t offset = 16; offset > 0; offset >>= 1) { + s += __shfl_down_sync(0xffffffffu, s, 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)s; + } + __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 v = (lane < ONPAIR_WARPS_PER_BLOCK) ? warp_sums[lane] : 0; + uint64_t incl = v; +#pragma unroll + for (uint8_t offset = 1; offset < 32; offset <<= 1) { + const uint64_t y = __shfl_up_sync(0xffffffffu, incl, offset); + if (lane >= offset) { + incl += y; + } + } + const uint64_t aggregate = __shfl_sync(0xffffffffu, incl, 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. + OnPairPrefixOp prefix_op(tile_state, prefix_storage, SumU64(), tile_idx); + prefix = prefix_op(aggregate); + } + if (lane < ONPAIR_WARPS_PER_BLOCK) { + warp_excl[lane] = prefix + incl - v; + } + } + __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. 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 uint16_t *codes, + 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; + } + + onpair_batch_offsets_sweep<<>>(codes, + lens, + dict_size, + total_tokens, + chunk_offsets, + status, + tile_state, + num_batches); + 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..0c92f1e34ed --- /dev/null +++ b/vortex-cuda/cub/src/onpair.rs @@ -0,0 +1,68 @@ +// 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. +/// 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` `u16` values, 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 u16, + 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, + 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 index 08377bea86c..37e1765f100 100644 --- a/vortex-cuda/kernels/src/onpair.cu +++ b/vortex-cuda/kernels/src/onpair.cu @@ -5,61 +5,10 @@ #include -// Support kernels for OnPair GPU decompression. -// -// The decode kernel (`onpair_shmem_4tpt_split8read.cu`) consumes per-batch output -// offsets: `chunk_offsets[b]` is the count of decoded bytes preceding the b-th -// 128-token batch. Vortex does not store those offsets; they are regenerated on -// the GPU at decode time. `onpair_batch_sizes` reduces each batch's decoded size -// from the codes and the per-token length LUT, and a CUB exclusive scan over the -// result yields `chunk_offsets`. This scans only the compressed codes — it -// touches neither the dictionary bytes nor the output. - -// Tokens per decode batch: one warp of the decode kernel emits 128 tokens -// (4 tokens/thread). Must match the decode kernel's layout. -constexpr uint32_t ONPAIR_TOKENS_PER_BATCH = 128; - -// One warp per 128-token batch: sums `lens[codes[t]]` over the batch's (up to) -// 128 tokens and writes the total to `batch_sizes[b]`. Code reads are -// lane-consecutive (coalesced); the length LUT is small and cache-resident. -// -// A code outside the dictionary raises `status` to 1 and contributes zero -// bytes: the host must check the flag before trusting `batch_sizes` and before -// launching the decode kernel, whose dictionary gathers are unchecked. -extern "C" __global__ void onpair_batch_sizes(const uint16_t *__restrict codes, - const uint8_t *__restrict lens, - uint32_t dict_size, - uint64_t total_tokens, - uint64_t *__restrict batch_sizes, - uint32_t *__restrict status) { - const int lane = threadIdx.x & 31; - const uint32_t warp = threadIdx.x >> 5; - const uint64_t b = (uint64_t)blockIdx.x * (uint64_t)(blockDim.x >> 5) + (uint64_t)warp; - const uint64_t base = b * (uint64_t)ONPAIR_TOKENS_PER_BATCH; - if (base >= total_tokens) { - return; - } - uint32_t s = 0; -#pragma unroll - for (int k = 0; k < 4; ++k) { - const uint64_t i = base + (uint64_t)lane + (uint64_t)(k * 32); - if (i < total_tokens) { - const uint32_t code = (uint32_t)codes[i]; - if (code < dict_size) { - s += (uint32_t)lens[code]; - } else { - atomicMax(status, 1u); - } - } - } -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - s += __shfl_down_sync(0xffffffffu, s, offset); - } - if (lane == 0) { - batch_sizes[b] = (uint64_t)s; - } -} +// 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 +// view-construction kernels launched after the decode. // 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 @@ -79,7 +28,7 @@ __device__ inline void onpair_write_view(const uint64_t *__restrict row_offsets, if (len <= MAX_INLINED_SIZE) { uint32_t words[3] = {0, 0, 0}; #pragma unroll - for (uint32_t i = 0; i < MAX_INLINED_SIZE; i++) { + 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)); } diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu index 17db978923a..e3fa94b51c6 100644 --- a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -39,10 +39,10 @@ #endif #define WARP_BUF_BYTES 2080u -__device__ inline uint32_t onpair_warp_inclusive_scan_u32(uint32_t x, int lane) { +__device__ inline uint32_t onpair_warp_inclusive_scan_u32(uint32_t x, uint32_t lane) { constexpr unsigned mask = 0xffffffffu; #pragma unroll - for (int offset = 1; offset < 32; offset <<= 1) { + for (uint8_t offset = 1; offset < 32; offset <<= 1) { uint32_t y = __shfl_up_sync(mask, x, offset); if (lane >= offset) { x += y; @@ -70,7 +70,7 @@ __device__ inline OnPairTokens onpair_load_tokens(const uint16_t *__restrict cod uint64_t total_tokens) { OnPairTokens t; #pragma unroll - for (int k = 0; k < 4; ++k) { + 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]; @@ -89,11 +89,11 @@ __device__ inline OnPairTokens onpair_load_tokens(const uint16_t *__restrict cod // 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], int lane, uint32_t (&excl)[4]) { +__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 (int k = 0; k < 4; ++k) { + 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); @@ -110,7 +110,7 @@ __device__ inline void onpair_stage_tokens(const OnPairTokens &t, const uint8_t *__restrict dict_padded, uint8_t *__restrict s_buf) { #pragma unroll - for (int k = 0; k < 4; ++k) { + for (uint8_t k = 0; k < 4; ++k) { const uint32_t len = t.len[k]; if (len == 0u) { continue; @@ -119,8 +119,8 @@ __device__ inline void onpair_stage_tokens(const OnPairTokens &t, const uint8_t *lob = reinterpret_cast(&t.lo[k]); const uint32_t nlo = len < 8u ? len : 8u; #pragma unroll - for (int j = 0; j < 8; ++j) { - if (j < (int)nlo) { + for (uint8_t j = 0; j < 8; ++j) { + if (j < nlo) { s_buf[base + j] = lob[j]; } } @@ -129,8 +129,8 @@ __device__ inline void onpair_stage_tokens(const OnPairTokens &t, const uint2 hi = *reinterpret_cast(dict_padded + (size_t)t.code[k] * 16u + 8u); const uint8_t *hib = reinterpret_cast(&hi); #pragma unroll - for (int j = 0; j < 8; ++j) { - if (8 + j < (int)len) { + for (uint8_t j = 0; j < 8; ++j) { + if (8u + j < len) { s_buf[base + 8 + j] = hib[j]; } } @@ -147,9 +147,9 @@ __device__ inline void onpair_drain(const uint8_t *__restrict s_buf, uint64_t out_start, uint32_t head_pre, uint32_t warp_total, - int lane) { + uint32_t lane) { const uint32_t head = head_pre < warp_total ? head_pre : warp_total; - if ((uint32_t)lane < head) { + if (lane < head) { output_bytes[out_start + (uint64_t)lane] = s_buf[lane]; } if (head >= warp_total) { @@ -157,14 +157,14 @@ __device__ inline void onpair_drain(const uint8_t *__restrict s_buf, } const uint32_t body_chunks = (warp_total - head) >> 4; - for (uint32_t k = (uint32_t)lane; k < body_chunks; k += 32u) { + 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 ((uint32_t)lane < warp_total - tail_start) { + if (lane < warp_total - tail_start) { output_bytes[out_start + (uint64_t)tail_start + (uint64_t)lane] = s_buf[tail_start + lane]; } } @@ -177,7 +177,7 @@ onpair_shmem_4tpt_split8read(const uint16_t *__restrict codes, const uint8_t *__restrict lens, uint8_t *__restrict output_bytes, uint64_t total_tokens) { - const int lane = threadIdx.x & 31; + 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) { diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 2b2ee76cfac..4910455cdce 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -8,49 +8,64 @@ 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; -/// CUB `DeviceScan::ExclusiveSum` over device-resident `u64` values. -/// -/// Runs through the `i64` CUB instantiation: callers pass non-negative counts -/// whose prefix sums stay below `i64::MAX`, where two's complement `i64` and -/// `u64` addition produce identical bit patterns. -pub(crate) fn exclusive_sum_u64( - input: &CudaSlice, - len: usize, +/// 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. 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. +pub(crate) fn onpair_batch_offsets( + codes: &BufferHandle, + lens: &BufferHandle, + dict_size: u32, + num_tokens: usize, + num_batches: usize, + status: &mut CudaSlice, ctx: &mut CudaExecutionCtx, ) -> VortexResult> { - let len_i64 = i64::try_from(len)?; - let temp_bytes = scan::exclusive_sum_i64_temp_size(len_i64) - .map_err(|err| vortex_err!("CUB scan_exclusive_sum_i64_temp_size failed: {err}"))?; + 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 output = ctx.device_alloc::(len)?; + 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 (input_ptr, record_input) = input.device_ptr(stream); - let (output_ptr, record_output) = output.device_ptr_mut(stream); + 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(len, || unsafe { - scan::exclusive_sum_i64( + ctx.launch_external(num_tokens, || unsafe { + onpair::batch_offsets( temp_ptr as *mut c_void, temp_bytes, - input_ptr as *const i64, - output_ptr as *mut i64, - len_i64, + codes_ptr as *const u16, + 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 scan_exclusive_sum_i64 failed: {err}")) + .map_err(|err| vortex_err!("CUB onpair_batch_offsets failed: {err}")) })?; - drop((record_input, record_output, record_temp)); + drop((record_status, record_offsets, record_temp)); - Ok(output) + Ok(chunk_offsets) } pub(crate) fn exclusive_sum_i32( diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index 2ac05d593c2..8451d18714f 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -5,10 +5,10 @@ //! //! Decoding runs on the GPU over the flat token stream: //! -//! 1. `onpair_batch_sizes` reduces the decoded byte size of every 128-token -//! batch from the codes and the per-token length LUT, and a CUB exclusive -//! scan regenerates the per-batch output offsets (`chunk_offsets`) the -//! decode kernel positions its writes with. +//! 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_shmem_4tpt_split8read` gathers each token's bytes from the //! split dictionary layout and scatters them to the output byte stream. //! @@ -62,7 +62,7 @@ use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::arrow::I32Offsets; use crate::arrow::i32_offsets_from_lengths; -use crate::cub::exclusive_sum_u64; +use crate::cub::onpair_batch_offsets; use crate::executor::CudaExecute; use crate::executor::CudaExecutionCtx; use crate::kernel::encodings::DecodedVarBin; @@ -71,7 +71,7 @@ use crate::kernel::encodings::DecodedVarBin; 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 `kernels/src/onpair.cu`. +/// 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; @@ -162,13 +162,13 @@ struct OnPairDecoded { } /// Stage this array's code window and dictionary on the device and regenerate -/// the decode kernel's per-batch output offsets from them (steps 1–2 of the -/// pipeline: `onpair_batch_sizes` + CUB exclusive scan). +/// the decode kernel's per-batch output offsets from them in one fused sweep +/// (see [`onpair_batch_offsets`]). async fn stage_codes( onpair: &OnPairArray, code_start: usize, code_end: usize, - status: &CudaSlice, + status: &mut CudaSlice, ctx: &mut CudaExecutionCtx, ) -> VortexResult { // Widen this array's code window to the decode kernel's u16 ABI. @@ -179,7 +179,6 @@ async fn stage_codes( .execute::(ctx.execution_ctx())? .into_buffer::(); let num_tokens = codes.len(); - let num_tokens_u64 = u64::try_from(num_tokens)?; // 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 @@ -209,27 +208,15 @@ async fn stage_codes( let num_batches = num_tokens.div_ceil(TOKENS_PER_BATCH); let launch_config = batch_launch_config(num_batches)?; - - // Per-batch decoded sizes. One extra zeroed slot makes the exclusive - // scan's last element the total decoded byte count. - let mut batch_sizes = ctx.device_alloc::(num_batches + 1)?; - ctx.stream() - .memset_zeros(&mut batch_sizes) - .map_err(|e| vortex_err!("Failed to zero OnPair batch sizes: {e}"))?; - - let codes_view = codes_dev.cuda_view::()?; - let lens_view = lens_dev.cuda_view::()?; - let batch_sizes_fn = ctx.load_function_with_suffixes("onpair", &["batch_sizes"])?; - ctx.launch_kernel_config(&batch_sizes_fn, launch_config, num_tokens, |args| { - args.arg(&codes_view) - .arg(&lens_view) - .arg(&dict_size_u32) - .arg(&num_tokens_u64) - .arg(&batch_sizes) - .arg(status); - })?; - - let chunk_offsets = exclusive_sum_u64(&batch_sizes, num_batches + 1, ctx)?; + let chunk_offsets = onpair_batch_offsets( + &codes_dev, + &lens_dev, + dict_size_u32, + num_tokens, + num_batches, + status, + ctx, + )?; Ok(StagedCodes { codes: codes_dev, @@ -261,19 +248,16 @@ async fn decode_onpair_bytes( .uncompressed_lengths() .clone() .execute::(ctx.execution_ctx())?; - let total_size: u64 = match_each_integer_ptype!(lengths.ptype(), |P| { - let mut acc = 0u64; - #[allow(clippy::unnecessary_cast)] - for &length in lengths.as_slice::

() { - let length = u64::try_from(length as i128) - .map_err(|_| vortex_err!("OnPair uncompressed length cannot be negative"))?; - acc = acc - .checked_add(length) - .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; - } - VortexResult::Ok(acc) - })?; - let total_size = usize::try_from(total_size)?; + // The conversion is fallible only for the wide/signed ptype instantiations + // of the macro; u8/u16 make it infallible, so allow the lint wholesale. + #[allow(clippy::unnecessary_fallible_conversions)] + let total_size = match_each_integer_ptype!(lengths.ptype(), |P| { + lengths + .as_slice::

() + .iter() + .map(|&v| usize::try_from(v).vortex_expect("length must fit in usize")) + .sum() + }); // `codes_offsets` may be a sliced view of the original; its first and last // boundaries bound the contiguous run of `codes` belonging to this array's @@ -311,7 +295,7 @@ async fn decode_onpair_bytes( .memset_zeros(&mut status) .map_err(|e| vortex_err!("Failed to zero OnPair status flag: {e}"))?; - let staged = stage_codes(onpair, code_start, code_end, &status, ctx).await?; + let staged = stage_codes(onpair, code_start, code_end, &mut status, ctx).await?; // One synchronizing readback validates the compressed stream before the // decode kernel — whose dictionary gathers and output scatters are From f71366748a53d94c889b24c8bdda459033e447f3 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 15:17:12 +0100 Subject: [PATCH 08/11] fixes Signed-off-by: Robert Kruszewski --- encodings/onpair/src/decode.rs | 5 +- encodings/onpair/src/lib.rs | 1 + vortex-cuda/cub/kernels/onpair.cu | 30 ++-- vortex-cuda/kernels/src/onpair.cu | 2 +- .../src/onpair_shmem_4tpt_split8read.cu | 4 + vortex-cuda/src/cub.rs | 108 +++++++++++++++ vortex-cuda/src/kernel/encodings/onpair.rs | 130 ++++++++++-------- 7 files changed, 206 insertions(+), 74 deletions(-) diff --git a/encodings/onpair/src/decode.rs b/encodings/onpair/src/decode.rs index 96add100e36..3c16485c15a 100644 --- a/encodings/onpair/src/decode.rs +++ b/encodings/onpair/src/decode.rs @@ -35,8 +35,9 @@ pub(crate) fn collect_widened( /// 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( +/// start/end bounds, the CUDA decoder's code window) don't pay to materialise +/// every boundary. +pub fn code_boundary_at( codes_offsets: &ArrayRef, index: usize, ctx: &mut ExecutionCtx, diff --git a/encodings/onpair/src/lib.rs b/encodings/onpair/src/lib.rs index f9d8be5ff0d..bd6b150304a 100644 --- a/encodings/onpair/src/lib.rs +++ b/encodings/onpair/src/lib.rs @@ -22,6 +22,7 @@ mod tests; pub use array::*; pub use compress::*; +pub use decode::code_boundary_at; pub use onpair::CompactDictionaryView; pub use onpair::Config; pub use onpair::DEFAULT_CONFIG; diff --git a/vortex-cuda/cub/kernels/onpair.cu b/vortex-cuda/cub/kernels/onpair.cu index b95eefba6b5..4a11176055b 100644 --- a/vortex-cuda/cub/kernels/onpair.cu +++ b/vortex-cuda/cub/kernels/onpair.cu @@ -62,16 +62,16 @@ __launch_bounds__(ONPAIR_BLOCK_THREADS) void onpair_batch_offsets_sweep(const ui // 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 s = 0; + 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 k = 0; k < 4; ++k) { - const uint64_t i = base + (uint64_t)lane + (uint64_t)(k * 32); + 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) { - s += (uint32_t)lens[code]; + batch_bytes += (uint32_t)lens[code]; } else { atomicMax(status, 1u); } @@ -80,30 +80,30 @@ __launch_bounds__(ONPAIR_BLOCK_THREADS) void onpair_batch_offsets_sweep(const ui } #pragma unroll for (uint8_t offset = 16; offset > 0; offset >>= 1) { - s += __shfl_down_sync(0xffffffffu, s, offset); + 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)s; + 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 v = (lane < ONPAIR_WARPS_PER_BLOCK) ? warp_sums[lane] : 0; - uint64_t incl = v; + 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 y = __shfl_up_sync(0xffffffffu, incl, offset); + const uint64_t shifted = __shfl_up_sync(0xffffffffu, inclusive, offset); if (lane >= offset) { - incl += y; + inclusive += shifted; } } - const uint64_t aggregate = __shfl_sync(0xffffffffu, incl, 31); + const uint64_t aggregate = __shfl_sync(0xffffffffu, inclusive, 31); uint64_t prefix = 0; if (tile_idx == 0) { @@ -112,11 +112,15 @@ __launch_bounds__(ONPAIR_BLOCK_THREADS) void onpair_batch_offsets_sweep(const ui } } 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); - prefix = prefix_op(aggregate); + 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 + incl - v; + warp_excl[lane] = prefix + inclusive - lane_sum; } } __syncthreads(); diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu index 37e1765f100..061fff760c4 100644 --- a/vortex-cuda/kernels/src/onpair.cu +++ b/vortex-cuda/kernels/src/onpair.cu @@ -13,7 +13,7 @@ // 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 uint32_t MAX_INLINED_SIZE = 12; +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 Rust caller only diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu index e3fa94b51c6..be1161a501c 100644 --- a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -34,6 +34,10 @@ #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 diff --git a/vortex-cuda/src/cub.rs b/vortex-cuda/src/cub.rs index 4910455cdce..53c58095bac 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -100,3 +100,111 @@ 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. + async fn batch_offsets_roundtrip( + codes: Vec, + lens: Vec, + ) -> VortexResult<(Vec, u32)> { + 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 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, + &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(()) + } + + /// 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/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index 8451d18714f..44fd6c8af85 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -12,9 +12,10 @@ //! 2. `onpair_shmem_4tpt_split8read` gathers each token's bytes from the //! split dictionary layout and scatters them to the output byte stream. //! -//! The per-row lengths are only summed on the host (sizing the output and -//! cross-checking the code stream); per-row offsets are built solely by the -//! output path that needs them. The result is exposed either as a canonical +//! The total decoded size comes from the GPU offsets regeneration; the +//! lengths child is materialised on host once, and each output path derives +//! its row offsets from it — validating them against that total before they +//! index the decoded heap. The result is exposed either as a canonical //! `VarBinView` (views built on-device by `onpair_build_views` from //! host-prefix-summed offsets, or on host for heaps that exceed a single //! backing buffer) or as Arrow-compatible i32 offsets plus values via @@ -56,6 +57,7 @@ use vortex_onpair::OnPair; use vortex_onpair::OnPairArray; use vortex_onpair::OnPairArrayExt; use vortex_onpair::OnPairArraySlotsExt; +use vortex_onpair::code_boundary_at; use vortex_onpair::dict_view; use crate::CudaBufferExt; @@ -106,18 +108,20 @@ impl CudaExecute for OnPairExecutor { } } -/// Read one `codes_offsets` boundary by point lookup, so a sliced array never -/// materialises the whole per-row offsets child just to bound its codes. -fn code_boundary( - codes_offsets: &ArrayRef, - index: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult { - codes_offsets - .execute_scalar(index, ctx.execution_ctx())? - .as_primitive() - .as_::() - .ok_or_else(|| vortex_err!("OnPair codes_offsets[{index}] is null")) +/// Host sum of the per-row decoded lengths, rejecting negatives. +fn sum_lengths(lengths: &PrimitiveArray) -> VortexResult { + match_each_integer_ptype!(lengths.ptype(), |P| { + let mut acc = 0u64; + #[allow(clippy::unnecessary_cast)] + for &length in lengths.as_slice::

() { + let length = u64::try_from(length as i128) + .map_err(|_| vortex_err!("OnPair uncompressed length cannot be negative"))?; + acc = acc + .checked_add(length) + .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; + } + VortexResult::Ok(acc) + }) } /// All-empty output: `num_rows` inline empty views and no backing buffers. @@ -152,7 +156,8 @@ struct StagedCodes { struct OnPairDecoded { /// The flat decoded byte stream. bytes: CudaSlice, - /// Total decoded byte count. + /// Total decoded byte count, computed on device by the offsets + /// regeneration. total_size: usize, /// Host-resident per-row lengths. Each output path derives what it needs: /// the views fast path prefix-sums them into row offsets, the varbin path @@ -240,30 +245,18 @@ async fn decode_onpair_bytes( ) -> VortexResult> { let num_rows = onpair.len(); - // Sum the per-row decoded lengths on host — they are materialised there - // anyway. The total sizes the output allocation and cross-checks the code - // stream; per-row offsets are only built by the output paths that need - // them. + // Materialise the lengths child once; each output path derives its row + // offsets from it and validates them against the GPU-computed total. let lengths = onpair .uncompressed_lengths() .clone() .execute::(ctx.execution_ctx())?; - // The conversion is fallible only for the wide/signed ptype instantiations - // of the macro; u8/u16 make it infallible, so allow the lint wholesale. - #[allow(clippy::unnecessary_fallible_conversions)] - let total_size = match_each_integer_ptype!(lengths.ptype(), |P| { - lengths - .as_slice::

() - .iter() - .map(|&v| usize::try_from(v).vortex_expect("length must fit in usize")) - .sum() - }); // `codes_offsets` may be a sliced view of the original; its first and last // boundaries bound the contiguous run of `codes` belonging to this array's // rows (`slice` keeps the full `codes` child and only narrows the offsets). - let code_start = code_boundary(onpair.codes_offsets(), 0, ctx)?; - let code_end = code_boundary(onpair.codes_offsets(), num_rows, ctx)?; + let code_start = code_boundary_at(onpair.codes_offsets(), 0, ctx.execution_ctx())?; + let code_end = code_boundary_at(onpair.codes_offsets(), num_rows, ctx.execution_ctx())?; vortex_ensure!( code_start <= code_end, "OnPair codes_offsets must be nondecreasing" @@ -275,18 +268,15 @@ async fn decode_onpair_bytes( onpair.codes().len() ); - if total_size == 0 { - // Every token decodes to at least one byte. + if code_start == code_end { + // No codes: the array must decode to zero bytes. + let total = sum_lengths(&lengths)?; vortex_ensure!( - code_start == code_end, - "OnPair records zero decoded bytes but has codes" + total == 0, + "OnPair records {total} decoded bytes but has no codes" ); return Ok(None); } - vortex_ensure!( - code_start < code_end, - "OnPair records {total_size} decoded bytes but has no codes" - ); // Corruption flag raised by the batch-sizes kernel for a code outside the // dictionary; checked before the unchecked decode kernel is allowed to run. @@ -297,10 +287,10 @@ async fn decode_onpair_bytes( let staged = stage_codes(onpair, code_start, code_end, &mut status, ctx).await?; - // One synchronizing readback validates the compressed stream before the - // decode kernel — whose dictionary gathers and output scatters are - // unchecked — is allowed to run: every code indexed the dictionary, and - // the codes decode to exactly the byte count the lengths record. + // One synchronizing readback gates the decode kernel — whose dictionary + // gathers and output scatters are unchecked — and yields the GPU-computed + // total that sizes the output. The lengths child is validated against it + // by whichever output path materialises row offsets. let status = ctx .stream() .clone_dtoh(&status) @@ -319,14 +309,14 @@ async fn decode_onpair_bytes( .first() .copied() .ok_or_else(|| vortex_err!("OnPair batch offset scan returned no total"))?; - vortex_ensure!( - chunk_total == total_size as u64, - "OnPair codes decode to {chunk_total} bytes but uncompressed_lengths records {total_size}" - ); + let total_size = usize::try_from(chunk_total)?; + // 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 kernel's drain gates 16-byte stores on `out_start % 16` // relative to the buffer base, so the base must be 16-aligned. - let bytes = ctx.device_alloc::(total_size)?; + let mut bytes = ctx.device_alloc::(total_size)?; let (bytes_base_ptr, _) = bytes.device_ptr(ctx.stream()); assert_eq!( bytes_base_ptr % 16, @@ -350,7 +340,7 @@ async fn decode_onpair_bytes( .arg(&s8_view) .arg(&padded_view) .arg(&lens_view) - .arg(&bytes) + .arg(&mut bytes) .arg(&num_tokens_u64); }, )?; @@ -386,8 +376,7 @@ async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> Vorte // Fast path: the decoded heap fits a single BinaryView backing buffer, so // the per-row views build on-device. Only this path needs the u64 row - // offsets: prefix-sum the lengths here (negatives were already rejected - // by the total sum) and stage them on device. + // offsets: prefix-sum the lengths here and stage them on device. if total_size <= MAX_BUFFER_LEN { let row_offsets: Vec = match_each_integer_ptype!(lengths.ptype(), |P| { let mut offsets = Vec::with_capacity(lengths.len() + 1); @@ -395,20 +384,33 @@ async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> Vorte offsets.push(0u64); #[allow(clippy::unnecessary_cast)] for &length in lengths.as_slice::

() { - acc += length as u64; + let length = u64::try_from(length as i128) + .map_err(|_| vortex_err!("OnPair uncompressed length cannot be negative"))?; + acc = acc + .checked_add(length) + .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; offsets.push(acc); } - offsets - }); + VortexResult::Ok(offsets) + })?; + // The views index the decoded heap, so the lengths must account for + // exactly the bytes the codes decoded to. + let row_total = *row_offsets + .last() + .vortex_expect("row_offsets has at least one entry"); + vortex_ensure!( + row_total == total_size as u64, + "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" + ); let row_offsets_dev = ctx.copy_to_device(row_offsets)?.await?; let row_offsets_view = row_offsets_dev.cuda_view::()?; - let device_views = ctx.device_alloc::(num_rows)?; + 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) - .arg(&device_views) + .arg(&mut device_views) .arg(&num_rows_u64); })?; @@ -421,6 +423,12 @@ async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> Vorte // BinaryView offsets are u32. Heaps that need multiple backing buffers // roll the decoded bytes over on host, mirroring the CPU canonical path. + // The host views index the copied heap, so validate the lengths first. + let row_total = sum_lengths(&lengths)?; + vortex_ensure!( + row_total == total_size as u64, + "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" + ); let host_bytes = CudaDeviceBuffer::new(bytes) .copy_to_host(Alignment::new(1))? .await?; @@ -476,7 +484,13 @@ pub(crate) async fn decode_onpair_varbin( buffer: offsets, total, } = i32_offsets_from_lengths(decoded.lengths.clone(), ctx).await?; - debug_assert_eq!(total, decoded.total_size); + // The Arrow offsets index the decoded heap, so the lengths must account + // for exactly the bytes the codes decoded to. + vortex_ensure!( + total == decoded.total_size, + "OnPair codes decode to {} bytes but uncompressed_lengths records {total}", + decoded.total_size + ); Ok(DecodedVarBin { dtype, From f2818dfe3a7e7648cab7268eab9ed117ed5b7fda Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 21:06:01 +0100 Subject: [PATCH 09/11] fix more Signed-off-by: Robert Kruszewski --- encodings/onpair/src/decode.rs | 7 +- encodings/onpair/src/lib.rs | 1 - vortex-cuda/cub/kernels/filter.h | 6 +- vortex-cuda/cub/kernels/onpair.cu | 57 +- vortex-cuda/cub/src/onpair.rs | 11 +- vortex-cuda/kernels/src/onpair.cu | 115 +++- .../src/onpair_shmem_4tpt_split8read.cu | 50 +- vortex-cuda/src/cub.rs | 47 +- vortex-cuda/src/kernel/encodings/onpair.rs | 539 ++++++++++++------ 9 files changed, 615 insertions(+), 218 deletions(-) diff --git a/encodings/onpair/src/decode.rs b/encodings/onpair/src/decode.rs index 3c16485c15a..7aac8c3963f 100644 --- a/encodings/onpair/src/decode.rs +++ b/encodings/onpair/src/decode.rs @@ -32,12 +32,7 @@ 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, the CUDA decoder's code window) don't pay to materialise -/// every boundary. -pub fn code_boundary_at( +pub(crate) fn code_boundary_at( codes_offsets: &ArrayRef, index: usize, ctx: &mut ExecutionCtx, diff --git a/encodings/onpair/src/lib.rs b/encodings/onpair/src/lib.rs index bd6b150304a..f9d8be5ff0d 100644 --- a/encodings/onpair/src/lib.rs +++ b/encodings/onpair/src/lib.rs @@ -22,7 +22,6 @@ mod tests; pub use array::*; pub use compress::*; -pub use decode::code_boundary_at; pub use onpair::CompactDictionaryView; pub use onpair::Config; pub use onpair::DEFAULT_CONFIG; diff --git a/vortex-cuda/cub/kernels/filter.h b/vortex-cuda/cub/kernels/filter.h index 384dfa7bcb7..0fae1b639a8 100644 --- a/vortex-cuda/cub/kernels/filter.h +++ b/vortex-cuda/cub/kernels/filter.h @@ -102,13 +102,15 @@ cudaError_t scan_exclusive_sum_i64(void *d_temp, // 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. A code outside the dictionary +// 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 uint16_t *codes, + const void *codes, + uint32_t code_width, const uint8_t *lens, uint32_t dict_size, uint64_t total_tokens, diff --git a/vortex-cuda/cub/kernels/onpair.cu b/vortex-cuda/cub/kernels/onpair.cu index 4a11176055b..65336ec0886 100644 --- a/vortex-cuda/cub/kernels/onpair.cu +++ b/vortex-cuda/cub/kernels/onpair.cu @@ -41,12 +41,14 @@ __global__ void onpair_batch_offsets_init(OnPairTileState tile_state, int num_ti tile_state.InitializeStatus(num_tiles); } -// The fused sweep. 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. +// 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 uint16_t *__restrict codes, +__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, @@ -151,12 +153,15 @@ extern "C" cudaError_t onpair_batch_offsets_temp_size(size_t *temp_bytes, int64_ // 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. A code outside the dictionary -// raises `*status` to 1 and contributes zero bytes; the caller must check the -// flag before trusting the offsets. +// `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 uint16_t *codes, + const void *codes, + uint32_t code_width, const uint8_t *lens, uint32_t dict_size, uint64_t total_tokens, @@ -183,13 +188,31 @@ extern "C" cudaError_t onpair_batch_offsets(void *d_temp, return err; } - onpair_batch_offsets_sweep<<>>(codes, - lens, - dict_size, - total_tokens, - chunk_offsets, - status, - tile_state, - num_batches); + 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/onpair.rs b/vortex-cuda/cub/src/onpair.rs index 0c92f1e34ed..614b69c42b2 100644 --- a/vortex-cuda/cub/src/onpair.rs +++ b/vortex-cuda/cub/src/onpair.rs @@ -24,15 +24,16 @@ pub fn batch_offsets_temp_size(num_batches: i64) -> Result { /// 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. -/// A code outside the dictionary raises `*status` to 1 and contributes zero +/// `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` `u16` values, with -/// `total_tokens <= num_batches * 128`. +/// - `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`. @@ -40,7 +41,8 @@ pub fn batch_offsets_temp_size(num_batches: i64) -> Result { pub unsafe fn batch_offsets( d_temp: *mut c_void, temp_bytes: usize, - codes: *const u16, + codes: *const c_void, + code_width: u32, lens: *const u8, dict_size: u32, total_tokens: u64, @@ -55,6 +57,7 @@ pub unsafe fn batch_offsets( d_temp, temp_bytes, codes, + code_width, lens, dict_size, total_tokens, diff --git a/vortex-cuda/kernels/src/onpair.cu b/vortex-cuda/kernels/src/onpair.cu index 061fff760c4..e15cfec7391 100644 --- a/vortex-cuda/kernels/src/onpair.cu +++ b/vortex-cuda/kernels/src/onpair.cu @@ -8,23 +8,118 @@ // 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 -// view-construction kernels launched after the decode. +// 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 Rust caller only -// launches this when every offset fits the view's u32 fields and the decoded -// heap is exposed as backing buffer zero. -__device__ inline void onpair_write_view(const uint64_t *__restrict row_offsets, +// 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 uint64_t start = row_offsets[rid]; - const uint32_t len = (uint32_t)(row_offsets[rid + 1] - start); + 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 @@ -40,10 +135,10 @@ __device__ inline void onpair_write_view(const uint64_t *__restrict row_offsets, 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, (uint32_t)start); + views[rid] = make_uint4(len, prefix, 0, start); } -extern "C" __global__ void onpair_build_views(const uint64_t *__restrict row_offsets, +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) { diff --git a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu index be1161a501c..4143c75c4e9 100644 --- a/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu +++ b/vortex-cuda/kernels/src/onpair_shmem_4tpt_split8read.cu @@ -66,8 +66,10 @@ struct OnPairTokens { }; // Phase 1 — load this lane's 4 (code, dict_s8 bytes, length) triples. Tokens -// past the end of the stream load as empty. -__device__ inline OnPairTokens onpair_load_tokens(const uint16_t *__restrict codes, +// 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, @@ -173,14 +175,14 @@ __device__ inline void onpair_drain(const uint8_t *__restrict s_buf, } } -extern "C" __global__ ONPAIR_LAUNCH_BOUNDS void -onpair_shmem_4tpt_split8read(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) { +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; @@ -208,3 +210,31 @@ onpair_shmem_4tpt_split8read(const uint16_t *__restrict codes, 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/cub.rs b/vortex-cuda/src/cub.rs index 53c58095bac..0612a13fc07 100644 --- a/vortex-cuda/src/cub.rs +++ b/vortex-cuda/src/cub.rs @@ -21,11 +21,14 @@ 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. 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. +/// 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, @@ -52,7 +55,8 @@ pub(crate) fn onpair_batch_offsets( onpair::batch_offsets( temp_ptr as *mut c_void, temp_bytes, - codes_ptr as *const u16, + codes_ptr as *const c_void, + code_width, lens_ptr as *const u8, dict_size, total_tokens, @@ -109,15 +113,25 @@ mod tests { use crate::session::CudaSession; /// Upload synthetic codes and lengths, regenerate the chunk offsets, and - /// read them back together with the status flag. - async fn batch_offsets_roundtrip( - codes: Vec, + /// 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)> { + ) -> 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?; @@ -128,6 +142,7 @@ mod tests { let offsets = onpair_batch_offsets( &codes_dev, + code_width, &lens_dev, dict_size, num_tokens, @@ -194,6 +209,22 @@ mod tests { 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] diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index 44fd6c8af85..6fb3dc2302d 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -3,24 +3,37 @@ //! CUDA executor for OnPair decompression. //! -//! Decoding runs on the GPU over the flat token stream: +//! 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_shmem_4tpt_split8read` gathers each token's bytes from the -//! split dictionary layout and scatters them to the output byte stream. +//! 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. //! -//! The total decoded size comes from the GPU offsets regeneration; the -//! lengths child is materialised on host once, and each output path derives -//! its row offsets from it — validating them against that total before they -//! index the decoded heap. The result is exposed either as a canonical -//! `VarBinView` (views built on-device by `onpair_build_views` from -//! host-prefix-summed offsets, or on host for heaps that exceed a single -//! backing buffer) or as Arrow-compatible i32 offsets plus values via -//! [`decode_onpair_varbin`], which builds the offsets on device with -//! [`i32_offsets_from_lengths`] — mirroring the FSST varbin path. +//! 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; @@ -28,43 +41,46 @@ 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::builtins::ArrayBuiltins; use vortex::array::match_each_integer_ptype; use vortex::array::validity::Validity; -use vortex::buffer::Alignment; use vortex::dtype::DType; -use vortex::dtype::Nullability; +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::code_boundary_at; 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; @@ -102,28 +118,37 @@ impl CudaExecute for OnPairExecutor { ctx: &mut CudaExecutionCtx, ) -> VortexResult { let onpair = array - .try_downcast::() - .map_err(|_| vortex_err!("Expected OnPairArray"))?; + .as_typed::() + .ok_or_else(|| vortex_err!("Expected OnPairArray"))?; decode_onpair(onpair, ctx).await } } -/// Host sum of the per-row decoded lengths, rejecting negatives. +/// 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; - #[allow(clippy::unnecessary_cast)] for &length in lengths.as_slice::

() { - let length = u64::try_from(length as i128) - .map_err(|_| vortex_err!("OnPair uncompressed length cannot be negative"))?; acc = acc - .checked_add(length) + .checked_add(AsPrimitive::::as_(length)) .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; } - VortexResult::Ok(acc) + 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, @@ -137,15 +162,17 @@ async fn empty_views( })) } -/// The device-staged compressed token stream: the u16 code window, the split -/// dictionary layout, and the regenerated per-batch output offsets. +/// 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 code window. + /// is the total decoded byte count of the full code stream. chunk_offsets: CudaSlice, num_batches: usize, num_tokens: usize, @@ -154,41 +181,38 @@ struct StagedCodes { /// The shared result of the OnPair GPU decode pipeline. struct OnPairDecoded { - /// The flat decoded byte stream. - bytes: CudaSlice, - /// Total decoded byte count, computed on device by the offsets - /// regeneration. + /// 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, - /// Host-resident per-row lengths. Each output path derives what it needs: - /// the views fast path prefix-sums them into row offsets, the varbin path - /// builds Arrow i32 offsets from them on device, and the rollover path - /// consumes them directly. + /// 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 code window and dictionary on the device and regenerate +/// 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`]). +/// (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: &OnPairArray, - code_start: usize, - code_end: usize, + onpair: ArrayView<'_, OnPair>, + codes: PrimitiveArray, status: &mut CudaSlice, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - // Widen this array's code window to the decode kernel's u16 ABI. - let codes = onpair - .codes() - .slice(code_start..code_end)? - .cast(DType::Primitive(PType::U16, Nullability::NonNullable))? - .execute::(ctx.execution_ctx())? - .into_buffer::(); 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.as_view(), ctx.execution_ctx())?; + 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]; @@ -205,7 +229,7 @@ async fn stage_codes( } let (codes_dev, s8_dev, padded_dev, lens_dev) = futures::try_join!( - ctx.copy_to_device(codes)?, + ctx.ensure_on_device(codes_buffer), ctx.copy_to_device(dict_s8)?, ctx.copy_to_device(dict_padded)?, ctx.copy_to_device(lens)?, @@ -215,6 +239,7 @@ async fn stage_codes( 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, @@ -228,6 +253,7 @@ async fn stage_codes( dict_s8: s8_dev, dict_padded: padded_dev, lens: lens_dev, + dict_size: dict_size_u32, chunk_offsets, num_batches, num_tokens, @@ -235,49 +261,98 @@ async fn stage_codes( }) } -/// Run the OnPair decode pipeline: sum the per-row lengths on host, -/// regenerate the per-batch output offsets on the device, validate the -/// compressed stream, and decode the flat byte stream. Returns `Ok(None)` -/// when the array decodes to zero bytes. +/// 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: &OnPairArray, + onpair: ArrayView<'_, OnPair>, ctx: &mut CudaExecutionCtx, ) -> VortexResult> { let num_rows = onpair.len(); - // Materialise the lengths child once; each output path derives its row - // offsets from it and validates them against the GPU-computed total. + 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::(ctx.execution_ctx())?; + .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); + } - // `codes_offsets` may be a sliced view of the original; its first and last - // boundaries bound the contiguous run of `codes` belonging to this array's - // rows (`slice` keeps the full `codes` child and only narrows the offsets). - let code_start = code_boundary_at(onpair.codes_offsets(), 0, ctx.execution_ctx())?; - let code_end = code_boundary_at(onpair.codes_offsets(), num_rows, ctx.execution_ctx())?; - vortex_ensure!( - code_start <= code_end, - "OnPair codes_offsets must be nondecreasing" - ); - vortex_ensure!( - code_end <= onpair.codes().len(), - "OnPair codes_offsets end {} exceeds codes len {}", - code_end, - onpair.codes().len() - ); + // 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(); - if code_start == code_end { - // No codes: the array must decode to zero bytes. - let total = sum_lengths(&lengths)?; - vortex_ensure!( - total == 0, - "OnPair records {total} decoded bytes but has no codes" - ); - return Ok(None); + // 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)?; @@ -285,12 +360,76 @@ async fn decode_onpair_bytes( .memset_zeros(&mut status) .map_err(|e| vortex_err!("Failed to zero OnPair status flag: {e}"))?; - let staged = stage_codes(onpair, code_start, code_end, &mut status, ctx).await?; + 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 - // total that sizes the output. The lengths child is validated against it - // by whichever output path materialises row offsets. + // 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) @@ -309,14 +448,42 @@ async fn decode_onpair_bytes( .first() .copied() .ok_or_else(|| vortex_err!("OnPair batch offset scan returned no total"))?; - let total_size = usize::try_from(chunk_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 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::(total_size)?; + // 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, @@ -324,12 +491,7 @@ async fn decode_onpair_bytes( "output base not 16-aligned: {bytes_base_ptr:#x}", ); - 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::()?; - let decode_fn = ctx.load_function_with_suffixes("onpair_shmem_4tpt_split8read", &[])?; + let decode_fn = ctx.load_function_with_suffixes("onpair_shmem_4tpt_split8read", &[&ptype])?; ctx.launch_kernel_config( &decode_fn, staged.launch_config, @@ -345,14 +507,19 @@ async fn decode_onpair_bytes( }, )?; + // This array's rows as a zero-copy device slice of the decoded heap. + let heap = CudaDeviceBuffer::new(bytes); Ok(Some(OnPairDecoded { - bytes, + bytes: BufferHandle::new_device(heap.slice(byte_start..byte_end)), total_size, lengths, })) } -async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> VortexResult { +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(); @@ -361,11 +528,7 @@ async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> Vorte return Ok(Canonical::empty(&dtype)); } - if validity.definitely_all_null() { - return empty_views(num_rows, dtype, validity, ctx).await; - } - - let Some(decoded) = decode_onpair_bytes(&onpair, ctx).await? else { + let Some(decoded) = decode_onpair_bytes(onpair, ctx).await? else { return empty_views(num_rows, dtype, validity, ctx).await; }; let OnPairDecoded { @@ -374,65 +537,44 @@ async fn decode_onpair(onpair: OnPairArray, ctx: &mut CudaExecutionCtx) -> Vorte lengths, } = decoded; - // Fast path: the decoded heap fits a single BinaryView backing buffer, so - // the per-row views build on-device. Only this path needs the u64 row - // offsets: prefix-sum the lengths here and stage them on device. + // 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 row_offsets: Vec = match_each_integer_ptype!(lengths.ptype(), |P| { - let mut offsets = Vec::with_capacity(lengths.len() + 1); - let mut acc = 0u64; - offsets.push(0u64); - #[allow(clippy::unnecessary_cast)] - for &length in lengths.as_slice::

() { - let length = u64::try_from(length as i128) - .map_err(|_| vortex_err!("OnPair uncompressed length cannot be negative"))?; - acc = acc - .checked_add(length) - .ok_or_else(|| vortex_err!("OnPair decoded size overflow"))?; - offsets.push(acc); - } - VortexResult::Ok(offsets) - })?; - // The views index the decoded heap, so the lengths must account for - // exactly the bytes the codes decoded to. - let row_total = *row_offsets - .last() - .vortex_expect("row_offsets has at least one entry"); - vortex_ensure!( - row_total == total_size as u64, - "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" - ); - let row_offsets_dev = ctx.copy_to_device(row_offsets)?.await?; - let row_offsets_view = row_offsets_dev.cuda_view::()?; + 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) + .arg(&bytes_view) .arg(&mut device_views) .arg(&num_rows_u64); })?; let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(device_views))); - let bytes = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(bytes))); return Ok(Canonical::VarBinView(unsafe { VarBinViewArray::new_handle_unchecked(views, Arc::from([bytes]), dtype, validity) })); } - // BinaryView offsets are u32. Heaps that need multiple backing buffers - // roll the decoded bytes over on host, mirroring the CPU canonical path. - // The host views index the copied heap, so validate the lengths first. - let row_total = sum_lengths(&lengths)?; - vortex_ensure!( - row_total == total_size as u64, - "OnPair codes decode to {total_size} bytes but uncompressed_lengths records {row_total}" - ); - let host_bytes = CudaDeviceBuffer::new(bytes) - .copy_to_host(Alignment::new(1))? - .await?; - let host_bytes = host_bytes.slice(0..total_size); + // 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( @@ -458,13 +600,7 @@ pub(crate) async fn decode_onpair_varbin( let validity = onpair.array_validity(); let len = onpair.len(); - let decoded = if onpair.is_empty() || validity.definitely_all_null() { - None - } else { - decode_onpair_bytes(&onpair, ctx).await? - }; - - let Some(decoded) = decoded else { + 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)?); @@ -478,25 +614,25 @@ pub(crate) async fn decode_onpair_varbin( }); }; + let OnPairDecoded { + bytes, + total_size, + lengths, + } = decoded; + // Build the Arrow i32 offsets from the lengths on device; this also - // rejects heaps beyond Arrow's i32 offset range. + // rejects windows beyond Arrow's i32 offset range. let I32Offsets { buffer: offsets, total, - } = i32_offsets_from_lengths(decoded.lengths.clone(), ctx).await?; - // The Arrow offsets index the decoded heap, so the lengths must account - // for exactly the bytes the codes decoded to. - vortex_ensure!( - total == decoded.total_size, - "OnPair codes decode to {} bytes but uncompressed_lengths records {total}", - decoded.total_size - ); + } = i32_offsets_from_lengths(lengths, ctx).await?; + ensure_lengths_match(u64::try_from(total)?, total_size)?; Ok(DecodedVarBin { dtype, len, offsets, - values: BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(decoded.bytes))), + values: bytes, validity, }) } @@ -509,7 +645,9 @@ mod tests { use vortex::array::IntoArray; use vortex::array::arrays::VarBinArray; use vortex::array::assert_arrays_eq; + use vortex::array::builtins::ArrayBuiltins; use vortex::buffer::Buffer; + use vortex::dtype::Nullability; use vortex::error::VortexExpect; use vortex_array::VortexSessionExecute; use vortex_onpair::DEFAULT_DICT12_CONFIG; @@ -682,6 +820,87 @@ mod tests { 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_DICT12_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. + #[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 values: Vec> = [ + &b"tokenized token stream"[..], + b"tokenized", + b"token stream", + b"stream of tokens", + ] + .into_iter() + .cycle() + .take(64) + .map(Some) + .collect(); + let original = + compress_onpair(values, DType::Utf8(Nullability::NonNullable), &mut cuda_ctx)?; + let onpair = original + .clone() + .try_downcast::() + .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; + // A tiny corpus trains a tiny dictionary, so every code fits u8. + vortex_ensure!( + onpair.dict_offsets().len() <= 256, + "test corpus unexpectedly trained {} tokens", + onpair.dict_offsets().len() - 1 + ); + let narrowed = OnPair::try_new( + onpair.dtype().clone(), + onpair.dict_bytes_handle().clone(), + onpair.dict_offsets().clone(), + onpair + .codes() + .clone() + .cast(DType::Primitive(PType::U8, Nullability::NonNullable))?, + onpair.codes_offsets().clone(), + onpair.uncompressed_lengths().clone(), + onpair.array_validity(), + )?; + + let gpu_result = OnPairExecutor + .execute(narrowed.into_array(), &mut cuda_ctx) + .await?; + assert_device_resident(&gpu_result); + let host_result = gpu_result.into_host().await?.into_array(); + assert_arrays_eq!(original, 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)?; From c708ed0052bb21cffc416491baf5b230ec604957 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 21:41:27 +0100 Subject: [PATCH 10/11] fix Signed-off-by: Robert Kruszewski --- vortex-cuda/src/kernel/encodings/onpair.rs | 70 +++++++++++++--------- 1 file changed, 41 insertions(+), 29 deletions(-) diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index 6fb3dc2302d..ee1cd69f9fc 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -645,7 +645,6 @@ mod tests { use vortex::array::IntoArray; use vortex::array::arrays::VarBinArray; use vortex::array::assert_arrays_eq; - use vortex::array::builtins::ArrayBuiltins; use vortex::buffer::Buffer; use vortex::dtype::Nullability; use vortex::error::VortexExpect; @@ -851,12 +850,17 @@ mod tests { } /// 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 values: Vec> = [ + + let strings: Vec<&[u8]> = [ &b"tokenized token stream"[..], b"tokenized", b"token stream", @@ -864,40 +868,48 @@ mod tests { ] .into_iter() .cycle() - .take(64) - .map(Some) + .take(800) .collect(); - let original = - compress_onpair(values, DType::Utf8(Nullability::NonNullable), &mut cuda_ctx)?; - let onpair = original - .clone() - .try_downcast::() - .map_err(|array| vortex_err!("expected OnPair array, got {}", array.encoding_id()))?; - // A tiny corpus trains a tiny dictionary, so every code fits u8. - vortex_ensure!( - onpair.dict_offsets().len() <= 256, - "test corpus unexpectedly trained {} tokens", - onpair.dict_offsets().len() - 1 - ); - let narrowed = OnPair::try_new( - onpair.dtype().clone(), - onpair.dict_bytes_handle().clone(), - onpair.dict_offsets().clone(), - onpair - .codes() - .clone() - .cast(DType::Primitive(PType::U8, Nullability::NonNullable))?, - onpair.codes_offsets().clone(), - onpair.uncompressed_lengths().clone(), - onpair.array_validity(), + + // 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(narrowed.into_array(), &mut cuda_ctx) + .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!(original, host_result, &mut ctx); + assert_arrays_eq!(expected, host_result, &mut ctx); Ok(()) } From 1732f2a63311ff0cd96385d48479da1a406e4d78 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 4 Aug 2026 11:57:19 +0100 Subject: [PATCH 11/11] fixes Signed-off-by: Robert Kruszewski --- vortex-cuda/benches/onpair_cuda.rs | 4 ++-- vortex-cuda/src/kernel/encodings/onpair.rs | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-cuda/benches/onpair_cuda.rs b/vortex-cuda/benches/onpair_cuda.rs index 51e9d079351..0aa716f0c0a 100644 --- a/vortex-cuda/benches/onpair_cuda.rs +++ b/vortex-cuda/benches/onpair_cuda.rs @@ -27,7 +27,7 @@ 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_DICT12_CONFIG; +use vortex_onpair::DEFAULT_CONFIG; use vortex_onpair::onpair_compress; use crate::timed_launch_strategy::TimedLaunchStrategy; @@ -55,7 +55,7 @@ fn make_fixture(n: usize) -> OnPairBenchFixture { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let array = onpair_compress(&varbin, DEFAULT_DICT12_CONFIG, setup_ctx.execution_ctx()) + let array = onpair_compress(&varbin, DEFAULT_CONFIG, setup_ctx.execution_ctx()) .vortex_expect("OnPair compression failed"); OnPairBenchFixture { diff --git a/vortex-cuda/src/kernel/encodings/onpair.rs b/vortex-cuda/src/kernel/encodings/onpair.rs index ee1cd69f9fc..38cc51a3afb 100644 --- a/vortex-cuda/src/kernel/encodings/onpair.rs +++ b/vortex-cuda/src/kernel/encodings/onpair.rs @@ -649,7 +649,7 @@ mod tests { use vortex::dtype::Nullability; use vortex::error::VortexExpect; use vortex_array::VortexSessionExecute; - use vortex_onpair::DEFAULT_DICT12_CONFIG; + use vortex_onpair::DEFAULT_CONFIG; use vortex_onpair::onpair_compress; use super::*; @@ -683,7 +683,7 @@ mod tests { ctx: &mut CudaExecutionCtx, ) -> VortexResult { let varbin = VarBinArray::from_iter(strings, dtype).into_array(); - let onpair = onpair_compress(&varbin, DEFAULT_DICT12_CONFIG, ctx.execution_ctx())?; + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, ctx.execution_ctx())?; vortex_ensure!( onpair.as_opt::().is_some(), "expected OnPair array, got {}", @@ -806,7 +806,7 @@ mod tests { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let onpair = onpair_compress(&varbin, DEFAULT_DICT12_CONFIG, cuda_ctx.execution_ctx())?; + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, cuda_ctx.execution_ctx())?; let gpu_result = OnPairExecutor .execute(onpair.clone(), &mut cuda_ctx) @@ -837,7 +837,7 @@ mod tests { DType::Utf8(Nullability::NonNullable), ) .into_array(); - let onpair = onpair_compress(&varbin, DEFAULT_DICT12_CONFIG, cuda_ctx.execution_ctx())?; + let onpair = onpair_compress(&varbin, DEFAULT_CONFIG, cuda_ctx.execution_ctx())?; let sliced = onpair.slice(19_997..20_101)?; let gpu_result = OnPairExecutor