Skip to content

Commit 67b04d9

Browse files
0ax10ax1
committed
feat(cuda): export FSST directly to Arrow varbin
Decode FSST directly into device-resident Arrow offsets and values while retaining opt-in VarBinView export. Add session-level layout selection, coverage for both layouts, and CUDA benchmarks. Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>
1 parent 159207e commit 67b04d9

8 files changed

Lines changed: 672 additions & 119 deletions

File tree

‎vortex-cuda/benches/fsst_cuda.rs‎

Lines changed: 106 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
//! CUDA benchmarks for FSST decompression.
55
66
#![expect(clippy::unwrap_used)]
7-
#![expect(clippy::cast_possible_truncation)]
87

98
#[allow(dead_code)]
109
mod bench_config;
@@ -18,12 +17,20 @@ use criterion::BenchmarkId;
1817
use criterion::Criterion;
1918
use criterion::Throughput;
2019
use futures::executor::block_on;
20+
use vortex::array::ArrayRef;
2121
use vortex::array::IntoArray;
2222
use vortex::array::arrays::PrimitiveArray;
2323
use vortex::array::match_each_integer_ptype;
24+
use vortex::dtype::DType;
25+
use vortex::dtype::Nullability;
26+
use vortex::encodings::fsst::FSST;
2427
use vortex::encodings::fsst::FSSTArrayExt;
2528
use vortex::error::VortexExpect;
29+
use vortex_cuda::CudaDispatchMode;
2630
use vortex_cuda::CudaSession;
31+
use vortex_cuda::VarBinExportLayout;
32+
use vortex_cuda::arrow::DeviceArrayExt;
33+
use vortex_cuda::arrow::release_device_array;
2734
use vortex_cuda::executor::CudaArrayExt;
2835
use vortex_cuda_macros::cuda_available;
2936
use vortex_cuda_macros::cuda_not_available;
@@ -36,30 +43,63 @@ use crate::timed_launch_strategy::TimedLaunchStrategy;
3643
// other kernels benchmark.
3744
const BENCH_SIZES: &[(usize, &str)] = &[(10_000_000, "10M")];
3845

