Skip to content

feat(cuda): export FSST directly to Arrow varbin #8787

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
0ax1 merged 13 commits into from
Jul 16, 2026
Merged

feat(cuda): export FSST directly to Arrow varbin #8787

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
97452ec
feat(cuda): export FSST directly to Arrow varbin
0ax1 Jul 16, 2026
30da392
fix(cuda): apply varbin export layout to list element schema fields
0ax1 Jul 16, 2026
2fa8f3a
fix(cuda): fall back to view layout for FSST exports over i32 offset …
0ax1 Jul 16, 2026
e278c5d
chore(cuda): pair FSST varbin decode with the wide-store alignment as…
0ax1 Jul 16, 2026
d3c4445
fix(buffer): shift unaligned bits in BitBuffer::sliced instead of cop…
0ax1 Jul 16, 2026
0897e72
test(cuda): cover FSST varbin export contents, validity, slices, and …
0ax1 Jul 16, 2026
72bbbc0
chore(cuda): move ChunkedArray import to the test module top
0ax1 Jul 16, 2026
499579d
style(cuda): clang-format FSST kernel
0ax1 Jul 16, 2026
2325faf
fix(buffer): make unaligned BitBuffer slicing Miri-safe
0ax1 Jul 16, 2026
bab3cb6
clean
0ax1 Jul 16, 2026
80266a2
test(buffer): cover word-boundary-crossing slice in BitBuffer::sliced
0ax1 Jul 16, 2026
35b0295
fix(cuda): pass FSST validity bit offset
0ax1 Jul 16, 2026
2cc117e
docs(cuda): explain FSST view fast path
0ax1 Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 106 additions & 20 deletions vortex-cuda/benches/fsst_cuda.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
//! CUDA benchmarks for FSST decompression.

#![expect(clippy::unwrap_used)]
#![expect(clippy::cast_possible_truncation)]

#[allow(dead_code)]
mod bench_config;
Expand All @@ -18,12 +17,20 @@ 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::PrimitiveArray;
use vortex::array::match_each_integer_ptype;
use vortex::dtype::DType;
use vortex::dtype::Nullability;
use vortex::encodings::fsst::FSST;
use vortex::encodings::fsst::FSSTArrayExt;
use vortex::error::VortexExpect;
use vortex_cuda::CudaDispatchMode;
use vortex_cuda::CudaSession;
use vortex_cuda::VarBinExportLayout;
use vortex_cuda::arrow::DeviceArrayExt;
use vortex_cuda::arrow::release_device_array;
use vortex_cuda::executor::CudaArrayExt;
use vortex_cuda_macros::cuda_available;
use vortex_cuda_macros::cuda_not_available;
Expand All @@ -36,30 +43,63 @@ use crate::timed_launch_strategy::TimedLaunchStrategy;
// other kernels benchmark.
const BENCH_SIZES: &[(usize, &str)] = &[(10_000_000, "10M")];

fn session_with_varbin_layout(layout: VarBinExportLayout) -> vortex::session::VortexSession {
vortex::array::array_session().with_some(
CudaSession::try_default()
.vortex_expect("failed to create CUDA session")
.with_varbin_export_layout(layout),
)
}

struct FSSTBenchFixture {
utf8: ArrayRef,
binary: ArrayRef,
uncompressed_size: u64,
}

fn make_fixture(n: usize) -> FSSTBenchFixture {
let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session())
.vortex_expect("failed to create execution context");
let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx());

let lens = fsst
.uncompressed_lengths()
.clone()
.execute::<PrimitiveArray>(setup_ctx.execution_ctx())
.vortex_expect("canonicalize uncompressed_lengths");
#[allow(clippy::unnecessary_cast)]
let uncompressed_size = match_each_integer_ptype!(lens.ptype(), |P| {
lens.as_slice::<P>().iter().map(|x| *x as u64).sum()
});

let binary = FSST::try_new(
DType::Binary(Nullability::NonNullable),
fsst.symbols().clone(),
fsst.symbol_lengths().clone(),
fsst.codes(),
fsst.uncompressed_lengths().clone(),
setup_ctx.execution_ctx(),
)
.vortex_expect("rebuild FSST fixture with Binary dtype")
.into_array();

FSSTBenchFixture {
utf8: fsst.into_array(),
binary,
uncompressed_size,
}
}

fn benchmark_fsst_cuda_decompress(c: &mut Criterion) {
let mut group = c.benchmark_group("cuda");

for &(n, len_str) in BENCH_SIZES {
let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session())
.vortex_expect("failed to create execution context");
let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx());

let lens = fsst
.uncompressed_lengths()
.clone()
.execute::<PrimitiveArray>(setup_ctx.execution_ctx())
.vortex_expect("canonicalize uncompressed_lengths");
let total_size: usize = match_each_integer_ptype!(lens.ptype(), |P| {
lens.as_slice::<P>().iter().map(|x| *x as usize).sum()
});
let uncompressed_size = total_size as u64;

