Skip to content

Commit f742e8a

Browse files
0ax10ax1
authored
feat(cuda): export FSST directly to Arrow varbin (#8787)
## Summary Decode CUDA FSST arrays directly into Arrow’s standard offset-based `Utf8`/`Binary` layout: `i32` offsets plus one contiguous device values buffer. Previously, FSST decoded to `Utf8View`/`BinaryView`. Since cuDF does not consume view arrays, it then had to materialize the same offset-based representation itself. Direct varbin export avoids that additional conversion and enables BenchPress GPU reads for Vortex string columns. ## Behavior changes - CUDA sessions now default variable-length exports to standard Arrow `Utf8`/`Binary`. - Consumers that support views can opt in with: ```rust CudaSession::with_varbin_export_layout(VarBinExportLayout::VarBinView) --------- Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>
1 parent 54b045b commit f742e8a

8 files changed

Lines changed: 1091 additions & 176 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: 107 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,9 @@
5353
// the next add (≤ 8 bytes) within the 24-byte capacity.
5454
//
5555
// `codes_offsets` is templated over the four unsigned integer widths
56-
// (u8/u16/u32/u64). `output_offsets` is uint64_t.
56+
// (u8/u16/u32/u64). `output_offsets` is uint64_t for the view kernels
57+
// (`fsst_*`, which also take an optional views pointer) and int32_t Arrow
58+
// varbin offsets for the bytes-only varbin kernels (`fsst_varbin_*`).
5759

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

126-
template <typename OffT>
128+
// Arrow/Vortex variable-length view records are 16 bytes. Values up to 12
129+
// bytes are stored inline after the u32 length. Longer values store their
130+
// first four bytes, backing-buffer index, and byte offset.
131+
constexpr uint32_t MAX_INLINED_SIZE = 12;
132+
133+
template <typename CodeOffsetT, typename OutputOffsetT>
127134
struct FSSTArgs {
128135
// Compressed FSST code stream, contiguous across all strings. String
129136
// `sid`'s codes live in `[codes_offsets[sid], codes_offsets[sid + 1])`.
130137
const uint8_t *__restrict codes_bytes;
131138
// Per-string offsets into `codes_bytes`, length `num_strings + 1`.
132-
const OffT *__restrict codes_offsets;
139+
const CodeOffsetT *__restrict codes_offsets;
133140
// FSST symbol table.
134141
const uint64_t *__restrict symbols;
135142
// Length in bytes (1..=8) of each entry in `symbols`. The remaining bits
@@ -138,19 +145,57 @@ struct FSSTArgs {
138145
// Buffer to write decoded data into.
139146
uint8_t *__restrict output_bytes;
140147
// Per-string offsets into `output_bytes`, length `num_strings + 1`.
141-
const uint64_t *__restrict output_offsets;
148+
const OutputOffsetT *__restrict output_offsets;
142149
// Validity of each string.
143150
const uint8_t *__restrict validity_bits;
151+
// Bit offset of string zero within the first validity byte (0..7).
152+
uint64_t validity_bit_offset;
153+
// Optional output views, one 16-byte uint4 per string. A null pointer
154+
// requests bytes-only decoding for heaps that need the host rollover path.
155+
uint4 *__restrict output_views;
144156
};
145157

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

152-
OffT in_pos = args.codes_offsets[sid];
153-
const OffT in_end = args.codes_offsets[sid + 1];
167+
const uint64_t start = args.output_offsets[sid];
168+
const uint32_t len = (uint32_t)(args.output_offsets[sid + 1] - start);
169+
if (len <= MAX_INLINED_SIZE) {
170+
uint32_t words[3] = {0, 0, 0};
171+
#pragma unroll
172+
for (uint32_t i = 0; i < MAX_INLINED_SIZE; i++) {
173+
if (i < len) {
174+
words[i >> 2] |= (uint32_t)args.output_bytes[start + i] << (8u * (i & 3u));
175+
}
176+
}
177+
args.output_views[sid] = make_uint4(len, words[0], words[1], words[2]);
178+
return;
179+
}
180+
181+
const uint32_t prefix =
182+
(uint32_t)args.output_bytes[start] | ((uint32_t)args.output_bytes[start + 1] << 8u) |
183+
((uint32_t)args.output_bytes[start + 2] << 16u) | ((uint32_t)args.output_bytes[start + 3] << 24u);
184+
args.output_views[sid] = make_uint4(len, prefix, 0, (uint32_t)start);
185+
}
186+
187+
template <typename CodeOffsetT, typename OutputOffsetT>
188+
__device__ inline void fsst_decode_string(const FSSTArgs<CodeOffsetT, OutputOffsetT> &args, uint64_t sid) {
189+
const uint64_t validity_index = sid + args.validity_bit_offset;
190+
if (((args.validity_bits[validity_index >> 3] >> (validity_index & 7u)) & 1u) == 0u) {
191+
if (args.output_views != nullptr) {
192+
args.output_views[sid] = make_uint4(0, 0, 0, 0);
193+
}
194+
return;
195+
}
196+
197+
CodeOffsetT in_pos = args.codes_offsets[sid];
198+
const CodeOffsetT in_end = args.codes_offsets[sid + 1];
154199
uint64_t out_pos = args.output_offsets[sid];
155200
const uint64_t out_end = args.output_offsets[sid + 1];
156201

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

183228
scratch.push(sym, len);
184-
in_pos += (OffT)consumed;
229+
in_pos += (CodeOffsetT)consumed;
185230
}
186231

187232
// Epilogue: drain everything that's left.
188233
while (scratch.cursor > 0) {
189234
scratch.drain(args.output_bytes, out_pos, out_end);
190235
}
236+
237+
fsst_write_view(args, sid);
191238
}
192239

