diff --git a/.github/workflows/pr-bench-gpu-compress.yml b/.github/workflows/pr-bench-gpu-compress.yml index cd97755e636..505a89a2a75 100644 --- a/.github/workflows/pr-bench-gpu-compress.yml +++ b/.github/workflows/pr-bench-gpu-compress.yml @@ -32,6 +32,25 @@ 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" + - 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: | @@ -44,8 +63,68 @@ 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 + id: verify + shell: bash + continue-on-error: true + env: + 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. 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 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 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 + 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: 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 @@ -56,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: | @@ -68,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 @@ -85,3 +190,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/Cargo.lock b/Cargo.lock index fee1e3ebd34..43b8042e63e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1572,12 +1572,15 @@ dependencies = [ "async-trait", "bytes", "clap", + "cudarc", "futures", "indicatif", "itertools 0.14.0", "lance-bench", "parquet 58.4.0", "regex", + "serde", + "serde_json", "tempfile", "tokio", "tracing", diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 4046a12d42e..1a3d2b6ac60 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,12 +21,15 @@ 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 } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { 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,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"] @@ -45,7 +48,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..7a4c3e71e45 100644 --- a/benchmarks/compress-bench/README.md +++ b/benchmarks/compress-bench/README.md @@ -15,13 +15,138 @@ 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 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. +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 \ --features cuda,unstable_encodings -- --gpu-decompress + +# 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 +``` + +### 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 +benchmark and never enters the Rust build: + +```bash +uv pip install --extra-index-url https://pypi.nvidia.com cudf-cu12 pandas pyarrow +``` + +`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. + +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. + +### GPU-friendly Parquet writer settings + +Set in `src/gpu_writer.rs`: + +| Setting | Value | Why | +| --- | --- | --- | +| 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-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 + +`--gpu-verify` cross-checks device output against the CPU decoders on every iteration: + +- 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. + +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 ``` -On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure -storage bandwidth rather than page-cache hits. +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. diff --git a/benchmarks/compress-bench/src/gpu_parquet.rs b/benchmarks/compress-bench/src/gpu_parquet.rs new file mode 100644 index 00000000000..3ee0297f1a0 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_parquet.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! GPU Parquet decompression backend, timed through cuDF. +//! +//! 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. +//! +//! 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::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::bail; +use anyhow::ensure; +use arrow_array::RecordBatch; +use async_trait::async_trait; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use serde::Deserialize; +use tempfile::NamedTempFile; +use vortex_bench::Format; +use vortex_bench::compress::Compressor; + +use crate::gpu_writer::GpuCodec; +use crate::gpu_writer::gpu_writer_properties; + +/// 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 times cuDF reading them back. + /// + /// 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(std::fs::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)) + } +} + +#[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 report = run_cudf_read(gpu_file.path(), self.verify)?; + + ensure!( + report.rows > 0 && report.columns > 0, + "cuDF read {} rows and {} columns, expected a non-empty table", + report.rows, + report.columns + ); + + Ok(Duration::from_nanos(report.min_ns)) + } +} + +/// 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"); + } + + let output = command.output().with_context(|| { + format!("failed to run {CUDF_SCRIPT}; is cudf-cu12 installed on this host?") + })?; + + if !output.status.success() { + bail!( + "{CUDF_SCRIPT} exited with {}:\n{}", + output.status, + String::from_utf8_lossy(&output.stderr).trim() + ); + } + + 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_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 3dbb68bc7a8..b87811ecba6 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/benchmarks/compress-bench/src/gpu_vortex.rs @@ -1,37 +1,109 @@ // 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; 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 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; use vortex::array::IntoArray; +use vortex::array::VortexSessionExecute; use vortex::array::arrays::StructArray; 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; + +/// 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; +pub struct GpuVortexCompressor { + verify: bool, + direct_io: bool, + profile: Option, + dataset: Arc, + iteration: AtomicUsize, +} + +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, + direct_io: bool, + profile: Option, + dataset: &str, + ) -> Self { + Self { + verify, + direct_io, + profile, + dataset: dataset.into(), + iteration: AtomicUsize::new(0), + } + } +} #[async_trait] impl Compressor for GpuVortexCompressor { @@ -46,13 +118,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) @@ -61,25 +146,501 @@ impl Compressor for GpuVortexCompressor { output.sync_all().await?; drop(output); + if self.verify { + return verify_against_host_scan(gpu_file.path(), self.direct_io).await; + } + 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", + 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(); - 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 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() { - black_box(field.clone().execute_cuda(&mut cuda_ctx).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()? + .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 decoded_rows = 0usize; + let mut fields_per_batch = 0usize; + let mut field_count = 0usize; + 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(); + let Some(batch) = batches.next().await else { + read_time += read_start.elapsed(); + break; + }; + read_time += read_start.elapsed(); + + let struct_start = Instant::now(); + 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{}", + 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.map(|_| { + ( + 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 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(FieldTiming { + field_name: field_name.to_string(), + rows: field.len(), + encoding, + tree, + wall_time, + gpu_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(); + 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!( + phase, + batch_count, + field_count, + ?open_time, + ?scan_time, + ?read_time, + ?struct_time, + ?execute_time, + ?sync_time, + "GPU Vortex decode stages" + ); + 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 +/// 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 = 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?) +} + +/// 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, 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 + // 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 + // 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, direct_io).await?; + let mut gpu_batches = gpu_file + .scan()? + .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE)) + .into_array_stream()?; + let host_file = SESSION.open_options().open_path(host_path.path()).await?; + let mut host_batches = host_file + .scan()? + .with_split_by(SplitBy::RowCount(GPU_ROW_GROUP_SIZE)) + .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) { + (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::(&mut host_ctx)?; + ensure!( + gpu_record.len() == host_record.len(), + "batch {batch_index} 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(), + "batch {batch_index} field count differs between the GPU and CPU scans" + ); + + 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, + &mut host_ctx, + batch_index, + field_index, + )?; + fields_checked += 1; + } + + batch_index += 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( + 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)?; + + 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(), + ); + + 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 } diff --git a/benchmarks/compress-bench/src/gpu_writer.rs b/benchmarks/compress-bench/src/gpu_writer.rs new file mode 100644 index 00000000000..bd66b005c73 --- /dev/null +++ b/benchmarks/compress-bench/src/gpu_writer.rs @@ -0,0 +1,99 @@ +// 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::DEFAULT_MAX_ROW_GROUP_ROW_COUNT; +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; + +/// Rows per independently readable partition in both GPU benchmark formats. +/// +/// Parquet calls these row groups and the CUDA Vortex layout stores one chunk per partition. +pub const GPU_ROW_GROUP_SIZE: usize = DEFAULT_MAX_ROW_GROUP_ROW_COUNT; + +/// 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) + .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() +} + +#[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); + assert_eq!( + properties.max_row_group_row_count(), + Some(GPU_ROW_GROUP_SIZE) + ); + } +} diff --git a/benchmarks/compress-bench/src/lib.rs b/benchmarks/compress-bench/src/lib.rs index 68039996605..76e34172f6a 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 gpu_writer; pub mod parquet; pub mod vortex; diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1b1603e52c8..ae2389dd52e 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -1,16 +1,25 @@ // 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; +use anyhow::Context; 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; +#[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; +use futures::FutureExt; use indicatif::ProgressBar; use itertools::Itertools; use regex::Regex; @@ -18,6 +27,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; @@ -35,6 +46,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; @@ -44,6 +56,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)] @@ -67,11 +81,37 @@ 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 file is written with. + /// + /// 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. + /// + /// 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, + /// 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, + /// 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)] @@ -96,9 +136,36 @@ 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, + }); - let (formats, ops) = if args.gpu_decompress { - (vec![Format::OnDiskVortex], vec![CompressOp::Decompress]) + #[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!( + 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) }; @@ -108,7 +175,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 +183,39 @@ async fn main() -> anyhow::Result<()> { .await } +/// Settings for the GPU decompression mode. +#[derive(Clone, Copy, Debug)] +struct GpuOptions { + /// Parquet page codec to write the GPU file with. + 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, + /// 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_decompress: bool) -> Box { - if gpu_decompress { +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, + gpu.vortex_profile, + _dataset, + )) 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 +242,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 +269,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 +308,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 { @@ -225,18 +332,37 @@ async fn run_compress( let mut measurements = vec![]; let mut v3_records: Vec = Vec::new(); + // 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() { - let (m, mut records) = run_benchmark_for_dataset( - &progress, - &formats, - &ops, - iterations, - dataset_handle, - gpu_decompress, - ) - .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); @@ -249,6 +375,8 @@ async fn run_compress( 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)?; @@ -260,13 +388,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( @@ -275,7 +423,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 +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_decompress); + let compressor = get_compressor(*format, gpu, bench_name); for op in ops { let time = match op { @@ -302,7 +450,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 @@ -333,7 +482,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() @@ -342,7 +492,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 +511,53 @@ 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, + }); +} + +/// 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() + } +} diff --git a/scripts/cudf-parquet-read.py b/scripts/cudf-parquet-read.py new file mode 100644 index 00000000000..39e435c3a77 --- /dev/null +++ b/scripts/cudf-parquet-read.py @@ -0,0 +1,120 @@ +#!/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 +from datetime import date + + +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 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 = 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. + 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-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-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/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/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/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/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/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/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/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/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/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/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) })) 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/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index 86b7a88b276..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))); @@ -234,6 +229,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 +303,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/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/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(()) + } } diff --git a/vortex-cuda/src/kernel/encodings/fsst.rs b/vortex-cuda/src/kernel/encodings/fsst.rs index f9b90c8ac35..dd174f550d4 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; @@ -30,6 +33,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; @@ -42,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; @@ -121,42 +126,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 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 + } + } +} - // 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 - }) +/// 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, + } = fsst_i32_offsets(&fsst, lens, 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 = 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), + )?; + + 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. @@ -183,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)?); @@ -229,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), )?; @@ -277,6 +357,84 @@ 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, + 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, 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/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; 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); 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,