let fsst_array = fsst.into_array();

group.throughput(Throughput::Bytes(uncompressed_size));
let fixture = make_fixture(n);

group.throughput(Throughput::Bytes(fixture.uncompressed_size));
group.bench_with_input(
BenchmarkId::new("cuda/fsst/decompress", len_str),
&fsst_array,
BenchmarkId::new("cuda/fsst/decompress_to_varbinview", len_str),
&fixture.utf8,
|b, fsst_array| {
b.iter_custom(|iters| {
let timed = TimedLaunchStrategy::default();
Expand All @@ -68,6 +108,7 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) {
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 {
Expand All @@ -77,6 +118,51 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) {
});
},
);

for (name, array, layout) in [
(
"cuda/fsst/decompress_to_varbin",
&fixture.utf8,
VarBinExportLayout::VarBin,
),
(
"cuda/fsst/export_binary",
&fixture.binary,
VarBinExportLayout::VarBin,
),
(
"cuda/fsst/export_utf8_view",
&fixture.utf8,
VarBinExportLayout::VarBinView,
),
(
"cuda/fsst/export_binary_view",
&fixture.binary,
VarBinExportLayout::VarBinView,
),
] {
group.bench_with_input(BenchmarkId::new(name, len_str), array, |b, array| {
b.iter_custom(|iters| {
let timed = TimedLaunchStrategy::default();
let timer = timed.timer();

let session = session_with_varbin_layout(layout);
let mut cuda_ctx = CudaSession::create_execution_ctx(&session)
.vortex_expect("failed to create execution context")
.with_dispatch_mode(CudaDispatchMode::StandaloneOnly)
.with_launch_strategy(Arc::new(timed));

for _ in 0..iters {
let mut exported =
block_on((*array).clone().export_device_array(&mut cuda_ctx))
.vortex_expect("export FSST device array");
release_device_array(&mut exported);
}

Duration::from_nanos(timer.load(Ordering::Relaxed))
});
});
}
}

group.finish();
Expand Down
134 changes: 107 additions & 27 deletions vortex-cuda/kernels/src/fsst.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,9 @@
// the next add (≤ 8 bytes) within the 24-byte capacity.
//
// `codes_offsets` is templated over the four unsigned integer widths
// (u8/u16/u32/u64). `output_offsets` is uint64_t.
// (u8/u16/u32/u64). `output_offsets` is uint64_t for the view kernels
// (`fsst_*`, which also take an optional views pointer) and int32_t Arrow
// varbin offsets for the bytes-only varbin kernels (`fsst_varbin_*`).

// 24-byte scratch buffer split across three u64 lanes. `cursor` is the
// number of bytes currently buffered and the next-push offset.
Expand Down Expand Up @@ -123,13 +125,18 @@ struct Scratch {
}
};

template <typename OffT>
// 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;

template <typename CodeOffsetT, typename OutputOffsetT>
struct FSSTArgs {
// Compressed FSST code stream, contiguous across all strings. String
// `sid`'s codes live in `[codes_offsets[sid], codes_offsets[sid + 1])`.
const uint8_t *__restrict codes_bytes;
// Per-string offsets into `codes_bytes`, length `num_strings + 1`.
const OffT *__restrict codes_offsets;
const CodeOffsetT *__restrict codes_offsets;
// FSST symbol table.
const uint64_t *__restrict symbols;
// Length in bytes (1..=8) of each entry in `symbols`. The remaining bits
Expand All @@ -138,19 +145,57 @@ struct FSSTArgs {
// Buffer to write decoded data into.
uint8_t *__restrict output_bytes;
// Per-string offsets into `output_bytes`, length `num_strings + 1`.
const uint64_t *__restrict output_offsets;
const OutputOffsetT *__restrict output_offsets;
// Validity of each string.
const uint8_t *__restrict validity_bits;
// Bit offset of string zero within the first validity byte (0..7).
uint64_t validity_bit_offset;
// Optional output views, one 16-byte uint4 per string. A null pointer
// requests bytes-only decoding for heaps that need the host rollover path.
uint4 *__restrict output_views;
};

template <typename OffT>
__device__ inline void fsst_decode_string(const FSSTArgs<OffT> &args, uint64_t sid) {
if (((args.validity_bits[sid >> 3] >> (sid & 7u)) & 1u) == 0u) {
// Build one BinaryView from the bytes this thread just decoded. The Rust
// caller only provides output_views when every offset fits in the view's u32
// fields and the decoded heap is exposed as backing buffer zero.
template <typename CodeOffsetT, typename OutputOffsetT>
__device__ inline void fsst_write_view(const FSSTArgs<CodeOffsetT, OutputOffsetT> &args, uint64_t sid) {
if (args.output_views == nullptr) {
return;
}

OffT in_pos = args.codes_offsets[sid];
const OffT in_end = args.codes_offsets[sid + 1];
const uint64_t start = args.output_offsets[sid];
const uint32_t len = (uint32_t)(args.output_offsets[sid + 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)args.output_bytes[start + i] << (8u * (i & 3u));
}
}
args.output_views[sid] = make_uint4(len, words[0], words[1], words[2]);
return;
}

const uint32_t prefix =
(uint32_t)args.output_bytes[start] | ((uint32_t)args.output_bytes[start + 1] << 8u) |
((uint32_t)args.output_bytes[start + 2] << 16u) | ((uint32_t)args.output_bytes[start + 3] << 24u);
args.output_views[sid] = make_uint4(len, prefix, 0, (uint32_t)start);
}

template <typename CodeOffsetT, typename OutputOffsetT>
__device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffsetT> &args, uint64_t sid) {
const uint64_t validity_index = sid + args.validity_bit_offset;
if (((args.validity_bits[validity_index >> 3] >> (validity_index & 7u)) & 1u) == 0u) {
if (args.output_views != nullptr) {
args.output_views[sid] = make_uint4(0, 0, 0, 0);
}
return;
}

CodeOffsetT in_pos = args.codes_offsets[sid];
const CodeOffsetT in_end = args.codes_offsets[sid + 1];
uint64_t out_pos = args.output_offsets[sid];
const uint64_t out_end = args.output_offsets[sid + 1];

Expand Down Expand Up @@ -181,46 +226,81 @@ __device__ inline void fsst_decode_string(const FSSTArgs<OffT> &args, uint64_t s
sym &= mask;

scratch.push(sym, len);
in_pos += (OffT)consumed;
in_pos += (CodeOffsetT)consumed;
}

// Epilogue: drain everything that's left.
while (scratch.cursor > 0) {
scratch.drain(args.output_bytes, out_pos, out_end);
}

fsst_write_view(args, sid);
}

