feat: add a ci gpu compress benchmark

.github/workflows/bench-dispatch.yml

@@ -32,6 +32,22 @@ jobs:
     uses: ./.github/workflows/bench-pr.yml
     secrets: inherit
 
+  remove-gpu-compress-bench-label:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    if: github.event.label.name == 'action/benchmark-gpu-compress'
+    steps:
+      - uses: actions-ecosystem/action-remove-labels@2ce5d41b4b6aa8503e285553f75ed56e0a40bae0  # v1
+        if: github.event.pull_request.head.repo.full_name == 'vortex-data/vortex'
+        with:
+          labels: action/benchmark-gpu-compress
+          fail_on_error: true
+
+  gpu-compress-bench:
+    needs: remove-gpu-compress-bench-label
+    uses: ./.github/workflows/gpu-compress-bench-pr.yml
+    secrets: inherit
+
   remove-sql-label:
     runs-on: ubuntu-latest
     timeout-minutes: 10

.github/workflows/gpu-compress-bench-pr.yml

@@ -0,0 +1,87 @@
+# Runs the GPU-enabled compression decompression benchmarks for a pull request.
+# Called from bench-dispatch.yml when the `action/benchmark-gpu-compress` label is added.
+
+name: GPU Compression Benchmarks
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: false
+
+on:
+  workflow_call: { }
+  workflow_dispatch: { }
+
+permissions:
+  contents: read
+  pull-requests: write
+
+jobs:
+  gpu-compress-bench:
+    name: GPU Compression
+    timeout-minutes: 120
+    runs-on: >-
+      runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/extras=s3-cache/tag=gpu-compress-bench
+    steps:
+      - uses: runs-on/action@v2
+        with:
+          sccache: s3
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0  # v7
+        with:
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
+      - name: Setup benchmark environment
+        run: sudo bash scripts/setup-benchmark.sh
+      - uses: ./.github/actions/setup-rust
+        with:
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          enable-sccache: "true"
+      - uses: ./.github/actions/system-info
+      - name: Display NVIDIA GPU details
+        run: |
+          nvidia-smi
+          nvidia-smi -L
+          nvidia-smi -q -d Memory
+      - name: Build GPU compression benchmark
+        shell: bash
+        run: |
+          cargo build --locked --package compress-bench --profile release_debug --features cuda,unstable_encodings
+      - name: Run GPU compression benchmark
+        shell: bash
+        env:
+          RUST_BACKTRACE: full
+          # Do not enable VORTEX_EXPERIMENTAL_PATCHED_ARRAY here: it rewrites interior
+          # `Patches` into a top-level `vortex.patched` array, which has no CUDA decode
+          # kernel and makes GPU decompression fail with "No CUDA kernel for encoding
+          # vortex.patched". Leaving it unset keeps patches inline in the bitpacked array,
+          # which the CUDA bitpacked kernel decodes. This matches the GPU test workflow
+          # (cuda.yaml), which also sets only FLAT_LAYOUT_INLINE_ARRAY_NODE.
+          FLAT_LAYOUT_INLINE_ARRAY_NODE: "1"
+        run: |
+          bash scripts/bench-taskset.sh target/release_debug/compress-bench \
+            --gpu-decompress -d table > gpu-compress.txt
+          cat gpu-compress.txt
+      - name: Publish results
+        shell: bash
+        run: |
+          {
+            echo "# GPU Compression"
+            echo
+            echo '```text'
+            cat gpu-compress.txt
+            echo '```'
+          } > comment.md
+          cat comment.md >> "$GITHUB_STEP_SUMMARY"
+      - name: Comment PR
+        if: github.event.pull_request.head.repo.fork == false
+        uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b  # v3
+        with:
+          file-path: comment.md
+          comment-tag: bench-pr-comment-gpu-compress
+      - name: Comment PR on failure
+        if: failure() && github.event.pull_request.head.repo.fork == false
+        uses: thollander/actions-comment-pull-request@24bffb9b452ba05a4f3f77933840a6a841d1b32b  # v3
+        with:
+          message: |
+            # BENCHMARK FAILED
+
+            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

benchmarks/compress-bench/Cargo.toml

