Skip to content

GPU: pack FSST metadata into one upload #9437

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

Draft
joseph-isaacs wants to merge 1 commit into
base: joe/gpu-fsst-no-metadata-callbacks
Choose a base branch
from
Draft

GPU: pack FSST metadata into one upload #9437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
38 changes: 19 additions & 19 deletions vortex-cuda/kernels/src/fsst.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 <typename CodeOffsetT, typename OutputOffsetT>
struct FSSTArgs {
// Compressed FSST code stream, contiguous across all strings. String
Expand Down Expand Up @@ -249,8 +252,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffs
#define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \
extern "C" __global__ void fsst_##suffix(const uint8_t *__restrict codes_bytes, \
const CodeOffsetT *__restrict codes_offsets, \
const uint64_t *__restrict symbols, \
const uint8_t *__restrict symbol_lengths, \
const FSSTSymbolTable *__restrict symbol_table, \
const uint64_t *__restrict output_offsets, \
const uint8_t *__restrict validity_bits, \
uint64_t validity_bit_offset, \
Expand All @@ -260,8 +262,8 @@ __device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffs
const FSSTArgs<CodeOffsetT, uint64_t> args = { \
codes_bytes, \
codes_offsets, \
symbols, \
symbol_lengths, \
symbol_table->symbols, \
symbol_table->symbol_lengths, \
output_bytes, \
output_offsets, \
validity_bits, \
Expand All @@ -274,8 +276,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffs
#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 FSSTSymbolTable *__restrict symbol_table, \
const int32_t *__restrict output_offsets, \
const uint8_t *__restrict validity_bits, \
uint64_t validity_bit_offset, \
Expand All @@ -284,8 +285,8 @@ __device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffs
const FSSTArgs<CodeOffsetT, int32_t> args = { \
codes_bytes, \
codes_offsets, \
symbols, \
symbol_lengths, \
symbol_table->symbols, \
symbol_table->symbol_lengths, \
output_bytes, \
output_offsets, \
validity_bits, \
Expand All @@ -298,8 +299,7 @@ __device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffs
#define GENERATE_FSST_VARBINVIEW_KERNEL(suffix, CodeOffsetT) \
extern "C" __global__ void fsst_varbinview_##suffix(const uint8_t *__restrict codes_bytes, \
const CodeOffsetT *__restrict codes_offsets, \
const uint64_t *__restrict symbols, \
const uint8_t *__restrict symbol_lengths, \
const FSSTSymbolTable *__restrict symbol_table, \
const int32_t *__restrict output_offsets, \
const uint8_t *__restrict validity_bits, \
uint64_t validity_bit_offset, \
Expand All @@ -309,8 +309,8 @@ __device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffs
const FSSTArgs<CodeOffsetT, int32_t> args = { \
codes_bytes, \
codes_offsets, \
symbols, \
symbol_lengths, \
symbol_table->symbols, \
symbol_table->symbol_lengths, \
output_bytes, \
output_offsets, \
validity_bits, \
Expand Down
85 changes: 48 additions & 37 deletions vortex-cuda/src/kernel/encodings/fsst.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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::<FSSTSymbolTable>() == 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
Expand Down Expand Up @@ -183,21 +212,17 @@ 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::<Vec<_>>();
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,
..
} = 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),
Expand All @@ -208,8 +233,7 @@ where
let mut views = ctx.device_alloc::<i128>(num_strings)?;
let codes_bytes_view = codes_bytes.cuda_view::<u8>()?;
let codes_offsets_view = codes_offsets.cuda_view::<U>()?;
let symbols_view = symbols.cuda_view::<u64>()?;
let symbol_lengths_view = symbol_lengths.cuda_view::<u8>()?;
let symbol_table_view = symbol_table.cuda_view::<FSSTSymbolTable>()?;
let output_offsets_view = output_offsets.cuda_view::<i32>()?;
let validity_view = validity_device.cuda_view::<u8>()?;
let ptype = U::PTYPE.to_string();
Expand All @@ -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)
Expand Down Expand Up @@ -292,21 +315,17 @@ 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::<Vec<_>>();
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,
..
} = 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),
Expand All @@ -325,8 +344,7 @@ where

let codes_bytes_view = codes_bytes.cuda_view::<u8>()?;
let codes_offsets_view = codes_offsets.cuda_view::<U>()?;
let symbols_view = symbols.cuda_view::<u64>()?;
let symbol_lengths_view = symbol_lengths.cuda_view::<u8>()?;
let symbol_table_view = symbol_table.cuda_view::<FSSTSymbolTable>()?;
let output_offsets_view = output_offsets.cuda_view::<i32>()?;
let validity_view = validity_device.cuda_view::<u8>()?;
let ptype = U::PTYPE.to_string();
Expand All @@ -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)
Expand Down Expand Up @@ -431,12 +448,10 @@ where
}));
}

let symbols_u64 = fsst
.symbols()
.iter()
.map(|s| s.to_u64())
.collect::<Vec<_>>();
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,
Expand All @@ -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),
Expand All @@ -469,8 +482,7 @@ where

let codes_bytes_view = codes_bytes.cuda_view::<u8>()?;
let codes_offsets_view = codes_offsets.cuda_view::<U>()?;
let symbols_view = symbols.cuda_view::<u64>()?;
let symbol_lengths_view = symbol_lengths.cuda_view::<u8>()?;
let symbol_table_view = symbol_table.cuda_view::<FSSTSymbolTable>()?;
let output_offsets_view = output_offsets.cuda_view::<u64>()?;
let validity_view = validity_device.cuda_view::<u8>()?;

Expand All @@ -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)
Expand Down
Loading