diff --git a/benchmarks/compress-bench/src/main.rs b/benchmarks/compress-bench/src/main.rs index 1b1603e52c8..405d45ea603 100644 --- a/benchmarks/compress-bench/src/main.rs +++ b/benchmarks/compress-bench/src/main.rs @@ -18,6 +18,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; @@ -44,6 +46,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)] @@ -97,6 +101,11 @@ async fn main() -> anyhow::Result<()> { anyhow::bail!("--gpu-decompress requires building compress-bench with --features cuda"); } + #[cfg(feature = "cuda")] + if args.gpu_decompress { + SESSION.register(CudaSession::try_single_stream()?); + } + let (formats, ops) = if args.gpu_decompress { (vec![Format::OnDiskVortex], vec![CompressOp::Decompress]) } else { 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,