diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 12bc584dad6..ffcd8105e35 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -19,13 +19,11 @@ use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteStrategyBuilder; -use vortex::io::VortexReadAt; use vortex::io::runtime::BlockingRuntime; -use vortex::io::session::RuntimeSessionExt; use vortex::session::SessionExt; use vortex::session::VortexSession; +use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; -use vortex_cuda::PooledFileReadAt; use vortex_cuda::arrow::ArrowDeviceArray; use vortex_cuda::arrow::ArrowDeviceArrayStream; use vortex_cuda::arrow::DeviceArrayExt; @@ -110,7 +108,10 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file( }) } -/// Scan a local Vortex file through pinned host buffers and export an Arrow C Device stream. +/// Scan a local Vortex file and export an Arrow C Device stream. +/// +/// Footer and zone-map reads remain on the host. Data segments are staged through pinned host +/// buffers and transferred directly to the GPU. /// /// The file must use encodings and layouts supported by the CUDA execution path, such as files /// written by [`vx_cuda_array_sink_open_file`]. Pinned staging buffers are reused across scans made @@ -140,19 +141,8 @@ pub unsafe extern "C-unwind" fn vx_cuda_scan_path_arrow_device_stream( let path = unsafe { path.as_str() }?.to_owned(); let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?; - let cuda_session = session.get::(); - let stream = cuda_session.stream()?; - let pool = Arc::clone(cuda_session.pinned_buffer_pool()); - drop(cuda_session); - - let reader: Arc = Arc::new(PooledFileReadAt::open( - path, - session.handle(), - pool, - stream, - )?); let array_stream = ffi_runtime().block_on(async { - let file = session.open_options().open(reader).await?; + let file = session.open_options().with_cuda().open_path(path).await?; Ok::<_, vortex::error::VortexError>(file.scan()?.into_array_stream()?.boxed()) })?; let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?; diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 00f6cd1ac1c..eb712bc2b41 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -31,9 +31,9 @@ use vortex::file::WriteStrategyBuilder; use vortex::io::session::RuntimeSessionExt; use vortex::session::SessionExt; use vortex::session::VortexSession; +use vortex_cuda::CudaOpenOptionsExt; use vortex_cuda::CudaSession; use vortex_cuda::PooledByteBufferReadAt; -use vortex_cuda::PooledFileReadAt; use vortex_cuda::TracingLaunchStrategy; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda::layout::CudaFlatLayoutStrategy; @@ -158,11 +158,9 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes let pool = Arc::clone(cuda_session.pinned_buffer_pool()); let cuda_stream = cuda_session.stream()?; drop(cuda_session); - let handle = session.handle(); let gpu_file_handle = if gpu_file { - let reader = PooledFileReadAt::open(&path, handle, Arc::clone(&pool), cuda_stream)?; - session.open_options().open(Arc::new(reader)).await? + session.open_options().with_cuda().open_path(&path).await? } else { let (recompressed, footer) = recompress_for_gpu(&path, &session).await?; let reader = PooledByteBufferReadAt::new(recompressed, Arc::clone(&pool), cuda_stream); diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs new file mode 100644 index 00000000000..9a41f51f9b9 --- /dev/null +++ b/vortex-cuda/src/file.rs @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! CUDA-aware Vortex file opening. + +use std::path::Path; +use std::sync::Arc; + +use vortex::error::VortexResult; +use vortex::error::vortex_bail; +use vortex::file::VortexFile; +use vortex::file::VortexOpenOptions; +use vortex::io::session::RuntimeSessionExt; +use vortex::layout::Layout; +use vortex::layout::layouts::zoned::LegacyStats; +use vortex::layout::layouts::zoned::Zoned; +use vortex::layout::segments::SegmentFuture; +use vortex::layout::segments::SegmentId; +use vortex::layout::segments::SegmentSource; + +use crate::CudaSessionExt; +use crate::PooledFileReadAt; +use crate::layout::register_cuda_layout; + +/// Extension trait for opening CUDA-readable files from [`VortexOpenOptions`]. +pub trait CudaOpenOptionsExt { + /// Configure the opener to keep control-plane segments on the host and read data-plane + /// segments into CUDA device memory. + fn with_cuda(self) -> CudaOpenOptions; +} + +impl CudaOpenOptionsExt for VortexOpenOptions { + fn with_cuda(self) -> CudaOpenOptions { + CudaOpenOptions { inner: self } + } +} + +/// CUDA-aware Vortex file open options. +/// +/// Construct this adapter with [`CudaOpenOptionsExt::with_cuda`]. Standard file options should be +/// configured before calling `with_cuda`. +pub struct CudaOpenOptions { + inner: VortexOpenOptions, +} + +impl CudaOpenOptions { + /// Open a local Vortex file for CUDA execution. + /// + /// The footer and zone-map segments are read through the ordinary host path. All other file + /// segments are staged through the session-scoped pinned buffer pool and transferred to the + /// device. + pub async fn open_path(self, path: impl AsRef) -> VortexResult { + let path = path.as_ref().to_owned(); + register_cuda_layout(self.inner.session()); + + // A host cache cannot front the data source: cached host buffers would violate the CUDA + // source's device-placement contract. It remains enabled on the control-plane source. + let data_options = self.inner.clone().without_segment_cache(); + let control_file = self.inner.open_path(&path).await?; + let footer = control_file.footer().clone(); + + let session = control_file.session().clone(); + let cuda_session = session.cuda_session(); + let stream = cuda_session.stream()?; + let pool = Arc::clone(cuda_session.pinned_buffer_pool()); + drop(cuda_session); + + let reader = Arc::new(PooledFileReadAt::open( + &path, + session.handle(), + pool, + stream, + )?); + let data_file = data_options + .with_footer(footer.clone()) + .open(reader) + .await?; + + let control_segments = + control_plane_segments(footer.layout().as_ref(), footer.segment_map().len())?; + let routed = Arc::new(RoutingSegmentSource { + control_segments: control_segments.into(), + control: control_file.segment_source(), + data: data_file.segment_source(), + }); + + Ok(data_file.with_segment_source(routed)) + } +} + +struct RoutingSegmentSource { + control_segments: Arc<[bool]>, + control: Arc, + data: Arc, +} + +impl SegmentSource for RoutingSegmentSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + let Ok(idx) = usize::try_from(*id) else { + return self.data.request(id); + }; + if self.control_segments.get(idx).copied().unwrap_or(false) { + self.control.request(id) + } else { + self.data.request(id) + } + } +} + +fn control_plane_segments(layout: &dyn Layout, segment_count: usize) -> VortexResult> { + let mut control_segments = vec![false; segment_count]; + mark_zone_map_segments(layout, &mut control_segments)?; + Ok(control_segments) +} + +fn mark_zone_map_segments(layout: &dyn Layout, control_segments: &mut [bool]) -> VortexResult<()> { + if layout.is::() || layout.is::() { + mark_layout_segments(layout.child(1)?.as_ref(), control_segments)?; + mark_zone_map_segments(layout.child(0)?.as_ref(), control_segments)?; + return Ok(()); + } + + for child in layout.children()? { + mark_zone_map_segments(child.as_ref(), control_segments)?; + } + Ok(()) +} + +fn mark_layout_segments(layout: &dyn Layout, control_segments: &mut [bool]) -> VortexResult<()> { + for id in layout.segment_ids() { + let idx = usize::try_from(*id)?; + let Some(is_control) = control_segments.get_mut(idx) else { + vortex_bail!("layout references out-of-range segment {id}"); + }; + *is_control = true; + } + + for child in layout.children()? { + mark_layout_segments(child.as_ref(), control_segments)?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + + use futures::FutureExt; + use parking_lot::Mutex; + use vortex::array::buffer::BufferHandle; + use vortex::array::dtype::DType; + use vortex::array::dtype::Nullability::NonNullable; + use vortex::array::dtype::PType; + use vortex::array::dtype::StructFields; + use vortex::buffer::ByteBuffer; + use vortex::layout::IntoLayout; + use vortex::layout::layouts::flat::FlatLayout; + use vortex::layout::layouts::zoned::ZonedLayout; + use vortex::session::registry::ReadContext; + + use super::*; + + #[derive(Clone)] + struct RecordingSource { + requests: Arc>>, + value: u8, + } + + impl RecordingSource { + fn new(value: u8) -> Self { + Self { + requests: Arc::new(Mutex::new(Vec::new())), + value, + } + } + } + + impl SegmentSource for RecordingSource { + fn request(&self, id: SegmentId) -> SegmentFuture { + self.requests.lock().push(id); + let value = self.value; + async move { Ok(BufferHandle::new_host(ByteBuffer::from(vec![value]))) }.boxed() + } + } + + #[tokio::test] + async fn routes_control_and_data_segments() { + let control = RecordingSource::new(1); + let data = RecordingSource::new(2); + let source = RoutingSegmentSource { + control_segments: Arc::from([false, true, false]), + control: Arc::new(control.clone()), + data: Arc::new(data.clone()), + }; + + let control_buffer = source + .request(SegmentId::from(1)) + .await + .unwrap() + .unwrap_host(); + let data_buffer = source + .request(SegmentId::from(2)) + .await + .unwrap() + .unwrap_host(); + + assert_eq!(control_buffer.as_ref(), &[1]); + assert_eq!(data_buffer.as_ref(), &[2]); + assert_eq!(*control.requests.lock(), [SegmentId::from(1)]); + assert_eq!(*data.requests.lock(), [SegmentId::from(2)]); + } + + #[test] + fn marks_only_zone_map_subtree_segments_as_control_plane() -> VortexResult<()> { + let read_ctx = ReadContext::new([]); + let data_dtype = DType::Primitive(PType::I32, NonNullable); + let data = + FlatLayout::new(16, data_dtype, SegmentId::from(0), read_ctx.clone()).into_layout(); + let zones = FlatLayout::new( + 1, + DType::Struct(StructFields::empty(), NonNullable), + SegmentId::from(1), + read_ctx, + ) + .into_layout(); + let layout = + ZonedLayout::try_new(data, zones, NonZeroUsize::new(16).unwrap(), Arc::new([]))? + .into_layout(); + + let control_segments = control_plane_segments(layout.as_ref(), 2)?; + + assert_eq!(control_segments, [false, true]); + Ok(()) + } +} diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 5817e15436b..69a2f7263eb 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -14,6 +14,7 @@ mod device_buffer; mod device_read_at; pub mod dynamic_dispatch; pub mod executor; +mod file; pub mod hybrid_dispatch; mod kernel; pub mod layout; @@ -34,6 +35,8 @@ pub use device_read_at::CopyDeviceReadAt; pub use executor::CudaDispatchMode; pub use executor::CudaExecutionCtx; pub use executor::CudaKernelEvents; +pub use file::CudaOpenOptions; +pub use file::CudaOpenOptionsExt; use kernel::ALPExecutor; use kernel::BitPackedExecutor; use kernel::ConstantNumericExecutor; diff --git a/vortex-cuda/src/pooled_read_at.rs b/vortex-cuda/src/pooled_read_at.rs index 1b5ec4fb31a..d3024140c96 100644 --- a/vortex-cuda/src/pooled_read_at.rs +++ b/vortex-cuda/src/pooled_read_at.rs @@ -40,6 +40,9 @@ pub const DEFAULT_OBJECT_STORE_CONCURRENCY: usize = 192; /// /// Reads into a pooled pinned (page-locked) buffer, then submits a non-blocking /// H2D DMA transfer and returns a device `BufferHandle`. +/// +/// This is a data-plane reader. To open a complete local Vortex file, prefer +/// [`crate::CudaOpenOptionsExt::with_cuda`], which keeps the footer and zone maps on the host. #[derive(Clone)] pub struct PooledFileReadAt { uri: Arc, diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index c9a71f85c55..2fc27039c06 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -125,6 +125,18 @@ impl VortexFile { Arc::clone(&self.segment_source) } + /// Replace the segment source used by this file. + /// + /// Any cached layout reader is cleared so that subsequent scans construct readers over the + /// replacement source. + pub fn with_segment_source(mut self, segment_source: Arc) -> Self { + self.segment_source = segment_source; + if self.layout_reader_cache.is_some() { + self.layout_reader_cache = Some(OnceLock::new()); + } + self + } + /// Returns a reference to the Vortex session used to open this file. pub fn session(&self) -> &VortexSession { &self.session diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index 86df2b6d4e8..d5d60b700e5 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -71,6 +71,23 @@ pub struct VortexOpenOptions { cache_layout_reader: bool, } +impl Clone for VortexOpenOptions { + fn clone(&self) -> Self { + Self { + session: self.session.clone(), + segment_cache: self.segment_cache.clone(), + initial_read_size: self.initial_read_size, + file_size: self.file_size, + dtype: self.dtype.clone(), + footer: self.footer.clone(), + initial_read_segments: RwLock::new(self.initial_read_segments.read().clone()), + metrics_registry: self.metrics_registry.clone(), + labels: self.labels.clone(), + cache_layout_reader: self.cache_layout_reader, + } + } +} + /// Extension trait for constructing [`VortexOpenOptions`] from a session. pub trait OpenOptionsSessionExt: ArraySessionExt + LayoutSessionExt + RuntimeSessionExt + MemorySessionExt @@ -97,6 +114,11 @@ impl &VortexSession { + &self.session + } + /// Configure how many bytes to read from the end of the file before parsing the footer. /// /// The actual read is at least large enough to contain the maximum postscript and EOF marker, @@ -125,6 +147,15 @@ impl VortexOpenOptions { self } + /// Disable the configured segment cache. + /// + /// This is useful when deriving an opener for a source whose buffers have different memory + /// placement requirements from the configured host cache. + pub fn without_segment_cache(mut self) -> Self { + self.segment_cache = None; + self + } + /// Configure a known file size. /// /// This helps to prevent an I/O request to discover the size of the file.