diff --git a/Cargo.lock b/Cargo.lock index a471337bea4..c52c1b921c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11741,6 +11741,7 @@ dependencies = [ "futures", "itertools 0.14.0", "kanal", + "memmap2", "mimalloc", "num-traits", "object_store 0.13.2", @@ -11748,6 +11749,7 @@ dependencies = [ "prost 0.14.4", "rstest", "rustix", + "tempfile", "tokio", "tracing", "vortex", diff --git a/vortex-cuda/Cargo.toml b/vortex-cuda/Cargo.toml index eaffffd1251..d1e54b7a399 100644 --- a/vortex-cuda/Cargo.toml +++ b/vortex-cuda/Cargo.toml @@ -28,6 +28,7 @@ cudarc = { workspace = true, features = ["f16"] } futures = { workspace = true, features = ["executor"] } itertools = { workspace = true } kanal = { workspace = true } +memmap2 = { workspace = true } num-traits = { workspace = true } object_store = { workspace = true, features = ["fs"] } parking_lot = { workspace = true } @@ -51,6 +52,7 @@ criterion = { workspace = true } futures = { workspace = true, features = ["executor"] } mimalloc = { workspace = true } rstest = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["rt", "macros"] } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-cuda = { path = ".", features = ["_test-harness"] } diff --git a/vortex-cuda/kernels/src/mapped_copy.cu b/vortex-cuda/kernels/src/mapped_copy.cu new file mode 100644 index 00000000000..2075f21554d --- /dev/null +++ b/vortex-cuda/kernels/src/mapped_copy.cu @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#include + +// Copies host memory that the device reads through the host page tables (for example a mapped +// file on a coherent CPU-GPU system) into device memory. Callers align source and target to +// the same offset within 16 bytes, so the body moves 16-byte words. +extern "C" __global__ void +mapped_copy(const uint8_t *__restrict__ source, uint8_t *__restrict__ target, uint64_t len) { + const uint64_t thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const uint64_t stride = static_cast(gridDim.x) * blockDim.x; + + const uint64_t misalignment = reinterpret_cast(source) % 16; + const uint64_t head = misalignment == 0 ? 0 : (16 - misalignment < len ? 16 - misalignment : len); + for (uint64_t i = thread; i < head; i += stride) { + target[i] = source[i]; + } + + const uint64_t words = (len - head) / 16; + const uint4 *const source_words = reinterpret_cast(source + head); + uint4 *const target_words = reinterpret_cast(target + head); + for (uint64_t i = thread; i < words; i += stride) { + target_words[i] = source_words[i]; + } + + for (uint64_t i = head + words * 16 + thread; i < len; i += stride) { + target[i] = source[i]; + } +} diff --git a/vortex-cuda/src/file.rs b/vortex-cuda/src/file.rs index 2a9806729f7..fa8f9398373 100644 --- a/vortex-cuda/src/file.rs +++ b/vortex-cuda/src/file.rs @@ -76,15 +76,21 @@ impl CudaOpenOptions { 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_with_options( + let reader = PooledFileReadAt::open_with_options( &path, session.handle(), pool, stream, self.read_at_options, - )?); + )?; + // Direct I/O bypasses the page cache, so there are no cached pages to copy from. + let reader = if self.read_at_options.is_direct_io() { + reader + } else { + reader.with_mapped_copy(&path, &cuda_session)? + }; + drop(cuda_session); + let reader = Arc::new(reader); let data_file = data_options .with_footer(footer.clone()) .open(reader) diff --git a/vortex-cuda/src/pooled_read_at/file.rs b/vortex-cuda/src/pooled_read_at/file.rs index 6d97f1c6064..c212610b3c6 100644 --- a/vortex-cuda/src/pooled_read_at/file.rs +++ b/vortex-cuda/src/pooled_read_at/file.rs @@ -5,19 +5,32 @@ mod direct; use std::fs::File; +use std::mem::ManuallyDrop; use std::ops::Range; use std::path::Path; use std::sync::Arc; +use cudarc::driver::CudaFunction; +use cudarc::driver::CudaSlice; +use cudarc::driver::DevicePtrMut; +use cudarc::driver::LaunchConfig; +use cudarc::driver::PushKernelArg; +use cudarc::driver::sys::CUdevice_attribute; use futures::FutureExt; use futures::StreamExt; use futures::TryFutureExt; use futures::TryStreamExt; use futures::future::BoxFuture; use futures::stream; +#[cfg(target_os = "linux")] +use memmap2::Advice; +use memmap2::Mmap; +#[cfg(target_os = "linux")] +use parking_lot::Mutex; use tokio::sync::Semaphore; use vortex::array::buffer::BufferHandle; use vortex::buffer::Alignment; +use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; @@ -29,6 +42,7 @@ use vortex::io::std_file::read_exact_at; #[cfg(target_os = "linux")] use self::direct::DirectFileReadBackend; use crate::CudaDeviceBuffer; +use crate::CudaSession; use crate::pinned::PinnedByteBufferPool; use crate::pinned::PooledPinnedBuffer; use crate::stream::VortexCudaStream; @@ -40,6 +54,15 @@ pub const DEFAULT_FILE_CONCURRENCY: usize = 32; // to limit blocking read sizes and start transfers before the whole segment is read. const FILE_READ_CHUNK_BYTES: usize = 4 << 20; +/// Upper bound on blocks per mapped copy launch: a few per SM saturates host-memory reads. +const MAPPED_COPY_BLOCKS_MAX: u64 = 1024; +const MAPPED_COPY_THREADS: u32 = 256; + +/// Filling page tables from several threads contends on the process memory map and measured +/// slower than a single thread, so fills are serialized across all mapped files. +#[cfg(target_os = "linux")] +static MAPPED_FILL: Mutex<()> = Mutex::new(()); + /// Options controlling how [`PooledFileReadAt`] opens and reads a local file. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct PooledFileReadAtOptions { @@ -58,6 +81,10 @@ impl PooledFileReadAtOptions { self } + pub(crate) fn is_direct_io(self) -> bool { + self.direct_io + } + fn open(self, path: &Path) -> VortexResult> { #[cfg(target_os = "linux")] if self.direct_io { @@ -103,6 +130,164 @@ impl FileReadBackend for File { } } +unsafe extern "C" { + fn mincore(addr: *mut std::ffi::c_void, length: usize, vec: *mut u8) -> std::ffi::c_int; + fn getpagesize() -> std::ffi::c_int; +} + +/// A buffered file whose page-cache pages the device copies directly. +/// +/// A buffered `pread` copies cached pages into pinned staging on the CPU, the slowest step of a +/// warm read. Devices that read pageable memory through the host page tables can instead copy +/// from a mapping of those pages with a kernel. The CPU then only fills page tables, which is +/// several times faster than copying the data. +struct MappedFile { + /// Present until drop; unmapped only after the copy stream has drained. + mapping: Option, + copy: CudaFunction, + stream: VortexCudaStream, +} + +impl MappedFile { + fn open( + path: &Path, + stream: &VortexCudaStream, + session: &CudaSession, + ) -> VortexResult> { + if !cfg!(target_os = "linux") { + return Ok(None); + } + let uses_host_page_tables = stream + .context() + .attribute( + CUdevice_attribute::CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES, + ) + .map_err(|error| vortex_err!("Failed to query pageable memory access: {error}"))?; + if uses_host_page_tables != 1 { + return Ok(None); + } + let file = File::open(path)?; + // SAFETY: Vortex files are immutable while open. As with any file mapping, truncation + // by another process would fault on access. + let mapping = unsafe { Mmap::map(&file)? }; + let copy = session.load_function_with_suffixes("mapped_copy", &[])?; + Ok(Some(Self { + mapping: Some(mapping), + copy, + stream: stream.clone(), + })) + } + + fn mapping(&self) -> &Mmap { + self.mapping + .as_ref() + .vortex_expect("mapping is present until drop") + } + + /// Return the in-bounds file range if every overlapping page is in the page cache. + /// + /// Filling page tables for uncached pages reads them from storage one fault at a time on the + /// single fill thread, about 13x slower than parallel `pread`s on a cold Q6 scan. Uncached + /// ranges therefore take the staged path, which also caches them for later reads. + fn cached_range(&self, offset: u64, length: usize) -> Option> { + let mapping = self.mapping(); + let start = usize::try_from(offset).ok()?; + let range = start..start.checked_add(length)?; + let bytes = mapping.get(range.clone())?; + if bytes.is_empty() { + return Some(range); + } + // SAFETY: getpagesize has no preconditions. + let page = usize::try_from(unsafe { getpagesize() }).ok()?; + if !page.is_power_of_two() { + return None; + } + let begin = (bytes.as_ptr() as usize) & !(page - 1); + let end = bytes.as_ptr() as usize + bytes.len(); + let mut resident = vec![0u8; (end - begin).div_ceil(page)]; + // SAFETY: `begin..end` lies within the live mapping, rounded down to a page boundary, and + // `resident` holds one byte per page of that range. + let status = unsafe { + mincore( + begin as *mut std::ffi::c_void, + end - begin, + resident.as_mut_ptr(), + ) + }; + (status == 0 && resident.iter().all(|page| page & 1 == 1)).then_some(range) + } + + /// Copy a range from [`Self::cached_range`] into its padded device allocation. + /// + /// Works in chunks so filling the page tables of one chunk overlaps the device copy of the + /// previous one. Returns `output` with the copies recorded as its pending write. + fn copy_to_device( + &self, + range: Range, + mut output: CudaSlice, + ) -> VortexResult> { + let mapping = self.mapping(); + let start = range.start; + let length = range.len(); + let pad = start % 16; + self.stream + .context() + .bind_to_thread() + .map_err(|error| vortex_err!("Failed to bind CUDA context: {error}"))?; + // The guard records the queued copies as the output's write when this scope ends, + // including on error, so frees and readers on other streams wait for them. + let (base, write) = output.device_ptr_mut(&self.stream); + for chunk in (0..length).step_by(FILE_READ_CHUNK_BYTES) { + let size = (length - chunk).min(FILE_READ_CHUNK_BYTES); + #[cfg(target_os = "linux")] + { + let _fill = MAPPED_FILL.lock(); + mapping.advise_range(Advice::PopulateRead, start + chunk, size)?; + } + let source = mapping.as_ptr() as u64 + (start + chunk) as u64; + let target = base + (pad + chunk) as u64; + let len = size as u64; + let blocks = len + .div_ceil(u64::from(MAPPED_COPY_THREADS) * 16) + .clamp(1, MAPPED_COPY_BLOCKS_MAX); + let config = LaunchConfig { + grid_dim: (u32::try_from(blocks)?, 1, 1), + block_dim: (MAPPED_COPY_THREADS, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.copy); + launch.arg(&source).arg(&target).arg(&len); + // SAFETY: The source range is mapped and in bounds, the target covers `pad + length` + // bytes, and both share the same offset within 16 bytes, as the kernel requires. + unsafe { launch.launch(config) } + .map_err(|error| vortex_err!("Failed to launch mapped copy: {error}"))?; + } + drop(write); + Ok(output) + } +} + +impl Drop for MappedFile { + fn drop(&mut self) { + let Some(mapping) = self.mapping.take() else { + return; + }; + // Queued copies may still read the mapping. + if self.stream.synchronize().is_err() { + // Without proof the copies finished, never unmap memory the device may read. + let _leaked = ManuallyDrop::new(mapping); + return; + } + // Unmapping filled page tables costs milliseconds per GiB; keep it off the caller. If + // the thread cannot start, the returned closure unmaps here. + drop( + std::thread::Builder::new() + .name("vortex-unmap".to_string()) + .spawn(move || drop(mapping)), + ); + } +} + /// File reader that uses CUDA pinned host memory for I/O buffers and transfers /// directly to the GPU. /// @@ -120,6 +305,8 @@ pub struct PooledFileReadAt { pool: Arc, stream: VortexCudaStream, read_slots: Arc, + /// Set when buffered reads are copied by the device straight from the page cache. + mapped: Option>, } impl PooledFileReadAt { @@ -157,9 +344,48 @@ impl PooledFileReadAt { pool, stream, read_slots: Arc::new(Semaphore::new(DEFAULT_FILE_CONCURRENCY)), + mapped: None, }) } + /// Let the device copy buffered reads straight from the page cache when it reads pageable + /// host memory through the host page tables. Otherwise the reader is unchanged. + pub(crate) fn with_mapped_copy( + mut self, + path: &Path, + session: &CudaSession, + ) -> VortexResult { + self.mapped = MappedFile::open(path, &self.stream, session)?.map(Arc::new); + Ok(self) + } + + /// Copy a cached range, or return `None` to use staged reads. + /// + /// Residency checks and allocation stay on the submitting future; only page-table fills and + /// copy launches run on a blocking worker. + async fn read_mapped(&self, offset: u64, length: usize) -> VortexResult> { + let Some(mapped) = &self.mapped else { + return Ok(None); + }; + let Some(range) = mapped.cached_range(offset, length) else { + return Ok(None); + }; + // Start the target at the source's offset within 16 bytes so the copy moves aligned + // words. The slice keeps the file's alignment, which the caller's alignment follows. + let pad = range.start % 16; + let output = self.stream.device_alloc::(pad + length)?; + // The job owns the mapping and output until its copies are queued, even if cancelled. + let mapped = Arc::clone(mapped); + let output = self + .handle + .spawn_blocking(move || mapped.copy_to_device(range, output)) + .await?; + Ok(Some( + BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output))) + .slice(pad..pad + length), + )) + } + /// Read into pinned memory under the file's shared concurrency limit. /// The returned `requested_range` excludes any backend alignment padding. async fn read_host(&self, offset: u64, length: usize) -> VortexResult { @@ -210,6 +436,9 @@ impl VortexReadAt for PooledFileReadAt { offset.checked_add(u64::try_from(length)?).is_some(), "file read range overflow: offset={offset}, length={length}" ); + if let Some(buffer) = reader.read_mapped(offset, length).await? { + return Ok(buffer); + } if length <= FILE_READ_CHUNK_BYTES { let read = reader.read_host(offset, length).await?; let cuda_buf = read.buffer.transfer_to_device(&reader.stream)?; diff --git a/vortex-cuda/src/pooled_read_at/file/tests.rs b/vortex-cuda/src/pooled_read_at/file/tests.rs index c3eb9d4cdc6..2c389aff022 100644 --- a/vortex-cuda/src/pooled_read_at/file/tests.rs +++ b/vortex-cuda/src/pooled_read_at/file/tests.rs @@ -136,6 +136,7 @@ fn reader( pool: Arc::clone(cuda.pinned_buffer_pool()), stream: cuda.stream()?, read_slots: Arc::new(Semaphore::new(DEFAULT_FILE_CONCURRENCY)), + mapped: None, }; Ok((runtime, reader)) } @@ -309,3 +310,133 @@ async fn cancelled_reads_keep_slots_and_buffers_until_blocking_io_finishes() -> assert_eq!(backend.requests.lock().len(), 3); Ok(()) } + +const MAPPED_FILE_BYTES: usize = 2 * FILE_READ_CHUNK_BYTES + 4096; + +/// Opens a real buffered file of `file_byte` content with the device-side mapped copy enabled. +fn mapped_reader() -> VortexResult<( + CurrentThreadRuntime, + PooledFileReadAt, + tempfile::NamedTempFile, +)> { + let mut file = tempfile::NamedTempFile::new()?; + let bytes: Vec = (0..MAPPED_FILE_BYTES as u64).map(file_byte).collect(); + io::Write::write_all(&mut file, &bytes)?; + let runtime = CurrentThreadRuntime::new(); + let session = crate::cuda_session().with_handle(runtime.handle()); + let cuda = session.cuda_session(); + let reader = PooledFileReadAt::open( + file.path(), + session.handle(), + Arc::clone(cuda.pinned_buffer_pool()), + cuda.stream()?, + )? + .with_mapped_copy(file.path(), &cuda)?; + Ok((runtime, reader, file)) +} + +/// Mapped reads must return the file's bytes for any source offset within 16 bytes, across +/// chunk boundaries and up to the end of the file. +#[rstest] +#[case::empty(37, 0)] +#[case::single_byte_unaligned(5, 1)] +#[case::short_unaligned(15, 17)] +#[case::chunk_aligned(0, FILE_READ_CHUNK_BYTES - 1)] +#[case::chunks_with_tail(4091, 2 * FILE_READ_CHUNK_BYTES + 3)] +#[case::through_file_end(0, MAPPED_FILE_BYTES)] +#[crate::test] +async fn mapped_reads_preserve_bytes_and_offsets( + #[case] offset: u64, + #[case] length: usize, +) -> VortexResult<()> { + let (_runtime, reader, _file) = mapped_reader()?; + let Some(mapped) = &reader.mapped else { + // The device does not read pageable memory; the staged path is covered above. + return Ok(()); + }; + #[cfg(target_os = "linux")] + mapped.mapping().advise(Advice::PopulateRead)?; + let start = usize::try_from(offset)?; + assert_eq!( + mapped.cached_range(offset, length), + Some(start..start + length) + ); + // A staged fallback must fail rather than silently bypass mapped-copy coverage. + reader.read_slots.close(); + let buffer = reader + .read_at(offset, length, Alignment::of::()) + .await?; + assert_eq!(buffer.len(), length); + if length > 0 { + assert_bytes(buffer, offset, length).await?; + } + Ok(()) +} + +#[rstest] +#[case::empty_at_start(0, 0, Some(0..0))] +#[case::empty_at_end( + MAPPED_FILE_BYTES as u64, + 0, + Some(MAPPED_FILE_BYTES..MAPPED_FILE_BYTES) +)] +#[case::empty_past_end(MAPPED_FILE_BYTES as u64 + 1, 0, None)] +#[case::past_end(MAPPED_FILE_BYTES as u64 - 10, 11, None)] +#[case::offset_overflow(u64::MAX, 1, None)] +#[case::length_overflow(1, usize::MAX, None)] +#[crate::test] +async fn mapped_ranges_validate_bounds( + #[case] offset: u64, + #[case] length: usize, + #[case] expected: Option>, +) -> VortexResult<()> { + let (_runtime, reader, _file) = mapped_reader()?; + let Some(mapped) = &reader.mapped else { + return Ok(()); + }; + assert_eq!(mapped.cached_range(offset, length), expected); + Ok(()) +} + +#[crate::test] +async fn mapped_read_past_end_fails() -> VortexResult<()> { + let (_runtime, reader, _file) = mapped_reader()?; + let Some(mapped) = &reader.mapped else { + return Ok(()); + }; + let offset = MAPPED_FILE_BYTES as u64 - 10; + assert!(mapped.cached_range(offset, 11).is_none()); + // Out-of-range reads take the staged path, which reports the short read. + reader + .read_at(offset, 11, Alignment::of::()) + .await + .expect_err("a read past the end must fail"); + Ok(()) +} + +/// Uncached ranges must not be filled through the mapping, where cold page faults are +/// serialized; the staged path reads them and leaves them cached for the mapped path. +#[cfg(target_os = "linux")] +#[crate::test] +async fn uncached_ranges_take_the_staged_path() -> VortexResult<()> { + let (_runtime, reader, file) = mapped_reader()?; + let Some(mapped) = reader.mapped.as_ref() else { + return Ok(()); + }; + file.as_file().sync_all()?; + rustix::fs::fadvise(file.as_file(), 0, None, rustix::fs::Advice::DontNeed) + .map_err(io::Error::from)?; + let (offset, length) = (37, FILE_READ_CHUNK_BYTES + 5); + assert!(mapped.cached_range(offset, length).is_none()); + + let buffer = reader + .read_at(offset, length, Alignment::of::()) + .await?; + assert_bytes(buffer, offset, length).await?; + let start = usize::try_from(offset)?; + assert_eq!( + mapped.cached_range(offset, length), + Some(start..start + length) + ); + Ok(()) +}