perf(cuda): copy cached file pages directly on the GPU

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"] }

vortex-cuda/kernels/src/mapped_copy.cu

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#include <cstdint>
+
+// 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<uint64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    const uint64_t stride = static_cast<uint64_t>(gridDim.x) * blockDim.x;
+
+    const uint64_t misalignment = reinterpret_cast<uintptr_t>(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<const uint4 *>(source + head);
+    uint4 *const target_words = reinterpret_cast<uint4 *>(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];
+    }
+}

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)

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<Arc<dyn FileReadBackend>> {
         #[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<Mmap>,
+    copy: CudaFunction,
+    stream: VortexCudaStream,
+}
+
+impl MappedFile {
+    fn open(
+        path: &Path,
+        stream: &VortexCudaStream,
+        session: &CudaSession,
+    ) -> VortexResult<Option<Self>> {
+        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<Range<usize>> {
+        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<usize>,
+        mut output: CudaSlice<u8>,
+    ) -> VortexResult<CudaSlice<u8>> {
+        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<PinnedByteBufferPool>,
     stream: VortexCudaStream,
     read_slots: Arc<Semaphore>,
+    /// Set when buffered reads are copied by the device straight from the page cache.
+    mapped: Option<Arc<MappedFile>>,
 }
 
 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> {
+        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<Option<BufferHandle>> {
+        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::<u8>(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<PooledHostRead> {
@@ -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)?;