From 37af5b06bf27568475ee57648ed728718c2b83a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 12:02:29 +0000 Subject: [PATCH 01/22] Add a GPU Parquet decompression backend to the compression benchmark The GPU compression benchmark only measured Vortex, and only on a single dataset, so it could not say anything about how Vortex GPU decompression compares to Parquet, nor about encodings beyond FSST strings. Parquet compresses each page body independently, which is exactly the batch shape nvCOMP's device decompressors take and how cuDF's Parquet reader gets pages off the CPU. This adds a Parquet backend built on that: column chunks are staged on the device through the same pinned, direct-I/O reader the Vortex backend uses, then every page in a row group is decompressed in one batched nvCOMP launch. - vortex-nvcomp: bind the batched Snappy decompression entrypoints and the per-algorithm alignment queries, and share `DecompressBackend` between the Snappy and Zstd wrappers. - compress-bench: locate compressed page bodies by walking the per-page Thrift headers (`parquet::format::PageHeader` is deprecated and `parquet`'s own parser is crate-private), and write files with GPU-friendly settings: v1 pages, dictionary encoding, 1 MiB pages, Snappy by default. - Run both Vortex and Parquet under `--gpu-decompress`, and expand the GPU dataset set from one to nine so ALP, bit-packed, run-end, date/time-parts and null-heavy columns are covered alongside FSST strings. - Add `--gpu-verify`, which compares every GPU-decompressed page against the host codec and every GPU-decoded Vortex field against the CPU decode, and run it as a CI step before the timed benchmark. Independently of that flag, nvCOMP's per-page status and size arrays are checked on every iteration. Page decoding is not part of the Parquet measurement, so its numbers are an upper bound on a full GPU Parquet reader; the README states this. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 10 + Cargo.lock | 4 + Cargo.toml | 1 + benchmarks/compress-bench/Cargo.toml | 10 +- benchmarks/compress-bench/README.md | 73 ++- benchmarks/compress-bench/src/gpu_parquet.rs | 551 ++++++++++++++++++ benchmarks/compress-bench/src/gpu_vortex.rs | 47 +- benchmarks/compress-bench/src/lib.rs | 3 + benchmarks/compress-bench/src/main.rs | 164 ++++-- .../compress-bench/src/parquet_pages.rs | 530 +++++++++++++++++ vortex-cuda/nvcomp/build.rs | 7 + vortex-cuda/nvcomp/src/backend.rs | 52 ++ vortex-cuda/nvcomp/src/lib.rs | 4 + vortex-cuda/nvcomp/src/snappy.rs | 222 +++++++ vortex-cuda/nvcomp/src/zstd.rs | 47 +- 15 files changed, 1655 insertions(+), 70 deletions(-) create mode 100644 benchmarks/compress-bench/src/gpu_parquet.rs create mode 100644 benchmarks/compress-bench/src/parquet_pages.rs create mode 100644 vortex-cuda/nvcomp/src/backend.rs create mode 100644 vortex-cuda/nvcomp/src/snappy.rs diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 3abb5b54533..45bc59fb317 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -44,6 +44,16 @@ jobs: cargo build --locked --package compress-bench --profile release_debug --features cuda,unstable_encodings - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh + - name: Verify GPU decompression correctness + shell: bash + env: + RUST_BACKTRACE: full + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + # Cross-checks every GPU-decompressed page and field against the CPU decoders before + # any timings are taken. Verification runs inline, so this pass is not timed. + run: | + target/release_debug/compress-bench \ + --gpu-decompress --gpu-verify --iterations 1 -d table - name: Run GPU compression benchmark shell: bash env: diff --git a/Cargo.lock b/Cargo.lock index c764cd8eab7..e7c6dd2dfdb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1566,12 +1566,15 @@ dependencies = [ "async-trait", "bytes", "clap", + "cudarc", "futures", "indicatif", "itertools 0.14.0", "lance-bench", "parquet 58.4.0", "regex", + "rstest", + "snap", "tempfile", "tokio", "tracing", @@ -1579,6 +1582,7 @@ dependencies = [ "vortex-arrow", "vortex-bench", "vortex-cuda", + "zstd", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 997b3281d25..edf730760b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -245,6 +245,7 @@ similar = "3.0.0" sketches-ddsketch = "0.4.0" smallvec = "1.15.1" smol = "2.0.2" +snap = "1.1.2" spatialbench = "0.2" spatialbench-arrow = "0.2" # spatialbench still pins arrow 56, two majors behind the workspace arrow. Until upstream diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 4046a12d42e..c20e2861b2d 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,12 +21,14 @@ arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } +cudarc = { workspace = true, optional = true } futures = { workspace = true } indicatif = { workspace = true } itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } +snap = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } @@ -34,9 +36,14 @@ vortex = { workspace = true } vortex-arrow = { workspace = true } vortex-bench = { workspace = true } vortex-cuda = { workspace = true, optional = true } +zstd = { workspace = true } + +[dev-dependencies] +rstest = { workspace = true } +tempfile = { workspace = true } [features] -cuda = ["dep:tempfile", "dep:vortex-cuda"] +cuda = ["dep:cudarc", "dep:tempfile", "dep:vortex-cuda"] lance = ["dep:lance-bench"] unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"] @@ -45,7 +52,6 @@ name = "compress-bench" test = false [lib] -test = false [lints] workspace = true diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index bf2d3efc1db..132ce153450 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -15,13 +15,78 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format cargo run -p compress-bench --profile release_debug ``` -GPU decompression is opt-in and runs only the existing benchmark names allow-listed in -`src/main.rs`: +## GPU decompression + +`--gpu-decompress` is opt-in, requires the `cuda` feature, and restricts the suite to the +GPU dataset list in `src/main.rs`. It measures decompression only, for two backends: + +- **Vortex** — the file is written with CUDA-compatible BtrBlocks encodings only + (`only_cuda_compatible`) and a CUDA flat layout, then decoded on the device all the way to + canonical arrays. +- **Parquet** — the file is rewritten with GPU-friendly writer settings (see below) and every + page body is decompressed on the device with nvCOMP's batched Snappy or Zstd entrypoints, + the same decomposition cuDF's Parquet reader uses: column chunks are staged on the device, + then all pages go through a single batched launch. ```bash cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress + +# pick the Parquet page codec (default: snappy) +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --gpu-parquet-codec zstd ``` -On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure -storage bandwidth rather than page-cache hits. +On Linux both backends read through the same pinned, direct-I/O (`O_DIRECT`) reader, so +repeated iterations measure storage bandwidth rather than page-cache hits. + +### What the Parquet GPU number does and does not include + +Included: column chunk I/O, the host-to-device transfer, and the batched codec launch. + +Not included: page *decoding* — the dictionary, RLE and plain decoders that turn a +decompressed page into an Arrow array — because there is no Rust GPU Parquet page decoder to +call. The Vortex backend it is compared against decodes all the way to canonical arrays, so +the Parquet figure is an upper bound on what a full GPU Parquet reader could reach, and the +`vortex:parquet- gpu ratio decompress time` metric is biased in Parquet's favour. + +Walking the per-page Thrift headers also happens on the host, once per file, outside the +measurement; a real GPU Parquet reader decodes page headers on the device. + +### GPU-friendly Parquet writer settings + +Set in `src/parquet_pages.rs`: + +| Setting | Value | Why | +| --- | --- | --- | +| writer version | `PARQUET_1_0` | v1 pages compress the whole page body, which is the unit nvCOMP decompresses. v2 pages put uncompressed levels ahead of the compressed values in the same body. | +| compression | Snappy (default) or Zstd | The two Parquet codecs nvCOMP implements. Snappy has the higher device throughput and is the Parquet default. | +| dictionary | enabled | Keeps the decompressed payload small; the encoding GPU Parquet readers decode fastest. | +| data page size | 1 MiB | Large enough to amortize per-chunk setup, small enough to keep every SM fed. Matches the page size cuDF targets. | +| data page row limit | 1,000,000 | The 20k-row default caps narrow columns' pages far below 1 MiB. | +| statistics | chunk-level | Page statistics only inflate the headers that have to be walked on the host. | + +### Correctness + +`--gpu-verify` cross-checks device output against the CPU decoders on every iteration: + +- Parquet: each decompressed page is copied back and compared byte-for-byte against the host + Snappy/Zstd output for the same compressed bytes. +- Vortex: each GPU-decoded field is copied back and compared against the same field decoded + on the CPU, through Arrow with a pinned target type. + +Verification runs inline, so timings from a verifying run are not comparable to a plain one — +run it as its own pass: + +```bash +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1 +``` + +Independently of `--gpu-verify`, every Parquet GPU run checks nvCOMP's per-page status and +output-size arrays after the measurement, because a batched launch reports per-page failures +in device memory rather than by failing the launch. + +Page-header scanning is covered by CPU-only unit tests in `src/parquet_pages.rs`, which +assert the located page bodies decompress to exactly the bytes `parquet`'s own page reader +produces. diff --git a/benchmarks/compress-bench/src/gpu_parquet.rs b/benchmarks/compress-bench/src/gpu_parquet.rs new file mode 100644 index 00000000000..b8e6195d784 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_parquet.rs @@ -0,0 +1,551 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! GPU Parquet decompression backend. +//! +//! Parquet's compressed unit is the page, and every page body is an independent block of +//! Snappy or Zstd. That is precisely the batch shape nvCOMP's device decompressors take, and +//! it is how cuDF's Parquet reader gets pages off the CPU: read the column chunks to the +//! device, then decompress every page in one batched launch. +//! +//! What this measures is the decompression stage of a Parquet read — column chunk I/O, +//! host-to-device transfer, and the batched codec launch. Page *decoding* (the dictionary, +//! RLE and plain decoders that turn a decompressed page into an Arrow array) is not +//! included, because there is no Rust GPU Parquet page decoder to call. The Vortex GPU +//! backend it is compared against decodes all the way to canonical arrays, so the Parquet +//! numbers here are an upper bound on what a full GPU Parquet reader could achieve, and the +//! comparison is favourable to Parquet. + +use std::fs::File; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; + +use anyhow::Result; +use anyhow::anyhow; +use anyhow::ensure; +use arrow_array::RecordBatch; +use async_trait::async_trait; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtr; +use cudarc::driver::DevicePtrMut; +use futures::StreamExt; +use futures::TryStreamExt; +use futures::stream; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::metadata::ParquetMetaDataReader; +use tempfile::NamedTempFile; +use vortex::array::buffer::BufferHandle; +use vortex::array::buffer::DeviceBuffer; +use vortex::buffer::Alignment; +use vortex::buffer::Buffer; +use vortex::error::vortex_err; +use vortex::io::VortexReadAt; +use vortex::io::session::RuntimeSessionExt; +use vortex_bench::Format; +use vortex_bench::SESSION; +use vortex_bench::compress::Compressor; +use vortex_cuda::CudaBufferExt; +use vortex_cuda::CudaDeviceBuffer; +use vortex_cuda::CudaExecutionCtx; +use vortex_cuda::CudaSession; +use vortex_cuda::CudaSessionExt; +use vortex_cuda::PooledFileReadAt; +use vortex_cuda::PooledFileReadAtOptions; +use vortex_cuda::nvcomp::AlignmentRequirements; +use vortex_cuda::nvcomp::snappy; +use vortex_cuda::nvcomp::sys; +use vortex_cuda::nvcomp::sys::nvcompStatus_t; +use vortex_cuda::nvcomp::zstd as nvcomp_zstd; + +use crate::parquet_pages::ColumnChunkPages; +use crate::parquet_pages::GpuCodec; +use crate::parquet_pages::gpu_writer_properties; +use crate::parquet_pages::scan_compressed_pages; + +/// Parquet compressor whose decompression measurement runs the page codec on the GPU. +pub struct GpuParquetCompressor { + codec: GpuCodec, + verify: bool, +} + +impl GpuParquetCompressor { + /// Create a backend that writes pages with `codec` and decompresses them with nvCOMP. + /// + /// When `verify` is set, every decompressed page is copied back and compared against the + /// host codec's output before the measurement is reported. + pub fn new(codec: GpuCodec, verify: bool) -> Self { + Self { codec, verify } + } + + /// Rewrite the source Parquet file with GPU-friendly writer settings. + fn write_gpu_parquet(&self, parquet_path: &Path) -> Result<(NamedTempFile, u64)> { + let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(parquet_path)?)?; + let schema = Arc::clone(builder.schema()); + let batches: Vec = builder.build()?.collect::, _>>()?; + + let output = NamedTempFile::new()?; + let mut writer = ArrowWriter::try_new( + output.reopen()?, + schema, + Some(gpu_writer_properties(self.codec)), + )?; + for batch in batches { + writer.write(&batch)?; + } + writer.flush()?; + let size = writer.bytes_written() as u64; + writer.close()?; + Ok((output, size)) + } + + fn check_codec(&self, metadata: &ParquetMetaData) -> Result<()> { + for row_group in metadata.row_groups() { + for column in row_group.columns() { + ensure!( + self.codec.matches(column.compression()), + "column {} was written with {:?}, expected {}", + column.column_path(), + column.compression(), + self.codec.name() + ); + } + } + Ok(()) + } +} + +#[async_trait] +impl Compressor for GpuParquetCompressor { + fn format(&self) -> Format { + Format::Parquet + } + + async fn compress(&self, parquet_path: &Path) -> Result<(u64, Duration)> { + let start = Instant::now(); + let (_file, size) = self.write_gpu_parquet(parquet_path)?; + Ok((size, start.elapsed())) + } + + async fn decompress(&self, parquet_path: &Path) -> Result { + let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?; + + let file = File::open(gpu_file.path())?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; + self.check_codec(&metadata)?; + drop(file); + + // The page map is a property of the file, so it is resolved once rather than on every + // iteration. A GPU Parquet reader decodes page headers on the device as part of the + // read; this benchmark cannot, so the host walk is kept out of the measurement. + let file_bytes = std::fs::read(gpu_file.path())?; + let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; + let plan = DecompressPlan::build(&row_groups, self.codec)?; + // Only verification needs the host copy afterwards, to decompress the same bytes. + let file_bytes = self.verify.then_some(file_bytes); + + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + let reader = open_reader(gpu_file.path(), &cuda_ctx)?; + + // Row groups are staged and released one at a time so device residency stays bounded + // by the largest row group rather than by the whole file, matching how the Vortex + // backend streams one batch at a time. + let mut elapsed = Duration::ZERO; + let mut checks = Vec::with_capacity(plan.row_groups.len()); + for row_group in &plan.row_groups { + let start = Instant::now(); + let device_chunks = stage_row_group(row_group, &reader).await?; + let output = + decompress_pages(row_group, &device_chunks, self.codec, &mut cuda_ctx).await?; + cuda_ctx.synchronize_stream()?; + elapsed += start.elapsed(); + + let DecompressOutput { + output, + actual_sizes, + statuses, + } = output; + drop(device_chunks); + + if self.verify { + // Verified inline, since holding every row group's decompressed output would + // defeat the point of streaming them. This is why a verifying run's timings + // are not comparable. + let file_bytes = file_bytes + .as_deref() + .ok_or_else(|| anyhow!("verification requires the host file bytes"))?; + verify_against_host(row_group, file_bytes, self.codec, output).await?; + } + checks.push((actual_sizes, statuses)); + } + + // Checked outside the measurement, but on every iteration: a batched launch reports + // per-page failures in device memory rather than by failing the launch, so without + // this a broken run would simply look fast. + for (row_group, (actual_sizes, statuses)) in plan.row_groups.iter().zip(checks) { + check_statuses(row_group, actual_sizes, statuses).await?; + } + + Ok(elapsed) + } +} + +/// Stages a row group's column chunks on the device, bounding concurrent pinned staging. +async fn stage_row_group( + row_group: &RowGroupPlan, + reader: &PooledFileReadAt, +) -> Result> { + let reads = row_group + .chunks + .iter() + .map(|chunk| reader.read_at(chunk.offset, chunk.len, Alignment::of::())) + .collect::>(); + + stream::iter(reads) + .buffered(STAGING_CONCURRENCY) + .try_collect::>() + .await + .map_err(|e| anyhow!("failed to stage Parquet column chunks on the device: {e}")) +} + +/// Column chunks staged on the device at once. Matches `vortex-cuda`'s file read concurrency. +const STAGING_CONCURRENCY: usize = 32; + +/// Opens the GPU Parquet file through the same pinned, direct-I/O reader the Vortex GPU +/// backend uses, so both formats measure storage bandwidth rather than page-cache hits. +fn open_reader(path: &Path, cuda_ctx: &CudaExecutionCtx) -> Result { + let pool = Arc::clone(SESSION.cuda_session().pinned_buffer_pool()); + let options = PooledFileReadAtOptions::default(); + #[cfg(target_os = "linux")] + let options = options.with_direct_io(); + + Ok(PooledFileReadAt::open_with_options( + path, + SESSION.handle(), + pool, + cuda_ctx.stream().clone(), + options, + )?) +} + +/// A page to decompress, addressed relative to its column chunk's device buffer. +struct PlannedPage { + chunk: usize, + offset_in_chunk: usize, + compressed_len: usize, + uncompressed_len: usize, + output_offset: usize, +} + +/// The byte range of a column chunk to stage on the device. +struct PlannedChunk { + offset: u64, + len: usize, +} + +/// Everything needed to issue one batched nvCOMP launch over a row group's pages. +struct RowGroupPlan { + chunks: Vec, + pages: Vec, + output_len: usize, + max_uncompressed: usize, +} + +/// The per-row-group work for one file. +struct DecompressPlan { + row_groups: Vec, +} + +impl DecompressPlan { + fn build(row_groups: &[Vec], codec: GpuCodec) -> Result { + let alignment = decompress_alignments(codec)?; + ensure!( + alignment.output.is_power_of_two(), + "nvcomp reported a non-power-of-two output alignment of {}", + alignment.output + ); + + let row_groups = row_groups + .iter() + .map(|chunks| RowGroupPlan::build(chunks, codec, alignment)) + .collect::>>()?; + + ensure!( + row_groups.iter().any(|plan| !plan.pages.is_empty()), + "Parquet file contains no compressed pages" + ); + + Ok(Self { row_groups }) + } +} + +impl RowGroupPlan { + fn build( + chunks: &[ColumnChunkPages], + codec: GpuCodec, + alignment: AlignmentRequirements, + ) -> Result { + let mut planned_chunks = Vec::with_capacity(chunks.len()); + let mut pages = Vec::new(); + let mut output_len = 0usize; + let mut max_uncompressed = 0usize; + + for (index, chunk) in chunks.iter().enumerate() { + planned_chunks.push(PlannedChunk { + offset: chunk.offset, + len: chunk.len, + }); + + for page in &chunk.pages { + let offset_in_chunk = usize::try_from( + u64::try_from(page.offset)? + .checked_sub(chunk.offset) + .ok_or_else(|| anyhow!("page offset precedes its column chunk"))?, + )?; + // Pages are decompressed in place from their column chunk's device buffer. + // CUDA allocations are at least 256-byte aligned, so a page's device address + // meets nvcomp's requirement exactly when its chunk-relative offset does. + ensure!( + offset_in_chunk.is_multiple_of(alignment.input), + "page at file offset {} sits {offset_in_chunk} bytes into its column chunk, \ + which does not meet nvcomp's {} byte input alignment for {}; \ + use --gpu-parquet-codec snappy", + page.offset, + alignment.input, + codec.name() + ); + + let output_offset = output_len.next_multiple_of(alignment.output); + output_len = output_offset + page.uncompressed_len; + max_uncompressed = max_uncompressed.max(page.uncompressed_len); + + pages.push(PlannedPage { + chunk: index, + offset_in_chunk, + compressed_len: page.compressed_len, + uncompressed_len: page.uncompressed_len, + output_offset, + }); + } + } + + Ok(Self { + chunks: planned_chunks, + pages, + output_len, + max_uncompressed, + }) + } +} + +fn decompress_alignments(codec: GpuCodec) -> Result { + match codec { + GpuCodec::Snappy => { + snappy::decompress_alignment_requirements(snappy::SnappyDecompressOpts::default()) + } + GpuCodec::Zstd => nvcomp_zstd::decompress_alignment_requirements( + nvcomp_zstd::ZstdDecompressOpts::default(), + ), + } + .map_err(|e| anyhow!("nvcomp alignment query failed: {e}")) +} + +/// Device-side outputs of a batched decompression launch. +struct DecompressOutput { + output: CudaSlice, + actual_sizes: CudaSlice, + statuses: CudaSlice, +} + +/// Enqueues the batched decompression of every page in `plan` onto the context's stream. +async fn decompress_pages( + plan: &RowGroupPlan, + device_chunks: &[BufferHandle], + codec: GpuCodec, + ctx: &mut CudaExecutionCtx, +) -> Result { + let num_pages = plan.pages.len(); + + let temp_size = match codec { + GpuCodec::Snappy => { + snappy::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) + } + GpuCodec::Zstd => { + nvcomp_zstd::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) + } + } + .map_err(|e| anyhow!("nvcomp temp size query failed: {e}"))?; + + let chunk_bases = device_chunks + .iter() + .map(|handle| handle.cuda_device_ptr()) + .collect::, _>>()?; + + let mut output = ctx.device_alloc::(plan.output_len)?; + // Only the allocation address is needed to build the output pointer table; the device + // write itself is tracked by the guard taken around the launch below. + let output_base = { + let (base, _) = output.device_ptr(ctx.stream()); + base + }; + + let mut compressed_ptrs = Vec::with_capacity(num_pages); + let mut compressed_sizes = Vec::with_capacity(num_pages); + let mut uncompressed_sizes = Vec::with_capacity(num_pages); + let mut output_ptrs = Vec::with_capacity(num_pages); + for page in &plan.pages { + compressed_ptrs.push(chunk_bases[page.chunk] + page.offset_in_chunk as u64); + compressed_sizes.push(page.compressed_len); + uncompressed_sizes.push(page.uncompressed_len); + output_ptrs.push(output_base + page.output_offset as u64); + } + + let (compressed_ptrs, compressed_sizes, uncompressed_sizes, output_ptrs) = futures::try_join!( + ctx.copy_to_device(compressed_ptrs)?, + ctx.copy_to_device(compressed_sizes)?, + ctx.copy_to_device(uncompressed_sizes)?, + ctx.copy_to_device(output_ptrs)? + )?; + + let mut actual_sizes: CudaSlice = ctx.device_alloc(num_pages)?; + let mut statuses: CudaSlice = ctx.device_alloc(num_pages)?; + let mut temp: CudaSlice = ctx.device_alloc(temp_size)?; + + let stream = ctx.stream(); + let compressed_ptrs_view = compressed_ptrs.cuda_view::()?; + let compressed_sizes_view = compressed_sizes.cuda_view::()?; + let uncompressed_sizes_view = uncompressed_sizes.cuda_view::()?; + let output_ptrs_view = output_ptrs.cuda_view::()?; + + let (compressed_ptrs_ptr, record_compressed_ptrs) = compressed_ptrs_view.device_ptr(stream); + let (compressed_sizes_ptr, record_compressed_sizes) = compressed_sizes_view.device_ptr(stream); + let (uncompressed_sizes_ptr, record_uncompressed_sizes) = + uncompressed_sizes_view.device_ptr(stream); + let (output_ptrs_ptr, record_output_ptrs) = output_ptrs_view.device_ptr(stream); + let (_output_ptr, record_output) = output.device_ptr_mut(stream); + let (actual_sizes_ptr, record_actual_sizes) = actual_sizes.device_ptr_mut(stream); + let (statuses_ptr, record_statuses) = statuses.device_ptr_mut(stream); + let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); + + ctx.launch_external(plan.output_len, || { + // SAFETY: every pointer is derived from a live device allocation sized by the plan, + // and each batch metadata array holds exactly `num_pages` entries. + unsafe { + match codec { + GpuCodec::Snappy => snappy::decompress_async( + compressed_ptrs_ptr as _, + compressed_sizes_ptr as _, + uncompressed_sizes_ptr as _, + actual_sizes_ptr as _, + num_pages, + temp_ptr as _, + temp_size, + output_ptrs_ptr as _, + statuses_ptr as _, + stream.cu_stream().cast(), + ), + GpuCodec::Zstd => nvcomp_zstd::decompress_async( + compressed_ptrs_ptr as _, + compressed_sizes_ptr as _, + uncompressed_sizes_ptr as _, + actual_sizes_ptr as _, + num_pages, + temp_ptr as _, + temp_size, + output_ptrs_ptr as _, + statuses_ptr as _, + stream.cu_stream().cast(), + ), + } + .map_err(|e| vortex_err!("nvcomp decompress_async failed: {}", e)) + } + })?; + + drop(( + record_compressed_ptrs, + record_compressed_sizes, + record_uncompressed_sizes, + record_output_ptrs, + record_output, + record_actual_sizes, + record_statuses, + record_temp, + )); + // The temporary workspace must outlive the launch, which the stream ordering guarantees + // only while the allocation is alive. + drop(temp); + + Ok(DecompressOutput { + output, + actual_sizes, + statuses, + }) +} + +/// Copies the per-page status and size arrays back and fails on any mismatch. +async fn check_statuses( + plan: &RowGroupPlan, + actual_sizes: CudaSlice, + statuses: CudaSlice, +) -> Result<()> { + let statuses = CudaDeviceBuffer::new(statuses) + .copy_to_host(Alignment::of::())? + .await?; + let actual_sizes = CudaDeviceBuffer::new(actual_sizes) + .copy_to_host(Alignment::of::())? + .await?; + + let statuses = Buffer::::from_byte_buffer(statuses); + let actual_sizes = Buffer::::from_byte_buffer(actual_sizes); + + for (index, page) in plan.pages.iter().enumerate() { + let status = statuses.as_slice()[index]; + ensure!( + status == sys::nvcompStatus_t_nvcompSuccess, + "page {index} failed to decompress with nvcomp status {status}" + ); + let actual = actual_sizes.as_slice()[index]; + ensure!( + actual == page.uncompressed_len, + "page {index} decompressed to {actual} bytes, expected {}", + page.uncompressed_len + ); + } + + Ok(()) +} + +/// Compares every decompressed page against the host codec's output for the same bytes. +async fn verify_against_host( + plan: &RowGroupPlan, + file_bytes: &[u8], + codec: GpuCodec, + output: CudaSlice, +) -> Result<()> { + let device_output = CudaDeviceBuffer::new(output) + .copy_to_host(Alignment::of::())? + .await?; + let device_output = device_output.as_slice(); + + for (index, page) in plan.pages.iter().enumerate() { + let chunk = &plan.chunks[page.chunk]; + let start = usize::try_from(chunk.offset)? + page.offset_in_chunk; + let compressed = &file_bytes[start..start + page.compressed_len]; + let expected = codec.decompress_host(compressed, page.uncompressed_len)?; + let actual = &device_output[page.output_offset..page.output_offset + page.uncompressed_len]; + ensure!( + actual == expected.as_slice(), + "page {index} decompressed on the GPU differs from the host codec output" + ); + } + + tracing::info!( + "verified {} GPU-decompressed {} pages against the host codec", + plan.pages.len(), + codec.name() + ); + Ok(()) +} diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 3dbb68bc7a8..0e31d78ec00 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -8,9 +8,13 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use anyhow::ensure; +use arrow_schema::Field; use async_trait::async_trait; use futures::StreamExt; use tempfile::NamedTempFile; +use vortex::array::ArrayRef; +use vortex::array::ExecutionCtx; use vortex::array::IntoArray; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; @@ -18,10 +22,12 @@ use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteStrategyBuilder; +use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::compress::Compressor; use vortex_bench::conversions::parquet_to_vortex_chunks; +use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; #[cfg(target_os = "linux")] @@ -31,7 +37,20 @@ use vortex_cuda::layout::CudaFlatLayoutStrategy; use vortex_cuda::layout::register_cuda_layout; /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. -pub struct GpuVortexCompressor; +pub struct GpuVortexCompressor { + verify: bool, +} + +impl GpuVortexCompressor { + /// Create the backend. + /// + /// When `verify` is set, each GPU-decoded field is copied back to the host and compared + /// against the same field decoded on the CPU. Verification runs inline, so timings from a + /// verifying run are not comparable to a plain one. + pub fn new(verify: bool) -> Self { + Self { verify } + } +} #[async_trait] impl Compressor for GpuVortexCompressor { @@ -75,7 +94,13 @@ impl Compressor for GpuVortexCompressor { while let Some(batch) = batches.next().await { let record = batch?.execute::(cuda_ctx.execution_ctx())?; for field in record.iter_unmasked_fields() { - black_box(field.clone().execute_cuda(&mut cuda_ctx).await?); + let decoded = field.clone().execute_cuda(&mut cuda_ctx).await?; + if self.verify { + let host = decoded.into_host().await?.into_array(); + verify_field(field, host, cuda_ctx.execution_ctx())?; + } else { + black_box(decoded); + } } } cuda_ctx.synchronize_stream()?; @@ -83,3 +108,21 @@ impl Compressor for GpuVortexCompressor { Ok(start.elapsed()) } } + +/// Fails unless a GPU-decoded field matches the same field decoded on the CPU. +fn verify_field(compressed: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { + let expected = SESSION + .arrow() + .execute_arrow(compressed.clone(), None, ctx)?; + // Pin the Arrow target type so the two sides cannot land on different but equivalent + // encodings of the same logical values. + let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); + let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?; + + ensure!( + expected.to_data() == actual.to_data(), + "GPU decode of a {} field does not match the CPU decode", + compressed.encoding_id() + ); + Ok(()) +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index 68039996605..e82f8574dc6 100644 --- a/benchmarks/compress-bench/src/lib.rs +++ b/benchmarks/compress-bench/src/lib.rs @@ -4,6 +4,9 @@ #[cfg(feature = "lance")] pub use lance_bench::compress::LanceCompressor; #[cfg(feature = "cuda")] +pub mod gpu_parquet; +#[cfg(feature = "cuda")] pub mod gpu_vortex; pub mod parquet; +pub mod parquet_pages; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1b1603e52c8..a915112897f 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -8,8 +8,11 @@ use clap::Parser; #[cfg(feature = "lance")] use compress_bench::LanceCompressor; #[cfg(feature = "cuda")] +use compress_bench::gpu_parquet::GpuParquetCompressor; +#[cfg(feature = "cuda")] use compress_bench::gpu_vortex::GpuVortexCompressor; use compress_bench::parquet::ParquetCompressor; +use compress_bench::parquet_pages::GpuCodec; use compress_bench::vortex::VortexCompressor; use indicatif::ProgressBar; use itertools::Itertools; @@ -35,6 +38,7 @@ use vortex_bench::display::DisplayFormat; use vortex_bench::display::print_measurements_json; use vortex_bench::display::render_table; use vortex_bench::downloadable_dataset::DownloadableDataset; +use vortex_bench::measurements::CustomUnitMeasurement; use vortex_bench::public_bi::PBI_DATASETS; use vortex_bench::public_bi::PBIDataset::Arade; use vortex_bench::public_bi::PBIDataset::Bimbo; @@ -67,11 +71,23 @@ struct Args { ops: Vec, #[arg(long)] datasets: Option, - /// Run GPU decompression for the allow-listed benchmarks. + /// Run GPU decompression for the GPU-supported benchmarks. /// - /// This filters the suite to GPU-supported dataset names and runs only Vortex decompression. + /// Restricts the suite to datasets with verified CUDA decode support and measures + /// decompression only, for both Vortex and Parquet. #[arg(long)] gpu_decompress: bool, + /// Page codec the GPU Parquet backend writes and decompresses with. + /// + /// Snappy is the Parquet default and the codec nvCOMP decompresses fastest. + #[arg(long, value_enum, default_value_t)] + gpu_parquet_codec: GpuCodec, + /// Cross-check every GPU-decompressed result against the CPU decoder. + /// + /// Verification runs inline, so timings from a verifying run are not comparable to a + /// plain one. Intended to be run as its own pass. + #[arg(long)] + gpu_verify: bool, #[arg(short, long, default_value_t, value_enum)] display_format: DisplayFormat, #[arg(short, long)] @@ -97,8 +113,16 @@ async fn main() -> anyhow::Result<()> { anyhow::bail!("--gpu-decompress requires building compress-bench with --features cuda"); } - let (formats, ops) = if args.gpu_decompress { - (vec![Format::OnDiskVortex], vec![CompressOp::Decompress]) + let gpu = args.gpu_decompress.then_some(GpuOptions { + codec: args.gpu_parquet_codec, + verify: args.gpu_verify, + }); + + let (formats, ops) = if gpu.is_some() { + ( + vec![Format::Parquet, Format::OnDiskVortex], + vec![CompressOp::Decompress], + ) } else { (args.formats, args.ops) }; @@ -108,7 +132,7 @@ async fn main() -> anyhow::Result<()> { args.datasets.map(|d| Regex::new(&d)).transpose()?, formats, ops, - args.gpu_decompress, + gpu, args.display_format, args.output_path, args.ingest_output, @@ -116,15 +140,35 @@ async fn main() -> anyhow::Result<()> { .await } +/// Settings for the GPU decompression mode. +#[derive(Clone, Copy, Debug)] +struct GpuOptions { + /// Parquet page codec to write and decompress on the device. + codec: GpuCodec, + /// Cross-check decompressed output against the CPU decoders. + #[cfg_attr( + not(feature = "cuda"), + expect(dead_code, reason = "only the CUDA backends read this") + )] + verify: bool, +} + /// Get a compressor for the given format. -fn get_compressor(format: Format, gpu_decompress: bool) -> Box { - if gpu_decompress { +fn get_compressor(format: Format, gpu: Option) -> Box { + if let Some(gpu) = gpu { #[cfg(feature = "cuda")] + return match format { + Format::OnDiskVortex => { + Box::new(GpuVortexCompressor::new(gpu.verify)) as Box + } + Format::Parquet => Box::new(GpuParquetCompressor::new(gpu.codec, gpu.verify)), + _ => unimplemented!("GPU compress bench not implemented for {format}"), + }; + #[cfg(not(feature = "cuda"))] { - return Box::new(GpuVortexCompressor); + let _ = gpu; + unreachable!("GPU feature validation happens before selecting compressors"); } - #[cfg(not(feature = "cuda"))] - unreachable!("GPU feature validation happens before selecting compressors"); } match format { @@ -151,7 +195,7 @@ async fn run_compress( datasets_filter: Option, formats: Vec, ops: Vec, - gpu_decompress: bool, + gpu: Option, display_format: DisplayFormat, output_path: Option, ingest_output: Option, @@ -178,15 +222,26 @@ async fn run_compress( // ), ]; - // Add an existing benchmark name here only after its CUDA-compatible compression and - // decompression kernels have been verified end to end. - #[expect( - clippy::useless_vec, - reason = "this is an intentionally incremental allow-list of benchmark names" - )] - let gpu_decompress_benchmarks = vec!["TPC-H l_comment canonical"]; + // Datasets run in GPU mode. Add one only after its CUDA-compatible compression and + // decompression kernels have been verified end to end with `--gpu-verify`. Between them + // these cover FSST strings, ALP and bit-packed numerics, run-end and date/time-parts + // encodings, and columns with nulls. + // + // `StructListOfInts` is deliberately absent: its list layouts have no verified CUDA + // decode path yet. + let gpu_datasets: [&dyn Dataset; 9] = [ + &TPCHLCommentCanonical as &dyn Dataset, + &TPCHLCommentChunked, + &TaxiData, + PBI_DATASETS.get(Arade), + PBI_DATASETS.get(Bimbo), + PBI_DATASETS.get(CMSprovider), + PBI_DATASETS.get(Euro2016), + PBI_DATASETS.get(Food), + PBI_DATASETS.get(HashTags), + ]; - let datasets: Vec<&dyn Dataset> = [ + let all_datasets: Vec<&dyn Dataset> = [ &TaxiData as &dyn Dataset, PBI_DATASETS.get(Arade), PBI_DATASETS.get(Bimbo), @@ -206,10 +261,15 @@ async fn run_compress( ] .into_iter() .chain(structlistofints.iter().map(|d| d as &dyn Dataset)) + .collect(); + + let datasets: Vec<&dyn Dataset> = if gpu.is_some() { + gpu_datasets.to_vec() + } else { + all_datasets + } + .into_iter() .filter(|d| { - if gpu_decompress && !gpu_decompress_benchmarks.contains(&d.name()) { - return false; - } if let Some(filter) = datasets_filter.as_ref() { filter.is_match(d.name()) } else { @@ -226,15 +286,9 @@ async fn run_compress( let mut v3_records: Vec = Vec::new(); for dataset_handle in datasets.into_iter() { - let (m, mut records) = run_benchmark_for_dataset( - &progress, - &formats, - &ops, - iterations, - dataset_handle, - gpu_decompress, - ) - .await?; + let (m, mut records) = + run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, gpu) + .await?; measurements.push(m); v3_records.append(&mut records); } @@ -275,7 +329,7 @@ async fn run_benchmark_for_dataset( ops: &[CompressOp], iterations: usize, dataset_handle: &dyn Dataset, - gpu_decompress: bool, + gpu: Option, ) -> anyhow::Result<(CompressMeasurements, Vec)> { let bench_name = dataset_handle.name(); let (v3_dataset, v3_variant) = dataset_handle.v3_dataset_dims(); @@ -291,7 +345,7 @@ async fn run_benchmark_for_dataset( let mut v3_records: Vec = Vec::new(); for format in formats { - let compressor = get_compressor(*format, gpu_decompress); + let compressor = get_compressor(*format, gpu); for op in ops { let time = match op { @@ -342,7 +396,7 @@ async fn run_benchmark_for_dataset( v3_records.push(v3::compression_time_record( &result.timing, v3_dataset, - if gpu_decompress { + if gpu.is_some() { Some("gpu") } else { v3_variant @@ -361,12 +415,42 @@ async fn run_benchmark_for_dataset( } // Calculate cross-format ratios after all measurements. - calculate_ratios( - &measurements_map, - &compressed_sizes, - bench_name, - &mut ratios, - ); + match gpu { + // The shared ratio labels name the CPU suite's codec, which the GPU run does not + // necessarily use, so GPU mode emits its own correctly-labelled ratio. + Some(gpu) => push_gpu_ratio(&measurements_map, gpu, bench_name, &mut ratios), + None => calculate_ratios( + &measurements_map, + &compressed_sizes, + bench_name, + &mut ratios, + ), + } Ok((CompressMeasurements { timings, ratios }, v3_records)) } + +/// Emit the Vortex-versus-Parquet decompression ratio for a GPU run. +fn push_gpu_ratio( + measurements: &HashMap<(Format, CompressOp), Duration>, + gpu: GpuOptions, + bench_name: &str, + ratios: &mut Vec, +) { + let (Some(vortex_time), Some(parquet_time)) = ( + measurements.get(&(Format::OnDiskVortex, CompressOp::Decompress)), + measurements.get(&(Format::Parquet, CompressOp::Decompress)), + ) else { + return; + }; + + ratios.push(CustomUnitMeasurement { + name: format!( + "vortex:parquet-{} gpu ratio decompress time/{bench_name}", + gpu.codec.name() + ), + format: Format::OnDiskVortex, + unit: std::borrow::Cow::from("ratio"), + value: vortex_time.as_nanos() as f64 / parquet_time.as_nanos() as f64, + }); +} diff --git a/benchmarks/compress-bench/src/parquet_pages.rs b/benchmarks/compress-bench/src/parquet_pages.rs new file mode 100644 index 00000000000..d55598cd8dd --- /dev/null +++ b/benchmarks/compress-bench/src/parquet_pages.rs @@ -0,0 +1,530 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Locating the compressed page bodies inside a Parquet file. +//! +//! Parquet compresses each page body independently with a block codec, which is exactly the +//! shape nvCOMP's batched decompression entrypoints consume: an array of independent +//! compressed chunks with known uncompressed sizes. This module finds those chunks so the +//! GPU backend can hand the whole batch to the device in one launch, the same decomposition +//! cuDF's Parquet reader uses. +//! +//! Column chunk ranges come from the file footer; page boundaries within a chunk are only +//! discoverable by walking the per-page Thrift headers, so a minimal Thrift compact-protocol +//! reader lives here. `parquet::format::PageHeader` is deprecated and scheduled for removal, +//! and `parquet`'s own page-header parser is crate-private, so neither can be used. + +use anyhow::Result; +use anyhow::bail; +use anyhow::ensure; +use clap::ValueEnum; +use parquet::basic::Compression; +use parquet::basic::ZstdLevel; +use parquet::file::metadata::ParquetMetaData; +use parquet::file::properties::EnabledStatistics; +use parquet::file::properties::WriterProperties; +use parquet::file::properties::WriterVersion; + +/// Target size of a data page written for GPU decompression. +/// +/// nvCOMP decompresses one chunk per page, so pages must be large enough to amortize the +/// per-chunk setup yet numerous enough to fill the device. ~1 MiB is the page size cuDF's +/// Parquet reader is tuned around. +pub const GPU_DATA_PAGE_SIZE: usize = 1024 * 1024; + +/// Row cap per data page. +/// +/// `parquet`'s default caps pages at 20k rows, which produces pages far below +/// [`GPU_DATA_PAGE_SIZE`] for narrow columns and leaves the device underfed. +const GPU_DATA_PAGE_ROW_COUNT_LIMIT: usize = 1_000_000; + +/// Parquet page codecs that nvCOMP can decompress on the device. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)] +pub enum GpuCodec { + /// The Parquet default, and the codec with the highest device-side throughput. + #[default] + Snappy, + /// Matches the codec used by the CPU Parquet benchmark, at lower device throughput. + Zstd, +} + +impl GpuCodec { + /// The Parquet compression setting for this codec. + pub fn to_parquet(self) -> Compression { + match self { + GpuCodec::Snappy => Compression::SNAPPY, + GpuCodec::Zstd => Compression::ZSTD(ZstdLevel::default()), + } + } + + /// Short lowercase name, used in measurement labels. + pub fn name(self) -> &'static str { + match self { + GpuCodec::Snappy => "snappy", + GpuCodec::Zstd => "zstd", + } + } + + /// Whether a column chunk's codec matches this one. + pub fn matches(self, compression: Compression) -> bool { + matches!( + (self, compression), + (GpuCodec::Snappy, Compression::SNAPPY) | (GpuCodec::Zstd, Compression::ZSTD(_)) + ) + } + + /// Decompress a single page body on the host, for cross-checking device output. + pub fn decompress_host(self, compressed: &[u8], uncompressed_len: usize) -> Result> { + let decompressed = match self { + GpuCodec::Snappy => snap::raw::Decoder::new().decompress_vec(compressed)?, + GpuCodec::Zstd => zstd::bulk::decompress(compressed, uncompressed_len)?, + }; + ensure!( + decompressed.len() == uncompressed_len, + "page decompressed to {} bytes, page header declared {uncompressed_len}", + decompressed.len() + ); + Ok(decompressed) + } +} + +/// Writer properties tuned for GPU decompression. +pub fn gpu_writer_properties(codec: GpuCodec) -> WriterProperties { + WriterProperties::builder() + // V1 data pages compress the entire page body, which is the unit nvCOMP decompresses. + // V2 pages place uncompressed repetition/definition levels ahead of the compressed + // values inside one page body, which the batched entrypoints cannot express. + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_compression(codec.to_parquet()) + // Dictionary encoding keeps the decompressed payload small and is the encoding GPU + // Parquet readers decode fastest. + .set_dictionary_enabled(true) + .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) + .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) + // Per-page statistics only inflate the page headers that have to be walked on the host. + .set_statistics_enabled(EnabledStatistics::Chunk) + .build() +} + +/// A compressed page body located within a Parquet file. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CompressedPage { + /// Offset of the compressed body, i.e. just past the page header. + pub offset: usize, + /// Length of the compressed body in bytes. + pub compressed_len: usize, + /// Length of the body once decompressed. + pub uncompressed_len: usize, +} + +/// The pages of one column chunk, alongside the byte range the chunk occupies. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ColumnChunkPages { + /// Offset of the column chunk within the file. + pub offset: u64, + /// Length of the column chunk in bytes. + pub len: usize, + /// Compressed pages of this chunk, in file order, with file-absolute offsets. + pub pages: Vec, +} + +/// Walks every column chunk's page headers and returns the compressed page bodies. +/// +/// The outer `Vec` is one entry per row group, which is the unit a reader can stage on the +/// device and release before moving on. Chunks and the pages within them are in file order. +pub fn scan_compressed_pages( + file_bytes: &[u8], + metadata: &ParquetMetaData, +) -> Result>> { + let mut row_groups = Vec::with_capacity(metadata.row_groups().len()); + + for row_group in metadata.row_groups() { + let mut chunks = Vec::with_capacity(row_group.columns().len()); + for column in row_group.columns() { + let (chunk_offset, chunk_len) = column.byte_range(); + let mut pages = Vec::new(); + let (start, len) = (chunk_offset, chunk_len); + let start = usize::try_from(start)?; + let end = start + .checked_add(usize::try_from(len)?) + .filter(|end| *end <= file_bytes.len()) + .ok_or_else(|| { + anyhow::anyhow!( + "column chunk range {start}..+{len} extends past the {} byte file", + file_bytes.len() + ) + })?; + + let mut pos = start; + while pos < end { + let header = read_page_header(&file_bytes[pos..end])?; + let body = pos + header.header_len; + let body_end = body + .checked_add(header.compressed_len) + .filter(|body_end| *body_end <= end) + .ok_or_else(|| { + anyhow::anyhow!( + "page body at {body} of {} bytes overruns its column chunk", + header.compressed_len + ) + })?; + + match header.page_type { + PageType::Data | PageType::Dictionary => pages.push(CompressedPage { + offset: body, + compressed_len: header.compressed_len, + uncompressed_len: header.uncompressed_len, + }), + PageType::DataV2 => bail!( + "v2 data pages are not GPU-decompressible as a single chunk; \ + write the file with WriterVersion::PARQUET_1_0" + ), + // Index pages are not part of the column data and are never written by + // `parquet`; skip over the body rather than decompressing it. + PageType::Index => {} + } + + pos = body_end; + } + + chunks.push(ColumnChunkPages { + offset: chunk_offset, + len: usize::try_from(chunk_len)?, + pages, + }); + } + + row_groups.push(chunks); + } + + Ok(row_groups) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PageType { + Data, + Index, + Dictionary, + DataV2, +} + +struct PageHeaderInfo { + header_len: usize, + page_type: PageType, + compressed_len: usize, + uncompressed_len: usize, +} + +/// Thrift compact-protocol field types. +mod ttype { + pub(super) const STOP: u8 = 0x00; + pub(super) const BOOL_TRUE: u8 = 0x01; + pub(super) const BOOL_FALSE: u8 = 0x02; + pub(super) const I8: u8 = 0x03; + pub(super) const I16: u8 = 0x04; + pub(super) const I32: u8 = 0x05; + pub(super) const I64: u8 = 0x06; + pub(super) const DOUBLE: u8 = 0x07; + pub(super) const BINARY: u8 = 0x08; + pub(super) const LIST: u8 = 0x09; + pub(super) const SET: u8 = 0x0a; + pub(super) const MAP: u8 = 0x0b; + pub(super) const STRUCT: u8 = 0x0c; + pub(super) const UUID: u8 = 0x0d; +} + +/// Guards against unbounded recursion on malformed headers. +const MAX_STRUCT_DEPTH: u32 = 32; + +/// Reads the `PageHeader` at the start of `buf`, returning its fields and encoded length. +fn read_page_header(buf: &[u8]) -> Result { + let mut reader = CompactReader { buf, pos: 0 }; + let mut page_type = None; + let mut uncompressed_len = None; + let mut compressed_len = None; + let mut last_field_id = 0i16; + + while let Some((field_id, field_type)) = reader.read_field_header(&mut last_field_id)? { + match (field_id, field_type) { + (1, ttype::I32) => page_type = Some(reader.read_i32()?), + (2, ttype::I32) => uncompressed_len = Some(reader.read_i32()?), + (3, ttype::I32) => compressed_len = Some(reader.read_i32()?), + _ => reader.skip_value(field_type, 0)?, + } + } + + let page_type = match page_type { + Some(0) => PageType::Data, + Some(1) => PageType::Index, + Some(2) => PageType::Dictionary, + Some(3) => PageType::DataV2, + Some(other) => bail!("unknown Parquet page type {other}"), + None => bail!("page header is missing its page type"), + }; + + let uncompressed_len = + uncompressed_len.ok_or_else(|| anyhow::anyhow!("page header is missing its size"))?; + let compressed_len = compressed_len + .ok_or_else(|| anyhow::anyhow!("page header is missing its compressed size"))?; + + Ok(PageHeaderInfo { + header_len: reader.pos, + page_type, + compressed_len: usize::try_from(compressed_len)?, + uncompressed_len: usize::try_from(uncompressed_len)?, + }) +} + +/// Minimal reader for the subset of the Thrift compact protocol that page headers use. +struct CompactReader<'a> { + buf: &'a [u8], + pos: usize, +} + +impl CompactReader<'_> { + fn read_u8(&mut self) -> Result { + let byte = *self + .buf + .get(self.pos) + .ok_or_else(|| anyhow::anyhow!("page header ends mid-field"))?; + self.pos += 1; + Ok(byte) + } + + fn advance(&mut self, len: usize) -> Result<()> { + let end = self + .pos + .checked_add(len) + .filter(|end| *end <= self.buf.len()) + .ok_or_else(|| anyhow::anyhow!("page header ends mid-value"))?; + self.pos = end; + Ok(()) + } + + fn read_varint(&mut self) -> Result { + let mut value = 0u64; + for shift in (0..64).step_by(7) { + let byte = self.read_u8()?; + value |= u64::from(byte & 0x7f) << shift; + if byte & 0x80 == 0 { + return Ok(value); + } + } + bail!("varint in page header is not terminated") + } + + fn read_zigzag(&mut self) -> Result { + let encoded = self.read_varint()?; + Ok(((encoded >> 1) as i64) ^ -((encoded & 1) as i64)) + } + + fn read_i32(&mut self) -> Result { + Ok(i32::try_from(self.read_zigzag()?)?) + } + + /// Reads the next field header, or `None` at the struct's STOP byte. + fn read_field_header(&mut self, last_field_id: &mut i16) -> Result> { + let header = self.read_u8()?; + if header == ttype::STOP { + return Ok(None); + } + + let field_type = header & 0x0f; + let delta = header >> 4; + let field_id = if delta == 0 { + i16::try_from(self.read_zigzag()?)? + } else { + last_field_id + .checked_add(i16::from(delta)) + .ok_or_else(|| anyhow::anyhow!("field id overflow in page header"))? + }; + *last_field_id = field_id; + + Ok(Some((field_id, field_type))) + } + + fn skip_struct(&mut self, depth: u32) -> Result<()> { + ensure!( + depth < MAX_STRUCT_DEPTH, + "page header nests structs more than {MAX_STRUCT_DEPTH} deep" + ); + let mut last_field_id = 0i16; + while let Some((_, field_type)) = self.read_field_header(&mut last_field_id)? { + self.skip_value(field_type, depth + 1)?; + } + Ok(()) + } + + /// Skips a field value. Booleans carry their value in the field type, so consume nothing. + fn skip_value(&mut self, field_type: u8, depth: u32) -> Result<()> { + match field_type { + ttype::BOOL_TRUE | ttype::BOOL_FALSE => Ok(()), + ttype::I8 => self.advance(1), + ttype::I16 | ttype::I32 | ttype::I64 => self.read_varint().map(|_| ()), + ttype::DOUBLE => self.advance(8), + ttype::UUID => self.advance(16), + ttype::BINARY => { + let len = usize::try_from(self.read_varint()?)?; + self.advance(len) + } + ttype::LIST | ttype::SET => { + let (len, element_type) = self.read_collection_header()?; + for _ in 0..len { + self.skip_element(element_type, depth + 1)?; + } + Ok(()) + } + ttype::MAP => { + let len = usize::try_from(self.read_varint()?)?; + if len > 0 { + let types = self.read_u8()?; + let (key_type, value_type) = (types >> 4, types & 0x0f); + for _ in 0..len { + self.skip_element(key_type, depth + 1)?; + self.skip_element(value_type, depth + 1)?; + } + } + Ok(()) + } + ttype::STRUCT => self.skip_struct(depth), + other => bail!("unsupported Thrift compact type {other} in page header"), + } + } + + /// Skips one collection element. Unlike fields, booleans here occupy a byte of their own. + fn skip_element(&mut self, element_type: u8, depth: u32) -> Result<()> { + match element_type { + ttype::BOOL_TRUE | ttype::BOOL_FALSE => self.advance(1), + other => self.skip_value(other, depth), + } + } + + fn read_collection_header(&mut self) -> Result<(usize, u8)> { + let header = self.read_u8()?; + let element_type = header & 0x0f; + let len = match header >> 4 { + 0x0f => usize::try_from(self.read_varint()?)?, + short_len => usize::from(short_len), + }; + Ok((len, element_type)) + } +} + +#[cfg(test)] +mod tests { + use std::fs::File; + use std::sync::Arc; + + use arrow_array::Int64Array; + use arrow_array::RecordBatch; + use arrow_array::StringArray; + use arrow_schema::DataType; + use arrow_schema::Field; + use arrow_schema::Schema; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::ParquetMetaDataReader; + use parquet::file::reader::FileReader; + use parquet::file::reader::SerializedFileReader; + use rstest::rstest; + + use super::*; + + fn sample_batch() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("ints", DataType::Int64, false), + Field::new("strings", DataType::Utf8, false), + ])); + let ints = Int64Array::from_iter_values((0..50_000).map(|i| i % 977)); + let strings = + StringArray::from_iter_values((0..50_000).map(|i| format!("value-{}", i % 1_000))); + Ok(RecordBatch::try_new( + schema, + vec![Arc::new(ints), Arc::new(strings)], + )?) + } + + fn write_sample(path: &std::path::Path, codec: GpuCodec) -> Result<()> { + let batch = sample_batch()?; + let file = File::create(path)?; + let mut writer = + ArrowWriter::try_new(file, batch.schema(), Some(gpu_writer_properties(codec)))?; + writer.write(&batch)?; + writer.close()?; + Ok(()) + } + + /// The page bodies we locate must decompress to exactly the bytes `parquet` itself reads. + #[rstest] + #[case(GpuCodec::Snappy)] + #[case(GpuCodec::Zstd)] + fn scanned_pages_match_parquet_reader(#[case] codec: GpuCodec) -> Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("sample.parquet"); + write_sample(&path, codec)?; + + let file = File::open(&path)?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; + let file_bytes = std::fs::read(&path)?; + let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; + let pages = row_groups + .iter() + .flatten() + .flat_map(|chunk| chunk.pages.iter()) + .collect::>(); + + let reader = SerializedFileReader::new(File::open(&path)?)?; + let mut expected = Vec::new(); + for row_group in 0..reader.metadata().num_row_groups() { + let row_group_reader = reader.get_row_group(row_group)?; + for column in 0..row_group_reader.num_columns() { + let mut page_reader = row_group_reader.get_column_page_reader(column)?; + while let Some(page) = page_reader.get_next_page()? { + expected.push(page.buffer().to_vec()); + } + } + } + + assert_eq!(pages.len(), expected.len(), "page count mismatch"); + assert!(!pages.is_empty(), "expected the sample file to have pages"); + + for (page, expected) in pages.iter().zip(expected.iter()) { + let compressed = &file_bytes[page.offset..page.offset + page.compressed_len]; + let decompressed = codec.decompress_host(compressed, page.uncompressed_len)?; + assert_eq!(&decompressed, expected); + } + + Ok(()) + } + + /// Pages must tile their column chunks exactly, with no gap left unaccounted for. + #[test] + fn scanned_pages_cover_every_column_chunk() -> Result<()> { + let dir = tempfile::tempdir()?; + let path = dir.path().join("sample.parquet"); + write_sample(&path, GpuCodec::Snappy)?; + + let file = File::open(&path)?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; + let file_bytes = std::fs::read(&path)?; + let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; + let pages = row_groups + .iter() + .flatten() + .flat_map(|chunk| chunk.pages.iter()) + .collect::>(); + + let compressed: usize = pages.iter().map(|page| page.compressed_len).sum(); + let chunk_total: i64 = metadata + .row_groups() + .iter() + .flat_map(|rg| rg.columns()) + .map(|col| col.compressed_size()) + .sum(); + + // The chunk total includes the page headers, so the page bodies must be strictly + // smaller but within a header's worth per page. + assert!(compressed < usize::try_from(chunk_total)?); + assert!(compressed > usize::try_from(chunk_total)? - pages.len() * 256); + Ok(()) + } +} diff --git a/vortex-cuda/nvcomp/build.rs b/vortex-cuda/nvcomp/build.rs index 877fc68c58d..f34507378db 100644 --- a/vortex-cuda/nvcomp/build.rs +++ b/vortex-cuda/nvcomp/build.rs @@ -90,13 +90,20 @@ fn main() { let bindings = bindgen::Builder::default() .header(include_dir.join("nvcomp.h").to_string_lossy()) .header(include_dir.join("nvcomp/zstd.h").to_string_lossy()) + .header(include_dir.join("nvcomp/snappy.h").to_string_lossy()) .clang_arg(format!("-I{}", include_dir.display())) .clang_arg(format!("-I{}", cuda_stub_dir.display())) .allowlist_type("nvcompStatus_t") + .allowlist_type("nvcompAlignmentRequirements_t") .allowlist_type("nvcompBatchedZstdDecompressOpts_t") + .allowlist_type("nvcompBatchedSnappyDecompressOpts_t") .allowlist_type("nvcompDecompressBackend_t") .allowlist_function("nvcompBatchedZstdDecompressGetTempSizeAsync") + .allowlist_function("nvcompBatchedZstdDecompressGetRequiredAlignments") .allowlist_function("nvcompBatchedZstdDecompressAsync") + .allowlist_function("nvcompBatchedSnappyDecompressGetTempSizeAsync") + .allowlist_function("nvcompBatchedSnappyDecompressGetRequiredAlignments") + .allowlist_function("nvcompBatchedSnappyDecompressAsync") .dynamic_library_name("NvcompLibrary") .dynamic_link_require_all(true) .wrap_unsafe_ops(true) diff --git a/vortex-cuda/nvcomp/src/backend.rs b/vortex-cuda/nvcomp/src/backend.rs new file mode 100644 index 00000000000..5be1ebbde4c --- /dev/null +++ b/vortex-cuda/nvcomp/src/backend.rs @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Options and metadata shared by nvcomp's batched decompression APIs. + +use crate::sys; + +/// Backend selection for nvcomp decompression. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum DecompressBackend { + /// Let nvcomp auto-select the best backend for the hardware. + #[default] + Default, + /// Use hardware decompression + Hardware, + /// Use CUDA + Cuda, +} + +impl DecompressBackend { + pub(crate) fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { + match self { + Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, + Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, + Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, + } + } +} + +/// Minimum buffer alignments required by an nvcomp algorithm. +/// +/// Buffers passed to the batched decompression entrypoints must satisfy these alignments. +/// Exceeding them (for example 16- or 32-byte alignment) may improve throughput. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AlignmentRequirements { + /// Minimum alignment of each compressed input chunk. + pub input: usize, + /// Minimum alignment of each decompressed output chunk. + pub output: usize, + /// Minimum alignment of the temporary workspace buffer. + pub temp: usize, +} + +impl From for AlignmentRequirements { + fn from(value: sys::nvcompAlignmentRequirements_t) -> Self { + Self { + input: value.input, + output: value.output, + temp: value.temp, + } + } +} diff --git a/vortex-cuda/nvcomp/src/lib.rs b/vortex-cuda/nvcomp/src/lib.rs index 3ab502baa63..859eecbea22 100644 --- a/vortex-cuda/nvcomp/src/lib.rs +++ b/vortex-cuda/nvcomp/src/lib.rs @@ -30,9 +30,13 @@ use std::sync::OnceLock; )] pub mod sys; +mod backend; mod error; +pub mod snappy; pub mod zstd; +pub use backend::AlignmentRequirements; +pub use backend::DecompressBackend; pub use error::NvcompError; /// The loaded nvcomp library instance. diff --git a/vortex-cuda/nvcomp/src/snappy.rs b/vortex-cuda/nvcomp/src/snappy.rs new file mode 100644 index 00000000000..175c19d3a13 --- /dev/null +++ b/vortex-cuda/nvcomp/src/snappy.rs @@ -0,0 +1,222 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Wrappers around nvcomp's batched Snappy decompression API. +//! +//! Snappy is the default Parquet page codec and is the fastest of the codecs nvcomp +//! implements on device, which makes it the codec of choice when feeding Parquet pages +//! to the GPU. + +use std::ffi::c_void; + +use crate::backend::AlignmentRequirements; +pub use crate::backend::DecompressBackend; +use crate::error::NvcompError; +use crate::error::check_status; +use crate::nvcomp_library; +use crate::sys; + +/// The largest compressed chunk the Snappy decompressor accepts, in bytes. +pub const MAX_COMPRESSED_CHUNK_SIZE: usize = (1 << 31) - 1; + +/// Options for batched Snappy decompression. +#[derive(Debug, Clone, Copy, Default)] +pub struct SnappyDecompressOpts { + /// Which nvcomp backend performs the decompression. + pub backend: DecompressBackend, + /// Sort chunks by size before submitting them to the hardware decompression engine. + /// + /// Only used when `backend` selects the hardware engine. + pub sort_before_hw_decompress: bool, +} + +impl SnappyDecompressOpts { + fn to_nvcomp(self) -> sys::nvcompBatchedSnappyDecompressOpts_t { + sys::nvcompBatchedSnappyDecompressOpts_t { + backend: self.backend.to_nvcomp(), + sort_before_hw_decompress: i32::from(self.sort_before_hw_decompress), + reserved: [0; 56], + } + } +} + +/// Computes required temporary buffer size for batched Snappy decompression. +/// +/// # Arguments +/// +/// * `num_chunks` - Number of compressed chunks to decompress +/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk +/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks +/// +/// # Returns +/// +/// The required size in bytes for the temporary buffer. +pub fn get_decompress_temp_size( + num_chunks: usize, + max_uncompressed_chunk_bytes: usize, + max_total_uncompressed_bytes: usize, +) -> Result { + get_decompress_temp_size_with_opts( + num_chunks, + max_uncompressed_chunk_bytes, + max_total_uncompressed_bytes, + SnappyDecompressOpts::default(), + ) +} + +/// Computes required temporary buffer size with custom options. +/// +/// # Arguments +/// +/// * `num_chunks` - Number of compressed chunks to decompress +/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk +/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks +/// * `opts` - Decompression options +/// +/// # Returns +/// +/// The required size in bytes for the temporary buffer. +pub fn get_decompress_temp_size_with_opts( + num_chunks: usize, + max_uncompressed_chunk_bytes: usize, + max_total_uncompressed_bytes: usize, + opts: SnappyDecompressOpts, +) -> Result { + let library = nvcomp_library()?; + + let mut temp_bytes: usize = 0; + + let status = unsafe { + library.nvcompBatchedSnappyDecompressGetTempSizeAsync( + num_chunks, + max_uncompressed_chunk_bytes, + opts.to_nvcomp(), + &raw mut temp_bytes, + max_total_uncompressed_bytes, + ) + }; + + check_status(status)?; + Ok(temp_bytes) +} + +/// Returns the minimum buffer alignments required by batched Snappy decompression. +pub fn decompress_alignment_requirements( + opts: SnappyDecompressOpts, +) -> Result { + let library = nvcomp_library()?; + + let mut requirements = sys::nvcompAlignmentRequirements_t { + input: 0, + output: 0, + temp: 0, + }; + + let status = unsafe { + library.nvcompBatchedSnappyDecompressGetRequiredAlignments( + opts.to_nvcomp(), + &raw mut requirements, + ) + }; + + check_status(status)?; + Ok(requirements.into()) +} + +/// Launches batched Snappy decompression asynchronously on the GPU. +/// +/// This function decompresses multiple raw Snappy blocks in parallel on the GPU. All +/// pointer arguments must point to device memory, and the operation is executed +/// asynchronously on the provided CUDA stream. +/// +/// # Arguments +/// +/// * `device_compressed_ptrs` - Device pointer to array of pointers to compressed chunks +/// * `device_compressed_bytes` - Device pointer to array of compressed chunk sizes +/// * `device_uncompressed_bytes` - Device pointer to array of expected uncompressed sizes +/// * `device_actual_uncompressed_bytes` - Device pointer to array for actual uncompressed sizes (output) +/// * `num_chunks` - Number of chunks to decompress +/// * `device_temp_ptr` - Device pointer to temporary workspace buffer +/// * `temp_bytes` - Size of temporary buffer in bytes +/// * `device_uncompressed_ptrs` - Device pointer to array of pointers to output buffers +/// * `device_statuses` - Device pointer to array for per-chunk status codes (output) +/// * `stream` - CUDA stream to execute on +/// +/// # Safety +/// +/// - All device pointers must be valid and point to properly allocated device memory +/// - `device_compressed_ptrs` must point to valid device pointers +/// - `device_uncompressed_ptrs` must point to valid device pointers +/// - Each output buffer must have at least the corresponding `device_uncompressed_bytes` size +/// - `device_temp_ptr` must have at least `temp_bytes` allocated +/// - The stream must be valid +#[expect(clippy::too_many_arguments)] +pub unsafe fn decompress_async( + device_compressed_ptrs: *const *const c_void, + device_compressed_bytes: *const usize, + device_uncompressed_bytes: *const usize, + device_actual_uncompressed_bytes: *mut usize, + num_chunks: usize, + device_temp_ptr: *mut c_void, + temp_bytes: usize, + device_uncompressed_ptrs: *const *mut c_void, + device_statuses: *mut sys::nvcompStatus_t, + stream: sys::cudaStream_t, +) -> Result<(), NvcompError> { + // SAFETY: Caller has to ensure all pointers are valid. + unsafe { + decompress_async_with_opts( + device_compressed_ptrs, + device_compressed_bytes, + device_uncompressed_bytes, + device_actual_uncompressed_bytes, + num_chunks, + device_temp_ptr, + temp_bytes, + device_uncompressed_ptrs, + device_statuses, + stream, + SnappyDecompressOpts::default(), + ) + } +} + +/// Launches batched Snappy decompression asynchronously with custom options. +/// +/// # Safety +/// +/// Same requirements as [`decompress_async`]. +#[expect(clippy::too_many_arguments)] +pub unsafe fn decompress_async_with_opts( + device_compressed_ptrs: *const *const c_void, + device_compressed_bytes: *const usize, + device_uncompressed_bytes: *const usize, + device_actual_uncompressed_bytes: *mut usize, + num_chunks: usize, + device_temp_ptr: *mut c_void, + temp_bytes: usize, + device_uncompressed_ptrs: *const *mut c_void, + device_statuses: *mut sys::nvcompStatus_t, + stream: sys::cudaStream_t, + opts: SnappyDecompressOpts, +) -> Result<(), NvcompError> { + let library = nvcomp_library()?; + + let status = unsafe { + library.nvcompBatchedSnappyDecompressAsync( + device_compressed_ptrs, + device_compressed_bytes, + device_uncompressed_bytes, + device_actual_uncompressed_bytes, + num_chunks, + device_temp_ptr, + temp_bytes, + device_uncompressed_ptrs, + opts.to_nvcomp(), + device_statuses, + stream, + ) + }; + + check_status(status) +} diff --git a/vortex-cuda/nvcomp/src/zstd.rs b/vortex-cuda/nvcomp/src/zstd.rs index 44111901071..04769d2621a 100644 --- a/vortex-cuda/nvcomp/src/zstd.rs +++ b/vortex-cuda/nvcomp/src/zstd.rs @@ -5,33 +5,13 @@ use std::ffi::c_void; +use crate::backend::AlignmentRequirements; +pub use crate::backend::DecompressBackend; use crate::error::NvcompError; use crate::error::check_status; use crate::nvcomp_library; use crate::sys; -/// Backend selection for nvcomp decompression. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum DecompressBackend { - /// Let nvcomp auto-select the best backend for the hardware. - #[default] - Default, - /// Use hardware decompression - Hardware, - /// Use CUDA - Cuda, -} - -impl DecompressBackend { - fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { - match self { - Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, - Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, - Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, - } - } -} - /// Options for batched ZSTD decompression. #[derive(Debug, Clone, Copy, Default)] pub struct ZstdDecompressOpts { @@ -107,6 +87,29 @@ pub fn get_decompress_temp_size_with_opts( Ok(temp_bytes) } +/// Returns the minimum buffer alignments required by batched ZSTD decompression. +pub fn decompress_alignment_requirements( + opts: ZstdDecompressOpts, +) -> Result { + let library = nvcomp_library()?; + + let mut requirements = sys::nvcompAlignmentRequirements_t { + input: 0, + output: 0, + temp: 0, + }; + + let status = unsafe { + library.nvcompBatchedZstdDecompressGetRequiredAlignments( + opts.to_nvcomp(), + &raw mut requirements, + ) + }; + + check_status(status)?; + Ok(requirements.into()) +} + /// Launches batched ZSTD decompression asynchronously on the GPU. /// /// This function decompresses multiple ZSTD-compressed chunks in parallel on the GPU. From 5eb7a6d622bc3e1b956176e148d49db3ac3a45eb Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 12:50:11 +0000 Subject: [PATCH 02/22] Verify GPU Vortex decode against a separate host scan A CUDA scan hands back arrays whose buffers live in device memory, so decoding those same arrays through the host Arrow path panics rather than producing a CPU reference. Read the file a second time through the ordinary host reader and compare the two scans batch by batch instead. Signed-off-by: Claude --- benchmarks/compress-bench/src/gpu_vortex.rs | 101 ++++++++++++++++---- 1 file changed, 82 insertions(+), 19 deletions(-) diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 0e31d78ec00..d9be0b8730f 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -8,6 +8,7 @@ use std::time::Duration; use std::time::Instant; use anyhow::Result; +use anyhow::bail; use anyhow::ensure; use arrow_schema::Field; use async_trait::async_trait; @@ -80,27 +81,19 @@ impl Compressor for GpuVortexCompressor { output.sync_all().await?; drop(output); + if self.verify { + return verify_against_host_scan(gpu_file.path()).await; + } + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); - let open_options = SESSION.open_options().with_cuda(); - // Direct IO keeps repeated iterations measuring storage bandwidth rather than - // page-cache hits. It is only available on Linux. - #[cfg(target_os = "linux")] - let open_options = - open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); - let file = open_options.open_path(gpu_file.path()).await?; + let file = open_gpu(gpu_file.path()).await?; let mut batches = file.scan()?.into_array_stream()?; while let Some(batch) = batches.next().await { let record = batch?.execute::(cuda_ctx.execution_ctx())?; for field in record.iter_unmasked_fields() { - let decoded = field.clone().execute_cuda(&mut cuda_ctx).await?; - if self.verify { - let host = decoded.into_host().await?.into_array(); - verify_field(field, host, cuda_ctx.execution_ctx())?; - } else { - black_box(decoded); - } + black_box(field.clone().execute_cuda(&mut cuda_ctx).await?); } } cuda_ctx.synchronize_stream()?; @@ -109,11 +102,81 @@ impl Compressor for GpuVortexCompressor { } } +/// Opens a Vortex file for CUDA execution. +/// +/// On Linux direct IO keeps repeated iterations measuring storage bandwidth rather than +/// page-cache hits. +async fn open_gpu(path: &Path) -> Result { + let open_options = SESSION.open_options().with_cuda(); + #[cfg(target_os = "linux")] + let open_options = + open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); + Ok(open_options.open_path(path).await?) +} + +/// Decodes the same file on the GPU and on the CPU and fails on the first difference. +/// +/// The CPU reference comes from a second, host-only scan rather than from re-decoding the +/// GPU scan's arrays: a CUDA scan hands back arrays whose buffers live in device memory, +/// which the host decoders cannot read. +/// +/// Verification runs inline, so the returned duration is not comparable to a plain run. +async fn verify_against_host_scan(path: &Path) -> Result { + let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + let start = Instant::now(); + + let gpu_file = open_gpu(path).await?; + let mut gpu_batches = gpu_file.scan()?.into_array_stream()?; + let host_file = SESSION.open_options().open_path(path).await?; + let mut host_batches = host_file.scan()?.into_array_stream()?; + + let mut fields_checked = 0usize; + loop { + let (gpu_batch, host_batch) = (gpu_batches.next().await, host_batches.next().await); + let (gpu_batch, host_batch) = match (gpu_batch, host_batch) { + (Some(gpu_batch), Some(host_batch)) => (gpu_batch?, host_batch?), + (None, None) => break, + _ => bail!("the GPU and CPU scans of the same file produced different batch counts"), + }; + + let gpu_record = gpu_batch.execute::(cuda_ctx.execution_ctx())?; + let host_record = host_batch.execute::(cuda_ctx.execution_ctx())?; + ensure!( + gpu_record.len() == host_record.len(), + "batch length differs between the GPU and CPU scans: {} vs {}", + gpu_record.len(), + host_record.len() + ); + + let gpu_fields = gpu_record + .iter_unmasked_fields() + .cloned() + .collect::>(); + let host_fields = host_record + .iter_unmasked_fields() + .cloned() + .collect::>(); + ensure!( + gpu_fields.len() == host_fields.len(), + "field count differs between the GPU and CPU scans" + ); + + for (gpu_field, host_field) in gpu_fields.into_iter().zip(host_fields) { + let decoded = gpu_field.execute_cuda(&mut cuda_ctx).await?; + let decoded = decoded.into_host().await?.into_array(); + verify_field(&host_field, decoded, cuda_ctx.execution_ctx())?; + fields_checked += 1; + } + } + cuda_ctx.synchronize_stream()?; + + tracing::info!("verified {fields_checked} GPU-decoded Vortex fields against the CPU decode"); + Ok(start.elapsed()) +} + /// Fails unless a GPU-decoded field matches the same field decoded on the CPU. -fn verify_field(compressed: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { - let expected = SESSION - .arrow() - .execute_arrow(compressed.clone(), None, ctx)?; +fn verify_field(host: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { + let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?; // Pin the Arrow target type so the two sides cannot land on different but equivalent // encodings of the same logical values. let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); @@ -122,7 +185,7 @@ fn verify_field(compressed: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> ensure!( expected.to_data() == actual.to_data(), "GPU decode of a {} field does not match the CPU decode", - compressed.encoding_id() + host.encoding_id() ); Ok(()) } From e7a6b4189c3a4d1d4fe8e9f345f5a1d55c684118 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:01:28 +0000 Subject: [PATCH 03/22] Isolate the verification scans and report how a GPU decode differs Two changes to the Vortex GPU verification, after CI reported a `fastlanes.for` mismatch with no detail: - Read the CPU reference from a copy of the file. The session segment cache is keyed by URI and the CUDA reader deliberately bypasses it because its buffers are device-resident, so pointing both scans at one URI risks them sharing entries. - Synchronize the stream before copying a decoded field back, and report the Arrow types, lengths, null counts and the first differing row when the two decodes disagree. Signed-off-by: Claude --- benchmarks/compress-bench/src/gpu_vortex.rs | 92 ++++++++++++++++++--- 1 file changed, 81 insertions(+), 11 deletions(-) diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index d9be0b8730f..9b25b2f6dd4 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -10,6 +10,7 @@ use std::time::Instant; use anyhow::Result; use anyhow::bail; use anyhow::ensure; +use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::Field; use async_trait::async_trait; use futures::StreamExt; @@ -125,12 +126,19 @@ async fn verify_against_host_scan(path: &Path) -> Result { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); + // The host scan reads a copy rather than the same path. The session's segment cache is + // keyed by URI, and the CUDA reader deliberately bypasses it because its buffers are + // device-resident; running both scans against one URI risks the two sharing entries. + let host_path = NamedTempFile::new()?; + std::fs::copy(path, host_path.path())?; + let gpu_file = open_gpu(path).await?; let mut gpu_batches = gpu_file.scan()?.into_array_stream()?; - let host_file = SESSION.open_options().open_path(path).await?; + let host_file = SESSION.open_options().open_path(host_path.path()).await?; let mut host_batches = host_file.scan()?.into_array_stream()?; let mut fields_checked = 0usize; + let mut batch_index = 0usize; loop { let (gpu_batch, host_batch) = (gpu_batches.next().await, host_batches.next().await); let (gpu_batch, host_batch) = match (gpu_batch, host_batch) { @@ -143,7 +151,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { let host_record = host_batch.execute::(cuda_ctx.execution_ctx())?; ensure!( gpu_record.len() == host_record.len(), - "batch length differs between the GPU and CPU scans: {} vs {}", + "batch {batch_index} length differs between the GPU and CPU scans: {} vs {}", gpu_record.len(), host_record.len() ); @@ -158,15 +166,28 @@ async fn verify_against_host_scan(path: &Path) -> Result { .collect::>(); ensure!( gpu_fields.len() == host_fields.len(), - "field count differs between the GPU and CPU scans" + "batch {batch_index} field count differs between the GPU and CPU scans" ); - for (gpu_field, host_field) in gpu_fields.into_iter().zip(host_fields) { + for (field_index, (gpu_field, host_field)) in + gpu_fields.into_iter().zip(host_fields).enumerate() + { let decoded = gpu_field.execute_cuda(&mut cuda_ctx).await?; + // The decode is enqueued, not complete: make the writes visible before reading + // the buffers back to the host. + cuda_ctx.synchronize_stream()?; let decoded = decoded.into_host().await?.into_array(); - verify_field(&host_field, decoded, cuda_ctx.execution_ctx())?; + verify_field( + &host_field, + decoded, + cuda_ctx.execution_ctx(), + batch_index, + field_index, + )?; fields_checked += 1; } + + batch_index += 1; } cuda_ctx.synchronize_stream()?; @@ -175,17 +196,66 @@ async fn verify_against_host_scan(path: &Path) -> Result { } /// Fails unless a GPU-decoded field matches the same field decoded on the CPU. -fn verify_field(host: &ArrayRef, gpu: ArrayRef, ctx: &mut ExecutionCtx) -> Result<()> { +fn verify_field( + host: &ArrayRef, + gpu: ArrayRef, + ctx: &mut ExecutionCtx, + batch_index: usize, + field_index: usize, +) -> Result<()> { let expected = SESSION.arrow().execute_arrow(host.clone(), None, ctx)?; // Pin the Arrow target type so the two sides cannot land on different but equivalent // encodings of the same logical values. let target = Field::new("", expected.data_type().clone(), gpu.dtype().is_nullable()); let actual = SESSION.arrow().execute_arrow(gpu, Some(&target), ctx)?; - ensure!( - expected.to_data() == actual.to_data(), - "GPU decode of a {} field does not match the CPU decode", - host.encoding_id() + if expected.to_data() == actual.to_data() { + return Ok(()); + } + + bail!( + "GPU decode of a {} field does not match the CPU decode \ + (batch {batch_index}, field {field_index}){}", + host.encoding_id(), + describe_mismatch(&expected, &actual) + ) +} + +/// Builds a human-readable description of how two Arrow arrays differ. +fn describe_mismatch(expected: &ArrowArrayRef, actual: &ArrowArrayRef) -> String { + let mut description = format!( + "\n cpu: type={:?} len={} nulls={}\n gpu: type={:?} len={} nulls={}", + expected.data_type(), + expected.len(), + expected.null_count(), + actual.data_type(), + actual.len(), + actual.null_count(), ); - Ok(()) + + if expected.data_type() != actual.data_type() || expected.len() != actual.len() { + return description; + } + + // Binary search for the shortest prefix that already differs; its last element is the + // first mismatching row. + let (mut low, mut high) = (0usize, expected.len()); + while low < high { + let mid = low + (high - low) / 2 + 1; + if expected.slice(0, mid).to_data() == actual.slice(0, mid).to_data() { + low = mid; + } else { + high = mid - 1; + } + } + + if low < expected.len() { + description.push_str(&format!( + "\n first difference at row {low}:\n cpu: {:?}\n gpu: {:?}", + expected.slice(low, 1), + actual.slice(low, 1), + )); + } + + description } From 335d14af3e5cf5460af4167a9070bc2d23f519b3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:19:11 +0000 Subject: [PATCH 04/22] fix(cuda): add the frame of reference to bit-packed patch values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The bit-unpack kernel writes patch values straight into the output while the lane decoder adds the frame of reference to every unpacked value. Bit-packing exceptions are stored in the same reference-relative domain as the packed values, so under `FoR(BitPacked)` every patched position came out short by exactly the reference. The existing kernel tests could not catch this: they exercise `BitPacked` directly, where the reference is zero. The new `FoRExecutor` case bit-packs to 8 bits with values that overflow into patches and a non-zero reference. Found by the compression benchmark's new `--gpu-verify` pass, which reported a `fastlanes.for` field decoding row 8038 as 131072 where the CPU produced 393061 — a difference of exactly the 261989 reference. Also thread the dataset name through compress-bench failures, so a benchmark error says which dataset it came from. Signed-off-by: Claude --- benchmarks/compress-bench/src/main.rs | 7 +++- vortex-cuda/kernels/src/bit_unpack_16.cu | 9 ++++- vortex-cuda/kernels/src/bit_unpack_32.cu | 9 ++++- vortex-cuda/kernels/src/bit_unpack_64.cu | 9 ++++- vortex-cuda/kernels/src/bit_unpack_8.cu | 9 ++++- vortex-cuda/src/bit_unpack_gen.rs | 9 ++++- vortex-cuda/src/kernel/encodings/for_.rs | 47 ++++++++++++++++++++++++ 7 files changed, 87 insertions(+), 12 deletions(-) diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index a915112897f..63bcfac7025 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -4,6 +4,7 @@ use std::path::PathBuf; use std::time::Duration; +use anyhow::Context; use clap::Parser; #[cfg(feature = "lance")] use compress_bench::LanceCompressor; @@ -356,7 +357,8 @@ async fn run_benchmark_for_dataset( iterations, bench_name, ) - .await?; + .await + .with_context(|| format!("compressing {bench_name} as {format}"))?; compressed_sizes.insert(*format, result.compressed_size); let all_runs_ns: Vec = result .all_runs @@ -387,7 +389,8 @@ async fn run_benchmark_for_dataset( iterations, bench_name, ) - .await?; + .await + .with_context(|| format!("decompressing {bench_name} as {format}"))?; let all_runs_ns: Vec = result .all_runs .iter() diff --git a/vortex-cuda/kernels/src/bit_unpack_16.cu b/vortex-cuda/kernels/src/bit_unpack_16.cu index a784df201d3..43acff3d08d 100644 --- a/vortex-cuda/kernels/src/bit_unpack_16.cu +++ b/vortex-cuda/kernels/src/bit_unpack_16.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_16_device(const uint16_t *__restrict in, uint16_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_32.cu b/vortex-cuda/kernels/src/bit_unpack_32.cu index 3f8fcb5c227..b9f29c72c3a 100644 --- a/vortex-cuda/kernels/src/bit_unpack_32.cu +++ b/vortex-cuda/kernels/src/bit_unpack_32.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_32_device(const uint32_t *__restrict in, uint32_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_64.cu b/vortex-cuda/kernels/src/bit_unpack_64.cu index ebe0b125369..aebe500c653 100644 --- a/vortex-cuda/kernels/src/bit_unpack_64.cu +++ b/vortex-cuda/kernels/src/bit_unpack_64.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_64_device(const uint64_t *__restrict in, uint64_t *_ } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 16); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/kernels/src/bit_unpack_8.cu b/vortex-cuda/kernels/src/bit_unpack_8.cu index b2fcfd26f04..cabea862e59 100644 --- a/vortex-cuda/kernels/src/bit_unpack_8.cu +++ b/vortex-cuda/kernels/src/bit_unpack_8.cu @@ -13,11 +13,16 @@ __device__ void _bit_unpack_8_device(const uint8_t *__restrict in, uint8_t *__re } __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, 32); auto patch = cursor.next(); while (patch.index != FL_CHUNK) { - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); } __syncwarp(); diff --git a/vortex-cuda/src/bit_unpack_gen.rs b/vortex-cuda/src/bit_unpack_gen.rs index 2482c0996b8..7b3ce148753 100644 --- a/vortex-cuda/src/bit_unpack_gen.rs +++ b/vortex-cuda/src/bit_unpack_gen.rs @@ -152,11 +152,16 @@ __device__ void _bit_unpack_{bits}_device(const uint{bits}_t *__restrict in, uin }} __syncwarp(); - // Step 2: Apply patches to shared memory in parallel + // Step 2: Apply patches to shared memory in parallel. + // + // Patch values are stored in the same frame-of-reference domain as the packed values, so + // they take the same `+ reference` the lane decoder applies to every unpacked value. For a + // plain bit-packed array the reference is zero and this is a no-op; for FoR-over-bit-packed + // it is what keeps patched positions from coming out short by the reference. PatchesCursor cursor(patches, blockIdx.x, thread_idx, {thread_count}); auto patch = cursor.next(); while (patch.index != FL_CHUNK) {{ - shared_out[patch.index] = patch.value; + shared_out[patch.index] = patch.value + reference; patch = cursor.next(); }} __syncwarp(); diff --git a/vortex-cuda/src/kernel/encodings/for_.rs b/vortex-cuda/src/kernel/encodings/for_.rs index fefa4ff5ef6..cc3ceed2869 100644 --- a/vortex-cuda/src/kernel/encodings/for_.rs +++ b/vortex-cuda/src/kernel/encodings/for_.rs @@ -158,6 +158,7 @@ mod tests { use vortex::buffer::Buffer; use vortex::dtype::NativePType; use vortex::encodings::fastlanes::BitPacked; + use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::FoR; use vortex::encodings::fastlanes::FoRArray; use vortex::error::VortexExpect; @@ -228,4 +229,50 @@ mod tests { assert_arrays_eq!(for_array, gpu_result, &mut ctx); } + + /// Patched positions must pick up the frame of reference, exactly like unpacked ones. + /// + /// The bit-packed exceptions are stored reference-relative, so a decoder that writes them + /// straight into the output leaves every patched value short by the reference. A plain + /// bit-packed array cannot catch that: its reference is zero. + #[rstest] + #[case::u32(100_000u32)] + #[case::u64(1_000_000u64)] + #[crate::test] + async fn test_ffor_patched_values_include_reference(#[case] reference: T) -> VortexResult<()> + where + T: NativePType + Into + From, + { + let mut ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + // Values that fit in 8 bits, with a handful that do not and so become patches. + let mut values = (0..2048u32) + .map(|i| >::from(i % 200)) + .collect::>(); + for index in [7, 1023, 1024, 2047] { + values[index] = >::from(1u32 << 17); + } + + let values = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); + let packed = BitPacked::encode(&values, 8, &mut array_session().create_execution_ctx())?; + assert!( + packed.patches().is_some(), + "test setup expects the exceptions to be stored as patches" + ); + let for_array = FoR::try_new(packed.into_array(), reference.into())?; + + let gpu_result = FoRExecutor + .execute(for_array.clone().into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(for_array, gpu_result, &mut ctx); + + Ok(()) + } } From 5f297d4a716f4eff28a84b6178638ad600b57718 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:29:19 +0000 Subject: [PATCH 05/22] Survey every dataset in a GPU verification pass A verification run stopped at the first dataset that failed, so finding the GPU-clean set took one CI cycle per dataset. Run every dataset instead, recording failures and reporting them together at the end, then exit non-zero. Missing CUDA kernel support surfaces as a panic rather than an error, so the survey catches those too. Signed-off-by: Claude --- benchmarks/compress-bench/src/main.rs | 67 +++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 5 deletions(-) diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 63bcfac7025..6fb003f97e3 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::any::Any; +use std::panic::AssertUnwindSafe; use std::path::PathBuf; use std::time::Duration; @@ -15,6 +17,7 @@ use compress_bench::gpu_vortex::GpuVortexCompressor; use compress_bench::parquet::ParquetCompressor; use compress_bench::parquet_pages::GpuCodec; use compress_bench::vortex::VortexCompressor; +use futures::FutureExt; use indicatif::ProgressBar; use itertools::Itertools; use regex::Regex; @@ -286,18 +289,61 @@ async fn run_compress( let mut measurements = vec![]; let mut v3_records: Vec = Vec::new(); + // A verification pass reports on every dataset rather than stopping at the first failure: + // one run then says exactly which datasets decode correctly on the GPU and which do not. + let survey_all = gpu.is_some_and(|gpu| gpu.verify); + let mut failures: Vec<(String, anyhow::Error)> = Vec::new(); + for dataset_handle in datasets.into_iter() { - let (m, mut records) = - run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, gpu) - .await?; - measurements.push(m); - v3_records.append(&mut records); + let run = + run_benchmark_for_dataset(&progress, &formats, &ops, iterations, dataset_handle, gpu); + + // Missing CUDA kernel support surfaces as a panic rather than an error, so the survey + // has to catch those too or the first unsupported dataset ends the run. + let result = if survey_all { + match AssertUnwindSafe(run).catch_unwind().await { + Ok(result) => result, + Err(panic) => Err(anyhow::anyhow!("panicked: {}", panic_message(&panic))), + } + } else { + run.await + }; + + match result { + Ok((m, mut records)) => { + measurements.push(m); + v3_records.append(&mut records); + } + Err(error) if survey_all => { + tracing::error!("{}: {error:#}", dataset_handle.name()); + failures.push((dataset_handle.name().to_string(), error)); + } + Err(error) => return Err(error), + } } let measurements = CompressMeasurements::from_iter(measurements); progress.finish(); + if !failures.is_empty() { + eprintln!( + "\nGPU verification failed for {} dataset(s):", + failures.len() + ); + for (dataset, error) in &failures { + eprintln!(" - {dataset}: {error:#}"); + } + anyhow::bail!( + "GPU verification failed for: {}", + failures + .iter() + .map(|(dataset, _)| dataset.as_str()) + .collect::>() + .join(", ") + ); + } + if let Some(path) = ingest_output { v3::write_jsonl_to_path(&path, &v3_records)?; } @@ -457,3 +503,14 @@ fn push_gpu_ratio( value: vortex_time.as_nanos() as f64 / parquet_time.as_nanos() as f64, }); } + +/// Extracts the message from a caught panic payload. +fn panic_message(panic: &Box) -> String { + if let Some(message) = panic.downcast_ref::<&str>() { + (*message).to_string() + } else if let Some(message) = panic.downcast_ref::() { + message.clone() + } else { + "non-string panic payload".to_string() + } +} From b0eda650245bbde23d5e644ee5c9c7308386f8ea Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:40:24 +0000 Subject: [PATCH 06/22] CI: publish the GPU verification matrix to the PR The per-dataset verification verdicts were only visible by digging through a multi-thousand-line job log. Capture the verification output, publish the per-dataset results to the step summary and a PR comment, and keep failing the job through a separate gate step. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 44 +++++++++++++++++++-- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 45bc59fb317..94fb085eb31 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -45,15 +45,53 @@ jobs: - name: Setup benchmark environment run: sudo bash scripts/setup-benchmark.sh - name: Verify GPU decompression correctness + id: verify shell: bash + continue-on-error: true env: - RUST_BACKTRACE: full + RUST_BACKTRACE: "1" FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" # Cross-checks every GPU-decompressed page and field against the CPU decoders before - # any timings are taken. Verification runs inline, so this pass is not timed. + # any timings are taken. Verification runs inline, so this pass is not timed. It runs + # every dataset rather than stopping at the first failure, so one run reports the whole + # matrix; the gate step below still fails the job. run: | + set -o pipefail target/release_debug/compress-bench \ - --gpu-decompress --gpu-verify --iterations 1 -d table + --gpu-decompress --gpu-verify --iterations 1 -d table 2>&1 | tee gpu-verify.txt + - name: Publish verification result + shell: bash + # The per-dataset verdicts are worth surfacing on the PR whether or not they all pass: + # digging them out of a multi-thousand-line job log is otherwise the only way to see + # which encodings decode correctly on the GPU. + run: | + { + echo "# GPU decompression verification" + echo + if [ "${{ steps.verify.outcome }}" = "success" ]; then + echo "All GPU datasets matched the CPU decode." + else + echo "Verification failed. Per-dataset results:" + fi + echo + echo '```text' + grep -E "verified [0-9]+|GPU verification failed|^ - " gpu-verify.txt | tail -40 \ + || tail -40 gpu-verify.txt + echo '```' + } > verify-comment.md + cat verify-comment.md >> "$GITHUB_STEP_SUMMARY" + - name: Comment PR with verification result + if: github.event.pull_request.head.repo.fork == false + uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b # v3 + with: + file-path: verify-comment.md + comment-tag: bench-pr-comment-gpu-verify + - name: Fail if verification failed + if: steps.verify.outcome == 'failure' + shell: bash + run: | + echo "GPU decompression verification failed; see the verification comment." >&2 + exit 1 - name: Run GPU compression benchmark shell: bash env: From 66c75a50aa7148aa9a6628a77524e8dc010db9fd Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 3 Aug 2026 13:56:54 +0000 Subject: [PATCH 07/22] CI: install DuckDB for the GPU compression benchmark The Public BI datasets build their Parquet fixture through the DuckDB CLI, as in bench-pr.yml. The GPU job never installed it, so all six failed with ENOENT before reaching the GPU at all. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 94fb085eb31..7e60b9e1970 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -32,6 +32,13 @@ jobs: with: repo-token: ${{ secrets.GITHUB_TOKEN }} enable-sccache: "true" + - name: Install DuckDB + # The Public BI datasets are converted from CSV to Parquet through the DuckDB CLI, as + # in bench-pr.yml. Without it those datasets fail to materialise their fixture. + run: | + wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.5/duckdb_cli-linux-amd64.zip | funzip > duckdb + chmod +x duckdb + echo "$PWD" >> "$GITHUB_PATH" - uses: ./.github/actions/system-info - name: Display NVIDIA GPU details run: | From 64d481c424429539932e0ab1aa557248db7f89b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 14:50:15 +0000 Subject: [PATCH 08/22] Time the GPU Parquet number with cuDF instead of nvCOMP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The nvCOMP backend only ran the codec stage on the device: page decoding stayed on the CPU and was excluded from the measurement, so the Parquet figure was an upper bound and the comparison against Vortex was not like-for-like. cuDF's `read_parquet` does the whole read on the device — page header decode, decompression, dictionary/RLE/plain decoding and column assembly — which is the same amount of work the Vortex backend does when it decodes to canonical arrays. It is reached through the prebuilt `cudf-cu12` wheel, so it stays a runtime dependency and never enters the Rust build. Timing is taken inside scripts/cudf-parquet-read.py, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up read runs first. `--gpu-verify` now compares the cuDF frame against a CPU Parquet read. This removes the page scanner, the batched nvCOMP launch path and the nvCOMP Snappy bindings, all of which existed only to serve the codec-stage backend. What remains of the Parquet side is the GPU-friendly writer settings, now in gpu_writer.rs. Signed-off-by: Claude --- .github/workflows/gpu-compress-bench-pr.yml | 12 + Cargo.lock | 6 +- Cargo.toml | 1 - benchmarks/compress-bench/Cargo.toml | 11 +- benchmarks/compress-bench/README.md | 58 +- benchmarks/compress-bench/src/gpu_parquet.rs | 525 ++--------------- benchmarks/compress-bench/src/gpu_writer.rs | 88 +++ benchmarks/compress-bench/src/lib.rs | 2 +- benchmarks/compress-bench/src/main.rs | 12 +- .../compress-bench/src/parquet_pages.rs | 530 ------------------ scripts/cudf-parquet-read.py | 101 ++++ vortex-cuda/nvcomp/build.rs | 7 - vortex-cuda/nvcomp/src/backend.rs | 52 -- vortex-cuda/nvcomp/src/lib.rs | 4 - vortex-cuda/nvcomp/src/snappy.rs | 222 -------- vortex-cuda/nvcomp/src/zstd.rs | 47 +- 16 files changed, 318 insertions(+), 1360 deletions(-) create mode 100644 benchmarks/compress-bench/src/gpu_writer.rs delete mode 100644 benchmarks/compress-bench/src/parquet_pages.rs create mode 100644 scripts/cudf-parquet-read.py delete mode 100644 vortex-cuda/nvcomp/src/backend.rs delete mode 100644 vortex-cuda/nvcomp/src/snappy.rs diff --git a/.github/workflows/gpu-compress-bench-pr.yml b/.github/workflows/gpu-compress-bench-pr.yml index 7e60b9e1970..36f16223105 100644 --- a/.github/workflows/gpu-compress-bench-pr.yml +++ b/.github/workflows/gpu-compress-bench-pr.yml @@ -39,6 +39,18 @@ jobs: wget -qO- https://github.com/duckdb/duckdb/releases/download/v1.5.5/duckdb_cli-linux-amd64.zip | funzip > duckdb chmod +x duckdb echo "$PWD" >> "$GITHUB_PATH" + - name: Install uv + uses: spiraldb/actions/.github/actions/setup-uv@a746510eafaa926484c354541cfc49b2ec06cc63 # 0.18.6 + - name: Install cuDF + # The GPU Parquet number is a full cuDF `read_parquet`. cuDF ships prebuilt manylinux + # wheels on NVIDIA's index, so it stays a runtime dependency and never enters the Rust + # build. pandas and pyarrow back the --gpu-verify cross-check. + run: | + uv venv --python 3.12 .venv-cudf + uv pip install --python .venv-cudf \ + --extra-index-url https://pypi.nvidia.com \ + cudf-cu12 pandas pyarrow + echo "$PWD/.venv-cudf/bin" >> "$GITHUB_PATH" - uses: ./.github/actions/system-info - name: Display NVIDIA GPU details run: | diff --git a/Cargo.lock b/Cargo.lock index e7c6dd2dfdb..d599793b8c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1566,15 +1566,14 @@ dependencies = [ "async-trait", "bytes", "clap", - "cudarc", "futures", "indicatif", "itertools 0.14.0", "lance-bench", "parquet 58.4.0", "regex", - "rstest", - "snap", + "serde", + "serde_json", "tempfile", "tokio", "tracing", @@ -1582,7 +1581,6 @@ dependencies = [ "vortex-arrow", "vortex-bench", "vortex-cuda", - "zstd", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index edf730760b4..997b3281d25 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -245,7 +245,6 @@ similar = "3.0.0" sketches-ddsketch = "0.4.0" smallvec = "1.15.1" smol = "2.0.2" -snap = "1.1.2" spatialbench = "0.2" spatialbench-arrow = "0.2" # spatialbench still pins arrow 56, two majors behind the workspace arrow. Until upstream diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index c20e2861b2d..7b1dadc3209 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,14 +21,14 @@ arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } -cudarc = { workspace = true, optional = true } futures = { workspace = true } indicatif = { workspace = true } itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } -snap = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } tempfile = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } @@ -36,14 +36,9 @@ vortex = { workspace = true } vortex-arrow = { workspace = true } vortex-bench = { workspace = true } vortex-cuda = { workspace = true, optional = true } -zstd = { workspace = true } - -[dev-dependencies] -rstest = { workspace = true } -tempfile = { workspace = true } [features] -cuda = ["dep:cudarc", "dep:tempfile", "dep:vortex-cuda"] +cuda = ["dep:tempfile", "dep:vortex-cuda"] lance = ["dep:lance-bench"] unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"] diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 132ce153450..e14c117d5a6 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -23,55 +23,58 @@ GPU dataset list in `src/main.rs`. It measures decompression only, for two backe - **Vortex** — the file is written with CUDA-compatible BtrBlocks encodings only (`only_cuda_compatible`) and a CUDA flat layout, then decoded on the device all the way to canonical arrays. -- **Parquet** — the file is rewritten with GPU-friendly writer settings (see below) and every - page body is decompressed on the device with nvCOMP's batched Snappy or Zstd entrypoints, - the same decomposition cuDF's Parquet reader uses: column chunks are staged on the device, - then all pages go through a single batched launch. +- **Parquet** — the file is rewritten with GPU-friendly writer settings (see below) and read + back with [cuDF](https://github.com/rapidsai/cudf)'s `read_parquet`, which performs the + whole read on the device: page header decode, codec decompression, dictionary/RLE/plain + decoding and column assembly. + +Both sides therefore decode all the way to device-resident arrays, which is what makes the +`vortex:parquet- gpu ratio decompress time` metric a like-for-like comparison. ```bash cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress -# pick the Parquet page codec (default: snappy) +# pick the Parquet page codec the GPU file is written with (default: snappy) cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --gpu-parquet-codec zstd ``` -On Linux both backends read through the same pinned, direct-I/O (`O_DIRECT`) reader, so -repeated iterations measure storage bandwidth rather than page-cache hits. +### cuDF -### What the Parquet GPU number does and does not include +cuDF is reached through its prebuilt `cudf-cu12` wheel, so it is a runtime dependency of the +benchmark and never enters the Rust build: -Included: column chunk I/O, the host-to-device transfer, and the batched codec launch. +```bash +uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarrow +``` -Not included: page *decoding* — the dictionary, RLE and plain decoders that turn a -decompressed page into an Arrow array — because there is no Rust GPU Parquet page decoder to -call. The Vortex backend it is compared against decodes all the way to canonical arrays, so -the Parquet figure is an upper bound on what a full GPU Parquet reader could reach, and the -`vortex:parquet- gpu ratio decompress time` metric is biased in Parquet's favour. +`scripts/cudf-parquet-read.py` performs and times the read. Timing is taken inside that +script, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up +read runs first for the same reason. -Walking the per-page Thrift headers also happens on the host, once per file, outside the -measurement; a real GPU Parquet reader decodes page headers on the device. +Note that the two backends do not share an I/O path: the Vortex reader uses pinned buffers and +direct I/O (`O_DIRECT`) on Linux, while cuDF does its own host read and host-to-device +transfer. ### GPU-friendly Parquet writer settings -Set in `src/parquet_pages.rs`: +Set in `src/gpu_writer.rs`: | Setting | Value | Why | | --- | --- | --- | -| writer version | `PARQUET_1_0` | v1 pages compress the whole page body, which is the unit nvCOMP decompresses. v2 pages put uncompressed levels ahead of the compressed values in the same body. | -| compression | Snappy (default) or Zstd | The two Parquet codecs nvCOMP implements. Snappy has the higher device throughput and is the Parquet default. | +| writer version | `PARQUET_1_0` | v1 pages compress the whole page body; v2 pages put uncompressed levels ahead of the compressed values in the same body. | +| compression | Snappy (default) or Zstd | Snappy is the Parquet default and has the higher device-side throughput. | | dictionary | enabled | Keeps the decompressed payload small; the encoding GPU Parquet readers decode fastest. | -| data page size | 1 MiB | Large enough to amortize per-chunk setup, small enough to keep every SM fed. Matches the page size cuDF targets. | +| data page size | 1 MiB | Large enough to amortize per-page setup, small enough to keep every SM fed. Matches the page size cuDF targets. | | data page row limit | 1,000,000 | The 20k-row default caps narrow columns' pages far below 1 MiB. | -| statistics | chunk-level | Page statistics only inflate the headers that have to be walked on the host. | +| statistics | chunk-level | Page statistics only inflate the headers a reader has to walk. | ### Correctness `--gpu-verify` cross-checks device output against the CPU decoders on every iteration: -- Parquet: each decompressed page is copied back and compared byte-for-byte against the host - Snappy/Zstd output for the same compressed bytes. +- Parquet: the cuDF-read frame is compared against a CPU Parquet read of the same file. - Vortex: each GPU-decoded field is copied back and compared against the same field decoded on the CPU, through Arrow with a pinned target type. @@ -83,10 +86,5 @@ cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1 ``` -Independently of `--gpu-verify`, every Parquet GPU run checks nvCOMP's per-page status and -output-size arrays after the measurement, because a batched launch reports per-page failures -in device memory rather than by failing the launch. - -Page-header scanning is covered by CPU-only unit tests in `src/parquet_pages.rs`, which -assert the located page bodies decompress to exactly the bytes `parquet`'s own page reader -produces. +A verifying run reports on every dataset rather than stopping at the first failure, so one run +shows which datasets decode correctly on the GPU and which do not. diff --git a/benchmarks/compress-bench/src/gpu_parquet.rs b/benchmarks/compress-bench/src/gpu_parquet.rs index b8e6195d784..3ee0297f1a0 100644 --- a/benchmarks/compress-bench/src/gpu_parquet.rs +++ b/benchmarks/compress-bench/src/gpu_parquet.rs @@ -1,89 +1,70 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! GPU Parquet decompression backend. +//! GPU Parquet decompression backend, timed through cuDF. //! -//! Parquet's compressed unit is the page, and every page body is an independent block of -//! Snappy or Zstd. That is precisely the batch shape nvCOMP's device decompressors take, and -//! it is how cuDF's Parquet reader gets pages off the CPU: read the column chunks to the -//! device, then decompress every page in one batched launch. +//! cuDF's `read_parquet` performs the whole read on the device — page header decode, +//! codec decompression, dictionary/RLE/plain decoding and column assembly — which makes it +//! the like-for-like opponent for the Vortex GPU backend, which likewise decodes all the way +//! to canonical arrays on device. //! -//! What this measures is the decompression stage of a Parquet read — column chunk I/O, -//! host-to-device transfer, and the batched codec launch. Page *decoding* (the dictionary, -//! RLE and plain decoders that turn a decompressed page into an Arrow array) is not -//! included, because there is no Rust GPU Parquet page decoder to call. The Vortex GPU -//! backend it is compared against decodes all the way to canonical arrays, so the Parquet -//! numbers here are an upper bound on what a full GPU Parquet reader could achieve, and the -//! comparison is favourable to Parquet. +//! cuDF is reached through its prebuilt `cudf-cu12` wheel rather than by linking libcudf, so +//! it stays a runtime dependency of this benchmark and never enters the Rust build. The +//! measurement is taken inside [`CUDF_SCRIPT`], so interpreter start, `import cudf` and CUDA +//! context creation are excluded; only the reads themselves are timed. -use std::fs::File; use std::path::Path; +use std::process::Command; use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use anyhow::Context; use anyhow::Result; -use anyhow::anyhow; +use anyhow::bail; use anyhow::ensure; use arrow_array::RecordBatch; use async_trait::async_trait; -use cudarc::driver::CudaSlice; -use cudarc::driver::DevicePtr; -use cudarc::driver::DevicePtrMut; -use futures::StreamExt; -use futures::TryStreamExt; -use futures::stream; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; -use parquet::file::metadata::ParquetMetaData; -use parquet::file::metadata::ParquetMetaDataReader; +use serde::Deserialize; use tempfile::NamedTempFile; -use vortex::array::buffer::BufferHandle; -use vortex::array::buffer::DeviceBuffer; -use vortex::buffer::Alignment; -use vortex::buffer::Buffer; -use vortex::error::vortex_err; -use vortex::io::VortexReadAt; -use vortex::io::session::RuntimeSessionExt; use vortex_bench::Format; -use vortex_bench::SESSION; use vortex_bench::compress::Compressor; -use vortex_cuda::CudaBufferExt; -use vortex_cuda::CudaDeviceBuffer; -use vortex_cuda::CudaExecutionCtx; -use vortex_cuda::CudaSession; -use vortex_cuda::CudaSessionExt; -use vortex_cuda::PooledFileReadAt; -use vortex_cuda::PooledFileReadAtOptions; -use vortex_cuda::nvcomp::AlignmentRequirements; -use vortex_cuda::nvcomp::snappy; -use vortex_cuda::nvcomp::sys; -use vortex_cuda::nvcomp::sys::nvcompStatus_t; -use vortex_cuda::nvcomp::zstd as nvcomp_zstd; -use crate::parquet_pages::ColumnChunkPages; -use crate::parquet_pages::GpuCodec; -use crate::parquet_pages::gpu_writer_properties; -use crate::parquet_pages::scan_compressed_pages; +use crate::gpu_writer::GpuCodec; +use crate::gpu_writer::gpu_writer_properties; -/// Parquet compressor whose decompression measurement runs the page codec on the GPU. +/// Repo-relative path of the script that performs and times the cuDF read. +const CUDF_SCRIPT: &str = "scripts/cudf-parquet-read.py"; + +/// Parquet compressor whose decompression measurement is a full cuDF GPU read. pub struct GpuParquetCompressor { codec: GpuCodec, verify: bool, } +/// What the cuDF script reports back. +#[derive(Debug, Deserialize)] +struct CudfReadReport { + /// Fastest timed read, in nanoseconds. + min_ns: u64, + rows: u64, + columns: u64, +} + impl GpuParquetCompressor { - /// Create a backend that writes pages with `codec` and decompresses them with nvCOMP. + /// Create a backend that writes pages with `codec` and times cuDF reading them back. /// - /// When `verify` is set, every decompressed page is copied back and compared against the - /// host codec's output before the measurement is reported. + /// When `verify` is set, the GPU read is cross-checked against a CPU Parquet read of the + /// same file before the measurement is reported. pub fn new(codec: GpuCodec, verify: bool) -> Self { Self { codec, verify } } /// Rewrite the source Parquet file with GPU-friendly writer settings. fn write_gpu_parquet(&self, parquet_path: &Path) -> Result<(NamedTempFile, u64)> { - let builder = ParquetRecordBatchReaderBuilder::try_new(File::open(parquet_path)?)?; + let builder = ParquetRecordBatchReaderBuilder::try_new(std::fs::File::open(parquet_path)?)?; let schema = Arc::clone(builder.schema()); let batches: Vec = builder.build()?.collect::, _>>()?; @@ -101,21 +82,6 @@ impl GpuParquetCompressor { writer.close()?; Ok((output, size)) } - - fn check_codec(&self, metadata: &ParquetMetaData) -> Result<()> { - for row_group in metadata.row_groups() { - for column in row_group.columns() { - ensure!( - self.codec.matches(column.compression()), - "column {} was written with {:?}, expected {}", - column.column_path(), - column.compression(), - self.codec.name() - ); - } - } - Ok(()) - } } #[async_trait] @@ -132,420 +98,43 @@ impl Compressor for GpuParquetCompressor { async fn decompress(&self, parquet_path: &Path) -> Result { let (gpu_file, _) = self.write_gpu_parquet(parquet_path)?; - - let file = File::open(gpu_file.path())?; - let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; - self.check_codec(&metadata)?; - drop(file); - - // The page map is a property of the file, so it is resolved once rather than on every - // iteration. A GPU Parquet reader decodes page headers on the device as part of the - // read; this benchmark cannot, so the host walk is kept out of the measurement. - let file_bytes = std::fs::read(gpu_file.path())?; - let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; - let plan = DecompressPlan::build(&row_groups, self.codec)?; - // Only verification needs the host copy afterwards, to decompress the same bytes. - let file_bytes = self.verify.then_some(file_bytes); - - let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; - let reader = open_reader(gpu_file.path(), &cuda_ctx)?; - - // Row groups are staged and released one at a time so device residency stays bounded - // by the largest row group rather than by the whole file, matching how the Vortex - // backend streams one batch at a time. - let mut elapsed = Duration::ZERO; - let mut checks = Vec::with_capacity(plan.row_groups.len()); - for row_group in &plan.row_groups { - let start = Instant::now(); - let device_chunks = stage_row_group(row_group, &reader).await?; - let output = - decompress_pages(row_group, &device_chunks, self.codec, &mut cuda_ctx).await?; - cuda_ctx.synchronize_stream()?; - elapsed += start.elapsed(); - - let DecompressOutput { - output, - actual_sizes, - statuses, - } = output; - drop(device_chunks); - - if self.verify { - // Verified inline, since holding every row group's decompressed output would - // defeat the point of streaming them. This is why a verifying run's timings - // are not comparable. - let file_bytes = file_bytes - .as_deref() - .ok_or_else(|| anyhow!("verification requires the host file bytes"))?; - verify_against_host(row_group, file_bytes, self.codec, output).await?; - } - checks.push((actual_sizes, statuses)); - } - - // Checked outside the measurement, but on every iteration: a batched launch reports - // per-page failures in device memory rather than by failing the launch, so without - // this a broken run would simply look fast. - for (row_group, (actual_sizes, statuses)) in plan.row_groups.iter().zip(checks) { - check_statuses(row_group, actual_sizes, statuses).await?; - } - - Ok(elapsed) - } -} - -/// Stages a row group's column chunks on the device, bounding concurrent pinned staging. -async fn stage_row_group( - row_group: &RowGroupPlan, - reader: &PooledFileReadAt, -) -> Result> { - let reads = row_group - .chunks - .iter() - .map(|chunk| reader.read_at(chunk.offset, chunk.len, Alignment::of::())) - .collect::>(); - - stream::iter(reads) - .buffered(STAGING_CONCURRENCY) - .try_collect::>() - .await - .map_err(|e| anyhow!("failed to stage Parquet column chunks on the device: {e}")) -} - -/// Column chunks staged on the device at once. Matches `vortex-cuda`'s file read concurrency. -const STAGING_CONCURRENCY: usize = 32; - -/// Opens the GPU Parquet file through the same pinned, direct-I/O reader the Vortex GPU -/// backend uses, so both formats measure storage bandwidth rather than page-cache hits. -fn open_reader(path: &Path, cuda_ctx: &CudaExecutionCtx) -> Result { - let pool = Arc::clone(SESSION.cuda_session().pinned_buffer_pool()); - let options = PooledFileReadAtOptions::default(); - #[cfg(target_os = "linux")] - let options = options.with_direct_io(); - - Ok(PooledFileReadAt::open_with_options( - path, - SESSION.handle(), - pool, - cuda_ctx.stream().clone(), - options, - )?) -} - -/// A page to decompress, addressed relative to its column chunk's device buffer. -struct PlannedPage { - chunk: usize, - offset_in_chunk: usize, - compressed_len: usize, - uncompressed_len: usize, - output_offset: usize, -} - -/// The byte range of a column chunk to stage on the device. -struct PlannedChunk { - offset: u64, - len: usize, -} - -/// Everything needed to issue one batched nvCOMP launch over a row group's pages. -struct RowGroupPlan { - chunks: Vec, - pages: Vec, - output_len: usize, - max_uncompressed: usize, -} - -/// The per-row-group work for one file. -struct DecompressPlan { - row_groups: Vec, -} - -impl DecompressPlan { - fn build(row_groups: &[Vec], codec: GpuCodec) -> Result { - let alignment = decompress_alignments(codec)?; - ensure!( - alignment.output.is_power_of_two(), - "nvcomp reported a non-power-of-two output alignment of {}", - alignment.output - ); - - let row_groups = row_groups - .iter() - .map(|chunks| RowGroupPlan::build(chunks, codec, alignment)) - .collect::>>()?; + let report = run_cudf_read(gpu_file.path(), self.verify)?; ensure!( - row_groups.iter().any(|plan| !plan.pages.is_empty()), - "Parquet file contains no compressed pages" + report.rows > 0 && report.columns > 0, + "cuDF read {} rows and {} columns, expected a non-empty table", + report.rows, + report.columns ); - Ok(Self { row_groups }) + Ok(Duration::from_nanos(report.min_ns)) } } -impl RowGroupPlan { - fn build( - chunks: &[ColumnChunkPages], - codec: GpuCodec, - alignment: AlignmentRequirements, - ) -> Result { - let mut planned_chunks = Vec::with_capacity(chunks.len()); - let mut pages = Vec::new(); - let mut output_len = 0usize; - let mut max_uncompressed = 0usize; - - for (index, chunk) in chunks.iter().enumerate() { - planned_chunks.push(PlannedChunk { - offset: chunk.offset, - len: chunk.len, - }); - - for page in &chunk.pages { - let offset_in_chunk = usize::try_from( - u64::try_from(page.offset)? - .checked_sub(chunk.offset) - .ok_or_else(|| anyhow!("page offset precedes its column chunk"))?, - )?; - // Pages are decompressed in place from their column chunk's device buffer. - // CUDA allocations are at least 256-byte aligned, so a page's device address - // meets nvcomp's requirement exactly when its chunk-relative offset does. - ensure!( - offset_in_chunk.is_multiple_of(alignment.input), - "page at file offset {} sits {offset_in_chunk} bytes into its column chunk, \ - which does not meet nvcomp's {} byte input alignment for {}; \ - use --gpu-parquet-codec snappy", - page.offset, - alignment.input, - codec.name() - ); - - let output_offset = output_len.next_multiple_of(alignment.output); - output_len = output_offset + page.uncompressed_len; - max_uncompressed = max_uncompressed.max(page.uncompressed_len); - - pages.push(PlannedPage { - chunk: index, - offset_in_chunk, - compressed_len: page.compressed_len, - uncompressed_len: page.uncompressed_len, - output_offset, - }); - } - } - - Ok(Self { - chunks: planned_chunks, - pages, - output_len, - max_uncompressed, - }) - } -} - -fn decompress_alignments(codec: GpuCodec) -> Result { - match codec { - GpuCodec::Snappy => { - snappy::decompress_alignment_requirements(snappy::SnappyDecompressOpts::default()) - } - GpuCodec::Zstd => nvcomp_zstd::decompress_alignment_requirements( - nvcomp_zstd::ZstdDecompressOpts::default(), - ), +/// Runs the cuDF read script and returns the timing it measured. +fn run_cudf_read(path: &Path, verify: bool) -> Result { + let mut command = Command::new("python3"); + command.arg(CUDF_SCRIPT).arg(path); + if verify { + command.arg("--verify"); } - .map_err(|e| anyhow!("nvcomp alignment query failed: {e}")) -} - -/// Device-side outputs of a batched decompression launch. -struct DecompressOutput { - output: CudaSlice, - actual_sizes: CudaSlice, - statuses: CudaSlice, -} - -/// Enqueues the batched decompression of every page in `plan` onto the context's stream. -async fn decompress_pages( - plan: &RowGroupPlan, - device_chunks: &[BufferHandle], - codec: GpuCodec, - ctx: &mut CudaExecutionCtx, -) -> Result { - let num_pages = plan.pages.len(); - - let temp_size = match codec { - GpuCodec::Snappy => { - snappy::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) - } - GpuCodec::Zstd => { - nvcomp_zstd::get_decompress_temp_size(num_pages, plan.max_uncompressed, plan.output_len) - } - } - .map_err(|e| anyhow!("nvcomp temp size query failed: {e}"))?; - - let chunk_bases = device_chunks - .iter() - .map(|handle| handle.cuda_device_ptr()) - .collect::, _>>()?; - - let mut output = ctx.device_alloc::(plan.output_len)?; - // Only the allocation address is needed to build the output pointer table; the device - // write itself is tracked by the guard taken around the launch below. - let output_base = { - let (base, _) = output.device_ptr(ctx.stream()); - base - }; - - let mut compressed_ptrs = Vec::with_capacity(num_pages); - let mut compressed_sizes = Vec::with_capacity(num_pages); - let mut uncompressed_sizes = Vec::with_capacity(num_pages); - let mut output_ptrs = Vec::with_capacity(num_pages); - for page in &plan.pages { - compressed_ptrs.push(chunk_bases[page.chunk] + page.offset_in_chunk as u64); - compressed_sizes.push(page.compressed_len); - uncompressed_sizes.push(page.uncompressed_len); - output_ptrs.push(output_base + page.output_offset as u64); - } - - let (compressed_ptrs, compressed_sizes, uncompressed_sizes, output_ptrs) = futures::try_join!( - ctx.copy_to_device(compressed_ptrs)?, - ctx.copy_to_device(compressed_sizes)?, - ctx.copy_to_device(uncompressed_sizes)?, - ctx.copy_to_device(output_ptrs)? - )?; - - let mut actual_sizes: CudaSlice = ctx.device_alloc(num_pages)?; - let mut statuses: CudaSlice = ctx.device_alloc(num_pages)?; - let mut temp: CudaSlice = ctx.device_alloc(temp_size)?; - - let stream = ctx.stream(); - let compressed_ptrs_view = compressed_ptrs.cuda_view::()?; - let compressed_sizes_view = compressed_sizes.cuda_view::()?; - let uncompressed_sizes_view = uncompressed_sizes.cuda_view::()?; - let output_ptrs_view = output_ptrs.cuda_view::()?; - let (compressed_ptrs_ptr, record_compressed_ptrs) = compressed_ptrs_view.device_ptr(stream); - let (compressed_sizes_ptr, record_compressed_sizes) = compressed_sizes_view.device_ptr(stream); - let (uncompressed_sizes_ptr, record_uncompressed_sizes) = - uncompressed_sizes_view.device_ptr(stream); - let (output_ptrs_ptr, record_output_ptrs) = output_ptrs_view.device_ptr(stream); - let (_output_ptr, record_output) = output.device_ptr_mut(stream); - let (actual_sizes_ptr, record_actual_sizes) = actual_sizes.device_ptr_mut(stream); - let (statuses_ptr, record_statuses) = statuses.device_ptr_mut(stream); - let (temp_ptr, record_temp) = temp.device_ptr_mut(stream); - - ctx.launch_external(plan.output_len, || { - // SAFETY: every pointer is derived from a live device allocation sized by the plan, - // and each batch metadata array holds exactly `num_pages` entries. - unsafe { - match codec { - GpuCodec::Snappy => snappy::decompress_async( - compressed_ptrs_ptr as _, - compressed_sizes_ptr as _, - uncompressed_sizes_ptr as _, - actual_sizes_ptr as _, - num_pages, - temp_ptr as _, - temp_size, - output_ptrs_ptr as _, - statuses_ptr as _, - stream.cu_stream().cast(), - ), - GpuCodec::Zstd => nvcomp_zstd::decompress_async( - compressed_ptrs_ptr as _, - compressed_sizes_ptr as _, - uncompressed_sizes_ptr as _, - actual_sizes_ptr as _, - num_pages, - temp_ptr as _, - temp_size, - output_ptrs_ptr as _, - statuses_ptr as _, - stream.cu_stream().cast(), - ), - } - .map_err(|e| vortex_err!("nvcomp decompress_async failed: {}", e)) - } + let output = command.output().with_context(|| { + format!("failed to run {CUDF_SCRIPT}; is cudf-cu12 installed on this host?") })?; - drop(( - record_compressed_ptrs, - record_compressed_sizes, - record_uncompressed_sizes, - record_output_ptrs, - record_output, - record_actual_sizes, - record_statuses, - record_temp, - )); - // The temporary workspace must outlive the launch, which the stream ordering guarantees - // only while the allocation is alive. - drop(temp); - - Ok(DecompressOutput { - output, - actual_sizes, - statuses, - }) -} - -/// Copies the per-page status and size arrays back and fails on any mismatch. -async fn check_statuses( - plan: &RowGroupPlan, - actual_sizes: CudaSlice, - statuses: CudaSlice, -) -> Result<()> { - let statuses = CudaDeviceBuffer::new(statuses) - .copy_to_host(Alignment::of::())? - .await?; - let actual_sizes = CudaDeviceBuffer::new(actual_sizes) - .copy_to_host(Alignment::of::())? - .await?; - - let statuses = Buffer::::from_byte_buffer(statuses); - let actual_sizes = Buffer::::from_byte_buffer(actual_sizes); - - for (index, page) in plan.pages.iter().enumerate() { - let status = statuses.as_slice()[index]; - ensure!( - status == sys::nvcompStatus_t_nvcompSuccess, - "page {index} failed to decompress with nvcomp status {status}" - ); - let actual = actual_sizes.as_slice()[index]; - ensure!( - actual == page.uncompressed_len, - "page {index} decompressed to {actual} bytes, expected {}", - page.uncompressed_len - ); - } - - Ok(()) -} - -/// Compares every decompressed page against the host codec's output for the same bytes. -async fn verify_against_host( - plan: &RowGroupPlan, - file_bytes: &[u8], - codec: GpuCodec, - output: CudaSlice, -) -> Result<()> { - let device_output = CudaDeviceBuffer::new(output) - .copy_to_host(Alignment::of::())? - .await?; - let device_output = device_output.as_slice(); - - for (index, page) in plan.pages.iter().enumerate() { - let chunk = &plan.chunks[page.chunk]; - let start = usize::try_from(chunk.offset)? + page.offset_in_chunk; - let compressed = &file_bytes[start..start + page.compressed_len]; - let expected = codec.decompress_host(compressed, page.uncompressed_len)?; - let actual = &device_output[page.output_offset..page.output_offset + page.uncompressed_len]; - ensure!( - actual == expected.as_slice(), - "page {index} decompressed on the GPU differs from the host codec output" + if !output.status.success() { + bail!( + "{CUDF_SCRIPT} exited with {}:\n{}", + output.status, + String::from_utf8_lossy(&output.stderr).trim() ); } - tracing::info!( - "verified {} GPU-decompressed {} pages against the host codec", - plan.pages.len(), - codec.name() - ); - Ok(()) + serde_json::from_slice(&output.stdout).with_context(|| { + format!( + "could not parse the report from {CUDF_SCRIPT}: {}", + String::from_utf8_lossy(&output.stdout).trim() + ) + }) } diff --git a/benchmarks/compress-bench/src/gpu_writer.rs b/benchmarks/compress-bench/src/gpu_writer.rs new file mode 100644 index 00000000000..49052a5d6b4 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_writer.rs @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Parquet writer settings for the GPU benchmark. +//! +//! The GPU backend rewrites each dataset before reading it back with cuDF, so the file it +//! reads is written the way a GPU reader wants it rather than the way the CPU suite writes it. + +use clap::ValueEnum; +use parquet::basic::Compression; +use parquet::basic::ZstdLevel; +use parquet::file::properties::EnabledStatistics; +use parquet::file::properties::WriterProperties; +use parquet::file::properties::WriterVersion; + +/// Target size of a data page written for GPU reads. +/// +/// Pages are the unit a GPU reader decompresses in parallel, so they need to be large enough +/// to amortize per-page setup and numerous enough to fill the device. ~1 MiB is the page size +/// cuDF's Parquet reader is tuned around. +pub const GPU_DATA_PAGE_SIZE: usize = 1024 * 1024; + +/// Row cap per data page. +/// +/// `parquet`'s default caps pages at 20k rows, which produces pages far below +/// [`GPU_DATA_PAGE_SIZE`] for narrow columns and leaves the device underfed. +const GPU_DATA_PAGE_ROW_COUNT_LIMIT: usize = 1_000_000; + +/// Parquet page codecs the GPU benchmark can write. +#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)] +pub enum GpuCodec { + /// The Parquet default, and the codec with the highest device-side throughput. + #[default] + Snappy, + /// Matches the codec used by the CPU Parquet benchmark, at lower device throughput. + Zstd, +} + +impl GpuCodec { + /// The Parquet compression setting for this codec. + pub fn to_parquet(self) -> Compression { + match self { + GpuCodec::Snappy => Compression::SNAPPY, + GpuCodec::Zstd => Compression::ZSTD(ZstdLevel::default()), + } + } + + /// Short lowercase name, used in measurement labels. + pub fn name(self) -> &'static str { + match self { + GpuCodec::Snappy => "snappy", + GpuCodec::Zstd => "zstd", + } + } +} + +/// Writer properties tuned for a GPU read. +pub fn gpu_writer_properties(codec: GpuCodec) -> WriterProperties { + WriterProperties::builder() + // V1 data pages compress the whole page body. V2 pages place uncompressed + // repetition/definition levels ahead of the compressed values in the same body, which + // not every GPU reader path handles. + .set_writer_version(WriterVersion::PARQUET_1_0) + .set_compression(codec.to_parquet()) + // Dictionary encoding keeps the decompressed payload small and is the encoding GPU + // Parquet readers decode fastest. + .set_dictionary_enabled(true) + .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) + .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) + // Per-page statistics only inflate the page headers a reader has to walk. + .set_statistics_enabled(EnabledStatistics::Chunk) + .build() +} + +#[cfg(test)] +mod tests { + use parquet::file::properties::WriterVersion; + + use super::*; + + #[test] + fn gpu_properties_use_v1_pages_and_the_requested_codec() { + let properties = gpu_writer_properties(GpuCodec::Snappy); + assert_eq!(properties.writer_version(), WriterVersion::PARQUET_1_0); + assert_eq!(properties.compression(&"x".into()), Compression::SNAPPY); + assert_eq!(properties.data_page_size_limit(), GPU_DATA_PAGE_SIZE); + } +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index e82f8574dc6..76e34172f6a 100644 --- a/benchmarks/compress-bench/src/lib.rs +++ b/benchmarks/compress-bench/src/lib.rs @@ -7,6 +7,6 @@ pub use lance_bench::compress::LanceCompressor; pub mod gpu_parquet; #[cfg(feature = "cuda")] pub mod gpu_vortex; +pub mod gpu_writer; pub mod parquet; -pub mod parquet_pages; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 6fb003f97e3..34f67b13e28 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -14,8 +14,8 @@ use compress_bench::LanceCompressor; use compress_bench::gpu_parquet::GpuParquetCompressor; #[cfg(feature = "cuda")] use compress_bench::gpu_vortex::GpuVortexCompressor; +use compress_bench::gpu_writer::GpuCodec; use compress_bench::parquet::ParquetCompressor; -use compress_bench::parquet_pages::GpuCodec; use compress_bench::vortex::VortexCompressor; use futures::FutureExt; use indicatif::ProgressBar; @@ -81,9 +81,9 @@ struct Args { /// decompression only, for both Vortex and Parquet. #[arg(long)] gpu_decompress: bool, - /// Page codec the GPU Parquet backend writes and decompresses with. + /// Page codec the GPU Parquet file is written with. /// - /// Snappy is the Parquet default and the codec nvCOMP decompresses fastest. + /// Snappy is the Parquet default and the codec GPU readers decompress fastest. #[arg(long, value_enum, default_value_t)] gpu_parquet_codec: GpuCodec, /// Cross-check every GPU-decompressed result against the CPU decoder. @@ -147,13 +147,9 @@ async fn main() -> anyhow::Result<()> { /// Settings for the GPU decompression mode. #[derive(Clone, Copy, Debug)] struct GpuOptions { - /// Parquet page codec to write and decompress on the device. + /// Parquet page codec to write the GPU file with. codec: GpuCodec, /// Cross-check decompressed output against the CPU decoders. - #[cfg_attr( - not(feature = "cuda"), - expect(dead_code, reason = "only the CUDA backends read this") - )] verify: bool, } diff --git a/benchmarks/compress-bench/src/parquet_pages.rs b/benchmarks/compress-bench/src/parquet_pages.rs deleted file mode 100644 index d55598cd8dd..00000000000 --- a/benchmarks/compress-bench/src/parquet_pages.rs +++ /dev/null @@ -1,530 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Locating the compressed page bodies inside a Parquet file. -//! -//! Parquet compresses each page body independently with a block codec, which is exactly the -//! shape nvCOMP's batched decompression entrypoints consume: an array of independent -//! compressed chunks with known uncompressed sizes. This module finds those chunks so the -//! GPU backend can hand the whole batch to the device in one launch, the same decomposition -//! cuDF's Parquet reader uses. -//! -//! Column chunk ranges come from the file footer; page boundaries within a chunk are only -//! discoverable by walking the per-page Thrift headers, so a minimal Thrift compact-protocol -//! reader lives here. `parquet::format::PageHeader` is deprecated and scheduled for removal, -//! and `parquet`'s own page-header parser is crate-private, so neither can be used. - -use anyhow::Result; -use anyhow::bail; -use anyhow::ensure; -use clap::ValueEnum; -use parquet::basic::Compression; -use parquet::basic::ZstdLevel; -use parquet::file::metadata::ParquetMetaData; -use parquet::file::properties::EnabledStatistics; -use parquet::file::properties::WriterProperties; -use parquet::file::properties::WriterVersion; - -/// Target size of a data page written for GPU decompression. -/// -/// nvCOMP decompresses one chunk per page, so pages must be large enough to amortize the -/// per-chunk setup yet numerous enough to fill the device. ~1 MiB is the page size cuDF's -/// Parquet reader is tuned around. -pub const GPU_DATA_PAGE_SIZE: usize = 1024 * 1024; - -/// Row cap per data page. -/// -/// `parquet`'s default caps pages at 20k rows, which produces pages far below -/// [`GPU_DATA_PAGE_SIZE`] for narrow columns and leaves the device underfed. -const GPU_DATA_PAGE_ROW_COUNT_LIMIT: usize = 1_000_000; - -/// Parquet page codecs that nvCOMP can decompress on the device. -#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, ValueEnum)] -pub enum GpuCodec { - /// The Parquet default, and the codec with the highest device-side throughput. - #[default] - Snappy, - /// Matches the codec used by the CPU Parquet benchmark, at lower device throughput. - Zstd, -} - -impl GpuCodec { - /// The Parquet compression setting for this codec. - pub fn to_parquet(self) -> Compression { - match self { - GpuCodec::Snappy => Compression::SNAPPY, - GpuCodec::Zstd => Compression::ZSTD(ZstdLevel::default()), - } - } - - /// Short lowercase name, used in measurement labels. - pub fn name(self) -> &'static str { - match self { - GpuCodec::Snappy => "snappy", - GpuCodec::Zstd => "zstd", - } - } - - /// Whether a column chunk's codec matches this one. - pub fn matches(self, compression: Compression) -> bool { - matches!( - (self, compression), - (GpuCodec::Snappy, Compression::SNAPPY) | (GpuCodec::Zstd, Compression::ZSTD(_)) - ) - } - - /// Decompress a single page body on the host, for cross-checking device output. - pub fn decompress_host(self, compressed: &[u8], uncompressed_len: usize) -> Result> { - let decompressed = match self { - GpuCodec::Snappy => snap::raw::Decoder::new().decompress_vec(compressed)?, - GpuCodec::Zstd => zstd::bulk::decompress(compressed, uncompressed_len)?, - }; - ensure!( - decompressed.len() == uncompressed_len, - "page decompressed to {} bytes, page header declared {uncompressed_len}", - decompressed.len() - ); - Ok(decompressed) - } -} - -/// Writer properties tuned for GPU decompression. -pub fn gpu_writer_properties(codec: GpuCodec) -> WriterProperties { - WriterProperties::builder() - // V1 data pages compress the entire page body, which is the unit nvCOMP decompresses. - // V2 pages place uncompressed repetition/definition levels ahead of the compressed - // values inside one page body, which the batched entrypoints cannot express. - .set_writer_version(WriterVersion::PARQUET_1_0) - .set_compression(codec.to_parquet()) - // Dictionary encoding keeps the decompressed payload small and is the encoding GPU - // Parquet readers decode fastest. - .set_dictionary_enabled(true) - .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) - .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) - // Per-page statistics only inflate the page headers that have to be walked on the host. - .set_statistics_enabled(EnabledStatistics::Chunk) - .build() -} - -/// A compressed page body located within a Parquet file. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct CompressedPage { - /// Offset of the compressed body, i.e. just past the page header. - pub offset: usize, - /// Length of the compressed body in bytes. - pub compressed_len: usize, - /// Length of the body once decompressed. - pub uncompressed_len: usize, -} - -/// The pages of one column chunk, alongside the byte range the chunk occupies. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct ColumnChunkPages { - /// Offset of the column chunk within the file. - pub offset: u64, - /// Length of the column chunk in bytes. - pub len: usize, - /// Compressed pages of this chunk, in file order, with file-absolute offsets. - pub pages: Vec, -} - -/// Walks every column chunk's page headers and returns the compressed page bodies. -/// -/// The outer `Vec` is one entry per row group, which is the unit a reader can stage on the -/// device and release before moving on. Chunks and the pages within them are in file order. -pub fn scan_compressed_pages( - file_bytes: &[u8], - metadata: &ParquetMetaData, -) -> Result>> { - let mut row_groups = Vec::with_capacity(metadata.row_groups().len()); - - for row_group in metadata.row_groups() { - let mut chunks = Vec::with_capacity(row_group.columns().len()); - for column in row_group.columns() { - let (chunk_offset, chunk_len) = column.byte_range(); - let mut pages = Vec::new(); - let (start, len) = (chunk_offset, chunk_len); - let start = usize::try_from(start)?; - let end = start - .checked_add(usize::try_from(len)?) - .filter(|end| *end <= file_bytes.len()) - .ok_or_else(|| { - anyhow::anyhow!( - "column chunk range {start}..+{len} extends past the {} byte file", - file_bytes.len() - ) - })?; - - let mut pos = start; - while pos < end { - let header = read_page_header(&file_bytes[pos..end])?; - let body = pos + header.header_len; - let body_end = body - .checked_add(header.compressed_len) - .filter(|body_end| *body_end <= end) - .ok_or_else(|| { - anyhow::anyhow!( - "page body at {body} of {} bytes overruns its column chunk", - header.compressed_len - ) - })?; - - match header.page_type { - PageType::Data | PageType::Dictionary => pages.push(CompressedPage { - offset: body, - compressed_len: header.compressed_len, - uncompressed_len: header.uncompressed_len, - }), - PageType::DataV2 => bail!( - "v2 data pages are not GPU-decompressible as a single chunk; \ - write the file with WriterVersion::PARQUET_1_0" - ), - // Index pages are not part of the column data and are never written by - // `parquet`; skip over the body rather than decompressing it. - PageType::Index => {} - } - - pos = body_end; - } - - chunks.push(ColumnChunkPages { - offset: chunk_offset, - len: usize::try_from(chunk_len)?, - pages, - }); - } - - row_groups.push(chunks); - } - - Ok(row_groups) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum PageType { - Data, - Index, - Dictionary, - DataV2, -} - -struct PageHeaderInfo { - header_len: usize, - page_type: PageType, - compressed_len: usize, - uncompressed_len: usize, -} - -/// Thrift compact-protocol field types. -mod ttype { - pub(super) const STOP: u8 = 0x00; - pub(super) const BOOL_TRUE: u8 = 0x01; - pub(super) const BOOL_FALSE: u8 = 0x02; - pub(super) const I8: u8 = 0x03; - pub(super) const I16: u8 = 0x04; - pub(super) const I32: u8 = 0x05; - pub(super) const I64: u8 = 0x06; - pub(super) const DOUBLE: u8 = 0x07; - pub(super) const BINARY: u8 = 0x08; - pub(super) const LIST: u8 = 0x09; - pub(super) const SET: u8 = 0x0a; - pub(super) const MAP: u8 = 0x0b; - pub(super) const STRUCT: u8 = 0x0c; - pub(super) const UUID: u8 = 0x0d; -} - -/// Guards against unbounded recursion on malformed headers. -const MAX_STRUCT_DEPTH: u32 = 32; - -/// Reads the `PageHeader` at the start of `buf`, returning its fields and encoded length. -fn read_page_header(buf: &[u8]) -> Result { - let mut reader = CompactReader { buf, pos: 0 }; - let mut page_type = None; - let mut uncompressed_len = None; - let mut compressed_len = None; - let mut last_field_id = 0i16; - - while let Some((field_id, field_type)) = reader.read_field_header(&mut last_field_id)? { - match (field_id, field_type) { - (1, ttype::I32) => page_type = Some(reader.read_i32()?), - (2, ttype::I32) => uncompressed_len = Some(reader.read_i32()?), - (3, ttype::I32) => compressed_len = Some(reader.read_i32()?), - _ => reader.skip_value(field_type, 0)?, - } - } - - let page_type = match page_type { - Some(0) => PageType::Data, - Some(1) => PageType::Index, - Some(2) => PageType::Dictionary, - Some(3) => PageType::DataV2, - Some(other) => bail!("unknown Parquet page type {other}"), - None => bail!("page header is missing its page type"), - }; - - let uncompressed_len = - uncompressed_len.ok_or_else(|| anyhow::anyhow!("page header is missing its size"))?; - let compressed_len = compressed_len - .ok_or_else(|| anyhow::anyhow!("page header is missing its compressed size"))?; - - Ok(PageHeaderInfo { - header_len: reader.pos, - page_type, - compressed_len: usize::try_from(compressed_len)?, - uncompressed_len: usize::try_from(uncompressed_len)?, - }) -} - -/// Minimal reader for the subset of the Thrift compact protocol that page headers use. -struct CompactReader<'a> { - buf: &'a [u8], - pos: usize, -} - -impl CompactReader<'_> { - fn read_u8(&mut self) -> Result { - let byte = *self - .buf - .get(self.pos) - .ok_or_else(|| anyhow::anyhow!("page header ends mid-field"))?; - self.pos += 1; - Ok(byte) - } - - fn advance(&mut self, len: usize) -> Result<()> { - let end = self - .pos - .checked_add(len) - .filter(|end| *end <= self.buf.len()) - .ok_or_else(|| anyhow::anyhow!("page header ends mid-value"))?; - self.pos = end; - Ok(()) - } - - fn read_varint(&mut self) -> Result { - let mut value = 0u64; - for shift in (0..64).step_by(7) { - let byte = self.read_u8()?; - value |= u64::from(byte & 0x7f) << shift; - if byte & 0x80 == 0 { - return Ok(value); - } - } - bail!("varint in page header is not terminated") - } - - fn read_zigzag(&mut self) -> Result { - let encoded = self.read_varint()?; - Ok(((encoded >> 1) as i64) ^ -((encoded & 1) as i64)) - } - - fn read_i32(&mut self) -> Result { - Ok(i32::try_from(self.read_zigzag()?)?) - } - - /// Reads the next field header, or `None` at the struct's STOP byte. - fn read_field_header(&mut self, last_field_id: &mut i16) -> Result> { - let header = self.read_u8()?; - if header == ttype::STOP { - return Ok(None); - } - - let field_type = header & 0x0f; - let delta = header >> 4; - let field_id = if delta == 0 { - i16::try_from(self.read_zigzag()?)? - } else { - last_field_id - .checked_add(i16::from(delta)) - .ok_or_else(|| anyhow::anyhow!("field id overflow in page header"))? - }; - *last_field_id = field_id; - - Ok(Some((field_id, field_type))) - } - - fn skip_struct(&mut self, depth: u32) -> Result<()> { - ensure!( - depth < MAX_STRUCT_DEPTH, - "page header nests structs more than {MAX_STRUCT_DEPTH} deep" - ); - let mut last_field_id = 0i16; - while let Some((_, field_type)) = self.read_field_header(&mut last_field_id)? { - self.skip_value(field_type, depth + 1)?; - } - Ok(()) - } - - /// Skips a field value. Booleans carry their value in the field type, so consume nothing. - fn skip_value(&mut self, field_type: u8, depth: u32) -> Result<()> { - match field_type { - ttype::BOOL_TRUE | ttype::BOOL_FALSE => Ok(()), - ttype::I8 => self.advance(1), - ttype::I16 | ttype::I32 | ttype::I64 => self.read_varint().map(|_| ()), - ttype::DOUBLE => self.advance(8), - ttype::UUID => self.advance(16), - ttype::BINARY => { - let len = usize::try_from(self.read_varint()?)?; - self.advance(len) - } - ttype::LIST | ttype::SET => { - let (len, element_type) = self.read_collection_header()?; - for _ in 0..len { - self.skip_element(element_type, depth + 1)?; - } - Ok(()) - } - ttype::MAP => { - let len = usize::try_from(self.read_varint()?)?; - if len > 0 { - let types = self.read_u8()?; - let (key_type, value_type) = (types >> 4, types & 0x0f); - for _ in 0..len { - self.skip_element(key_type, depth + 1)?; - self.skip_element(value_type, depth + 1)?; - } - } - Ok(()) - } - ttype::STRUCT => self.skip_struct(depth), - other => bail!("unsupported Thrift compact type {other} in page header"), - } - } - - /// Skips one collection element. Unlike fields, booleans here occupy a byte of their own. - fn skip_element(&mut self, element_type: u8, depth: u32) -> Result<()> { - match element_type { - ttype::BOOL_TRUE | ttype::BOOL_FALSE => self.advance(1), - other => self.skip_value(other, depth), - } - } - - fn read_collection_header(&mut self) -> Result<(usize, u8)> { - let header = self.read_u8()?; - let element_type = header & 0x0f; - let len = match header >> 4 { - 0x0f => usize::try_from(self.read_varint()?)?, - short_len => usize::from(short_len), - }; - Ok((len, element_type)) - } -} - -#[cfg(test)] -mod tests { - use std::fs::File; - use std::sync::Arc; - - use arrow_array::Int64Array; - use arrow_array::RecordBatch; - use arrow_array::StringArray; - use arrow_schema::DataType; - use arrow_schema::Field; - use arrow_schema::Schema; - use parquet::arrow::ArrowWriter; - use parquet::file::metadata::ParquetMetaDataReader; - use parquet::file::reader::FileReader; - use parquet::file::reader::SerializedFileReader; - use rstest::rstest; - - use super::*; - - fn sample_batch() -> Result { - let schema = Arc::new(Schema::new(vec![ - Field::new("ints", DataType::Int64, false), - Field::new("strings", DataType::Utf8, false), - ])); - let ints = Int64Array::from_iter_values((0..50_000).map(|i| i % 977)); - let strings = - StringArray::from_iter_values((0..50_000).map(|i| format!("value-{}", i % 1_000))); - Ok(RecordBatch::try_new( - schema, - vec![Arc::new(ints), Arc::new(strings)], - )?) - } - - fn write_sample(path: &std::path::Path, codec: GpuCodec) -> Result<()> { - let batch = sample_batch()?; - let file = File::create(path)?; - let mut writer = - ArrowWriter::try_new(file, batch.schema(), Some(gpu_writer_properties(codec)))?; - writer.write(&batch)?; - writer.close()?; - Ok(()) - } - - /// The page bodies we locate must decompress to exactly the bytes `parquet` itself reads. - #[rstest] - #[case(GpuCodec::Snappy)] - #[case(GpuCodec::Zstd)] - fn scanned_pages_match_parquet_reader(#[case] codec: GpuCodec) -> Result<()> { - let dir = tempfile::tempdir()?; - let path = dir.path().join("sample.parquet"); - write_sample(&path, codec)?; - - let file = File::open(&path)?; - let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; - let file_bytes = std::fs::read(&path)?; - let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; - let pages = row_groups - .iter() - .flatten() - .flat_map(|chunk| chunk.pages.iter()) - .collect::>(); - - let reader = SerializedFileReader::new(File::open(&path)?)?; - let mut expected = Vec::new(); - for row_group in 0..reader.metadata().num_row_groups() { - let row_group_reader = reader.get_row_group(row_group)?; - for column in 0..row_group_reader.num_columns() { - let mut page_reader = row_group_reader.get_column_page_reader(column)?; - while let Some(page) = page_reader.get_next_page()? { - expected.push(page.buffer().to_vec()); - } - } - } - - assert_eq!(pages.len(), expected.len(), "page count mismatch"); - assert!(!pages.is_empty(), "expected the sample file to have pages"); - - for (page, expected) in pages.iter().zip(expected.iter()) { - let compressed = &file_bytes[page.offset..page.offset + page.compressed_len]; - let decompressed = codec.decompress_host(compressed, page.uncompressed_len)?; - assert_eq!(&decompressed, expected); - } - - Ok(()) - } - - /// Pages must tile their column chunks exactly, with no gap left unaccounted for. - #[test] - fn scanned_pages_cover_every_column_chunk() -> Result<()> { - let dir = tempfile::tempdir()?; - let path = dir.path().join("sample.parquet"); - write_sample(&path, GpuCodec::Snappy)?; - - let file = File::open(&path)?; - let metadata = ParquetMetaDataReader::new().parse_and_finish(&file)?; - let file_bytes = std::fs::read(&path)?; - let row_groups = scan_compressed_pages(&file_bytes, &metadata)?; - let pages = row_groups - .iter() - .flatten() - .flat_map(|chunk| chunk.pages.iter()) - .collect::>(); - - let compressed: usize = pages.iter().map(|page| page.compressed_len).sum(); - let chunk_total: i64 = metadata - .row_groups() - .iter() - .flat_map(|rg| rg.columns()) - .map(|col| col.compressed_size()) - .sum(); - - // The chunk total includes the page headers, so the page bodies must be strictly - // smaller but within a header's worth per page. - assert!(compressed < usize::try_from(chunk_total)?); - assert!(compressed > usize::try_from(chunk_total)? - pages.len() * 256); - Ok(()) - } -} diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py new file mode 100644 index 00000000000..d597762e114 --- /dev/null +++ b/scripts/cudf-parquet-read.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Times a full GPU Parquet read with cuDF. + +`cudf.read_parquet` performs the entire read on the device: page header decode, +codec decompression, dictionary/RLE/plain decoding and column assembly. That makes it +the like-for-like opponent for the Vortex GPU backend, which also decodes all the way +to canonical arrays on device. + +Timing excludes interpreter start, `import cudf`, CUDA context creation and any JIT +warm-up, all of which are paid once per process and are not part of a read. A warm-up +read runs first for exactly that reason. + +Emits one JSON object on stdout so the benchmark can parse it. +""" + +import argparse +import json +import sys +import time + + +def synchronize() -> None: + """Block until queued device work finishes. + + `cudf.read_parquet` returns a materialized DataFrame, but synchronizing explicitly + keeps the measurement honest if that ever stops being true. + """ + try: + import cupy + + cupy.cuda.runtime.deviceSynchronize() + except ImportError: + pass + + +def verify(path: str, frame) -> None: + """Fails unless the GPU read matches a CPU Parquet read of the same file.""" + import pandas as pd + from pandas.testing import assert_frame_equal + + expected = pd.read_parquet(path) + actual = frame.to_pandas() + + # cuDF and pyarrow can land on different-but-equivalent dtypes (nullable vs numpy + # backed, for instance), so compare values and leave dtype policy out of it. + assert_frame_equal(actual, expected, check_dtype=False) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", help="Parquet file to read") + parser.add_argument( + "--iterations", type=int, default=1, help="timed reads to perform" + ) + parser.add_argument( + "--verify", + action="store_true", + help="cross-check the GPU read against a CPU Parquet read", + ) + args = parser.parse_args() + + import cudf + + # Warm-up: pays CUDA context creation and any first-call JIT so they stay out of + # the timed reads below. + warmup = cudf.read_parquet(args.path) + synchronize() + + if args.verify: + verify(args.path, warmup) + + rows, columns = warmup.shape + del warmup + + runs_ns = [] + for _ in range(max(args.iterations, 1)): + start = time.perf_counter_ns() + frame = cudf.read_parquet(args.path) + synchronize() + runs_ns.append(time.perf_counter_ns() - start) + del frame + + json.dump( + { + "min_ns": min(runs_ns), + "runs_ns": runs_ns, + "rows": int(rows), + "columns": int(columns), + "verified": bool(args.verify), + }, + sys.stdout, + ) + sys.stdout.write("\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vortex-cuda/nvcomp/build.rs b/vortex-cuda/nvcomp/build.rs index f34507378db..877fc68c58d 100644 --- a/vortex-cuda/nvcomp/build.rs +++ b/vortex-cuda/nvcomp/build.rs @@ -90,20 +90,13 @@ fn main() { let bindings = bindgen::Builder::default() .header(include_dir.join("nvcomp.h").to_string_lossy()) .header(include_dir.join("nvcomp/zstd.h").to_string_lossy()) - .header(include_dir.join("nvcomp/snappy.h").to_string_lossy()) .clang_arg(format!("-I{}", include_dir.display())) .clang_arg(format!("-I{}", cuda_stub_dir.display())) .allowlist_type("nvcompStatus_t") - .allowlist_type("nvcompAlignmentRequirements_t") .allowlist_type("nvcompBatchedZstdDecompressOpts_t") - .allowlist_type("nvcompBatchedSnappyDecompressOpts_t") .allowlist_type("nvcompDecompressBackend_t") .allowlist_function("nvcompBatchedZstdDecompressGetTempSizeAsync") - .allowlist_function("nvcompBatchedZstdDecompressGetRequiredAlignments") .allowlist_function("nvcompBatchedZstdDecompressAsync") - .allowlist_function("nvcompBatchedSnappyDecompressGetTempSizeAsync") - .allowlist_function("nvcompBatchedSnappyDecompressGetRequiredAlignments") - .allowlist_function("nvcompBatchedSnappyDecompressAsync") .dynamic_library_name("NvcompLibrary") .dynamic_link_require_all(true) .wrap_unsafe_ops(true) diff --git a/vortex-cuda/nvcomp/src/backend.rs b/vortex-cuda/nvcomp/src/backend.rs deleted file mode 100644 index 5be1ebbde4c..00000000000 --- a/vortex-cuda/nvcomp/src/backend.rs +++ /dev/null @@ -1,52 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Options and metadata shared by nvcomp's batched decompression APIs. - -use crate::sys; - -/// Backend selection for nvcomp decompression. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum DecompressBackend { - /// Let nvcomp auto-select the best backend for the hardware. - #[default] - Default, - /// Use hardware decompression - Hardware, - /// Use CUDA - Cuda, -} - -impl DecompressBackend { - pub(crate) fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { - match self { - Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, - Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, - Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, - } - } -} - -/// Minimum buffer alignments required by an nvcomp algorithm. -/// -/// Buffers passed to the batched decompression entrypoints must satisfy these alignments. -/// Exceeding them (for example 16- or 32-byte alignment) may improve throughput. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AlignmentRequirements { - /// Minimum alignment of each compressed input chunk. - pub input: usize, - /// Minimum alignment of each decompressed output chunk. - pub output: usize, - /// Minimum alignment of the temporary workspace buffer. - pub temp: usize, -} - -impl From for AlignmentRequirements { - fn from(value: sys::nvcompAlignmentRequirements_t) -> Self { - Self { - input: value.input, - output: value.output, - temp: value.temp, - } - } -} diff --git a/vortex-cuda/nvcomp/src/lib.rs b/vortex-cuda/nvcomp/src/lib.rs index 859eecbea22..3ab502baa63 100644 --- a/vortex-cuda/nvcomp/src/lib.rs +++ b/vortex-cuda/nvcomp/src/lib.rs @@ -30,13 +30,9 @@ use std::sync::OnceLock; )] pub mod sys; -mod backend; mod error; -pub mod snappy; pub mod zstd; -pub use backend::AlignmentRequirements; -pub use backend::DecompressBackend; pub use error::NvcompError; /// The loaded nvcomp library instance. diff --git a/vortex-cuda/nvcomp/src/snappy.rs b/vortex-cuda/nvcomp/src/snappy.rs deleted file mode 100644 index 175c19d3a13..00000000000 --- a/vortex-cuda/nvcomp/src/snappy.rs +++ /dev/null @@ -1,222 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Wrappers around nvcomp's batched Snappy decompression API. -//! -//! Snappy is the default Parquet page codec and is the fastest of the codecs nvcomp -//! implements on device, which makes it the codec of choice when feeding Parquet pages -//! to the GPU. - -use std::ffi::c_void; - -use crate::backend::AlignmentRequirements; -pub use crate::backend::DecompressBackend; -use crate::error::NvcompError; -use crate::error::check_status; -use crate::nvcomp_library; -use crate::sys; - -/// The largest compressed chunk the Snappy decompressor accepts, in bytes. -pub const MAX_COMPRESSED_CHUNK_SIZE: usize = (1 << 31) - 1; - -/// Options for batched Snappy decompression. -#[derive(Debug, Clone, Copy, Default)] -pub struct SnappyDecompressOpts { - /// Which nvcomp backend performs the decompression. - pub backend: DecompressBackend, - /// Sort chunks by size before submitting them to the hardware decompression engine. - /// - /// Only used when `backend` selects the hardware engine. - pub sort_before_hw_decompress: bool, -} - -impl SnappyDecompressOpts { - fn to_nvcomp(self) -> sys::nvcompBatchedSnappyDecompressOpts_t { - sys::nvcompBatchedSnappyDecompressOpts_t { - backend: self.backend.to_nvcomp(), - sort_before_hw_decompress: i32::from(self.sort_before_hw_decompress), - reserved: [0; 56], - } - } -} - -/// Computes required temporary buffer size for batched Snappy decompression. -/// -/// # Arguments -/// -/// * `num_chunks` - Number of compressed chunks to decompress -/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk -/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks -/// -/// # Returns -/// -/// The required size in bytes for the temporary buffer. -pub fn get_decompress_temp_size( - num_chunks: usize, - max_uncompressed_chunk_bytes: usize, - max_total_uncompressed_bytes: usize, -) -> Result { - get_decompress_temp_size_with_opts( - num_chunks, - max_uncompressed_chunk_bytes, - max_total_uncompressed_bytes, - SnappyDecompressOpts::default(), - ) -} - -/// Computes required temporary buffer size with custom options. -/// -/// # Arguments -/// -/// * `num_chunks` - Number of compressed chunks to decompress -/// * `max_uncompressed_chunk_bytes` - Maximum uncompressed size of any single chunk -/// * `max_total_uncompressed_bytes` - Total uncompressed size across all chunks -/// * `opts` - Decompression options -/// -/// # Returns -/// -/// The required size in bytes for the temporary buffer. -pub fn get_decompress_temp_size_with_opts( - num_chunks: usize, - max_uncompressed_chunk_bytes: usize, - max_total_uncompressed_bytes: usize, - opts: SnappyDecompressOpts, -) -> Result { - let library = nvcomp_library()?; - - let mut temp_bytes: usize = 0; - - let status = unsafe { - library.nvcompBatchedSnappyDecompressGetTempSizeAsync( - num_chunks, - max_uncompressed_chunk_bytes, - opts.to_nvcomp(), - &raw mut temp_bytes, - max_total_uncompressed_bytes, - ) - }; - - check_status(status)?; - Ok(temp_bytes) -} - -/// Returns the minimum buffer alignments required by batched Snappy decompression. -pub fn decompress_alignment_requirements( - opts: SnappyDecompressOpts, -) -> Result { - let library = nvcomp_library()?; - - let mut requirements = sys::nvcompAlignmentRequirements_t { - input: 0, - output: 0, - temp: 0, - }; - - let status = unsafe { - library.nvcompBatchedSnappyDecompressGetRequiredAlignments( - opts.to_nvcomp(), - &raw mut requirements, - ) - }; - - check_status(status)?; - Ok(requirements.into()) -} - -/// Launches batched Snappy decompression asynchronously on the GPU. -/// -/// This function decompresses multiple raw Snappy blocks in parallel on the GPU. All -/// pointer arguments must point to device memory, and the operation is executed -/// asynchronously on the provided CUDA stream. -/// -/// # Arguments -/// -/// * `device_compressed_ptrs` - Device pointer to array of pointers to compressed chunks -/// * `device_compressed_bytes` - Device pointer to array of compressed chunk sizes -/// * `device_uncompressed_bytes` - Device pointer to array of expected uncompressed sizes -/// * `device_actual_uncompressed_bytes` - Device pointer to array for actual uncompressed sizes (output) -/// * `num_chunks` - Number of chunks to decompress -/// * `device_temp_ptr` - Device pointer to temporary workspace buffer -/// * `temp_bytes` - Size of temporary buffer in bytes -/// * `device_uncompressed_ptrs` - Device pointer to array of pointers to output buffers -/// * `device_statuses` - Device pointer to array for per-chunk status codes (output) -/// * `stream` - CUDA stream to execute on -/// -/// # Safety -/// -/// - All device pointers must be valid and point to properly allocated device memory -/// - `device_compressed_ptrs` must point to valid device pointers -/// - `device_uncompressed_ptrs` must point to valid device pointers -/// - Each output buffer must have at least the corresponding `device_uncompressed_bytes` size -/// - `device_temp_ptr` must have at least `temp_bytes` allocated -/// - The stream must be valid -#[expect(clippy::too_many_arguments)] -pub unsafe fn decompress_async( - device_compressed_ptrs: *const *const c_void, - device_compressed_bytes: *const usize, - device_uncompressed_bytes: *const usize, - device_actual_uncompressed_bytes: *mut usize, - num_chunks: usize, - device_temp_ptr: *mut c_void, - temp_bytes: usize, - device_uncompressed_ptrs: *const *mut c_void, - device_statuses: *mut sys::nvcompStatus_t, - stream: sys::cudaStream_t, -) -> Result<(), NvcompError> { - // SAFETY: Caller has to ensure all pointers are valid. - unsafe { - decompress_async_with_opts( - device_compressed_ptrs, - device_compressed_bytes, - device_uncompressed_bytes, - device_actual_uncompressed_bytes, - num_chunks, - device_temp_ptr, - temp_bytes, - device_uncompressed_ptrs, - device_statuses, - stream, - SnappyDecompressOpts::default(), - ) - } -} - -/// Launches batched Snappy decompression asynchronously with custom options. -/// -/// # Safety -/// -/// Same requirements as [`decompress_async`]. -#[expect(clippy::too_many_arguments)] -pub unsafe fn decompress_async_with_opts( - device_compressed_ptrs: *const *const c_void, - device_compressed_bytes: *const usize, - device_uncompressed_bytes: *const usize, - device_actual_uncompressed_bytes: *mut usize, - num_chunks: usize, - device_temp_ptr: *mut c_void, - temp_bytes: usize, - device_uncompressed_ptrs: *const *mut c_void, - device_statuses: *mut sys::nvcompStatus_t, - stream: sys::cudaStream_t, - opts: SnappyDecompressOpts, -) -> Result<(), NvcompError> { - let library = nvcomp_library()?; - - let status = unsafe { - library.nvcompBatchedSnappyDecompressAsync( - device_compressed_ptrs, - device_compressed_bytes, - device_uncompressed_bytes, - device_actual_uncompressed_bytes, - num_chunks, - device_temp_ptr, - temp_bytes, - device_uncompressed_ptrs, - opts.to_nvcomp(), - device_statuses, - stream, - ) - }; - - check_status(status) -} diff --git a/vortex-cuda/nvcomp/src/zstd.rs b/vortex-cuda/nvcomp/src/zstd.rs index 04769d2621a..44111901071 100644 --- a/vortex-cuda/nvcomp/src/zstd.rs +++ b/vortex-cuda/nvcomp/src/zstd.rs @@ -5,13 +5,33 @@ use std::ffi::c_void; -use crate::backend::AlignmentRequirements; -pub use crate::backend::DecompressBackend; use crate::error::NvcompError; use crate::error::check_status; use crate::nvcomp_library; use crate::sys; +/// Backend selection for nvcomp decompression. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum DecompressBackend { + /// Let nvcomp auto-select the best backend for the hardware. + #[default] + Default, + /// Use hardware decompression + Hardware, + /// Use CUDA + Cuda, +} + +impl DecompressBackend { + fn to_nvcomp(self) -> sys::nvcompDecompressBackend_t { + match self { + Self::Default => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_DEFAULT, + Self::Hardware => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_HARDWARE, + Self::Cuda => sys::nvcompDecompressBackend_t_NVCOMP_DECOMPRESS_BACKEND_CUDA, + } + } +} + /// Options for batched ZSTD decompression. #[derive(Debug, Clone, Copy, Default)] pub struct ZstdDecompressOpts { @@ -87,29 +107,6 @@ pub fn get_decompress_temp_size_with_opts( Ok(temp_bytes) } -/// Returns the minimum buffer alignments required by batched ZSTD decompression. -pub fn decompress_alignment_requirements( - opts: ZstdDecompressOpts, -) -> Result { - let library = nvcomp_library()?; - - let mut requirements = sys::nvcompAlignmentRequirements_t { - input: 0, - output: 0, - temp: 0, - }; - - let status = unsafe { - library.nvcompBatchedZstdDecompressGetRequiredAlignments( - opts.to_nvcomp(), - &raw mut requirements, - ) - }; - - check_status(status)?; - Ok(requirements.into()) -} - /// Launches batched ZSTD decompression asynchronously on the GPU. /// /// This function decompresses multiple ZSTD-compressed chunks in parallel on the GPU. From c5a74d17e3abb501c2f427eef447f0f39407b82c Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 15:07:10 +0000 Subject: [PATCH 09/22] Run the GPU verification reference through a host execution context The reference side of the Vortex verification was executing through the CUDA context: the host scan's batches and both Arrow conversions were handed `cuda_ctx.execution_ctx()`. A CUDA context allocates its outputs in device memory, so the Arrow conversion then read a device buffer from the host and panicked with "unwrap_host called for Device allocation" on the string-heavy Public BI datasets, where canonicalisation goes through the buffer directly. Signed-off-by: Claude --- benchmarks/compress-bench/src/gpu_vortex.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 9b25b2f6dd4..2b9e5399e94 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -18,6 +18,7 @@ use tempfile::NamedTempFile; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; use vortex::array::arrays::struct_::StructArrayExt; use vortex::compressor::BtrBlocksCompressorBuilder; @@ -124,6 +125,10 @@ async fn open_gpu(path: &Path) -> Result { /// Verification runs inline, so the returned duration is not comparable to a plain run. async fn verify_against_host_scan(path: &Path) -> Result { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + // Everything on the reference side — the host scan and both Arrow conversions — has to run + // through a plain host context. A CUDA context allocates its outputs in device memory, and + // the Arrow conversion then reads those buffers on the host. + let mut host_ctx = SESSION.create_execution_ctx(); let start = Instant::now(); // The host scan reads a copy rather than the same path. The session's segment cache is @@ -148,7 +153,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { }; let gpu_record = gpu_batch.execute::(cuda_ctx.execution_ctx())?; - let host_record = host_batch.execute::(cuda_ctx.execution_ctx())?; + let host_record = host_batch.execute::(&mut host_ctx)?; ensure!( gpu_record.len() == host_record.len(), "batch {batch_index} length differs between the GPU and CPU scans: {} vs {}", @@ -180,7 +185,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { verify_field( &host_field, decoded, - cuda_ctx.execution_ctx(), + &mut host_ctx, batch_index, field_index, )?; From c3a0ad928a93d27f3aecb07492c6f9bcb496b62b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 16:15:57 +0000 Subject: [PATCH 10/22] Publish GPU timings even when a dataset fails to decode The verification gate ran before the timed pass, so a single unsupported dataset meant the run produced no numbers at all. Two changes: - the per-dataset survey now applies to any GPU run, not just a verifying one, and the timing tables render before the failure summary, so datasets that do decode still publish their numbers; - the workflow runs the benchmark before the gate and fails the job at the end on either a failed verification or a failed benchmark. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 19 +++++--- benchmarks/compress-bench/src/main.rs | 50 +++++++++++---------- 2 files changed, 40 insertions(+), 29 deletions(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index d2975cbbb91..a4e477f3fe3 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -105,14 +105,13 @@ jobs: with: file-path: verify-comment.md comment-tag: bench-pr-comment-gpu-verify - - name: Fail if verification failed - if: steps.verify.outcome == 'failure' - shell: bash - run: | - echo "GPU decompression verification failed; see the verification comment." >&2 - exit 1 - name: Run GPU compression benchmark + id: bench shell: bash + # A dataset the GPU cannot decode is reported per dataset and still fails the run, but + # the timing tables are printed first, so the datasets that do decode publish their + # numbers. The gate at the end of the job turns either failure into a job failure. + continue-on-error: true env: RUST_BACKTRACE: full # Do not enable VORTEX_EXPERIMENTAL_PATCHED_ARRAY here: it rewrites interior @@ -152,3 +151,11 @@ jobs: GPU Compression failed. Check the [workflow run](${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}) for details. comment-tag: bench-pr-comment-gpu-compress + - name: Fail if verification or the benchmark failed + if: steps.verify.outcome == 'failure' || steps.bench.outcome == 'failure' + shell: bash + # Last, so both the verification matrix and the timing tables are already published. + run: | + echo "verify=${{ steps.verify.outcome }} bench=${{ steps.bench.outcome }}" >&2 + echo "GPU decompression failed; see the verification and results comments." >&2 + exit 1 diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 34f67b13e28..fefe61c44f3 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -285,9 +285,9 @@ async fn run_compress( let mut measurements = vec![]; let mut v3_records: Vec = Vec::new(); - // A verification pass reports on every dataset rather than stopping at the first failure: - // one run then says exactly which datasets decode correctly on the GPU and which do not. - let survey_all = gpu.is_some_and(|gpu| gpu.verify); + // A GPU pass reports on every dataset rather than stopping at the first failure, so one run + // says which datasets decode correctly on the GPU and still yields numbers for the rest. + let survey_all = gpu.is_some(); let mut failures: Vec<(String, anyhow::Error)> = Vec::new(); for dataset_handle in datasets.into_iter() { @@ -322,30 +322,14 @@ async fn run_compress( progress.finish(); - if !failures.is_empty() { - eprintln!( - "\nGPU verification failed for {} dataset(s):", - failures.len() - ); - for (dataset, error) in &failures { - eprintln!(" - {dataset}: {error:#}"); - } - anyhow::bail!( - "GPU verification failed for: {}", - failures - .iter() - .map(|(dataset, _)| dataset.as_str()) - .collect::>() - .join(", ") - ); - } - if let Some(path) = ingest_output { v3::write_jsonl_to_path(&path, &v3_records)?; } let mut writer = create_output_writer(&display_format, output_path, BENCHMARK_ID)?; + // The tables render before any failure is reported, so a partially failing GPU matrix still + // publishes the numbers for the datasets that did decode. match display_format { DisplayFormat::Table => { render_table(&mut writer, measurements.timings, &targets)?; @@ -357,13 +341,33 @@ async fn run_compress( } else { vec![] }, - ) + )?; } DisplayFormat::GhJson => { print_measurements_json(&mut writer, measurements.timings, DOC_PATH)?; - print_measurements_json(&mut writer, measurements.ratios, DOC_PATH) + print_measurements_json(&mut writer, measurements.ratios, DOC_PATH)?; } } + + if !failures.is_empty() { + eprintln!( + "\nGPU decompression failed for {} dataset(s):", + failures.len() + ); + for (dataset, error) in &failures { + eprintln!(" - {dataset}: {error:#}"); + } + anyhow::bail!( + "GPU decompression failed for: {}", + failures + .iter() + .map(|(dataset, _)| dataset.as_str()) + .collect::>() + .join(", ") + ); + } + + Ok(()) } async fn run_benchmark_for_dataset( From ca22f1883f46335888801a263186d2437b5002fd Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 16:17:53 +0000 Subject: [PATCH 11/22] Document that a partial GPU matrix still publishes its timings Signed-off-by: Claude --- benchmarks/compress-bench/README.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index e14c117d5a6..b1fd9bc6c7c 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -86,5 +86,7 @@ cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --gpu-verify --iterations 1 ``` -A verifying run reports on every dataset rather than stopping at the first failure, so one run -shows which datasets decode correctly on the GPU and which do not. +Any `--gpu-decompress` run reports on every dataset rather than stopping at the first failure, so +one run shows which datasets decode correctly on the GPU and which do not. The timing tables are +rendered before the failure summary, so a dataset the GPU cannot decode still leaves the rest of +the matrix with numbers — the process exits non-zero either way. From 2afc3b79cbbcdf6779908ecd6d4df957526799ce Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 16:58:34 +0000 Subject: [PATCH 12/22] CI: publish the full error detail from a failed GPU verification The per-dataset grep captures only the first line of each error, so a Python traceback from the cuDF script or a Rust backtrace never reached the comment. Attach the tail of the raw output in a collapsed block on failure, and match the summary line's current wording. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index a4e477f3fe3..9131518615b 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -94,9 +94,22 @@ jobs: fi echo echo '```text' - grep -E "verified [0-9]+|GPU verification failed|^ - " gpu-verify.txt | tail -40 \ + grep -E "verified [0-9]+|GPU decompression failed|^ - " gpu-verify.txt | tail -40 \ || tail -40 gpu-verify.txt echo '```' + # The per-dataset lines above carry only the first line of each error. Python + # tracebacks, Rust backtraces and first-differing-row dumps span several lines, so + # the tail of the raw output goes in a collapsed block rather than back in the log. + if [ "${{ steps.verify.outcome }}" != "success" ]; then + echo + echo "
Full error detail" + echo + echo '```text' + tail -200 gpu-verify.txt + echo '```' + echo + echo "
" + fi } > verify-comment.md cat verify-comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR with verification result From 6a0fd9dfc1e13836381100fab251d5debba30fe6 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 13 Aug 2026 17:23:56 +0000 Subject: [PATCH 13/22] Fix three defects the first cuDF comparison run exposed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(cuda): copy validity back to the host in `into_host` `CanonicalCudaExt::into_host` migrated a canonical array's values buffer but passed its validity through untouched, so a nullable array came back to the host half-migrated and the first host read of the mask panicked with "unwrap_host called for Device allocation" — via `Validity::execute_mask` -> `BoolArray::into_bit_buffer`. Non-nullable arrays were unaffected, which is why it only showed up on the Public BI tables. The `Bool` arm already carried a TODO for exactly this. Do not compare Parquet DATE columns across representations pyarrow materialises a DATE column as `datetime.date` objects and cuDF as `datetime64[s]`. The values agree, but `check_dtype=False` does not bridge object-vs-datetime64, so the comparison reported every row as different and failed both TPC-H datasets. Coerce both sides to datetime64 first. Read the Vortex GPU file through the page cache by default cuDF takes an untimed warm-up read, so its timed read is served from the page cache, while the Vortex reader used `O_DIRECT` on every iteration and paid real disk reads each time. That compared a read of the disk against a read of RAM. Direct IO is now off by default and available behind `--gpu-direct-io` for measuring storage bandwidth, which is not a decode comparison. Signed-off-by: Claude --- benchmarks/compress-bench/README.md | 11 +++++-- benchmarks/compress-bench/src/gpu_vortex.rs | 30 +++++++++++------- benchmarks/compress-bench/src/main.rs | 12 +++++++- scripts/cudf-parquet-read.py | 23 ++++++++++++-- vortex-cuda/src/canonical.rs | 34 ++++++++++++++++++--- 5 files changed, 88 insertions(+), 22 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index b1fd9bc6c7c..5033bca34a7 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -53,9 +53,14 @@ uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarro script, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up read runs first for the same reason. -Note that the two backends do not share an I/O path: the Vortex reader uses pinned buffers and -direct I/O (`O_DIRECT`) on Linux, while cuDF does its own host read and host-to-device -transfer. +Both backends read a warm file by default. cuDF runs an untimed warm-up read before the timed +one, so its timed read hits the page cache; the Vortex reader therefore does **not** use direct +I/O by default, because `O_DIRECT` would bypass the page cache and compare a Vortex read of the +disk against a cuDF read of RAM. `--gpu-direct-io` turns it back on to measure storage bandwidth +instead — a different question, and the resulting ratio is not a decode comparison. + +The remaining asymmetry is the transfer path: the Vortex reader uses pinned buffers, while cuDF +does its own host read and host-to-device copy. ### GPU-friendly Parquet writer settings diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 2b9e5399e94..832466978a0 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -42,6 +42,7 @@ use vortex_cuda::layout::register_cuda_layout; /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. pub struct GpuVortexCompressor { verify: bool, + direct_io: bool, } impl GpuVortexCompressor { @@ -50,8 +51,8 @@ impl GpuVortexCompressor { /// When `verify` is set, each GPU-decoded field is copied back to the host and compared /// against the same field decoded on the CPU. Verification runs inline, so timings from a /// verifying run are not comparable to a plain one. - pub fn new(verify: bool) -> Self { - Self { verify } + pub fn new(verify: bool, direct_io: bool) -> Self { + Self { verify, direct_io } } } @@ -84,12 +85,12 @@ impl Compressor for GpuVortexCompressor { drop(output); if self.verify { - return verify_against_host_scan(gpu_file.path()).await; + return verify_against_host_scan(gpu_file.path(), self.direct_io).await; } let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; let start = Instant::now(); - let file = open_gpu(gpu_file.path()).await?; + let file = open_gpu(gpu_file.path(), self.direct_io).await?; let mut batches = file.scan()?.into_array_stream()?; while let Some(batch) = batches.next().await { @@ -106,13 +107,20 @@ impl Compressor for GpuVortexCompressor { /// Opens a Vortex file for CUDA execution. /// -/// On Linux direct IO keeps repeated iterations measuring storage bandwidth rather than -/// page-cache hits. -async fn open_gpu(path: &Path) -> Result { +/// `direct_io` is off by default so this backend is comparable with the cuDF one: cuDF reads +/// through the page cache after an untimed warm-up read, so leaving direct IO on would compare +/// a Vortex read of the disk against a cuDF read of RAM. Turning it on measures storage +/// bandwidth instead, which is a different question and not comparable across the two. +async fn open_gpu(path: &Path, direct_io: bool) -> Result { let open_options = SESSION.open_options().with_cuda(); #[cfg(target_os = "linux")] - let open_options = - open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()); + let open_options = if direct_io { + open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io()) + } else { + open_options + }; + #[cfg(not(target_os = "linux"))] + let _ = direct_io; Ok(open_options.open_path(path).await?) } @@ -123,7 +131,7 @@ async fn open_gpu(path: &Path) -> Result { /// which the host decoders cannot read. /// /// Verification runs inline, so the returned duration is not comparable to a plain run. -async fn verify_against_host_scan(path: &Path) -> Result { +async fn verify_against_host_scan(path: &Path, direct_io: bool) -> Result { let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; // Everything on the reference side — the host scan and both Arrow conversions — has to run // through a plain host context. A CUDA context allocates its outputs in device memory, and @@ -137,7 +145,7 @@ async fn verify_against_host_scan(path: &Path) -> Result { let host_path = NamedTempFile::new()?; std::fs::copy(path, host_path.path())?; - let gpu_file = open_gpu(path).await?; + let gpu_file = open_gpu(path, direct_io).await?; let mut gpu_batches = gpu_file.scan()?.into_array_stream()?; let host_file = SESSION.open_options().open_path(host_path.path()).await?; let mut host_batches = host_file.scan()?.into_array_stream()?; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index fefe61c44f3..8232a9a6018 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -92,6 +92,13 @@ struct Args { /// plain one. Intended to be run as its own pass. #[arg(long)] gpu_verify: bool, + /// Read the Vortex GPU file with direct IO, bypassing the page cache. + /// + /// Off by default: cuDF reads through the page cache after an untimed warm-up, so direct + /// IO would compare a Vortex read of the disk against a cuDF read of RAM. Turn it on to + /// measure storage bandwidth instead, and do not read the ratio as a decode comparison. + #[arg(long)] + gpu_direct_io: bool, #[arg(short, long, default_value_t, value_enum)] display_format: DisplayFormat, #[arg(short, long)] @@ -120,6 +127,7 @@ async fn main() -> anyhow::Result<()> { let gpu = args.gpu_decompress.then_some(GpuOptions { codec: args.gpu_parquet_codec, verify: args.gpu_verify, + direct_io: args.gpu_direct_io, }); let (formats, ops) = if gpu.is_some() { @@ -151,6 +159,8 @@ struct GpuOptions { codec: GpuCodec, /// Cross-check decompressed output against the CPU decoders. verify: bool, + /// Read the Vortex file with direct IO instead of through the page cache. + direct_io: bool, } /// Get a compressor for the given format. @@ -159,7 +169,7 @@ fn get_compressor(format: Format, gpu: Option) -> Box { - Box::new(GpuVortexCompressor::new(gpu.verify)) as Box + Box::new(GpuVortexCompressor::new(gpu.verify, gpu.direct_io)) as Box } Format::Parquet => Box::new(GpuParquetCompressor::new(gpu.codec, gpu.verify)), _ => unimplemented!("GPU compress bench not implemented for {format}"), diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py index d597762e114..39e435c3a77 100644 --- a/scripts/cudf-parquet-read.py +++ b/scripts/cudf-parquet-read.py @@ -20,6 +20,7 @@ import json import sys import time +from datetime import date def synchronize() -> None: @@ -36,13 +37,31 @@ def synchronize() -> None: pass +def normalize(frame): + """Collapses representation differences that are not value differences. + + A Parquet DATE column comes back from pyarrow as a column of `datetime.date` + objects but from cuDF as `datetime64[s]`. Those hold the same instants, yet + `check_dtype=False` does not bridge them because one side is `object`, so the + comparison reports every row as different. Coercing both sides to datetime64 + compares the dates themselves. + """ + import pandas as pd + + for name in frame.columns: + column = frame[name] + if column.dtype == object and len(column) and isinstance(column.iloc[0], date): + frame[name] = pd.to_datetime(column) + return frame + + def verify(path: str, frame) -> None: """Fails unless the GPU read matches a CPU Parquet read of the same file.""" import pandas as pd from pandas.testing import assert_frame_equal - expected = pd.read_parquet(path) - actual = frame.to_pandas() + expected = normalize(pd.read_parquet(path)) + actual = normalize(frame.to_pandas()) # cuDF and pyarrow can land on different-but-equivalent dtypes (nullable vs numpy # backed, for instance), so compare values and leave dtype policy out of it. diff --git a/vortex-cuda/src/canonical.rs b/vortex-cuda/src/canonical.rs index 9f1fce7e68e..6da059f7cd5 100644 --- a/vortex-cuda/src/canonical.rs +++ b/vortex-cuda/src/canonical.rs @@ -23,6 +23,7 @@ use vortex::array::arrays::varbinview::BinaryView; use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::legacy_session; +use vortex::array::validity::Validity; use vortex::buffer::BitBuffer; use vortex::buffer::Buffer; use vortex::buffer::ByteBuffer; @@ -36,6 +37,26 @@ pub trait CanonicalCudaExt { Self: Sized; } +/// Copies an array-backed validity mask back to the host. +/// +/// Only [`Validity::Array`] owns a buffer; the other variants are metadata and pass through. +/// Migrating the values of a nullable array without its validity leaves the mask on the +/// device, and the first host read of it — canonicalising to Arrow, say — panics in +/// `BufferHandle::unwrap_host`. +#[allow(clippy::disallowed_methods)] +async fn validity_into_host(validity: Validity) -> VortexResult { + let Validity::Array(array) = validity else { + return Ok(validity); + }; + Ok(Validity::Array( + array + .execute::(&mut legacy_session().create_execution_ctx())? + .into_host() + .await? + .into_array(), + )) +} + #[async_trait] impl CanonicalCudaExt for Canonical { #[allow(clippy::disallowed_methods)] @@ -67,13 +88,11 @@ impl CanonicalCudaExt for Canonical { struct_fields.names().clone(), host_fields, len, - validity, + validity_into_host(validity).await?, ))) } n @ Canonical::Null(_) => Ok(n), Canonical::Bool(bool) => { - // NOTE: update to copy to host when adding buffer handle. - // Also update other method to copy validity to host. let len = bool.len(); let validity = bool.validity()?; let BoolDataParts { bits, meta } = bool.into_data().into_parts(len); @@ -83,7 +102,10 @@ impl CanonicalCudaExt for Canonical { meta.len(), meta.offset(), ); - Ok(Canonical::Bool(BoolArray::new(bits, validity))) + Ok(Canonical::Bool(BoolArray::new( + bits, + validity_into_host(validity).await?, + ))) } Canonical::Primitive(prim) => { let PrimitiveDataParts { @@ -95,7 +117,7 @@ impl CanonicalCudaExt for Canonical { Ok(Canonical::Primitive(PrimitiveArray::from_byte_buffer( buffer.try_into_host()?.await?, ptype, - validity, + validity_into_host(validity).await?, ))) } Canonical::Decimal(decimal) => { @@ -106,6 +128,7 @@ impl CanonicalCudaExt for Canonical { validity, .. } = decimal.into_data_parts(); + let validity = validity_into_host(validity).await?; Ok(Canonical::Decimal(unsafe { DecimalArray::new_unchecked_handle( BufferHandle::new_host(values.try_into_host()?.await?), @@ -136,6 +159,7 @@ impl CanonicalCudaExt for Canonical { let host_buffers = try_join_all(host_buffers).await?; let host_buffers: Arc<[ByteBuffer]> = Arc::from(host_buffers); + let validity = validity_into_host(validity).await?; Ok(Canonical::VarBinView(unsafe { VarBinViewArray::new_unchecked(host_views, host_buffers, dtype, validity) })) From b9fed6ad16f89d1218829355fd38c2155bc16348 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 14 Aug 2026 09:05:58 +0000 Subject: [PATCH 14/22] CI: measure a CPU decompression baseline on the same GPU runner The GPU ratio says which of two GPU readers is faster, not whether either beats the CPU decoders. Run the same binary over the same datasets with the CPU path on the same machine and publish it alongside, so the GPU numbers can be read against something. Also capture the benchmark's exit status rather than letting `shell: bash`'s -e skip the `cat`, which kept the timing tables out of the job log and left them only in the PR comment. Signed-off-by: Claude --- .github/workflows/pr-bench-gpu-compress.yml | 28 ++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index 9131518615b..505a89a2a75 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -135,9 +135,29 @@ jobs: # (cuda.yaml), which also sets only FLAT_LAYOUT_INLINE_ARRAY_NODE. FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" run: | + # `shell: bash` implies -e, so a non-zero benchmark exit would skip the cat and leave + # the timing tables out of the job log entirely. Capture the status instead. + status=0 bash scripts/bench-taskset.sh target/release_debug/compress-bench \ - --gpu-decompress -d table > gpu-compress.txt + --gpu-decompress -d table > gpu-compress.txt || status=$? cat gpu-compress.txt + exit "$status" + - name: Run CPU decompression baseline + id: cpu-bench + continue-on-error: true + shell: bash + env: + RUST_BACKTRACE: full + FLAT_LAYOUT_INLINE_ARRAY_NODE: "1" + # Same binary, same datasets, same machine, CPU decoders. Without this the GPU ratio + # says which of two GPU readers is faster but not whether either beats the CPU, which + # is the only number that says whether GPU decoding is worth doing at all. + run: | + bash scripts/bench-taskset.sh target/release_debug/compress-bench \ + --ops decompress \ + --datasets 'TPC-H l_comment|taxi|Arade|Bimbo|CMSprovider|Euro2016|Food|HashTags' \ + -d table > cpu-compress.txt + cat cpu-compress.txt - name: Publish results shell: bash run: | @@ -147,6 +167,12 @@ jobs: echo '```text' cat gpu-compress.txt echo '```' + echo + echo "## CPU baseline (same machine, same datasets)" + echo + echo '```text' + cat cpu-compress.txt || echo "CPU baseline did not produce output" + echo '```' } > comment.md cat comment.md >> "$GITHUB_STEP_SUMMARY" - name: Comment PR From f3bceedaa0f20f497febc1ee237baceda734f673 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 16:19:15 +0100 Subject: [PATCH 15/22] Keep GPU FSST decompression on device Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- Cargo.lock | 1 + benchmarks/compress-bench/Cargo.toml | 3 +- benchmarks/compress-bench/README.md | 16 +- benchmarks/compress-bench/src/gpu_vortex.rs | 191 +++++++++++++++++-- benchmarks/compress-bench/src/gpu_writer.rs | 11 ++ vortex-bench/src/conversions.rs | 72 ++++++-- vortex-cuda/kernels/src/fsst.cu | 30 +++ vortex-cuda/src/kernel/encodings/fsst.rs | 194 ++++++++++++++++---- 8 files changed, 444 insertions(+), 74 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a80d6f85f02..43b8042e63e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1572,6 +1572,7 @@ dependencies = [ "async-trait", "bytes", "clap", + "cudarc", "futures", "indicatif", "itertools 0.14.0", diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 7b1dadc3209..1a3d2b6ac60 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,6 +21,7 @@ arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } +cudarc = { workspace = true, features = ["nvtx"], optional = true } futures = { workspace = true } indicatif = { workspace = true } itertools = { workspace = true } @@ -38,7 +39,7 @@ vortex-bench = { workspace = true } vortex-cuda = { workspace = true, optional = true } [features] -cuda = ["dep:tempfile", "dep:vortex-cuda"] +cuda = ["dep:cudarc", "dep:tempfile", "dep:vortex-cuda"] lance = ["dep:lance-bench"] unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"] diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 5033bca34a7..749bf4ac1fd 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -30,6 +30,9 @@ GPU dataset list in `src/main.rs`. It measures decompression only, for two backe Both sides therefore decode all the way to device-resident arrays, which is what makes the `vortex:parquet- gpu ratio decompress time` metric a like-for-like comparison. +The generated files also use the same 1,048,576-row physical partition size: Parquet row groups +and Vortex root chunks. Vortex input batches are concatenated and sliced at those exact boundaries +before writing, so smaller source batches cannot leak into its on-disk layout. ```bash cargo run -p compress-bench --profile release_debug \ @@ -53,11 +56,13 @@ uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarro script, so interpreter start, `import cudf` and CUDA context creation are excluded; a warm-up read runs first for the same reason. -Both backends read a warm file by default. cuDF runs an untimed warm-up read before the timed -one, so its timed read hits the page cache; the Vortex reader therefore does **not** use direct -I/O by default, because `O_DIRECT` would bypass the page cache and compare a Vortex read of the -disk against a cuDF read of RAM. `--gpu-direct-io` turns it back on to measure storage bandwidth -instead — a different question, and the resulting ratio is not a decode comparison. +Both backends read a warm file by default. Each runs an untimed full read before a separately +opened timed read, warming the OS page cache, allocator, and CUDA modules. Neither reuses decoded +arrays, and the Vortex CUDA opener disables its data-segment cache. The Vortex reader therefore +does **not** use direct I/O by default, because `O_DIRECT` would bypass the page cache and compare +a Vortex read of the disk against a cuDF read of RAM. `--gpu-direct-io` turns it back on to measure +storage bandwidth instead — a different question, and the resulting ratio is not a decode +comparison. The remaining asymmetry is the transfer path: the Vortex reader uses pinned buffers, while cuDF does its own host read and host-to-device copy. @@ -73,6 +78,7 @@ Set in `src/gpu_writer.rs`: | dictionary | enabled | Keeps the decompressed payload small; the encoding GPU Parquet readers decode fastest. | | data page size | 1 MiB | Large enough to amortize per-page setup, small enough to keep every SM fed. Matches the page size cuDF targets. | | data page row limit | 1,000,000 | The 20k-row default caps narrow columns' pages far below 1 MiB. | +| row-group / root-chunk rows | 1,048,576 | Gives both formats the same independently readable physical partitions and amortizes GPU launch overhead. | | statistics | chunk-level | Page statistics only inflate the headers a reader has to walk. | ### Correctness diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 832466978a0..627eab23f0d 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -13,6 +13,8 @@ use anyhow::ensure; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::Field; use async_trait::async_trait; +use cudarc::driver::sys::CUevent_flags::CU_EVENT_DEFAULT; +use cudarc::nvtx::safe::scoped_range; use futures::StreamExt; use tempfile::NamedTempFile; use vortex::array::ArrayRef; @@ -24,21 +26,26 @@ use vortex::array::arrays::struct_::StructArrayExt; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; +use vortex::layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex::layout::layouts::compressed::CompressingStrategy; +use vortex::layout::scan::split_by::SplitBy; use vortex_arrow::ArrowSessionExt; use vortex_bench::Format; use vortex_bench::SESSION; use vortex_bench::compress::Compressor; -use vortex_bench::conversions::parquet_to_vortex_chunks; +use vortex_bench::conversions::parquet_to_vortex_chunks_with_batch_size; use vortex_cuda::CanonicalCudaExt; use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; #[cfg(target_os = "linux")] use vortex_cuda::PooledFileReadAtOptions; use vortex_cuda::executor::CudaArrayExt; +use vortex_cuda::executor::CudaExecutionCtx; use vortex_cuda::layout::CudaFlatLayoutStrategy; use vortex_cuda::layout::register_cuda_layout; +use crate::gpu_writer::GPU_ROW_GROUP_SIZE; + /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. pub struct GpuVortexCompressor { verify: bool, @@ -69,13 +76,26 @@ impl Compressor for GpuVortexCompressor { async fn decompress(&self, parquet_path: &Path) -> Result { register_cuda_layout(&SESSION); - let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).await?; + // Match the Parquet writer's row-group size so both formats expose the same number of + // independently readable row partitions to their GPU decoder. The default Arrow reader + // batch size is only 8K rows, which otherwise creates hundreds of tiny CUDA-flat layouts + // and thousands of tiny kernel launches for Vortex. + let uncompressed = parquet_to_vortex_chunks_with_batch_size( + parquet_path.to_path_buf(), + Some(GPU_ROW_GROUP_SIZE), + ) + .await?; let gpu_file = NamedTempFile::new()?; let mut output = tokio::fs::File::create(gpu_file.path()).await?; - let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(BtrBlocksCompressorBuilder::default().only_cuda_compatible()) - .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())) - .build(); + // Preserve the exact input partitions at the root. The general-purpose file strategy + // splits struct fields and may introduce field-specific chunk boundaries, which makes a + // fixed-row GPU scan yield ChunkedArray fields that have no CUDA execution kernel. + let strategy = Arc::new(ChunkedLayoutStrategy::new(CompressingStrategy::new( + CudaFlatLayoutStrategy::default(), + BtrBlocksCompressorBuilder::default() + .only_cuda_compatible() + .build(), + ))); SESSION .write_options() .with_strategy(strategy) @@ -89,20 +109,151 @@ impl Compressor for GpuVortexCompressor { } let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?; + // Match cuDF's untimed read: pay CUDA module loading, allocator initialization and page + // faults before measuring. The helper opens and scans the file afresh each time, and the + // CUDA opener disables its data-segment cache, so no decoded arrays are reused. + decode_gpu_file(gpu_file.path(), self.direct_io, &mut cuda_ctx, "warmup").await?; + let start = Instant::now(); - let file = open_gpu(gpu_file.path(), self.direct_io).await?; - let mut batches = file.scan()?.into_array_stream()?; + decode_gpu_file(gpu_file.path(), self.direct_io, &mut cuda_ctx, "timed").await?; + Ok(start.elapsed()) + } +} + +/// Decode every row and column from a fresh file open into device-resident canonical arrays. +async fn decode_gpu_file( + path: &Path, + direct_io: bool, + cuda_ctx: &mut CudaExecutionCtx, + phase: &'static str, +) -> Result<()> { + let open_start = Instant::now(); + let file = open_gpu(path, direct_io).await?; + let open_time = open_start.elapsed(); + + let scan_start = Instant::now(); + let mut batches = file + .scan()? + .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE)) + .into_array_stream()?; + let scan_time = scan_start.elapsed(); + + let mut read_time = Duration::ZERO; + let mut struct_time = Duration::ZERO; + let mut execute_time = Duration::ZERO; + let mut batch_count = 0usize; + let mut field_count = 0usize; + // Diagnostic modes: `wall` times the host call, `gpu` also brackets its stream work with + // CUDA events, and `nsys` labels every call with NVTX. All modes perturb the benchmark and + // print their records only after the final stream synchronization. + let profile_mode = (phase == "timed") + .then(|| std::env::var("VORTEX_GPU_PROFILE_FIELDS").ok()) + .flatten(); + let profile_fields = profile_mode.is_some(); + let profile_gpu_spans = profile_mode.as_deref() == Some("gpu"); + let profile_nsys = profile_mode.as_deref() == Some("nsys"); + let mut field_timings = Vec::new(); + loop { + let read_start = Instant::now(); + let Some(batch) = batches.next().await else { + read_time += read_start.elapsed(); + break; + }; + read_time += read_start.elapsed(); - while let Some(batch) = batches.next().await { - let record = batch?.execute::(cuda_ctx.execution_ctx())?; - for field in record.iter_unmasked_fields() { - black_box(field.clone().execute_cuda(&mut cuda_ctx).await?); + let struct_start = Instant::now(); + let record = batch?.execute::(cuda_ctx.execution_ctx())?; + struct_time += struct_start.elapsed(); + batch_count += 1; + if phase == "warmup" && std::env::var_os("VORTEX_GPU_DUMP_ARRAY_TREES").is_some() { + eprintln!( + "VORTEX_GPU_ARRAY_TREE batch={}\n{}", + batch_count - 1, + record.clone().into_array().display_tree() + ); + } + + for (field_index, (field, field_name)) in record + .iter_unmasked_fields() + .zip(record.struct_fields().names().iter()) + .enumerate() + { + let metadata = profile_fields.then(|| { + ( + field.encoding_id().to_string(), + field + .display_tree_encodings_only() + .to_string() + .replace('\n', " | "), + ) + }); + let before = profile_gpu_spans + .then(|| cuda_ctx.stream().record_event(Some(CU_EVENT_DEFAULT))) + .transpose()?; + let nsys_range = profile_nsys.then(|| { + scoped_range(format!( + "vortex_field batch={} field={field_index} name={field_name} encoding={}", + batch_count - 1, + field.encoding_id(), + )) + }); + let execute_start = Instant::now(); + black_box(field.clone().execute_cuda(cuda_ctx).await?); + let wall_time = execute_start.elapsed(); + drop(nsys_range); + execute_time += wall_time; + let timing_events = if let Some(before) = before { + let after = cuda_ctx.stream().record_event(Some(CU_EVENT_DEFAULT))?; + Some((before, after)) + } else { + None + }; + if let Some((encoding, tree)) = metadata { + field_timings.push(( + batch_count - 1, + field_index, + field_name.to_string(), + field.len(), + encoding, + tree, + wall_time, + timing_events, + )); } + field_count += 1; } - cuda_ctx.synchronize_stream()?; + } - Ok(start.elapsed()) + let sync_start = Instant::now(); + cuda_ctx.synchronize_stream()?; + let sync_time = sync_start.elapsed(); + + for (batch, field, name, rows, encoding, tree, wall_time, timing_events) in field_timings { + let gpu_span_us = timing_events + .map(|(before, after)| before.elapsed_ms(&after)) + .transpose()? + .map(|milliseconds| Duration::from_secs_f32(milliseconds / 1000.0)) + .map(|duration| duration.as_micros().to_string()) + .unwrap_or_else(|| "NA".to_string()); + eprintln!( + "VORTEX_GPU_FIELD_TIMING\tbatch={batch}\tfield={field}\tname={name}\trows={rows}\tencoding={encoding}\twall_us={}\tgpu_span_us={gpu_span_us}\ttree={tree}", + wall_time.as_micros(), + ); } + + tracing::debug!( + phase, + batch_count, + field_count, + ?open_time, + ?scan_time, + ?read_time, + ?struct_time, + ?execute_time, + ?sync_time, + "GPU Vortex decode stages" + ); + Ok(()) } /// Opens a Vortex file for CUDA execution. @@ -146,9 +297,15 @@ async fn verify_against_host_scan(path: &Path, direct_io: bool) -> Result WriterProperties { .set_dictionary_enabled(true) .set_data_page_size_limit(GPU_DATA_PAGE_SIZE) .set_data_page_row_count_limit(GPU_DATA_PAGE_ROW_COUNT_LIMIT) + .set_max_row_group_row_count(Some(GPU_ROW_GROUP_SIZE)) // Per-page statistics only inflate the page headers a reader has to walk. .set_statistics_enabled(EnabledStatistics::Chunk) .build() @@ -84,5 +91,9 @@ mod tests { assert_eq!(properties.writer_version(), WriterVersion::PARQUET_1_0); assert_eq!(properties.compression(&"x".into()), Compression::SNAPPY); assert_eq!(properties.data_page_size_limit(), GPU_DATA_PAGE_SIZE); + assert_eq!( + properties.max_row_group_row_count(), + Some(GPU_ROW_GROUP_SIZE) + ); } } diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 69cba42c6b0..6ce87e26722 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -6,6 +6,8 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use arrow_array::RecordBatch; +use arrow_select::concat::concat_batches; use futures::StreamExt; use futures::TryStreamExt; use parquet::arrow::AsyncArrowWriter; @@ -96,14 +98,46 @@ fn calculate_concurrency() -> usize { /// Note: This loads the entire file into memory. For large files, use the streaming conversion like /// in [`parquet_to_vortex_stream`] instead. pub async fn parquet_to_vortex_chunks(parquet_path: PathBuf) -> anyhow::Result { + parquet_to_vortex_chunks_with_batch_size(parquet_path, None).await +} + +/// Read a Parquet file into Vortex chunks, using `batch_size` as the maximum rows per chunk. +/// +/// This is useful when the physical row partitions of a generated Vortex file need to match +/// another format's row groups. `None` preserves the Parquet reader's default batch size. +pub async fn parquet_to_vortex_chunks_with_batch_size( + parquet_path: PathBuf, + batch_size: Option, +) -> anyhow::Result { let file = File::open(parquet_path).await?; let builder = ParquetRecordBatchStreamBuilder::new(file).await?; + let builder = if let Some(batch_size) = batch_size { + builder.with_batch_size(batch_size) + } else { + builder + }; let reader = builder.build()?; - let chunks: Vec = parquet_to_vortex_stream(reader) - .map(|r| r.map_err(anyhow::Error::from)) - .try_collect() - .await?; + let chunks: Vec = if let Some(batch_size) = batch_size { + let batches: Vec = reader.map_err(anyhow::Error::from).try_collect().await?; + let schema = batches + .first() + .map(RecordBatch::schema) + .ok_or_else(|| anyhow::anyhow!("cannot convert an empty Parquet file"))?; + let combined = concat_batches(&schema, &batches)?; + + let mut chunks = Vec::with_capacity(combined.num_rows().div_ceil(batch_size)); + for start in (0..combined.num_rows()).step_by(batch_size) { + let len = batch_size.min(combined.num_rows() - start); + chunks.push(record_batch_to_vortex(combined.slice(start, len))?); + } + chunks + } else { + parquet_to_vortex_stream(reader) + .map(|r| r.map_err(anyhow::Error::from)) + .try_collect() + .await? + }; Ok(ChunkedArray::from_iter(chunks)) } @@ -116,22 +150,26 @@ pub fn parquet_to_vortex_stream( reader: ParquetRecordBatchStream, ) -> impl futures::Stream> { reader.map(move |result| { - result.map_err(|e| vortex_err!(External: e)).and_then(|rb| { - let schema = rb.schema(); - let chunk = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; - let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); - - // Canonicalize the chunk. - chunk.append_to_builder( - builder.as_mut(), - &mut VortexSession::default().create_execution_ctx(), - )?; - - Ok(builder.finish()) - }) + result + .map_err(|e| vortex_err!(External: e)) + .and_then(record_batch_to_vortex) }) } +fn record_batch_to_vortex(rb: RecordBatch) -> VortexResult { + let schema = rb.schema(); + let chunk = SESSION.arrow().from_arrow_record_batch(rb, &schema)?; + let mut builder = builder_with_capacity(chunk.dtype(), chunk.len()); + + // Canonicalize the chunk. + chunk.append_to_builder( + builder.as_mut(), + &mut VortexSession::default().create_execution_ctx(), + )?; + + Ok(builder.finish()) +} + /// Convert a single Parquet file to Vortex format using streaming. /// /// Streams data directly from Parquet to Vortex without loading the entire file into memory. diff --git a/vortex-cuda/kernels/src/fsst.cu b/vortex-cuda/kernels/src/fsst.cu index 22e8790a479..13c17a3d533 100644 --- a/vortex-cuda/kernels/src/fsst.cu +++ b/vortex-cuda/kernels/src/fsst.cu @@ -295,6 +295,31 @@ __device__ inline void fsst_decode_string(const FSSTArgs args = { \ + codes_bytes, \ + codes_offsets, \ + symbols, \ + symbol_lengths, \ + output_bytes, \ + output_offsets, \ + validity_bits, \ + validity_bit_offset, \ + output_views, \ + }; \ + 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) @@ -304,3 +329,8 @@ 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) + +GENERATE_FSST_VARBINVIEW_KERNEL(u8, uint8_t) +GENERATE_FSST_VARBINVIEW_KERNEL(u16, uint16_t) +GENERATE_FSST_VARBINVIEW_KERNEL(u32, uint32_t) +GENERATE_FSST_VARBINVIEW_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index f9b90c8ac35..5a5ed630a1d 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -30,6 +30,7 @@ use vortex::buffer::Alignment; use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::NativePType; +use vortex::dtype::PType; use vortex::encodings::fsst::FSST; use vortex::encodings::fsst::FSSTArray; use vortex::encodings::fsst::FSSTArrayExt; @@ -121,42 +122,117 @@ impl CudaExecute for FSSTExecutor { })); } - let lens = fsst - .uncompressed_lengths() - .clone() - .execute_cuda(ctx) - .await? - .into_host() - .await? - .into_primitive(); - let codes_offsets = fsst - .codes() - .offsets() - .clone() - .execute_cuda(ctx) - .await? - .into_primitive(); - - // Prefix-sum lens to per-string u64 output offsets so the kernel - // knows where to write each decoded string. - let output_offsets: Vec = match_each_integer_ptype!(lens.ptype(), |P| { - let mut out = Vec::with_capacity(lens.len() + 1); - let mut acc: u64 = 0; - out.push(0u64); - #[allow(clippy::unnecessary_cast)] - for &l in lens.as_slice::

() { - acc += l as u64; - out.push(acc); - } - out - }); + if can_build_i32_offsets(&fsst) { + decode_fsst_varbinview(fsst, ctx).await + } else { + decode_fsst_host_varbinview(fsst, ctx).await + } + } +} + +/// Decode FSST directly into a device-resident canonical `VarBinView` array. +async fn decode_fsst_varbinview( + fsst: FSSTArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let dtype = fsst.dtype().clone(); + let validity = fsst.codes().validity()?; + let len = fsst.len(); + let lens = fsst + .uncompressed_lengths() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + let codes_offsets = fsst + .codes() + .offsets() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + let I32Offsets { + buffer: output_offsets, + total: total_size, + } = i32_offsets_from_lengths(lens, ctx).await?; - // Dispatch on the unsigned width; signed and unsigned offsets of the - // same width share an identical byte representation. - match_each_unsigned_integer_ptype!(codes_offsets.ptype().to_unsigned(), |U| { - decode_fsst::(fsst, codes_offsets, lens, output_offsets, ctx).await - }) + if total_size == 0 { + let views = ctx.copy_to_device(vec![0i128; len])?.await?; + return Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([]), dtype, validity) + })); } + + match_each_unsigned_integer_ptype!(codes_offsets.ptype().to_unsigned(), |U| { + decode_fsst_varbinview_typed::(fsst, codes_offsets, output_offsets, total_size, ctx) + .await + }) +} + +async fn decode_fsst_varbinview_typed( + fsst: FSSTArray, + codes_offsets: PrimitiveArray, + output_offsets: BufferHandle, + total_size: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + U: NativePType + DeviceRepr + Send + Sync + 'static, +{ + let dtype = fsst.dtype().clone(); + 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 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, symbol_lengths, validity_device, codes_bytes, codes_offsets) = futures::try_join!( + ctx.copy_to_device(symbols_u64)?, + ctx.copy_to_device(symbol_lengths)?, + ctx.ensure_on_device(validity_bits), + ctx.ensure_on_device(codes_bytes_handle), + ctx.ensure_on_device(codes_offsets_buffer), + )?; + + let mut output = ctx.device_alloc::(total_size)?; + 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 output_offsets_view = output_offsets.cuda_view::()?; + let validity_view = validity_device.cuda_view::()?; + let ptype = U::PTYPE.to_string(); + let cuda_function = ctx.load_function_with_suffixes("fsst", &["varbinview", &ptype])?; + + 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(&output_offsets_view) + .arg(&validity_view) + .arg(&validity_bit_offset) + .arg(&mut output) + .arg(&mut views) + .arg(&num_strings_u64); + })?; + + let views = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(views))); + let values = BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))); + Ok(Canonical::VarBinView(unsafe { + VarBinViewArray::new_handle_unchecked(views, Arc::from([values]), dtype, validity) + })) } /// Decode FSST directly into Arrow-compatible i32 offsets and contiguous values on device. @@ -277,6 +353,56 @@ where }) } +fn can_build_i32_offsets(fsst: &FSSTArray) -> bool { + let max_length = match fsst.uncompressed_lengths().dtype().as_ptype() { + PType::U8 => u8::MAX as usize, + PType::U16 => u16::MAX as usize, + PType::U32 => u32::MAX as usize, + PType::U64 => usize::MAX, + _ => return false, + }; + fsst.len() + .checked_mul(max_length) + .is_some_and(|max_total| max_total <= i32::MAX as usize) +} + +async fn decode_fsst_host_varbinview( + fsst: FSSTArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let lens = fsst + .uncompressed_lengths() + .clone() + .execute_cuda(ctx) + .await? + .into_host() + .await? + .into_primitive(); + let codes_offsets = fsst + .codes() + .offsets() + .clone() + .execute_cuda(ctx) + .await? + .into_primitive(); + + let output_offsets: Vec = match_each_integer_ptype!(lens.ptype(), |P| { + let mut out = Vec::with_capacity(lens.len() + 1); + let mut acc: u64 = 0; + out.push(0u64); + #[allow(clippy::unnecessary_cast)] + for &length in lens.as_slice::

() { + acc += length as u64; + out.push(acc); + } + out + }); + + match_each_unsigned_integer_ptype!(codes_offsets.ptype().to_unsigned(), |U| { + decode_fsst::(fsst, codes_offsets, lens, output_offsets, ctx).await + }) +} + async fn decode_fsst( fsst: FSSTArray, codes_offsets: PrimitiveArray, From b9b95267e40bf7412c579dcbdbe4bf9ee9a1186e Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 17:26:31 +0100 Subject: [PATCH 16/22] Remove FSST scalar waits and support unsigned datetime parts Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- benchmarks/compress-bench/README.md | 4 ++ benchmarks/compress-bench/src/main.rs | 11 +++-- vortex-btrblocks/src/schemes/string/fsst.rs | 8 ++++ vortex-cuda/kernels/src/arrow_offsets.cu | 25 +++++++++++ vortex-cuda/kernels/src/date_time_parts.cu | 20 +++++++-- vortex-cuda/src/arrow/mod.rs | 1 + vortex-cuda/src/arrow/offsets.rs | 44 +++++++++++++++++++ .../src/kernel/encodings/date_time_parts.rs | 44 +++++++++++++++++-- vortex-cuda/src/kernel/encodings/fsst.rs | 38 ++++++++++++++-- 9 files changed, 180 insertions(+), 15 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 749bf4ac1fd..4dcdd049d88 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -41,6 +41,10 @@ cargo run -p compress-bench --profile release_debug \ # pick the Parquet page codec the GPU file is written with (default: snappy) cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --gpu-parquet-codec zstd + +# isolate one backend for diagnostics; the default remains parquet,vortex +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --formats vortex ``` ### cuDF diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 8232a9a6018..7a45053fb7d 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -131,10 +131,13 @@ async fn main() -> anyhow::Result<()> { }); let (formats, ops) = if gpu.is_some() { - ( - vec![Format::Parquet, Format::OnDiskVortex], - vec![CompressOp::Decompress], - ) + for format in &args.formats { + anyhow::ensure!( + matches!(format, Format::Parquet | Format::OnDiskVortex), + "GPU decompression supports only parquet and vortex, found {format}" + ); + } + (args.formats, vec![CompressOp::Decompress]) } else { (args.formats, args.ops) }; diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 7250be26207..9bb7367ebd4 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -16,6 +16,7 @@ use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArraySlotsExt; +use vortex_array::expr::stats::Stat; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -84,6 +85,10 @@ impl Scheme for FSSTScheme { .clone() .execute::(exec_ctx)? .narrow(exec_ctx)?; + let length_stats = uncompressed_lengths_primitive + .as_ref() + .statistics() + .compute_all(&[Stat::Min, Stat::Sum], exec_ctx)?; let compressed_original_lengths = compressor.compress_child( &uncompressed_lengths_primitive.into_array(), &compress_ctx, @@ -91,6 +96,9 @@ impl Scheme for FSSTScheme { 0, exec_ctx, )?; + compressed_original_lengths + .statistics() + .set_iter(length_stats.into_iter()); let codes_offsets_primitive = fsst .codes() diff --git a/vortex-cuda/kernels/src/arrow_offsets.cu b/vortex-cuda/kernels/src/arrow_offsets.cu index 79bbb2435c7..076cc2f732d 100644 --- a/vortex-cuda/kernels/src/arrow_offsets.cu +++ b/vortex-cuda/kernels/src/arrow_offsets.cu @@ -61,3 +61,28 @@ GENERATE_FROM_LENGTHS_KERNEL(u8, uint8_t, false) GENERATE_FROM_LENGTHS_KERNEL(u16, uint16_t, false) GENERATE_FROM_LENGTHS_KERNEL(u32, uint32_t, false) GENERATE_FROM_LENGTHS_KERNEL(u64, uint64_t, false) + +// Convert lengths that have already been proven nonnegative with an i32-representable sum. This +// avoids allocating and copying a status word when trusted exact Min/Sum statistics provide the +// same proof on the host. +#define GENERATE_FROM_KNOWN_LENGTHS_KERNEL(suffix, LengthT) \ + extern "C" __global__ void arrow_offsets_from_known_lengths_##suffix(const LengthT *__restrict lengths, \ + int32_t *__restrict scan, \ + uint64_t len) { \ + const uint64_t scan_len = len + 1; \ + 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_stop = min(block_start + elements_per_block, scan_len); \ + for (uint64_t idx = block_start + threadIdx.x; idx < block_stop; idx += blockDim.x) { \ + scan[idx] = idx == len ? 0 : (int32_t)lengths[idx]; \ + } \ + } + +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(i8, int8_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(i16, int16_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(i32, int32_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(i64, int64_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(u8, uint8_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(u16, uint16_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(u32, uint32_t) +GENERATE_FROM_KNOWN_LENGTHS_KERNEL(u64, uint64_t) diff --git a/vortex-cuda/kernels/src/date_time_parts.cu b/vortex-cuda/kernels/src/date_time_parts.cu index ccb3e614991..af625046583 100644 --- a/vortex-cuda/kernels/src/date_time_parts.cu +++ b/vortex-cuda/kernels/src/date_time_parts.cu @@ -46,19 +46,31 @@ __device__ void date_time_parts(const DaysT *__restrict days, X(i8, int8_t) \ X(i16, int16_t) \ X(i32, int32_t) \ - X(i64, int64_t) + X(i64, int64_t) \ + X(u8, uint8_t) \ + X(u16, uint16_t) \ + X(u32, uint32_t) \ + X(u64, uint64_t) #define EXPAND_SUBSECONDS(d, DT, s, ST) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i8, int8_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i16, int16_t) \ GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i32, int32_t) \ - GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i64, int64_t) + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, i64, int64_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u8, uint8_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u16, uint16_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u32, uint32_t) \ + GENERATE_DATE_TIME_PARTS_KERNEL(d, DT, s, ST, u64, uint64_t) #define EXPAND_SECONDS(d, DT) \ EXPAND_SUBSECONDS(d, DT, i8, int8_t) \ EXPAND_SUBSECONDS(d, DT, i16, int16_t) \ EXPAND_SUBSECONDS(d, DT, i32, int32_t) \ - EXPAND_SUBSECONDS(d, DT, i64, int64_t) + EXPAND_SUBSECONDS(d, DT, i64, int64_t) \ + EXPAND_SUBSECONDS(d, DT, u8, uint8_t) \ + EXPAND_SUBSECONDS(d, DT, u16, uint16_t) \ + EXPAND_SUBSECONDS(d, DT, u32, uint32_t) \ + EXPAND_SUBSECONDS(d, DT, u64, uint64_t) -// Generate all 64 kernels (4³) +// Components are narrowed independently, so generate every signed/unsigned integer combination. EXPAND_DAYS(EXPAND_SECONDS) diff --git a/vortex-cuda/src/arrow/mod.rs b/vortex-cuda/src/arrow/mod.rs index f606e623d4b..e976851d43b 100644 --- a/vortex-cuda/src/arrow/mod.rs +++ b/vortex-cuda/src/arrow/mod.rs @@ -33,6 +33,7 @@ use cudarc::driver::CudaStream; use cudarc::driver::DevicePtr; use cudarc::runtime::sys::cudaEvent_t; pub(crate) use offsets::I32Offsets; +pub(crate) use offsets::i32_offsets_from_known_lengths; pub(crate) use offsets::i32_offsets_from_lengths; use vortex::array::ArrayRef; use vortex::array::arrays::Dict; diff --git a/vortex-cuda/src/arrow/offsets.rs b/vortex-cuda/src/arrow/offsets.rs index 4eda482cb41..5ee592918d3 100644 --- a/vortex-cuda/src/arrow/offsets.rs +++ b/vortex-cuda/src/arrow/offsets.rs @@ -46,6 +46,50 @@ pub(crate) async fn i32_offsets_from_lengths( }) } +/// Build offsets without status validation when the caller has already proven that every length +/// is nonnegative and their exact sum fits in an Arrow `i32` offset. +pub(crate) async fn i32_offsets_from_known_lengths( + lengths: PrimitiveArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + let len = lengths.len(); + let ptype = lengths.ptype(); + let PrimitiveDataParts { buffer, .. } = lengths.into_data_parts(); + let lengths = ctx.ensure_on_device(buffer).await?; + + match_each_integer_ptype!(ptype, |L| { + i32_offsets_from_known_lengths_typed::(&lengths, len, ctx) + }) +} + +fn i32_offsets_from_known_lengths_typed( + lengths: &BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + L: NativePType + DeviceRepr + Send + Sync + 'static, +{ + let scan_len = len + .checked_add(1) + .ok_or_else(|| vortex_err!("Arrow offset count overflow"))?; + let lengths_view = lengths.cuda_view::()?; + let mut scan_input = ctx.device_alloc::(scan_len)?; + let ptype = L::PTYPE.to_string(); + let scan_kernel = + ctx.load_function_with_suffixes("arrow_offsets", &["from", "known", "lengths", &ptype])?; + let len_u64 = u64::try_from(len)?; + + ctx.launch_kernel(&scan_kernel, scan_len, |args| { + args.arg(&lengths_view).arg(&mut scan_input).arg(&len_u64); + })?; + + let offsets = exclusive_sum_i32(&scan_input, scan_len, ctx)?; + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + offsets, + )))) +} + async fn i32_offsets_from_lengths_typed( lengths: &BufferHandle, len: usize, diff --git a/vortex-cuda/src/kernel/encodings/date_time_parts.rs b/vortex-cuda/src/kernel/encodings/date_time_parts.rs index bff691e262b..4dfcad7bf51 100644 --- a/vortex-cuda/src/kernel/encodings/date_time_parts.rs +++ b/vortex-cuda/src/kernel/encodings/date_time_parts.rs @@ -15,7 +15,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::TemporalArray; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::buffer::BufferHandle; -use vortex::array::match_each_signed_integer_ptype; +use vortex::array::match_each_integer_ptype; use vortex::array::validity::Validity; use vortex::dtype::DType; use vortex::dtype::NativePType; @@ -110,9 +110,9 @@ impl CudaExecute for DateTimePartsExecutor { let seconds_ptype = seconds_prim.ptype(); let subseconds_ptype = subseconds_prim.ptype(); - match_each_signed_integer_ptype!(days_ptype, |DaysT| { - match_each_signed_integer_ptype!(seconds_ptype, |SecondsT| { - match_each_signed_integer_ptype!(subseconds_ptype, |SubsecondsT| { + match_each_integer_ptype!(days_ptype, |DaysT| { + match_each_integer_ptype!(seconds_ptype, |SecondsT| { + match_each_integer_ptype!(subseconds_ptype, |SubsecondsT| { decode_datetimeparts_typed::( days_prim, seconds_prim, @@ -297,6 +297,42 @@ mod tests { Ok(()) } + #[crate::test] + async fn test_cuda_datetimeparts_unsigned_components() -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + + let len = 3; + let temporal = TemporalArray::new_timestamp( + PrimitiveArray::new(buffer![0i64; len], Validity::NonNullable).into_array(), + TimeUnit::Nanoseconds, + None, + ); + let dtp_array = DateTimeParts::try_new( + temporal.dtype().clone(), + PrimitiveArray::new(buffer![20_000u16, 20_001, 20_002], Validity::NonNullable) + .into_array(), + PrimitiveArray::new(buffer![80_000u32, 80_001, 80_002], Validity::NonNullable) + .into_array(), + PrimitiveArray::new( + buffer![900_000_000u32, 900_000_001, 900_000_002], + Validity::NonNullable, + ) + .into_array(), + )?; + + let gpu_result = DateTimePartsExecutor + .execute(dtp_array.clone().into_array(), &mut cuda_ctx) + .await? + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(dtp_array, gpu_result, &mut ctx); + Ok(()) + } + #[crate::test] async fn test_cuda_datetimeparts_large_array() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index 5a5ed630a1d..d49ad40452d 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -23,6 +23,9 @@ use vortex::array::arrays::varbinview::build_views::MAX_BUFFER_LEN; use vortex::array::arrays::varbinview::build_views::build_views; use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBuffer; +use vortex::array::expr::stats::Precision; +use vortex::array::expr::stats::Stat; +use vortex::array::expr::stats::StatsProvider; use vortex::array::match_each_integer_ptype; use vortex::array::match_each_unsigned_integer_ptype; use vortex::array::validity::Validity; @@ -43,6 +46,7 @@ use crate::CanonicalCudaExt; use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::arrow::I32Offsets; +use crate::arrow::i32_offsets_from_known_lengths; use crate::arrow::i32_offsets_from_lengths; use crate::executor::CudaArrayExt; use crate::executor::CudaExecute; @@ -122,7 +126,7 @@ impl CudaExecute for FSSTExecutor { })); } - if can_build_i32_offsets(&fsst) { + if exact_nonnegative_length_sum(&fsst).is_some() || can_build_i32_offsets(&fsst) { decode_fsst_varbinview(fsst, ctx).await } else { decode_fsst_host_varbinview(fsst, ctx).await @@ -154,7 +158,7 @@ async fn decode_fsst_varbinview( let I32Offsets { buffer: output_offsets, total: total_size, - } = i32_offsets_from_lengths(lens, ctx).await?; + } = fsst_i32_offsets(&fsst, lens, ctx).await?; if total_size == 0 { let views = ctx.copy_to_device(vec![0i128; len])?.await?; @@ -259,7 +263,7 @@ pub(crate) async fn decode_fsst_varbin( let I32Offsets { buffer: output_offsets, total: total_size, - } = i32_offsets_from_lengths(lens, ctx).await?; + } = fsst_i32_offsets(&fsst, lens, ctx).await?; if total_size == 0 { let allocation = CudaDeviceBuffer::new(ctx.device_alloc::(1)?); @@ -353,6 +357,34 @@ where }) } +async fn fsst_i32_offsets( + fsst: &FSSTArray, + lengths: PrimitiveArray, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + if let Some(total) = exact_nonnegative_length_sum(fsst) { + return Ok(I32Offsets { + buffer: i32_offsets_from_known_lengths(lengths, ctx).await?, + total, + }); + } + + i32_offsets_from_lengths(lengths, ctx).await +} + +fn exact_nonnegative_length_sum(fsst: &FSSTArray) -> Option { + let stats = fsst.uncompressed_lengths().statistics(); + let Precision::Exact(min) = stats.get(Stat::Min) else { + return None; + }; + let Precision::Exact(sum) = stats.get(Stat::Sum) else { + return None; + }; + let min = i64::try_from(&min).ok()?; + let total = usize::try_from(&sum).ok()?; + (min >= 0 && total <= i32::MAX as usize).then_some(total) +} + fn can_build_i32_offsets(fsst: &FSSTArray) -> bool { let max_length = match fsst.uncompressed_lengths().dtype().as_ptype() { PType::U8 => u8::MAX as usize, From f0819143299b15e944dd112e2d839a6f54a4e3aa Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 17:53:55 +0100 Subject: [PATCH 17/22] Decode nullable RunEnd arrays on CUDA Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/kernels/src/runend.cu | 94 ++++++++++++++++++++++ vortex-cuda/src/kernel/arrays/masked.rs | 94 ++++++++++++++++++++++ vortex-cuda/src/kernel/arrays/mod.rs | 2 + vortex-cuda/src/kernel/encodings/runend.rs | 70 +++++++++++----- vortex-cuda/src/kernel/mod.rs | 1 + vortex-cuda/src/lib.rs | 3 + 6 files changed, 246 insertions(+), 18 deletions(-) create mode 100644 vortex-cuda/src/kernel/arrays/masked.rs diff --git a/vortex-cuda/kernels/src/runend.cu b/vortex-cuda/kernels/src/runend.cu index a3f1d245dbe..74a3aa50e90 100644 --- a/vortex-cuda/kernels/src/runend.cu +++ b/vortex-cuda/kernels/src/runend.cu @@ -155,3 +155,97 @@ GENERATE_RUNEND_KERNELS_FOR_VALUE(i64, int64_t) GENERATE_RUNEND_KERNELS_FOR_VALUE(f16, __half) GENERATE_RUNEND_KERNELS_FOR_VALUE(f32, float) GENERATE_RUNEND_KERNELS_FOR_VALUE(f64, double) + +template +__device__ void runend_decode_validity_kernel(const EndsT *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values_validity, + uint64_t values_validity_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output_validity) { + const uint64_t output_bytes = (output_len + 7) / 8; + const uint64_t bytes_per_block = static_cast(blockDim.x) * ELEMENTS_PER_THREAD; + const uint64_t block_start = static_cast(blockIdx.x) * bytes_per_block; + const uint64_t block_stop = min(block_start + bytes_per_block, output_bytes); + + for (uint64_t output_byte = block_start + threadIdx.x; output_byte < block_stop; + output_byte += blockDim.x) { + const uint64_t output_start = output_byte * 8; + const uint64_t output_stop = min(output_start + 8, output_len); + uint8_t bits = 0; + for (uint64_t idx = output_start; idx < output_stop; ++idx) { + uint64_t run_idx = upper_bound(ends, num_runs, idx + offset); + if (run_idx >= num_runs) { + run_idx = num_runs - 1; + } + const uint64_t validity_idx = values_validity_offset + run_idx; + bits |= ((values_validity[validity_idx / 8] >> (validity_idx % 8)) & 1) << (idx % 8); + } + output_validity[output_byte] = bits; + } +} + +extern "C" __global__ void runend_validity_u8(const uint8_t *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values_validity, + uint64_t values_validity_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output_validity) { + runend_decode_validity_kernel(ends, + num_runs, + values_validity, + values_validity_offset, + offset, + output_len, + output_validity); +} + +extern "C" __global__ void runend_validity_u16(const uint16_t *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values_validity, + uint64_t values_validity_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output_validity) { + runend_decode_validity_kernel(ends, + num_runs, + values_validity, + values_validity_offset, + offset, + output_len, + output_validity); +} + +extern "C" __global__ void runend_validity_u32(const uint32_t *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values_validity, + uint64_t values_validity_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output_validity) { + runend_decode_validity_kernel(ends, + num_runs, + values_validity, + values_validity_offset, + offset, + output_len, + output_validity); +} + +extern "C" __global__ void runend_validity_u64(const uint64_t *const __restrict ends, + uint64_t num_runs, + const uint8_t *const __restrict values_validity, + uint64_t values_validity_offset, + uint64_t offset, + uint64_t output_len, + uint8_t *const __restrict output_validity) { + runend_decode_validity_kernel(ends, + num_runs, + values_validity, + values_validity_offset, + offset, + output_len, + output_validity); +} diff --git a/vortex-cuda/src/kernel/arrays/masked.rs b/vortex-cuda/src/kernel/arrays/masked.rs new file mode 100644 index 00000000000..a551e952f50 --- /dev/null +++ b/vortex-cuda/src/kernel/arrays/masked.rs @@ -0,0 +1,94 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use async_trait::async_trait; +use tracing::instrument; +use vortex::array::ArrayRef; +use vortex::array::Canonical; +use vortex::array::arrays::Masked; +use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::masked::MaskedArrayExt; +use vortex::array::arrays::masked::MaskedArraySlotsExt; +use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::error::vortex_err; + +use crate::executor::CudaArrayExt; +use crate::executor::CudaExecute; +use crate::executor::CudaExecutionCtx; +use crate::executor::execute_validity_cuda; + +#[derive(Debug)] +pub(crate) struct MaskedExecutor; + +#[async_trait] +impl CudaExecute for MaskedExecutor { + #[instrument(level = "trace", skip_all, fields(executor = ?self))] + async fn execute( + &self, + array: ArrayRef, + ctx: &mut CudaExecutionCtx, + ) -> VortexResult { + let array = array + .try_downcast::() + .map_err(|_| vortex_err!("Expected MaskedArray"))?; + let len = array.len(); + let validity = execute_validity_cuda(array.masked_validity(), len, ctx).await?; + + match array.child().clone().execute_cuda(ctx).await? { + Canonical::Primitive(primitive) => { + let PrimitiveDataParts { ptype, buffer, .. } = primitive.into_data_parts(); + Ok(Canonical::Primitive(PrimitiveArray::from_buffer_handle( + buffer, ptype, validity, + ))) + } + canonical => vortex_bail!( + "CUDA Masked execution currently supports primitive children, got {}", + canonical.dtype() + ), + } + } +} + +#[cfg(test)] +mod tests { + use vortex::array::IntoArray; + use vortex::array::VortexSessionExecute; + use vortex::array::arrays::BoolArray; + use vortex::array::arrays::MaskedArray; + use vortex::array::arrays::PrimitiveArray; + use vortex::array::assert_arrays_eq; + use vortex::array::validity::Validity; + use vortex::error::VortexExpect; + use vortex::error::VortexResult; + + use super::*; + use crate::CanonicalCudaExt; + use crate::session::CudaSession; + + #[crate::test] + async fn test_cuda_masked_primitive() -> VortexResult<()> { + let mut ctx = vortex::array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let validity = + BoolArray::from_iter([false, false, true, false, true, true, false, false, false]) + .into_array() + .slice(2..7)?; + let masked = MaskedArray::try_new( + PrimitiveArray::from_iter([10i64, 20, 30, 40, 50]).into_array(), + Validity::Array(validity), + )?; + + let actual = MaskedExecutor + .execute(masked.clone().into_array(), &mut cuda_ctx) + .await? + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(masked, actual, &mut ctx); + Ok(()) + } +} diff --git a/vortex-cuda/src/kernel/arrays/mod.rs b/vortex-cuda/src/kernel/arrays/mod.rs index ab81934bb27..c4df15873a0 100644 --- a/vortex-cuda/src/kernel/arrays/mod.rs +++ b/vortex-cuda/src/kernel/arrays/mod.rs @@ -3,8 +3,10 @@ mod constant; mod dict; +mod masked; mod shared; pub(crate) use constant::ConstantNumericExecutor; pub(crate) use dict::DictExecutor; +pub(crate) use masked::MaskedExecutor; pub(crate) use shared::SharedExecutor; diff --git a/vortex-cuda/src/kernel/encodings/runend.rs b/vortex-cuda/src/kernel/encodings/runend.rs index 36ceb8c7b8b..bf5c295c54b 100644 --- a/vortex-cuda/src/kernel/encodings/runend.rs +++ b/vortex-cuda/src/kernel/encodings/runend.rs @@ -10,8 +10,10 @@ use tracing::instrument; use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::IntoArray; +use vortex::array::arrays::BoolArray; use vortex::array::arrays::ConstantArray; use vortex::array::arrays::PrimitiveArray; +use vortex::array::arrays::bool::BoolDataParts; use vortex::array::arrays::primitive::PrimitiveDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::match_each_native_ptype; @@ -149,10 +151,36 @@ async fn decode_runend_typed { unreachable!("AllInvalid should be handled by RunEndExecutor::execute") } - Validity::Array(_) => { - vortex_bail!( - "RunEnd GPU decoding does not yet support per-element validity in values; falling back to CPU" - ); + Validity::Array(array) => { + let values_validity = array.execute_cuda(ctx).await?.into_bool(); + let BoolDataParts { bits, meta } = values_validity.into_data().into_parts(num_runs); + let values_validity_device = ctx.ensure_on_device(bits).await?; + let values_validity_view = values_validity_device.cuda_view::()?; + let output_bytes = output_len.div_ceil(8); + let mut output_validity = ctx.device_alloc::(output_bytes)?; + let validity_kernel = + ctx.load_function_with_suffixes("runend", &["validity", &E::PTYPE.to_string()])?; + let values_validity_offset = u64::try_from(meta.offset())?; + + ctx.launch_kernel(&validity_kernel, output_bytes, |args| { + args.arg(&ends_view) + .arg(&num_runs_u64) + .arg(&values_validity_view) + .arg(&values_validity_offset) + .arg(&offset_u64) + .arg(&output_len_u64) + .arg(&mut output_validity); + })?; + + Validity::Array( + BoolArray::new_handle( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output_validity))), + 0, + output_len, + Validity::NonNullable, + ) + .into_array(), + ) } }; @@ -182,7 +210,6 @@ mod tests { use super::*; use crate::CanonicalCudaExt; - use crate::executor::CudaArrayExt; use crate::session::CudaSession; fn make_runend_array(ends: Vec, values: Vec, ctx: &mut ExecutionCtx) -> RunEndArray @@ -303,7 +330,7 @@ mod tests { } #[crate::test] - async fn test_cuda_runend_nullable_values_falls_back_to_cpu() -> VortexResult<()> { + async fn test_cuda_runend_nullable_values() -> VortexResult<()> { let mut ctx = vortex_array::array_session().create_execution_ctx(); let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) .vortex_expect("failed to create execution context"); @@ -312,19 +339,26 @@ mod tests { let ends_array = PrimitiveArray::new(Buffer::from(vec![3u32, 6, 10]), Validity::NonNullable) .into_array(); - let validity = - Validity::Array(BoolArray::from_iter([true, false, true].into_iter()).into_array()); - let values_array = - PrimitiveArray::new(Buffer::from(vec![10i32, 0, 30]), validity).into_array(); - let runend_array = RunEnd::new(ends_array, values_array, cuda_ctx.execution_ctx()); - - // execute_cuda should fall back to CPU and still produce the correct result. - let gpu_result = runend_array - .clone() - .into_array() - .execute_cuda(&mut cuda_ctx) + // Slice the validity to exercise a non-zero input bit offset, then slice the RunEnd + // array to exercise a non-zero logical offset and a partial final output byte. + let values_validity = BoolArray::from_iter([ + false, false, false, true, false, true, false, false, false, false, + ]) + .into_array() + .slice(3..6)?; + let values_array = PrimitiveArray::new( + Buffer::from(vec![10i32, 0, 30]), + Validity::Array(values_validity), + ) + .into_array(); + // SAFETY: ends are increasing, ends/values have equal length, and [offset, offset + + // length) = [1, 10) is covered by the final run end. + let runend_array = unsafe { RunEnd::new_unchecked(ends_array, values_array, 1, 9) }; + + let gpu_result = RunEndExecutor + .execute(runend_array.clone().into_array(), &mut cuda_ctx) .await - .vortex_expect("GPU/CPU fallback should succeed") + .vortex_expect("GPU decompression failed") .into_host() .await? .into_array(); diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index 36735024c7f..b9b01714b2f 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -31,6 +31,7 @@ mod slice; pub(crate) use arrays::ConstantNumericExecutor; pub(crate) use arrays::DictExecutor; +pub(crate) use arrays::MaskedExecutor; pub(crate) use arrays::SharedExecutor; pub use encodings::ZstdKernelPrep; pub use encodings::zstd_kernel_prepare; diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 3c712d20fb8..b0f31863214 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -48,6 +48,7 @@ use kernel::FSSTExecutor; use kernel::FilterExecutor; use kernel::FoRExecutor; pub use kernel::LaunchStrategy; +use kernel::MaskedExecutor; use kernel::RunEndExecutor; use kernel::SharedExecutor; pub use kernel::TracingLaunchStrategy; @@ -73,6 +74,7 @@ use vortex::array::ArrayVTable; use vortex::array::arrays::Constant; use vortex::array::arrays::Dict; use vortex::array::arrays::Filter; +use vortex::array::arrays::Masked; use vortex::array::arrays::Shared; use vortex::array::arrays::Slice; use vortex::encodings::alp::ALP; @@ -118,6 +120,7 @@ pub fn initialize_cuda(session: &CudaSession) { session.register_kernel(Shared.id(), &SharedExecutor); session.register_kernel(FoR.id(), &FoRExecutor); session.register_kernel(FSST.id(), &FSSTExecutor); + session.register_kernel(Masked.id(), &MaskedExecutor); session.register_kernel(RunEnd.id(), &RunEndExecutor); session.register_kernel(Sequence.id(), &SequenceExecutor); session.register_kernel(ZigZag.id(), &ZigZagExecutor); From 25818d73f4b35ffe92badc239a608c1e36e54129 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 18:45:03 +0100 Subject: [PATCH 18/22] Remove FSST metadata upload callbacks Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/src/kernel/encodings/fsst.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index d49ad40452d..dd174f550d4 100644 --- a/vortex-cuda/src/kernel/encodings/fsst.rs +++ b/vortex-cuda/src/kernel/encodings/fsst.rs @@ -200,10 +200,10 @@ where } = codes_offsets.into_data_parts(); let (validity_bit_offset, validity_bits) = cuda_validity(&validity, num_strings, ctx).await?; - let (symbols, symbol_lengths, validity_device, codes_bytes, codes_offsets) = futures::try_join!( - ctx.copy_to_device(symbols_u64)?, - ctx.copy_to_device(symbol_lengths)?, - ctx.ensure_on_device(validity_bits), + 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), ctx.ensure_on_device(codes_offsets_buffer), )?; @@ -309,10 +309,10 @@ where } = codes_offsets.into_data_parts(); let (validity_bit_offset, validity_bits) = cuda_validity(&validity, len, ctx).await?; - let (symbols, symbol_lengths, validity_device, codes_bytes, codes_offsets) = futures::try_join!( - ctx.copy_to_device(symbols_u64)?, - ctx.copy_to_device(symbol_lengths)?, - ctx.ensure_on_device(validity_bits), + 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), ctx.ensure_on_device(codes_offsets_buffer), )?; From 3893bb7248c75d6309d9de04d89f307fb10161ff Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 18:45:22 +0100 Subject: [PATCH 19/22] Add single-stream CUDA benchmark fast path Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- benchmarks/compress-bench/src/main.rs | 9 +++++++++ vortex-cuda/src/session.rs | 24 ++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 7a45053fb7d..e9134dad1c9 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -25,6 +25,8 @@ use vortex::utils::aliases::hash_map::HashMap; use vortex_bench::Engine; use vortex_bench::Format; use vortex_bench::LogFormat; +#[cfg(feature = "cuda")] +use vortex_bench::SESSION; use vortex_bench::Target; use vortex_bench::compress::CompressMeasurements; use vortex_bench::compress::CompressOp; @@ -52,6 +54,8 @@ use vortex_bench::public_bi::PBIDataset::Food; use vortex_bench::public_bi::PBIDataset::HashTags; use vortex_bench::setup_logging_and_tracing_with_format; use vortex_bench::v3; +#[cfg(feature = "cuda")] +use vortex_cuda::CudaSession; #[derive(Parser, Debug)] #[command(version, about, long_about = None)] @@ -130,6 +134,11 @@ async fn main() -> anyhow::Result<()> { direct_io: args.gpu_direct_io, }); + #[cfg(feature = "cuda")] + if gpu.is_some() { + SESSION.register(CudaSession::try_single_stream()?); + } + let (formats, ops) = if gpu.is_some() { for format in &args.formats { anyhow::ensure!( diff --git a/vortex-cuda/src/session.rs b/vortex-cuda/src/session.rs index a5410db0c99..06edbdcce92 100644 --- a/vortex-cuda/src/session.rs +++ b/vortex-cuda/src/session.rs @@ -114,6 +114,30 @@ impl CudaSession { } } + /// Creates a single-stream CUDA session using device 0, with event tracking disabled. + /// + /// Every execution context created from this session shares the same stream. This avoids + /// cudarc's per-buffer events, which are only needed to synchronize buffer use across streams. + pub fn try_single_stream() -> VortexResult { + // cudarc panics rather than returning an error when the CUDA driver library cannot be + // loaded, so catch any unwind here to uphold this constructor's no-panic contract. + match catch_unwind(AssertUnwindSafe(|| -> VortexResult { + let context = CudaContext::new(0) + .map_err(|err| vortex_err!("failed to initialize CUDA device 0: {err}"))?; + // SAFETY: this context is private to a session whose pool contains exactly one stream, + // and event tracking is disabled before any device buffers can be allocated. + unsafe { context.disable_event_tracking() }; + let this = Self::with_stream_pool_capacity(context, 1); + initialize_cuda(&this); + Ok(this) + })) { + Ok(result) => result, + Err(_) => Err(vortex_err!( + "failed to initialize CUDA: the driver library is unavailable" + )), + } + } + /// Creates a new CUDA execution context. pub fn create_execution_ctx( vortex_session: &vortex::session::VortexSession, From a214287bb066bc01fa904530e54df97db77840f0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 19:25:10 +0100 Subject: [PATCH 20/22] Add reusable GPU Vortex profiling metrics Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- benchmarks/compress-bench/README.md | 45 ++++ benchmarks/compress-bench/src/gpu_vortex.rs | 283 +++++++++++++++++--- benchmarks/compress-bench/src/main.rs | 35 ++- 3 files changed, 324 insertions(+), 39 deletions(-) diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index 4dcdd049d88..7a4c3e71e45 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -47,6 +47,51 @@ cargo run -p compress-bench --profile release_debug \ --features cuda,unstable_encodings -- --gpu-decompress --formats vortex ``` +### Vortex GPU profiling + +`--gpu-vortex-profile wall|gpu|nsys` enables opt-in diagnostics for the Vortex backend. `wall` +records host timings, `gpu` also brackets every field dispatch with CUDA events, and `nsys` adds +per-field NVTX ranges. These modes perturb the measured run; use them to explain a result, then +rerun without the flag for the comparison number. + +After each timed stream synchronization, the benchmark writes one JSON record to stderr with +`record="vortex_gpu_decompress_profile"`. It includes file/layout sizes and counts, decoded rows, +batches and field dispatches; `stages` contains microsecond wall times; and `encodings` groups calls, +rows and wall time by full encoding tree and field name. In `gpu` mode, each encoding group also has +`gpu_us`; it is `null` in the other modes. + +```bash +cargo run -p compress-bench --profile release_debug \ + --features cuda,unstable_encodings -- --gpu-decompress --formats vortex \ + --datasets '^(Arade|Bimbo|CMSprovider)$' --iterations 3 \ + --gpu-vortex-profile gpu 2> /tmp/vortex-gpu-profile.log + +# Ignore non-JSON progress/log lines and average the main stages by dataset. +jq -Rs ' + [split("\n")[] | fromjson? | + select(.record == "vortex_gpu_decompress_profile")] + | group_by(.dataset) + | map({ + dataset: .[0].dataset, + runs: length, + total_us: (map(.stages.total_us) | add / length), + read_us: (map(.stages.read_us) | add / length), + dispatch_us: (map(.stages.field_dispatch_us) | add / length), + gpu_us: (map([.encodings[].gpu_us // empty] | add) | add / length) + }) +' /tmp/vortex-gpu-profile.log +``` + +`open_us` covers opening and footer metadata, `scan_plan_us` builds the scan stream, `read_us` +awaits batches, `struct_dispatch_us` materializes each struct batch, `field_dispatch_us` measures +CPU planning/enqueue time for field decodes, and `final_sync_us` is the remaining device tail. +`profile_overhead_us` is the remainder spent collecting diagnostics, primarily encoding-tree +formatting and CUDA-event bookkeeping; it makes the profiler's own perturbation explicit. +Because CUDA work is asynchronous, `read_us` can include I/O, layout execution, backpressure, and +waiting for earlier device work; it is not pure storage time. `gpu_us` is the device-stream time +between field events. Allocation/free, upload, wait, event, and callback counts are not available +from this record; use Nsight Systems for those runtime-wide counts. + ### cuDF cuDF is reached through its prebuilt `cudf-cu12` wheel, so it is a runtime dependency of the diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 627eab23f0d..b87811ecba6 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -1,9 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::BTreeMap; +use std::collections::BTreeSet; use std::hint::black_box; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; @@ -13,9 +17,12 @@ use anyhow::ensure; use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::Field; use async_trait::async_trait; +use clap::ValueEnum; +use cudarc::driver::CudaEvent; use cudarc::driver::sys::CUevent_flags::CU_EVENT_DEFAULT; use cudarc::nvtx::safe::scoped_range; use futures::StreamExt; +use serde::Serialize; use tempfile::NamedTempFile; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; @@ -46,10 +53,34 @@ use vortex_cuda::layout::register_cuda_layout; use crate::gpu_writer::GPU_ROW_GROUP_SIZE; +/// Optional diagnostics for the Vortex GPU decompression path. +#[derive(Clone, Copy, Debug, ValueEnum)] +pub enum GpuVortexProfile { + /// Record host wall time for each stage and encoding dispatch group. + Wall, + /// Also bracket each field dispatch with CUDA events and report device-stream elapsed time. + Gpu, + /// Add NVTX ranges for correlating field dispatches with an Nsight Systems capture. + Nsys, +} + +impl GpuVortexProfile { + fn records_gpu_spans(self) -> bool { + matches!(self, Self::Gpu) + } + + fn records_nsys_ranges(self) -> bool { + matches!(self, Self::Nsys) + } +} + /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU. pub struct GpuVortexCompressor { verify: bool, direct_io: bool, + profile: Option, + dataset: Arc, + iteration: AtomicUsize, } impl GpuVortexCompressor { @@ -58,8 +89,19 @@ impl GpuVortexCompressor { /// When `verify` is set, each GPU-decoded field is copied back to the host and compared /// against the same field decoded on the CPU. Verification runs inline, so timings from a /// verifying run are not comparable to a plain one. - pub fn new(verify: bool, direct_io: bool) -> Self { - Self { verify, direct_io } + pub fn new( + verify: bool, + direct_io: bool, + profile: Option, + dataset: &str, + ) -> Self { + Self { + verify, + direct_io, + profile, + dataset: dataset.into(), + iteration: AtomicUsize::new(0), + } } } @@ -112,25 +154,139 @@ impl Compressor for GpuVortexCompressor { // Match cuDF's untimed read: pay CUDA module loading, allocator initialization and page // faults before measuring. The helper opens and scans the file afresh each time, and the // CUDA opener disables its data-segment cache, so no decoded arrays are reused. - decode_gpu_file(gpu_file.path(), self.direct_io, &mut cuda_ctx, "warmup").await?; + decode_gpu_file( + gpu_file.path(), + self.direct_io, + &mut cuda_ctx, + "warmup", + None, + ) + .await?; + let profile = self.profile.map(|mode| ProfileRun { + mode, + dataset: self.dataset.as_ref(), + iteration: self.iteration.fetch_add(1, Ordering::Relaxed), + }); let start = Instant::now(); - decode_gpu_file(gpu_file.path(), self.direct_io, &mut cuda_ctx, "timed").await?; + decode_gpu_file( + gpu_file.path(), + self.direct_io, + &mut cuda_ctx, + "timed", + profile, + ) + .await?; Ok(start.elapsed()) } } +/// Runtime information attached to one profiled timed decode. +#[derive(Clone, Copy)] +struct ProfileRun<'a> { + mode: GpuVortexProfile, + dataset: &'a str, + iteration: usize, +} + +struct FieldTiming { + field_name: String, + rows: usize, + encoding: String, + tree: String, + wall_time: Duration, + gpu_events: Option<(CudaEvent, CudaEvent)>, +} + +#[derive(Default)] +struct EncodingAggregate { + calls: usize, + rows: usize, + fields: BTreeSet, + wall_time: Duration, + gpu_time_us: Option, +} + +#[derive(Serialize)] +struct EncodingProfileRecord { + encoding: String, + tree: String, + fields: Vec, + calls: usize, + rows: usize, + wall_us: u64, + gpu_us: Option, +} + +#[derive(Serialize)] +struct StageProfileRecord { + total_us: u64, + open_us: u64, + scan_plan_us: u64, + read_us: u64, + struct_dispatch_us: u64, + field_dispatch_us: u64, + final_sync_us: u64, + profile_overhead_us: u64, +} + +#[derive(Serialize)] +struct GpuProfileRecord<'a> { + record: &'static str, + version: u8, + dataset: &'a str, + iteration: usize, + mode: &'static str, + direct_io: bool, + file_bytes: u64, + data_segment_bytes: u64, + data_segments: usize, + root_layout: String, + root_layout_children: usize, + file_rows: u64, + decoded_rows: usize, + batches: usize, + fields_per_batch: usize, + field_dispatches: usize, + stages: StageProfileRecord, + encodings: Vec, +} + /// Decode every row and column from a fresh file open into device-resident canonical arrays. async fn decode_gpu_file( path: &Path, direct_io: bool, cuda_ctx: &mut CudaExecutionCtx, phase: &'static str, + profile: Option>, ) -> Result<()> { + let total_start = profile.map(|_| Instant::now()); let open_start = Instant::now(); let file = open_gpu(path, direct_io).await?; let open_time = open_start.elapsed(); + let file_bytes = profile + .is_some() + .then(|| std::fs::metadata(path).map(|metadata| metadata.len())) + .transpose()?; + let (data_segments, data_segment_bytes, root_layout, root_layout_children, file_rows) = + if profile.is_some() { + let footer = file.footer(); + ( + footer.segment_map().len(), + footer + .segment_map() + .iter() + .map(|segment| u64::from(segment.length)) + .sum(), + footer.layout().encoding_id().to_string(), + footer.layout().nchildren(), + file.row_count(), + ) + } else { + (0, 0, String::new(), 0, 0) + }; + let scan_start = Instant::now(); let mut batches = file .scan()? @@ -142,16 +298,11 @@ async fn decode_gpu_file( let mut struct_time = Duration::ZERO; let mut execute_time = Duration::ZERO; let mut batch_count = 0usize; + let mut decoded_rows = 0usize; + let mut fields_per_batch = 0usize; let mut field_count = 0usize; - // Diagnostic modes: `wall` times the host call, `gpu` also brackets its stream work with - // CUDA events, and `nsys` labels every call with NVTX. All modes perturb the benchmark and - // print their records only after the final stream synchronization. - let profile_mode = (phase == "timed") - .then(|| std::env::var("VORTEX_GPU_PROFILE_FIELDS").ok()) - .flatten(); - let profile_fields = profile_mode.is_some(); - let profile_gpu_spans = profile_mode.as_deref() == Some("gpu"); - let profile_nsys = profile_mode.as_deref() == Some("nsys"); + let profile_gpu_spans = profile.is_some_and(|run| run.mode.records_gpu_spans()); + let profile_nsys = profile.is_some_and(|run| run.mode.records_nsys_ranges()); let mut field_timings = Vec::new(); loop { let read_start = Instant::now(); @@ -165,6 +316,10 @@ async fn decode_gpu_file( let record = batch?.execute::(cuda_ctx.execution_ctx())?; struct_time += struct_start.elapsed(); batch_count += 1; + decoded_rows += record.len(); + if batch_count == 1 { + fields_per_batch = record.struct_fields().names().len(); + } if phase == "warmup" && std::env::var_os("VORTEX_GPU_DUMP_ARRAY_TREES").is_some() { eprintln!( "VORTEX_GPU_ARRAY_TREE batch={}\n{}", @@ -178,7 +333,7 @@ async fn decode_gpu_file( .zip(record.struct_fields().names().iter()) .enumerate() { - let metadata = profile_fields.then(|| { + let metadata = profile.map(|_| { ( field.encoding_id().to_string(), field @@ -202,23 +357,21 @@ async fn decode_gpu_file( let wall_time = execute_start.elapsed(); drop(nsys_range); execute_time += wall_time; - let timing_events = if let Some(before) = before { + let gpu_events = if let Some(before) = before { let after = cuda_ctx.stream().record_event(Some(CU_EVENT_DEFAULT))?; Some((before, after)) } else { None }; if let Some((encoding, tree)) = metadata { - field_timings.push(( - batch_count - 1, - field_index, - field_name.to_string(), - field.len(), + field_timings.push(FieldTiming { + field_name: field_name.to_string(), + rows: field.len(), encoding, tree, wall_time, - timing_events, - )); + gpu_events, + }); } field_count += 1; } @@ -227,18 +380,76 @@ async fn decode_gpu_file( let sync_start = Instant::now(); cuda_ctx.synchronize_stream()?; let sync_time = sync_start.elapsed(); - - for (batch, field, name, rows, encoding, tree, wall_time, timing_events) in field_timings { - let gpu_span_us = timing_events - .map(|(before, after)| before.elapsed_ms(&after)) - .transpose()? - .map(|milliseconds| Duration::from_secs_f32(milliseconds / 1000.0)) - .map(|duration| duration.as_micros().to_string()) - .unwrap_or_else(|| "NA".to_string()); - eprintln!( - "VORTEX_GPU_FIELD_TIMING\tbatch={batch}\tfield={field}\tname={name}\trows={rows}\tencoding={encoding}\twall_us={}\tgpu_span_us={gpu_span_us}\ttree={tree}", - wall_time.as_micros(), - ); + let total_time = total_start.map(|start| start.elapsed()); + + if let Some(profile) = profile { + let mut encodings: BTreeMap<(String, String), EncodingAggregate> = BTreeMap::new(); + for timing in field_timings { + let gpu_time_us = timing + .gpu_events + .map(|(before, after)| before.elapsed_ms(&after)) + .transpose()? + .map(|milliseconds| duration_us(Duration::from_secs_f32(milliseconds / 1_000.0))); + let aggregate = encodings.entry((timing.encoding, timing.tree)).or_default(); + aggregate.calls += 1; + aggregate.rows += timing.rows; + aggregate.fields.insert(timing.field_name); + aggregate.wall_time += timing.wall_time; + aggregate.gpu_time_us = match (aggregate.gpu_time_us, gpu_time_us) { + (Some(total), Some(elapsed)) => Some(total.saturating_add(elapsed)), + (None, Some(elapsed)) => Some(elapsed), + (total, None) => total, + }; + } + let encodings = encodings + .into_iter() + .map(|((encoding, tree), aggregate)| EncodingProfileRecord { + encoding, + tree, + fields: aggregate.fields.into_iter().collect(), + calls: aggregate.calls, + rows: aggregate.rows, + wall_us: duration_us(aggregate.wall_time), + gpu_us: aggregate.gpu_time_us, + }) + .collect(); + let accounted_time = + open_time + scan_time + read_time + struct_time + execute_time + sync_time; + let total_time = total_time.unwrap_or_default(); + let record = GpuProfileRecord { + record: "vortex_gpu_decompress_profile", + version: 1, + dataset: profile.dataset, + iteration: profile.iteration, + mode: match profile.mode { + GpuVortexProfile::Wall => "wall", + GpuVortexProfile::Gpu => "gpu", + GpuVortexProfile::Nsys => "nsys", + }, + direct_io, + file_bytes: file_bytes.unwrap_or_default(), + data_segment_bytes, + data_segments, + root_layout, + root_layout_children, + file_rows, + decoded_rows, + batches: batch_count, + fields_per_batch, + field_dispatches: field_count, + stages: StageProfileRecord { + total_us: duration_us(total_time), + open_us: duration_us(open_time), + scan_plan_us: duration_us(scan_time), + read_us: duration_us(read_time), + struct_dispatch_us: duration_us(struct_time), + field_dispatch_us: duration_us(execute_time), + final_sync_us: duration_us(sync_time), + profile_overhead_us: duration_us(total_time.saturating_sub(accounted_time)), + }, + encodings, + }; + eprintln!("{}", serde_json::to_string(&record)?); } tracing::debug!( @@ -256,6 +467,10 @@ async fn decode_gpu_file( Ok(()) } +fn duration_us(duration: Duration) -> u64 { + u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) +} + /// Opens a Vortex file for CUDA execution. /// /// `direct_io` is off by default so this backend is comparable with the cuDF one: cuDF reads diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index e9134dad1c9..ae2389dd52e 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -14,6 +14,8 @@ use compress_bench::LanceCompressor; use compress_bench::gpu_parquet::GpuParquetCompressor; #[cfg(feature = "cuda")] use compress_bench::gpu_vortex::GpuVortexCompressor; +#[cfg(feature = "cuda")] +use compress_bench::gpu_vortex::GpuVortexProfile; use compress_bench::gpu_writer::GpuCodec; use compress_bench::parquet::ParquetCompressor; use compress_bench::vortex::VortexCompressor; @@ -103,6 +105,13 @@ struct Args { /// measure storage bandwidth instead, and do not read the ratio as a decode comparison. #[arg(long)] gpu_direct_io: bool, + /// Emit machine-readable Vortex GPU decode metrics after each timed stream synchronization. + /// + /// `wall` reports stage and encoding dispatch wall time, `gpu` additionally records CUDA + /// event spans around each field, and `nsys` adds per-field NVTX ranges. + #[cfg(feature = "cuda")] + #[arg(long, value_enum)] + gpu_vortex_profile: Option, #[arg(short, long, default_value_t, value_enum)] display_format: DisplayFormat, #[arg(short, long)] @@ -127,11 +136,21 @@ async fn main() -> anyhow::Result<()> { if args.gpu_decompress && !cfg!(feature = "cuda") { anyhow::bail!("--gpu-decompress requires building compress-bench with --features cuda"); } + #[cfg(feature = "cuda")] + if args.gpu_vortex_profile.is_some() && !args.gpu_decompress { + anyhow::bail!("--gpu-vortex-profile requires --gpu-decompress"); + } + #[cfg(feature = "cuda")] + if args.gpu_vortex_profile.is_some() && args.gpu_verify { + anyhow::bail!("--gpu-vortex-profile cannot be combined with --gpu-verify"); + } let gpu = args.gpu_decompress.then_some(GpuOptions { codec: args.gpu_parquet_codec, verify: args.gpu_verify, direct_io: args.gpu_direct_io, + #[cfg(feature = "cuda")] + vortex_profile: args.gpu_vortex_profile, }); #[cfg(feature = "cuda")] @@ -173,16 +192,22 @@ struct GpuOptions { verify: bool, /// Read the Vortex file with direct IO instead of through the page cache. direct_io: bool, + /// Optional diagnostics for the Vortex GPU path. + #[cfg(feature = "cuda")] + vortex_profile: Option, } /// Get a compressor for the given format. -fn get_compressor(format: Format, gpu: Option) -> Box { +fn get_compressor(format: Format, gpu: Option, _dataset: &str) -> Box { if let Some(gpu) = gpu { #[cfg(feature = "cuda")] return match format { - Format::OnDiskVortex => { - Box::new(GpuVortexCompressor::new(gpu.verify, gpu.direct_io)) as Box - } + Format::OnDiskVortex => Box::new(GpuVortexCompressor::new( + gpu.verify, + gpu.direct_io, + gpu.vortex_profile, + _dataset, + )) as Box, Format::Parquet => Box::new(GpuParquetCompressor::new(gpu.codec, gpu.verify)), _ => unimplemented!("GPU compress bench not implemented for {format}"), }; @@ -414,7 +439,7 @@ async fn run_benchmark_for_dataset( let mut v3_records: Vec = Vec::new(); for format in formats { - let compressor = get_compressor(*format, gpu); + let compressor = get_compressor(*format, gpu, bench_name); for op in ops { let time = match op { From fccded92070d45d0004ebf330f4f94d7a06a0ee5 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 19:25:53 +0100 Subject: [PATCH 21/22] Keep GPU patch indices on device Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/benches/dynamic_dispatch_cuda.rs | 4 +- vortex-cuda/kernels/src/patches.cuh | 56 ++++++++++-------- vortex-cuda/kernels/src/patches.h | 9 +-- vortex-cuda/src/kernel/encodings/bitpacked.rs | 59 +++++++++++++++++++ vortex-cuda/src/kernel/patches/mod.rs | 30 +++++----- vortex-cuda/src/kernel/patches/types.rs | 38 +++++------- 6 files changed, 128 insertions(+), 68 deletions(-) diff --git a/vortex-cuda/benches/dynamic_dispatch_cuda.rs b/vortex-cuda/benches/dynamic_dispatch_cuda.rs index 2d9db5ddbbd..d9115ba08d8 100644 --- a/vortex-cuda/benches/dynamic_dispatch_cuda.rs +++ b/vortex-cuda/benches/dynamic_dispatch_cuda.rs @@ -410,7 +410,8 @@ mod standalone { struct NullGpuPatches { chunk_offsets: *mut c_void, chunk_offset_type: u32, - indices: *mut u32, + indices_type: u32, + indices: *mut c_void, values: *mut c_void, offset: u32, offset_within_chunk: u32, @@ -425,6 +426,7 @@ mod standalone { const NULL: Self = Self { chunk_offsets: ptr::null_mut(), chunk_offset_type: 2, + indices_type: 2, indices: ptr::null_mut(), values: ptr::null_mut(), offset: 0, diff --git a/vortex-cuda/kernels/src/patches.cuh b/vortex-cuda/kernels/src/patches.cuh index 076bf11e72c..8acb068a89b 100644 --- a/vortex-cuda/kernels/src/patches.cuh +++ b/vortex-cuda/kernels/src/patches.cuh @@ -6,21 +6,25 @@ #include "fastlanes_common.cuh" #include "patches.h" -/// Load a chunk offset value, dispatching on the runtime type. -__device__ inline uint32_t load_chunk_offset(const GPUPatches &patches, uint32_t idx) { - switch (patches.chunk_offset_type) { - case CO_U8: - return reinterpret_cast(patches.chunk_offsets)[idx]; - case CO_U16: - return reinterpret_cast(patches.chunk_offsets)[idx]; - case CO_U32: - return reinterpret_cast(patches.chunk_offsets)[idx]; - case CO_U64: - return static_cast(reinterpret_cast(patches.chunk_offsets)[idx]); +/// Load an unsigned integer value, dispatching on the runtime type. +__device__ inline uint32_t load_unsigned(const void *values, UnsignedType type, uint32_t idx) { + switch (type) { + case UNSIGNED_U8: + return reinterpret_cast(values)[idx]; + case UNSIGNED_U16: + return reinterpret_cast(values)[idx]; + case UNSIGNED_U32: + return reinterpret_cast(values)[idx]; + case UNSIGNED_U64: + return static_cast(reinterpret_cast(values)[idx]); } return 0; } +__device__ inline uint32_t load_chunk_offset(const GPUPatches &patches, uint32_t idx) { + return load_unsigned(patches.chunk_offsets, patches.chunk_offset_type, idx); +} + /// A single patch: a within-chunk index and its replacement value. /// A sentinel patch has index == FL_CHUNK, which can never match a valid /// within-chunk position (0–FL_CHUNK-1). @@ -49,17 +53,14 @@ public: /// Construct a cursor for this thread's portion of patches in the chunk. __device__ PatchesCursor(const GPUPatches &patches, uint32_t chunk, uint32_t thread_idx, uint32_t n_threads) { - if (patches.chunk_offsets == nullptr) { - indices = nullptr; - values = nullptr; - remaining = 0; - return; - } + indices = nullptr; + indices_type = UNSIGNED_U32; + index = 0; + values = nullptr; + remaining = 0; + chunk_base = 0; - if (chunk >= patches.n_chunks) { - indices = nullptr; - values = nullptr; - remaining = 0; + if (patches.chunk_offsets == nullptr || chunk >= patches.n_chunks) { return; } @@ -94,7 +95,9 @@ public: uint32_t start = patches_start_idx + my_start; remaining = my_end - my_start; - indices = patches.indices + start; + indices = patches.indices; + indices_type = patches.indices_type; + index = start; values = reinterpret_cast(patches.values) + start; // The iterator returns indices relative to the start of the chunk. @@ -110,16 +113,19 @@ public: if (remaining == 0) { return {FL_CHUNK, T {}}; } - uint16_t within_chunk = static_cast(*indices - chunk_base); + uint16_t within_chunk = + static_cast(load_unsigned(indices, indices_type, index) - chunk_base); Patch patch = {within_chunk, *values}; - indices++; + index++; values++; remaining--; return patch; } private: - const uint32_t *indices; + const void *indices; + UnsignedType indices_type; + uint32_t index; const T *values; uint32_t remaining; uint32_t chunk_base; diff --git a/vortex-cuda/kernels/src/patches.h b/vortex-cuda/kernels/src/patches.h index 32dfa0de2cc..6f520adc3ee 100644 --- a/vortex-cuda/kernels/src/patches.h +++ b/vortex-cuda/kernels/src/patches.h @@ -9,8 +9,8 @@ extern "C" { #endif -/// Type tag for chunk_offsets pointer. -typedef enum { CO_U8 = 0, CO_U16 = 1, CO_U32 = 2, CO_U64 = 3 } ChunkOffsetType; +/// Type tag for an unsigned integer pointer. +typedef enum { UNSIGNED_U8 = 0, UNSIGNED_U16 = 1, UNSIGNED_U32 = 2, UNSIGNED_U64 = 3 } UnsignedType; static const uint32_t PATCH_DERIVE_INDICES_BASE = UINT32_MAX; @@ -24,8 +24,9 @@ static const uint32_t PATCH_DERIVE_INDICES_BASE = UINT32_MAX; /// A NULL chunk_offsets pointer indicates no patches are present. typedef struct { void *chunk_offsets; - ChunkOffsetType chunk_offset_type; - uint32_t *indices; + UnsignedType chunk_offset_type; + UnsignedType indices_type; + void *indices; void *values; uint32_t offset; uint32_t offset_within_chunk; diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 86b7a88b276..3ec161a6223 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -234,6 +234,7 @@ mod tests { use vortex::array::validity::Validity::NonNullable; use vortex::buffer::Buffer; use vortex::buffer::buffer; + use vortex::dtype::PType; use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::error::VortexExpect; use vortex_array::VortexSessionExecute; @@ -307,6 +308,64 @@ mod tests { Ok(()) } + #[rstest] + #[case::u8(PrimitiveArray::from_iter([0u8, 100, 200]).into_array())] + #[case::u16(PrimitiveArray::from_iter([0u16, 100, 200]).into_array())] + #[case::u32(PrimitiveArray::from_iter([0u32, 100, 200]).into_array())] + #[case::u64(PrimitiveArray::from_iter([0u64, 100, 200]).into_array())] + #[crate::test] + fn test_cuda_bitunpack_native_patch_index_widths( + #[case] indices: ArrayRef, + ) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) + .vortex_expect("failed to create execution context"); + let mut values: Vec = (0..1024).map(|i| i % 16).collect(); + values[0] = 500; + values[100] = 600; + values[200] = 700; + let expected = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(); + + let encoded = BitPacked::encode(&expected, 4, &mut ctx)?; + let BitPackedDataParts { + offset, + bit_width, + len, + packed, + patches, + validity, + } = BitPacked::into_parts(encoded); + let original_patches = patches.vortex_expect("expected patches"); + let native_patches = vortex_array::patches::Patches::new( + len, + 0, + indices, + original_patches.values().clone(), + original_patches.chunk_offsets().clone(), + )?; + let encoded = BitPacked::try_new( + packed, + PType::U16, + validity, + Some(native_patches), + bit_width, + len, + offset, + )?; + + let gpu_result = block_on(async { + BitPackedExecutor + .execute(encoded.into_array(), &mut cuda_ctx) + .await + .vortex_expect("GPU decompression failed") + .into_host() + .await + .map(|array| array.into_array()) + })?; + assert_arrays_eq!(expected, gpu_result, &mut ctx); + Ok(()) + } + #[rstest] #[case::bw_1(1)] #[case::bw_2(2)] diff --git a/vortex-cuda/src/kernel/patches/mod.rs b/vortex-cuda/src/kernel/patches/mod.rs index 7c651b7e507..727ef4b7611 100644 --- a/vortex-cuda/src/kernel/patches/mod.rs +++ b/vortex-cuda/src/kernel/patches/mod.rs @@ -27,13 +27,13 @@ use crate::CudaBufferExt; use crate::CudaDeviceBuffer; use crate::CudaExecutionCtx; use crate::executor::CudaArrayExt; -use crate::kernel::patches::gpu::ChunkOffsetType; -use crate::kernel::patches::gpu::ChunkOffsetType_CO_U8; -use crate::kernel::patches::gpu::ChunkOffsetType_CO_U16; -use crate::kernel::patches::gpu::ChunkOffsetType_CO_U32; -use crate::kernel::patches::gpu::ChunkOffsetType_CO_U64; use crate::kernel::patches::gpu::GPUPatches; use crate::kernel::patches::gpu::PATCH_DERIVE_INDICES_BASE; +use crate::kernel::patches::gpu::UnsignedType; +use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U8; +use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U16; +use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U32; +use crate::kernel::patches::gpu::UnsignedType_UNSIGNED_U64; use crate::kernel::patches::types::DevicePatches; // Safe because `GPUPatches` contains only raw pointers, POD integers, and an enum. @@ -44,7 +44,8 @@ impl GPUPatches { /// `chunk_offsets` pointer is the signal `PatchesCursor` checks for. pub(crate) const NULL_PATCHES: Self = Self { chunk_offsets: std::ptr::null_mut(), - chunk_offset_type: ChunkOffsetType_CO_U32, + chunk_offset_type: UnsignedType_UNSIGNED_U32, + indices_type: UnsignedType_UNSIGNED_U32, indices: std::ptr::null_mut(), values: std::ptr::null_mut(), offset: 0, @@ -55,14 +56,14 @@ impl GPUPatches { }; } -/// Convert a [`PType`] to the corresponding [`ChunkOffsetType`] for GPU patches. -pub(crate) fn ptype_to_chunk_offset_type(ptype: PType) -> VortexResult { +/// Convert a [`PType`] to the corresponding [`UnsignedType`] for GPU patches. +pub(crate) fn ptype_to_unsigned_type(ptype: PType) -> VortexResult { match ptype { - PType::U8 => Ok(ChunkOffsetType_CO_U8), - PType::U16 => Ok(ChunkOffsetType_CO_U16), - PType::U32 => Ok(ChunkOffsetType_CO_U32), - PType::U64 => Ok(ChunkOffsetType_CO_U64), - _ => vortex_bail!("Invalid PType for chunk_offsets: {:?}", ptype), + PType::U8 => Ok(UnsignedType_UNSIGNED_U8), + PType::U16 => Ok(UnsignedType_UNSIGNED_U16), + PType::U32 => Ok(UnsignedType_UNSIGNED_U32), + PType::U64 => Ok(UnsignedType_UNSIGNED_U64), + _ => vortex_bail!("Invalid unsigned PType: {:?}", ptype), } } @@ -77,7 +78,8 @@ pub(crate) fn build_gpu_patches( match device_patches { Some(p) => Ok(GPUPatches { chunk_offsets: p.chunk_offsets.cuda_device_ptr()? as _, - chunk_offset_type: ptype_to_chunk_offset_type(p.chunk_offset_ptype)?, + chunk_offset_type: ptype_to_unsigned_type(p.chunk_offset_ptype)?, + indices_type: ptype_to_unsigned_type(p.indices_ptype)?, indices: p.indices.cuda_device_ptr()? as _, values: p.values.cuda_device_ptr()? as _, offset: p.offset as u32, diff --git a/vortex-cuda/src/kernel/patches/types.rs b/vortex-cuda/src/kernel/patches/types.rs index 3bfe2270b66..da0058a63fc 100644 --- a/vortex-cuda/src/kernel/patches/types.rs +++ b/vortex-cuda/src/kernel/patches/types.rs @@ -6,31 +6,29 @@ use std::mem::size_of; use std::ops::Range; -use num_traits::ToPrimitive; use vortex::array::buffer::BufferHandle; use vortex::buffer::Alignment; -use vortex::buffer::Buffer; -use vortex::buffer::BufferMut; use vortex::buffer::ByteBufferMut; use vortex::dtype::PType; -use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::patches::PATCH_CHUNK_SIZE; use vortex_array::patches::Patches; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use crate::CudaBufferExt; use crate::CudaExecutionCtx; use crate::executor::CudaArrayExt; use crate::kernel::patches::gpu::GPUPatches; use crate::kernel::patches::gpu::PATCH_DERIVE_INDICES_BASE; -use crate::kernel::patches::ptype_to_chunk_offset_type; +use crate::kernel::patches::ptype_to_unsigned_type; /// A set of device-resident patches. pub struct DevicePatches { pub(crate) chunk_offsets: BufferHandle, pub(crate) chunk_offset_ptype: PType, pub(crate) indices: BufferHandle, + pub(crate) indices_ptype: PType, pub(crate) values: BufferHandle, pub(crate) offset: usize, pub(crate) offset_within_chunk: usize, @@ -55,6 +53,12 @@ pub(crate) async fn load_device_patches( ctx: &mut CudaExecutionCtx, ) -> VortexResult { let offset = patches.offset(); + vortex_ensure!( + offset + .checked_add(patches.array_len()) + .is_some_and(|end| end <= u32::MAX as usize), + "CUDA patches require offset + array length to fit in u32" + ); let offset_within_chunk = patches.offset_within_chunk().unwrap_or_default(); // Get or compute chunk_offsets let Some(co) = patches.chunk_offsets() else { @@ -68,7 +72,7 @@ pub(crate) async fn load_device_patches( (co_canonical.buffer_handle().clone(), ptype, len) }; - // Load indices - must be converted to u32 for GPU use + // Load indices at their native width. let indices = patches .indices() .clone() @@ -76,23 +80,7 @@ pub(crate) async fn load_device_patches( .await? .into_primitive(); let indices_ptype = indices.ptype(); - #[expect(clippy::expect_used)] - let indices = if indices_ptype == PType::U32 { - indices.buffer_handle().clone() - } else { - // Convert indices to u32 - let indices_buf = indices.buffer_handle().to_host().await; - let indices_u32 = match_each_unsigned_integer_ptype!(indices_ptype, |I| { - let src: Buffer = Buffer::from_byte_buffer(indices_buf); - let mut dst: BufferMut = BufferMut::with_capacity(src.len()); - for &idx in src.as_slice() { - // Indices are limited to u32 range for GPU - dst.push(idx.to_u32().expect("index should fit in u32")); - } - dst.freeze() - }); - BufferHandle::new_host(indices_u32.into_byte_buffer()) - }; + let indices = indices.buffer_handle().clone(); // Load values let values = patches @@ -113,6 +101,7 @@ pub(crate) async fn load_device_patches( chunk_offsets, chunk_offset_ptype, indices, + indices_ptype, values, offset, offset_within_chunk, @@ -134,7 +123,8 @@ fn build_gpu_patches( // chunk_offset_type and indices) which would be UB when serialized. let mut gpu_patches: GPUPatches = unsafe { std::mem::zeroed() }; gpu_patches.chunk_offsets = dp.chunk_offsets.cuda_device_ptr()? as _; - gpu_patches.chunk_offset_type = ptype_to_chunk_offset_type(dp.chunk_offset_ptype)?; + gpu_patches.chunk_offset_type = ptype_to_unsigned_type(dp.chunk_offset_ptype)?; + gpu_patches.indices_type = ptype_to_unsigned_type(dp.indices_ptype)?; gpu_patches.indices = dp.indices.cuda_device_ptr()? as _; gpu_patches.values = dp.values.cuda_device_ptr()? as _; gpu_patches.offset = dp.offset as u32; From 5b7e3a3df6a3c3373e3f8b9bcd427ee0bc80f830 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 14 Aug 2026 19:26:06 +0100 Subject: [PATCH 22/22] Avoid synchronizing fused GPU patch decodes Signed-off-by: Joe Isaacs <2413449+joseph-isaacs@users.noreply.github.com> --- vortex-cuda/src/kernel/encodings/bitpacked.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 3ec161a6223..da67326e893 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -206,11 +206,6 @@ where .arg(&patches_arg); })?; - // Patch-free decodes need no host synchronization. - if device_patches.is_some() { - ctx.synchronize_stream()?; - } - let output_buf = CudaDeviceBuffer::new(output_slice); let output_handle = BufferHandle::new_device(output_buf.slice_typed::(offset..(offset + len)));