From b7bf156c0941f0367dac6769d705f10640eb6a79 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 17 Aug 2026 19:28:09 +0100 Subject: [PATCH] Pack GPU FSST metadata into one upload Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/kernels/src/fsst.cu | 38 +++++------ vortex-cuda/src/kernel/encodings/fsst.rs | 85 +++++++++++++----------- 2 files changed, 67 insertions(+), 56 deletions(-) diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index 13c17a3d533..ef7f119c29e 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -40,13 +40,9 @@ // 16-aligned) and the epilogue tail (< 16 bytes left, no room for u128). // In steady state out_pos stays 16-aligned and u128 fires repeatedly. // -// The 256-entry symbol table (≤ 2 KB) is read directly from global memory. -// Staging it into shared memory measured ~3% slower at 10M rows and ~15% -// slower at 1M rows (benchmarked on clickbench URLs). The hypothesis is that L1 -// already holds the table after a few iterations and the explicit shared copy -// adds bank-conflict latency on the warp-divergent `symbols[code]` reads; the -// gap is wider at 1M because the kernel is less bandwidth-bound there, so -// per-load latency shows up more. +// The 255-entry symbols and lengths arrays are packed into one global-memory +// table. This preserves fast cached global loads while requiring only one small +// allocation and upload for each FSST decode. // // Decoded symbols are masked to their valid byte length so the table's high // bits never leak. The main loop drains to `scratch.cursor ≤ 16`, keeping @@ -130,6 +126,13 @@ struct Scratch { // first four bytes, backing-buffer index, and byte offset. constexpr uint32_t MAX_INLINED_SIZE = 12; +struct alignas(8) FSSTSymbolTable { + uint64_t symbols[255]; + uint8_t symbol_lengths[255]; +}; + +static_assert(sizeof(FSSTSymbolTable) == 2296, "FSSTSymbolTable must match the Rust kernel argument"); + template struct FSSTArgs { // Compressed FSST code stream, contiguous across all strings. String @@ -249,8 +252,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ codes_bytes, \ codes_offsets, \ - symbols, \ - symbol_lengths, \ + symbol_table->symbols, \ + symbol_table->symbol_lengths, \ output_bytes, \ output_offsets, \ validity_bits, \ @@ -274,8 +276,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ codes_bytes, \ codes_offsets, \ - symbols, \ - symbol_lengths, \ + symbol_table->symbols, \ + symbol_table->symbol_lengths, \ output_bytes, \ output_offsets, \ validity_bits, \ @@ -298,8 +299,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ codes_bytes, \ codes_offsets, \ - symbols, \ - symbol_lengths, \ + symbol_table->symbols, \ + symbol_table->symbol_lengths, \ output_bytes, \ output_offsets, \ validity_bits, \ diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 8f17fadabe9..77faf508836 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -10,6 +10,7 @@ use async_trait::async_trait; use cudarc::driver::DevicePtr; use cudarc::driver::DeviceRepr; use cudarc::driver::PushKernelArg; +use cudarc::driver::ValidAsZeroBits; use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; @@ -32,6 +33,7 @@ use vortex::dtype::DType; use vortex::dtype::NativePType; use vortex::dtype::PType; use vortex::encodings::fsst::FSST; +use vortex::encodings::fsst::FSST_SYMBOL_TABLE_LEN; use vortex::encodings::fsst::FSSTArray; use vortex::encodings::fsst::FSSTArrayExt; use vortex::encodings::fsst::FSSTArraySlotsExt; @@ -58,6 +60,33 @@ pub(crate) struct FSSTVarBin { pub(crate) validity: Validity, } +/// The complete FSST table uploaded to a single device buffer. +/// +/// This is plain data with the same C layout as `FSSTSymbolTable` in `fsst.cu`. +/// Packing it avoids allocating and uploading two separate tiny device buffers for every decode. +#[repr(C)] +#[derive(Debug)] +struct FSSTSymbolTable { + symbols: [u64; FSST_SYMBOL_TABLE_LEN], + symbol_lengths: [u8; FSST_SYMBOL_TABLE_LEN], +} + +// SAFETY: `FSSTSymbolTable` has a stable C layout, contains only integer arrays, and all-zero +// bit patterns are valid for every field. +unsafe impl DeviceRepr for FSSTSymbolTable {} +unsafe impl ValidAsZeroBits for FSSTSymbolTable {} + +impl From<&FSSTArray> for FSSTSymbolTable { + fn from(fsst: &FSSTArray) -> Self { + Self { + symbols: std::array::from_fn(|index| fsst.padded_symbols()[index].to_u64()), + symbol_lengths: std::array::from_fn(|index| fsst.padded_symbol_lengths()[index]), + } + } +} + +const _: () = assert!(size_of::() == 2296); + /// Returns validity backing bytes and the bit offset of the first row for the FSST kernels. /// /// `BitBuffer` slices normalize whole-byte offsets but can retain a sub-byte offset. Passing that @@ -183,12 +212,10 @@ where let validity = fsst.codes().validity()?; let num_strings = fsst.len(); let num_strings_u64 = u64::try_from(num_strings)?; - let symbols_u64 = fsst - .symbols() - .iter() - .map(|symbol| symbol.to_u64()) - .collect::>(); - let symbol_lengths = fsst.padded_symbol_lengths().slice(0..fsst.n_symbols()); + let symbol_table = FSSTSymbolTable::from(&fsst); + let symbol_table = ctx + .stream() + .copy_to_device_sync(std::slice::from_ref(&symbol_table))?; let codes_bytes_handle = fsst.codes_bytes_handle().clone(); let PrimitiveDataParts { buffer: codes_offsets_buffer, @@ -196,8 +223,6 @@ where } = codes_offsets.into_data_parts(); let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx).await?; - let symbols = ctx.stream().copy_to_device_sync(&symbols_u64)?; - let symbol_lengths = ctx.stream().copy_to_device_sync(symbol_lengths.as_ref())?; let validity_device = ctx.ensure_on_device_sync(validity_bits)?; let (codes_bytes, codes_offsets) = futures::try_join!( ctx.ensure_on_device(codes_bytes_handle), @@ -208,8 +233,7 @@ where let mut views = ctx.device_alloc::(num_strings)?; let codes_bytes_view = codes_bytes.cuda_view::()?; let codes_offsets_view = codes_offsets.cuda_view::()?; - let symbols_view = symbols.cuda_view::()?; - let symbol_lengths_view = symbol_lengths.cuda_view::()?; + let symbol_table_view = symbol_table.cuda_view::()?; let output_offsets_view = output_offsets.cuda_view::()?; let validity_view = validity_device.cuda_view::()?; let ptype = U::PTYPE.to_string(); @@ -218,8 +242,7 @@ where ctx.launch_kernel(&cuda_function, num_strings, |args| { args.arg(&codes_bytes_view) .arg(&codes_offsets_view) - .arg(&symbols_view) - .arg(&symbol_lengths_view) + .arg(&symbol_table_view) .arg(&output_offsets_view) .arg(&validity_view) .arg(&validity_bit_offset) @@ -292,12 +315,10 @@ where let validity = fsst.codes().validity()?; let len = fsst.len(); let len_u64 = len as u64; - let symbols_u64 = fsst - .symbols() - .iter() - .map(|s| s.to_u64()) - .collect::>(); - let symbol_lengths = fsst.padded_symbol_lengths().slice(0..fsst.n_symbols()); + let symbol_table = FSSTSymbolTable::from(&fsst); + let symbol_table = ctx + .stream() + .copy_to_device_sync(std::slice::from_ref(&symbol_table))?; let codes_bytes_handle = fsst.codes_bytes_handle().clone(); let PrimitiveDataParts { buffer: codes_offsets_buffer, @@ -305,8 +326,6 @@ where } = codes_offsets.into_data_parts(); let (validity_bit_offset, validity_bits) = cuda_validity(&validity, len, ctx).await?; - let symbols = ctx.stream().copy_to_device_sync(&symbols_u64)?; - let symbol_lengths = ctx.stream().copy_to_device_sync(symbol_lengths.as_ref())?; let validity_device = ctx.ensure_on_device_sync(validity_bits)?; let (codes_bytes, codes_offsets) = futures::try_join!( ctx.ensure_on_device(codes_bytes_handle), @@ -325,8 +344,7 @@ where let codes_bytes_view = codes_bytes.cuda_view::()?; let codes_offsets_view = codes_offsets.cuda_view::()?; - let symbols_view = symbols.cuda_view::()?; - let symbol_lengths_view = symbol_lengths.cuda_view::()?; + let symbol_table_view = symbol_table.cuda_view::()?; let output_offsets_view = output_offsets.cuda_view::()?; let validity_view = validity_device.cuda_view::()?; let ptype = U::PTYPE.to_string(); @@ -335,8 +353,7 @@ where ctx.launch_kernel(&cuda_function, len, |args| { args.arg(&codes_bytes_view) .arg(&codes_offsets_view) - .arg(&symbols_view) - .arg(&symbol_lengths_view) + .arg(&symbol_table_view) .arg(&output_offsets_view) .arg(&validity_view) .arg(&validity_bit_offset) @@ -431,12 +448,10 @@ where })); } - let symbols_u64 = fsst - .symbols() - .iter() - .map(|s| s.to_u64()) - .collect::>(); - let symbol_lengths = fsst.padded_symbol_lengths().slice(0..fsst.n_symbols()); + let symbol_table = FSSTSymbolTable::from(&fsst); + let symbol_table = ctx + .stream() + .copy_to_device_sync(std::slice::from_ref(&symbol_table))?; let codes_bytes_handle = fsst.codes_bytes_handle().clone(); let PrimitiveDataParts { buffer: codes_offsets_buffer, @@ -445,9 +460,7 @@ where let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx).await?; - let (symbols, symbol_lengths, output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( - ctx.copy_to_device(symbols_u64)?, - ctx.copy_to_device(symbol_lengths)?, + let (output_offsets, validity_device, codes_bytes, codes_offsets) = futures::try_join!( ctx.copy_to_device(output_offsets)?, ctx.ensure_on_device(validity_bits), ctx.ensure_on_device(codes_bytes_handle), @@ -469,8 +482,7 @@ where let codes_bytes_view = codes_bytes.cuda_view::()?; let codes_offsets_view = codes_offsets.cuda_view::()?; - let symbols_view = symbols.cuda_view::()?; - let symbol_lengths_view = symbol_lengths.cuda_view::()?; + let symbol_table_view = symbol_table.cuda_view::()?; let output_offsets_view = output_offsets.cuda_view::()?; let validity_view = validity_device.cuda_view::()?; @@ -479,8 +491,7 @@ where ctx.launch_kernel(&cuda_function, num_strings, |args| { args.arg(&codes_bytes_view) .arg(&codes_offsets_view) - .arg(&symbols_view) - .arg(&symbol_lengths_view) + .arg(&symbol_table_view) .arg(&output_offsets_view) .arg(&validity_view) .arg(&validity_bit_offset)