@@ -27,15 +27,18 @@ itertools = { workspace = true }
 lance-bench = { path = "../lance-bench", optional = true }
 parquet = { workspace = true }
 regex = { workspace = true }
+tempfile = { workspace = true, optional = true }
 tokio = { workspace = true, features = ["full"] }
 tracing = { workspace = true }
 vortex = { workspace = true }
 vortex-arrow = { workspace = true }
 vortex-bench = { workspace = true }
+vortex-cuda = { workspace = true, optional = true }
 
 [features]
+cuda = ["dep:tempfile", "dep:vortex-cuda"]
 lance = ["dep:lance-bench"]
-unstable_encodings = ["vortex/unstable_encodings"]
+unstable_encodings = ["vortex/unstable_encodings", "vortex-cuda?/unstable_encodings"]
 
 [[bin]]
 name = "compress-bench"

benchmarks/compress-bench/README.md

@@ -14,3 +14,14 @@ See [`src/main.rs`](./src/main.rs) for the dataset list and CLI flags (`--format
 ```bash
 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`:
+
+```bash
+cargo run -p compress-bench --profile release_debug \
+  --features cuda,unstable_encodings -- --gpu-decompress
+```
+
+On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure
+storage bandwidth rather than page-cache hits.

benchmarks/compress-bench/src/gpu_vortex.rs

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::hint::black_box;
+use std::path::Path;
+use std::sync::Arc;
+use std::time::Duration;
+use std::time::Instant;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use futures::StreamExt;
+use tempfile::NamedTempFile;
+use vortex::array::IntoArray;
+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_bench::Format;
+use vortex_bench::SESSION;
+use vortex_bench::compress::Compressor;
+use vortex_bench::conversions::parquet_to_vortex_chunks;
+use vortex_cuda::CudaOpenOptionsExt;
+use vortex_cuda::CudaSession;
+#[cfg(target_os = "linux")]
+use vortex_cuda::PooledFileReadAtOptions;
+use vortex_cuda::executor::CudaArrayExt;
+use vortex_cuda::layout::CudaFlatLayoutStrategy;
+use vortex_cuda::layout::register_cuda_layout;
+
+/// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU.
+pub struct GpuVortexCompressor;
+
+#[async_trait]
+impl Compressor for GpuVortexCompressor {
+    fn format(&self) -> Format {
+        Format::OnDiskVortex
+    }
+
+    async fn compress(&self, _parquet_path: &Path) -> Result<(u64, Duration)> {
+        anyhow::bail!("GPU compress-bench only supports decompression measurements")
+    }
+
+    async fn decompress(&self, parquet_path: &Path) -> Result<Duration> {
+        register_cuda_layout(&SESSION);
+
+        let uncompressed = parquet_to_vortex_chunks(parquet_path.to_path_buf()).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();
+        SESSION
+            .write_options()
+            .with_strategy(strategy)
+            .write(&mut output, uncompressed.into_array().to_array_stream())
+            .await?;
+        output.sync_all().await?;
+        drop(output);
+
+        let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?;
+        let start = Instant::now();
+        let open_options = SESSION.open_options().with_cuda();
+        // Direct IO keeps repeated iterations measuring storage bandwidth rather than
+        // page-cache hits. It is only available on Linux.
+        #[cfg(target_os = "linux")]
+        let open_options =
+            open_options.with_read_at_options(PooledFileReadAtOptions::default().with_direct_io());
+        let file = open_options.open_path(gpu_file.path()).await?;
+        let mut batches = file.scan()?.into_array_stream()?;
+
+        while let Some(batch) = batches.next().await {
+            let record = batch?.execute::<StructArray>(cuda_ctx.execution_ctx())?;
+            for field in record.iter_unmasked_fields() {
+                black_box(field.clone().execute_cuda(&mut cuda_ctx).await?);
+            }
+        }
+        cuda_ctx.synchronize_stream()?;
+
+        Ok(start.elapsed())
+    }
+}

benchmarks/compress-bench/src/lib.rs

@@ -3,5 +3,7 @@
 
 #[cfg(feature = "lance")]
 pub use lance_bench::compress::LanceCompressor;
+#[cfg(feature = "cuda")]
+pub mod gpu_vortex;
 pub mod parquet;
 pub mod vortex;