diff --git a/Cargo.lock b/Cargo.lock index 0d55c729ae6..a006177e2a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1572,12 +1572,14 @@ dependencies = [ "async-trait", "bytes", "clap", + "cudarc", "futures", "indicatif", "itertools 0.14.0", "lance-bench", "parquet 58.4.0", "regex", + "serde_json", "tempfile", "tokio", "tracing", diff --git a/benchmarks/compress-bench/Cargo.toml b/benchmarks/compress-bench/Cargo.toml index 4046a12d42e..0f233a3a22c 100644 --- a/benchmarks/compress-bench/Cargo.toml +++ b/benchmarks/compress-bench/Cargo.toml @@ -21,12 +21,14 @@ arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } clap = { workspace = true, features = ["derive"] } +cudarc = { workspace = true, optional = true } futures = { workspace = true } indicatif = { workspace = true } itertools = { workspace = true } lance-bench = { path = "../lance-bench", optional = true } parquet = { workspace = true } regex = { workspace = true } +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"] diff --git a/benchmarks/compress-bench/README.md b/benchmarks/compress-bench/README.md index bf2d3efc1db..1fb1e87e3fa 100644 --- a/benchmarks/compress-bench/README.md +++ b/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 +``` diff --git a/benchmarks/compress-bench/src/gpu_vortex.rs b/benchmarks/compress-bench/src/gpu_vortex.rs index 3dbb68bc7a8..c86391f8454 100644 --- a/benchmarks/compress-bench/src/gpu_vortex.rs +++ b/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, +} + /// 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::(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::(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::>(); + 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) +}