GPU benchmark: add reusable decompression profiling metrics

benchmarks/compress-bench/Cargo.toml

@@ -21,12 +21,14 @@ arrow-schema = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 clap = { workspace = true, features = ["derive"] }
+cudarc = { workspace = true, optional = true }
 futures = { workspace = true }
 indicatif = { workspace = true }
 itertools = { workspace = true }
 lance-bench = { path = "../lance-bench", optional = true }
 parquet = { workspace = true }
 regex = { workspace = true }
+serde_json = { workspace = true }
 tempfile = { workspace = true, optional = true }
 tokio = { workspace = true, features = ["full"] }
 tracing = { workspace = true }
@@ -36,7 +38,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"]
 

benchmarks/compress-bench/README.md

@@ -25,3 +25,19 @@ cargo run -p compress-bench --profile release_debug \
 
 On Linux, GPU files are read with direct IO (`O_DIRECT`) so repeated iterations measure
 storage bandwidth rather than page-cache hits.
+
+Set `VORTEX_GPU_PROFILE=wall` to emit one JSON record per decompression with file/layout counts,
+decoded rows, batch and field-dispatch counts, host time for open, scan planning, reads, struct and
+field dispatch, and final synchronization, plus wall time grouped by full encoding tree. Set it to
+`gpu` to additionally bracket field dispatches with CUDA events and report each encoding group's
+device-stream time. Profiling perturbs the measurement; rerun without it for comparison numbers.
+
+```bash
+VORTEX_GPU_PROFILE=gpu cargo run -p compress-bench --profile release_debug \
+  --features cuda,unstable_encodings -- --gpu-decompress --iterations 1 \
+  2> /tmp/vortex-gpu-profile.log
+
+jq -Rs '[split("\n")[] | fromjson? |
+  select(.record == "vortex_gpu_decompress_profile")]' \
+  /tmp/vortex-gpu-profile.log
+```

benchmarks/compress-bench/src/gpu_vortex.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::collections::BTreeMap;
 use std::hint::black_box;
 use std::path::Path;
 use std::sync::Arc;
@@ -9,7 +10,10 @@ use std::time::Instant;
 
 use anyhow::Result;
 use async_trait::async_trait;
+use cudarc::driver::CudaEvent;
+use cudarc::driver::sys::CUevent_flags::CU_EVENT_DEFAULT;
 use futures::StreamExt;
+use serde_json::json;
 use tempfile::NamedTempFile;
 use vortex::array::IntoArray;
 use vortex::array::arrays::StructArray;
@@ -30,6 +34,22 @@ use vortex_cuda::executor::CudaArrayExt;
 use vortex_cuda::layout::CudaFlatLayoutStrategy;
 use vortex_cuda::layout::register_cuda_layout;
 
+struct FieldTiming {
+    encoding: String,
+    tree: String,
+    rows: usize,
+    wall: Duration,
+    events: Option<(CudaEvent, CudaEvent)>,
+}
+
+#[derive(Default)]
+struct EncodingTiming {
+    calls: usize,
+    rows: usize,
+    wall: Duration,
+    gpu_us: Option<u64>,
+}
+
 /// Vortex compressor whose decompression measurement executes CUDA-compatible files on the GPU.
 pub struct GpuVortexCompressor;
 