46+
fn session_with_varbin_layout(layout: VarBinExportLayout) -> vortex::session::VortexSession {
47+
vortex::array::array_session().with_some(
48+
CudaSession::try_default()
49+
.vortex_expect("failed to create CUDA session")
50+
.with_varbin_export_layout(layout),
51+
)
52+
}
53+
54+
struct FSSTBenchFixture {
55+
utf8: ArrayRef,
56+
binary: ArrayRef,
57+
uncompressed_size: u64,
58+
}
59+
60+
fn make_fixture(n: usize) -> FSSTBenchFixture {
61+
let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session())
62+
.vortex_expect("failed to create execution context");
63+
let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx());
64+
65+
let lens = fsst
66+
.uncompressed_lengths()
67+
.clone()
68+
.execute::<PrimitiveArray>(setup_ctx.execution_ctx())
69+
.vortex_expect("canonicalize uncompressed_lengths");
70+
#[allow(clippy::unnecessary_cast)]
71+
let uncompressed_size = match_each_integer_ptype!(lens.ptype(), |P| {
72+
lens.as_slice::<P>().iter().map(|x| *x as u64).sum()
73+
});
74+
75+
let binary = FSST::try_new(
76+
DType::Binary(Nullability::NonNullable),
77+
fsst.symbols().clone(),
78+
fsst.symbol_lengths().clone(),
79+
fsst.codes(),
80+
fsst.uncompressed_lengths().clone(),
81+
setup_ctx.execution_ctx(),
82+
)
83+
.vortex_expect("rebuild FSST fixture with Binary dtype")
84+
.into_array();
85+
86+
FSSTBenchFixture {
87+
utf8: fsst.into_array(),
88+
binary,
89+
uncompressed_size,
90+
}
91+
}
92+
3993
fn benchmark_fsst_cuda_decompress(c: &mut Criterion) {
4094
let mut group = c.benchmark_group("cuda");
4195

4296
for &(n, len_str) in BENCH_SIZES {
43-
let mut setup_ctx = CudaSession::create_execution_ctx(&vortex_cuda::cuda_session())
44-
.vortex_expect("failed to create execution context");
45-
let fsst = make_fsst_clickbench_urls(n, setup_ctx.execution_ctx());
46-
47-
let lens = fsst
48-
.uncompressed_lengths()
49-
.clone()
50-
.execute::<PrimitiveArray>(setup_ctx.execution_ctx())
51-
.vortex_expect("canonicalize uncompressed_lengths");
52-
let total_size: usize = match_each_integer_ptype!(lens.ptype(), |P| {
53-
lens.as_slice::<P>().iter().map(|x| *x as usize).sum()
54-
});
55-
let uncompressed_size = total_size as u64;
56-
57-
let fsst_array = fsst.into_array();
58-
59-
group.throughput(Throughput::Bytes(uncompressed_size));
97+
let fixture = make_fixture(n);
98+
99+
group.throughput(Throughput::Bytes(fixture.uncompressed_size));
60100
group.bench_with_input(
61-
BenchmarkId::new("cuda/fsst/decompress", len_str),
62-
&fsst_array,
101+
BenchmarkId::new("cuda/fsst/decompress_to_varbinview", len_str),
102+
&fixture.utf8,
63103
|b, fsst_array| {
64104
b.iter_custom(|iters| {
65105
let timed = TimedLaunchStrategy::default();
@@ -68,6 +108,7 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) {
68108
let mut cuda_ctx =
69109
CudaSession::create_execution_ctx(&vortex_cuda::cuda_session())
70110
.vortex_expect("failed to create execution context")
111+
.with_dispatch_mode(CudaDispatchMode::StandaloneOnly)
71112
.with_launch_strategy(Arc::new(timed));
72113

73114
for _ in 0..iters {
@@ -77,6 +118,51 @@ fn benchmark_fsst_cuda_decompress(c: &mut Criterion) {
77118
});
78119
},
79120
);
121+
122+
for (name, array, layout) in [
123+
(
124+
"cuda/fsst/decompress_to_varbin",
125+
&fixture.utf8,
126+
VarBinExportLayout::VarBin,
127+
),
128+
(
129+
"cuda/fsst/export_binary",
130+
&fixture.binary,
131+
VarBinExportLayout::VarBin,
132+
),
133+
(
134+
"cuda/fsst/export_utf8_view",
135+
&fixture.utf8,
136+
VarBinExportLayout::VarBinView,
137+
),
138+
(
139+
"cuda/fsst/export_binary_view",
140+
&fixture.binary,
141+
VarBinExportLayout::VarBinView,
142+
),
143+
] {
144+
group.bench_with_input(BenchmarkId::new(name, len_str), array, |b, array| {
145+
b.iter_custom(|iters| {
146+
let timed = TimedLaunchStrategy::default();
147+
let timer = timed.timer();
148+
149+
let session = session_with_varbin_layout(layout);
150+
let mut cuda_ctx = CudaSession::create_execution_ctx(&session)
151+
.vortex_expect("failed to create execution context")
152+
.with_dispatch_mode(CudaDispatchMode::StandaloneOnly)
153+
.with_launch_strategy(Arc::new(timed));
154+
155+
for _ in 0..iters {
156+
let mut exported =
157+
block_on((*array).clone().export_device_array(&mut cuda_ctx))
158+
.vortex_expect("export FSST device array");
159+
release_device_array(&mut exported);
160+
}
161+
162+
Duration::from_nanos(timer.load(Ordering::Relaxed))
163+
});
164+
});
165+
}
80166
}
81167

82168
group.finish();

‎vortex-cuda/kernels/src/fsst.cu‎

Lines changed: 93 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -123,13 +123,18 @@ struct Scratch {
123123
}
124124
};
125125

126-
template <typename OffT>
126+
// Arrow/Vortex variable-length view records are 16 bytes. Values up to 12
127+
// bytes are stored inline after the u32 length. Longer values store their
128+
// first four bytes, backing-buffer index, and byte offset.
129+
constexpr uint32_t MAX_INLINED_SIZE = 12;
130+
131+
template <typename CodeOffsetT, typename OutputOffsetT>
127132
struct FSSTArgs {
128133
// Compressed FSST code stream, contiguous across all strings. String
129134
// `sid`'s codes live in `[codes_offsets[sid], codes_offsets[sid + 1])`.
130135
const uint8_t *__restrict codes_bytes;
131136
// Per-string offsets into `codes_bytes`, length `num_strings + 1`.
132-
const OffT *__restrict codes_offsets;
137+
const CodeOffsetT *__restrict codes_offsets;
133138
// FSST symbol table.
134139
const uint64_t *__restrict symbols;
135140
// Length in bytes (1..=8) of each entry in `symbols`. The remaining bits
@@ -138,19 +143,55 @@ struct FSSTArgs {
138143
// Buffer to write decoded data into.
139144
uint8_t *__restrict output_bytes;
140145
// Per-string offsets into `output_bytes`, length `num_strings + 1`.
141-
const uint64_t *__restrict output_offsets;
146+
const OutputOffsetT *__restrict output_offsets;
142147
// Validity of each string.
143148
const uint8_t *__restrict validity_bits;
149+
// Optional output views, one 16-byte uint4 per string. A null pointer
150+
// requests bytes-only decoding for heaps that need the host rollover path.
151+
uint4 *__restrict output_views;
144152
};
145153

146-
template <typename OffT>
147-
__device__ inline void fsst_decode_string(const FSSTArgs<OffT> &args, uint64_t sid) {
154+
// Build one BinaryView from the bytes this thread just decoded. The Rust
155+
// caller only provides output_views when every offset fits in the view's u32
156+
// fields and the decoded heap is exposed as backing buffer zero.
157+
template <typename CodeOffsetT, typename OutputOffsetT>
158+
__device__ inline void fsst_write_view(const FSSTArgs<CodeOffsetT, OutputOffsetT> &args, uint64_t sid) {
159+
if (args.output_views == nullptr) {
160+
return;
161+
}
162+
163+
const uint64_t start = args.output_offsets[sid];
164+
const uint32_t len = (uint32_t)(args.output_offsets[sid + 1] - start);
165+
if (len <= MAX_INLINED_SIZE) {
166+
uint32_t words[3] = {0, 0, 0};
167+
#pragma unroll
168+
for (uint32_t i = 0; i < MAX_INLINED_SIZE; i++) {
169+
if (i < len) {
170+
words[i >> 2] |= (uint32_t)args.output_bytes[start + i] << (8u * (i & 3u));
171+
}
172+
}
173+
args.output_views[sid] = make_uint4(len, words[0], words[1], words[2]);
174+
return;
175+
}
176+
177+
const uint32_t prefix = (uint32_t)args.output_bytes[start] |
178+
((uint32_t)args.output_bytes[start + 1] << 8u) |
179+
((uint32_t)args.output_bytes[start + 2] << 16u) |
180+
((uint32_t)args.output_bytes[start + 3] << 24u);
181+
args.output_views[sid] = make_uint4(len, prefix, 0, (uint32_t)start);
182+
}
183+
184+
template <typename CodeOffsetT, typename OutputOffsetT>
185+
__device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffsetT> &args, uint64_t sid) {
148186
if (((args.validity_bits[sid >> 3] >> (sid & 7u)) & 1u) == 0u) {
187+
if (args.output_views != nullptr) {
188+
args.output_views[sid] = make_uint4(0, 0, 0, 0);
189+
}
149190
return;
150191
}
151192

152-
OffT in_pos = args.codes_offsets[sid];
153-
const OffT in_end = args.codes_offsets[sid + 1];
193+
CodeOffsetT in_pos = args.codes_offsets[sid];
194+
const CodeOffsetT in_end = args.codes_offsets[sid + 1];
154195
uint64_t out_pos = args.output_offsets[sid];
155196
const uint64_t out_end = args.output_offsets[sid + 1];
156197

@@ -181,46 +222,66 @@ __device__ inline void fsst_decode_string(const FSSTArgs<OffT> &args, uint64_t s
181222
sym &= mask;
182223

183224
scratch.push(sym, len);
184-
in_pos += (OffT)consumed;
225+
in_pos += (CodeOffsetT)consumed;
185226
}
186227

187228
// Epilogue: drain everything that's left.
188229
while (scratch.cursor > 0) {
189230
scratch.drain(args.output_bytes, out_pos, out_end);
190231
}
232+
233+
fsst_write_view(args, sid);
191234
}
192235

193-
#define GENERATE_FSST_KERNEL(suffix, OffT) \
236+
#define FSST_GRID_STRIDE_LOOP(CodeOffsetT, OutputOffsetT, args) \
237+
const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \
238+
const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \
239+
const uint64_t block_end = (block_start + elements_per_block < num_strings) \
240+
? (block_start + elements_per_block) \
241+
: num_strings; \
242+
for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \
243+
fsst_decode_string<CodeOffsetT, OutputOffsetT>(args, sid); \
244+
}
245+
246+
#define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \
194247
extern "C" __global__ void fsst_##suffix(const uint8_t *__restrict codes_bytes, \
195-
const OffT *__restrict codes_offsets, \
248+
const CodeOffsetT *__restrict codes_offsets, \
196249
const uint64_t *__restrict symbols, \
197250
const uint8_t *__restrict symbol_lengths, \
198251
const uint64_t *__restrict output_offsets, \
199252
const uint8_t *__restrict validity_bits, \
200253
uint8_t *__restrict output_bytes, \
254+
uint4 *__restrict output_views, \
201255
uint64_t num_strings) { \
202-
const FSSTArgs<OffT> args = { \
203-
codes_bytes, \
204-
codes_offsets, \
205-
symbols, \
206-
symbol_lengths, \
207-
output_bytes, \
208-
output_offsets, \
209-
validity_bits, \
256+
const FSSTArgs<CodeOffsetT, uint64_t> args = { \
257+
codes_bytes, codes_offsets, symbols, symbol_lengths, output_bytes, output_offsets, \
258+
validity_bits, output_views, \
210259
}; \
211-
\
212-
const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \
213-
const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \
214-
const uint64_t block_end = (block_start + elements_per_block < num_strings) \
215-
? (block_start + elements_per_block) \
216-
: num_strings; \
217-
\
218-
for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \
219-
fsst_decode_string<OffT>(args, sid); \
220-
} \
260+
FSST_GRID_STRIDE_LOOP(CodeOffsetT, uint64_t, args) \
221261
}
222262

223-
GENERATE_FSST_KERNEL(u8, uint8_t)
224-
GENERATE_FSST_KERNEL(u16, uint16_t)
225-
GENERATE_FSST_KERNEL(u32, uint32_t)
226-
GENERATE_FSST_KERNEL(u64, uint64_t)
263+
#define GENERATE_FSST_VARBIN_KERNEL(suffix, CodeOffsetT) \
264+
extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \
265+
const CodeOffsetT *__restrict codes_offsets, \
266+
const uint64_t *__restrict symbols, \
267+
const uint8_t *__restrict symbol_lengths, \
268+
const int32_t *__restrict output_offsets, \
269+
const uint8_t *__restrict validity_bits, \
270+
uint8_t *__restrict output_bytes, \
271+
uint64_t num_strings) { \
272+
const FSSTArgs<CodeOffsetT, int32_t> args = { \
273+
codes_bytes, codes_offsets, symbols, symbol_lengths, output_bytes, output_offsets, \
274+
validity_bits, nullptr, \
275+
}; \
276+
FSST_GRID_STRIDE_LOOP(CodeOffsetT, int32_t, args) \
277+
}
278+
279+
GENERATE_FSST_VIEW_KERNEL(u8, uint8_t)
280+
GENERATE_FSST_VIEW_KERNEL(u16, uint16_t)
281+
GENERATE_FSST_VIEW_KERNEL(u32, uint32_t)
282+
GENERATE_FSST_VIEW_KERNEL(u64, uint64_t)
283+
284+
GENERATE_FSST_VARBIN_KERNEL(u8, uint8_t)
285+
GENERATE_FSST_VARBIN_KERNEL(u16, uint16_t)
286+
GENERATE_FSST_VARBIN_KERNEL(u32, uint32_t)
287+
GENERATE_FSST_VARBIN_KERNEL(u64, uint64_t)

0 commit comments

Comments
 (0)