193-
#define GENERATE_FSST_KERNEL(suffix, OffT) \
240+
#define FSST_GRID_STRIDE_LOOP(CodeOffsetT, OutputOffsetT, args) \
241+
const uint64_t elements_per_block = (uint64_t)blockDim.x * ELEMENTS_PER_THREAD; \
242+
const uint64_t block_start = (uint64_t)blockIdx.x * elements_per_block; \
243+
const uint64_t block_end = \
244+
(block_start + elements_per_block < num_strings) ? (block_start + elements_per_block) : num_strings; \
245+
for (uint64_t sid = block_start + threadIdx.x; sid < block_end; sid += blockDim.x) { \
246+
fsst_decode_string<CodeOffsetT, OutputOffsetT>(args, sid); \
247+
}
248+
249+
#define GENERATE_FSST_VIEW_KERNEL(suffix, CodeOffsetT) \
194250
extern "C" __global__ void fsst_##suffix(const uint8_t *__restrict codes_bytes, \
195-
const OffT *__restrict codes_offsets, \
251+
const CodeOffsetT *__restrict codes_offsets, \
196252
const uint64_t *__restrict symbols, \
197253
const uint8_t *__restrict symbol_lengths, \
198254
const uint64_t *__restrict output_offsets, \
199255
const uint8_t *__restrict validity_bits, \
256+
uint64_t validity_bit_offset, \
200257
uint8_t *__restrict output_bytes, \
258+
uint4 *__restrict output_views, \
201259
uint64_t num_strings) { \
202-
const FSSTArgs<OffT> args = { \
260+
const FSSTArgs<CodeOffsetT, uint64_t> args = { \
203261
codes_bytes, \
204262
codes_offsets, \
205263
symbols, \
206264
symbol_lengths, \
207265
output_bytes, \
208266
output_offsets, \
209267
validity_bits, \
268+
validity_bit_offset, \
269+
output_views, \
210270
}; \
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-
} \
271+
FSST_GRID_STRIDE_LOOP(CodeOffsetT, uint64_t, args) \
221272
}
222273

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)
274+
#define GENERATE_FSST_VARBIN_KERNEL(suffix, CodeOffsetT) \
275+
extern "C" __global__ void fsst_varbin_##suffix(const uint8_t *__restrict codes_bytes, \
276+
const CodeOffsetT *__restrict codes_offsets, \
277+
const uint64_t *__restrict symbols, \
278+
const uint8_t *__restrict symbol_lengths, \
279+
const int32_t *__restrict output_offsets, \
280+
const uint8_t *__restrict validity_bits, \
281+
uint64_t validity_bit_offset, \
282+
uint8_t *__restrict output_bytes, \
283+
uint64_t num_strings) { \
284+
const FSSTArgs<CodeOffsetT, int32_t> args = { \
285+
codes_bytes, \
286+
codes_offsets, \
287+
symbols, \
288+
symbol_lengths, \
289+
output_bytes, \
290+
output_offsets, \
291+
validity_bits, \
292+
validity_bit_offset, \
293+
nullptr, \
294+
}; \
295+
FSST_GRID_STRIDE_LOOP(CodeOffsetT, int32_t, args) \
296+
}
297+
298+
GENERATE_FSST_VIEW_KERNEL(u8, uint8_t)
299+
GENERATE_FSST_VIEW_KERNEL(u16, uint16_t)
300+
GENERATE_FSST_VIEW_KERNEL(u32, uint32_t)
301+
GENERATE_FSST_VIEW_KERNEL(u64, uint64_t)
302+
303+
GENERATE_FSST_VARBIN_KERNEL(u8, uint8_t)
304+
GENERATE_FSST_VARBIN_KERNEL(u16, uint16_t)
305+
GENERATE_FSST_VARBIN_KERNEL(u32, uint32_t)
306+
GENERATE_FSST_VARBIN_KERNEL(u64, uint64_t)

0 commit comments

Comments
 (0)