@@ -62,24 +82,196 @@ impl Compressor for GpuVortexCompressor {
         drop(output);
 
         let mut cuda_ctx = CudaSession::create_execution_ctx(&SESSION)?;
+        let profile = std::env::var("VORTEX_GPU_PROFILE").ok();
+        anyhow::ensure!(
+            profile
+                .as_deref()
+                .is_none_or(|mode| matches!(mode, "wall" | "gpu")),
+            "VORTEX_GPU_PROFILE must be wall or gpu"
+        );
+        if profile.is_none() {
+            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()?;
+
+            return Ok(start.elapsed());
+        }
+
+        let profile_gpu = profile.as_deref() == Some("gpu");
         let start = Instant::now();
+        let open_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()?;
+        let open_time = open_start.elapsed();
+
+        let (file_bytes, data_segments, data_segment_bytes, root_layout, root_children, file_rows) =
+            if profile.is_some() {
+                let footer = file.footer();
+                (
+                    std::fs::metadata(gpu_file.path())?.len(),
+                    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, 0, String::new(), 0, 0)
+            };
 
-        while let Some(batch) = batches.next().await {
+        let scan_start = Instant::now();
+        let mut batches = file.scan()?.into_array_stream()?;
+        let scan_time = scan_start.elapsed();
+        let mut read_time = Duration::ZERO;
+        let mut struct_time = Duration::ZERO;
+        let mut field_time = Duration::ZERO;
+        let mut batch_count = 0usize;
+        let mut decoded_rows = 0usize;
+        let mut field_count = 0usize;
+        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::<StructArray>(cuda_ctx.execution_ctx())?;
+            struct_time += struct_start.elapsed();
+            batch_count += 1;
+            decoded_rows += record.len();
             for field in record.iter_unmasked_fields() {
+                let metadata = profile.as_ref().map(|_| {
+                    (
+                        field.encoding_id().to_string(),
+                        field
+                            .display_tree_encodings_only()
+                            .to_string()
+                            .replace('\n', " | "),
+                    )
+                });
+                let before = profile_gpu
+                    .then(|| cuda_ctx.stream().record_event(Some(CU_EVENT_DEFAULT)))
+                    .transpose()?;
+                let field_start = Instant::now();
                 black_box(field.clone().execute_cuda(&mut cuda_ctx).await?);
+                let wall = field_start.elapsed();
+                field_time += wall;
+                let events = if let Some(before) = before {
+                    Some((
+                        before,
+                        cuda_ctx.stream().record_event(Some(CU_EVENT_DEFAULT))?,
+                    ))
+                } else {
+                    None
+                };
+                if let Some((encoding, tree)) = metadata {
+                    field_timings.push(FieldTiming {
+                        encoding,
+                        tree,
+                        rows: field.len(),
+                        wall,
+                        events,
+                    });
+                }
+                field_count += 1;
             }
         }
+        let sync_start = Instant::now();
         cuda_ctx.synchronize_stream()?;
+        let sync_time = sync_start.elapsed();
+        let total = start.elapsed();
+
+        if let Some(mode) = profile {
+            let mut encodings = BTreeMap::<(String, String), EncodingTiming>::new();
+            for timing in field_timings {
+                let gpu_us = timing
+                    .events
+                    .map(|(before, after)| before.elapsed_ms(&after))
+                    .transpose()?
+                    .map(|ms| duration_us(Duration::from_secs_f32(ms / 1_000.0)));
+                let aggregate = encodings.entry((timing.encoding, timing.tree)).or_default();
+                aggregate.calls += 1;
+                aggregate.rows += timing.rows;
+                aggregate.wall += timing.wall;
+                aggregate.gpu_us = match (aggregate.gpu_us, gpu_us) {
+                    (Some(total), Some(value)) => Some(total.saturating_add(value)),
+                    (None, value) => value,
+                    (value, None) => value,
+                };
+            }
+            let accounted =
+                open_time + scan_time + read_time + struct_time + field_time + sync_time;
+            let encodings = encodings
+                .into_iter()
+                .map(|((encoding, tree), timing)| {
+                    json!({
+                        "encoding": encoding,
+                        "tree": tree,
+                        "calls": timing.calls,
+                        "rows": timing.rows,
+                        "wall_us": duration_us(timing.wall),
+                        "gpu_us": timing.gpu_us,
+                    })
+                })
+                .collect::<Vec<_>>();
+            eprintln!(
+                "{}",
+                json!({
+                    "record": "vortex_gpu_decompress_profile",
+                    "version": 1,
+                    "dataset_path": parquet_path,
+                    "mode": mode,
+                    "file_bytes": file_bytes,
+                    "data_segments": data_segments,
+                    "data_segment_bytes": data_segment_bytes,
+                    "root_layout": root_layout,
+                    "root_layout_children": root_children,
+                    "file_rows": file_rows,
+                    "decoded_rows": decoded_rows,
+                    "batches": batch_count,
+                    "field_dispatches": field_count,
+                    "stages": {
+                        "total_us": duration_us(total),
+                        "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(field_time),
+                        "final_sync_us": duration_us(sync_time),
+                        "profile_overhead_us": duration_us(total.saturating_sub(accounted)),
+                    },
+                    "encodings": encodings,
+                })
+            );
+        }
 
-        Ok(start.elapsed())
+        Ok(total)
     }
 }
+
+fn duration_us(duration: Duration) -> u64 {
+    u64::try_from(duration.as_micros()).unwrap_or(u64::MAX)
+}