#define GENERATE_FSST_KERNEL(suffix, OffT) \
#define FSST_GRID_STRIDE_LOOP(CodeOffsetT, OutputOffsetT, args) \
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_strings) ? (block_start + elements_per_block) : num_strings; \
for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \
fsst_decode_string<CodeOffsetT, OutputOffsetT>(args, sid); \
}

#define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \
extern "C" __global__ void fsst_##suffix(const uint8_t *__restrict codes_bytes, \
const OffT *__restrict codes_offsets, \
const CodeOffsetT *__restrict codes_offsets, \
const uint64_t *__restrict symbols, \
const uint8_t *__restrict symbol_lengths, \
const uint64_t *__restrict output_offsets, \
const uint8_t *__restrict validity_bits, \
uint64_t validity_bit_offset, \
uint8_t *__restrict output_bytes, \
uint4 *__restrict output_views, \
uint64_t num_strings) { \
const FSSTArgs<OffT> args = { \
const FSSTArgs<CodeOffsetT, uint64_t> args = { \
codes_bytes, \
codes_offsets, \
symbols, \
symbol_lengths, \
output_bytes, \
output_offsets, \
validity_bits, \
validity_bit_offset, \
output_views, \
}; \
\
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_strings) \
? (block_start + elements_per_block) \
: num_strings; \
\
for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \
fsst_decode_string<OffT>(args, sid); \
} \
FSST_GRID_STRIDE_LOOP(CodeOffsetT, uint64_t, args) \
}

GENERATE_FSST_KERNEL(u8, uint8_t)
GENERATE_FSST_KERNEL(u16, uint16_t)
GENERATE_FSST_KERNEL(u32, uint32_t)
GENERATE_FSST_KERNEL(u64, uint64_t)
#define GENERATE_FSST_VARBIN_KERNEL(suffix, CodeOffsetT) \
extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \
const CodeOffsetT *__restrict codes_offsets, \
const uint64_t *__restrict symbols, \
const uint8_t *__restrict symbol_lengths, \
const int32_t *__restrict output_offsets, \
const uint8_t *__restrict validity_bits, \
uint64_t validity_bit_offset, \
uint8_t *__restrict output_bytes, \
uint64_t num_strings) { \
const FSSTArgs<CodeOffsetT, int32_t> args = { \
codes_bytes, \
codes_offsets, \
symbols, \
symbol_lengths, \
output_bytes, \
output_offsets, \
validity_bits, \
validity_bit_offset, \
nullptr, \
}; \
FSST_GRID_STRIDE_LOOP(CodeOffsetT, int32_t, args) \
}

GENERATE_FSST_VIEW_KERNEL(u8, uint8_t)
GENERATE_FSST_VIEW_KERNEL(u16, uint16_t)
GENERATE_FSST_VIEW_KERNEL(u32, uint32_t)
GENERATE_FSST_VIEW_KERNEL(u64, uint64_t)

GENERATE_FSST_VARBIN_KERNEL(u8, uint8_t)
GENERATE_FSST_VARBIN_KERNEL(u16, uint16_t)
GENERATE_FSST_VARBIN_KERNEL(u32, uint32_t)
GENERATE_FSST_VARBIN_KERNEL(u64, uint64_t)
Loading
Loading