Skip to content

Commit 32390c0

Browse files
authored
perf(cuda): copy cached file pages directly on the GPU (#10716)
## Problem Cached file pages still pass through pinned staging buffers before reaching the GPU. ## Solution On supported Linux devices, copy resident mapped-file ranges directly with a CUDA kernel; retain staged reads for other cases. ## Assessment Hardware-dependent fast path with mapping lifetime and cancellation guards. Keeps 4 MiB chunks, so it is independent of #10715. Based on #10726. Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>
1 parent 3258a3b commit 32390c0

6 files changed

Lines changed: 404 additions & 4 deletions

File tree

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎vortex-cuda/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ cudarc = { workspace = true, features = ["f16"] }
2727
futures = { workspace = true, features = ["executor"] }
2828
itertools = { workspace = true }
2929
kanal = { workspace = true }
30+
memmap2 = { workspace = true }
3031
num-traits = { workspace = true }
3132
object_store = { workspace = true, features = ["fs"] }
3233
once_cell = { workspace = true }
@@ -51,6 +52,7 @@ criterion = { workspace = true }
5152
futures = { workspace = true, features = ["executor"] }
5253
mimalloc = { workspace = true }
5354
rstest = { workspace = true }
55+
tempfile = { workspace = true }
5456
tokio = { workspace = true, features = ["rt", "macros"] }
5557
vortex-array = { workspace = true, features = ["_test-harness"] }
5658
vortex-cuda = { path = ".", features = ["_test-harness"] }
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
#include <cstdint>
5+
6+
// Copies host memory that the device reads through the host page tables (for example a mapped
7+
// file on a coherent CPU-GPU system) into device memory. Callers align source and target to
8+
// the same offset within 16 bytes, so the body moves 16-byte words.
9+
extern "C" __global__ void
10+
mapped_copy(const uint8_t *__restrict__ source, uint8_t *__restrict__ target, uint64_t len) {
11+
const uint64_t thread = static_cast<uint64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
12+
const uint64_t stride = static_cast<uint64_t>(gridDim.x) * blockDim.x;
13+
14+
const uint64_t misalignment = reinterpret_cast<uintptr_t>(source) % 16;
15+
const uint64_t head = misalignment == 0 ? 0 : (16 - misalignment < len ? 16 - misalignment : len);
16+
for (uint64_t i = thread; i < head; i += stride) {
17+
target[i] = source[i];
18+
}
19+
20+
const uint64_t words = (len - head) / 16;
21+
const uint4 *const source_words = reinterpret_cast<const uint4 *>(source + head);
22+
uint4 *const target_words = reinterpret_cast<uint4 *>(target + head);
23+
for (uint64_t i = thread; i < words; i += stride) {
24+
target_words[i] = source_words[i];
25+
}
26+
27+
for (uint64_t i = head + words * 16 + thread; i < len; i += stride) {
28+
target[i] = source[i];
29+
}
30+
}

‎vortex-cuda/src/file.rs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -76,15 +76,21 @@ impl CudaOpenOptions {
7676
let cuda_session = session.cuda_session();
7777
let stream = cuda_session.stream()?;
7878
let pool = Arc::clone(cuda_session.pinned_buffer_pool());
79-
drop(cuda_session);
80-
81-
let reader = Arc::new(PooledFileReadAt::open_with_options(
79+
let reader = PooledFileReadAt::open_with_options(
8280
&path,
8381
session.handle(),
8482
pool,
8583
stream,
8684
self.read_at_options,
87-
)?);
85+
)?;
86+
// Direct I/O bypasses the page cache, so there are no cached pages to copy from.
87+
let reader = if self.read_at_options.is_direct_io() {
88+
reader
89+
} else {
90+
reader.with_mapped_copy(&path, &cuda_session)?
91+
};
92+
drop(cuda_session);
93+
let reader = Arc::new(reader);
8894
let data_file = data_options
8995
.with_footer(footer.clone())
9096
.open(reader)

‎vortex-cuda/src/pooled_read_at/file.rs‎

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,32 @@
55
mod direct;
66

77
use std::fs::File;
8+
use std::mem::ManuallyDrop;
89
use std::ops::Range;
910
use std::path::Path;
1011
use std::sync::Arc;
1112

13+
use cudarc::driver::CudaFunction;
14+
use cudarc::driver::CudaSlice;
15+
use cudarc::driver::DevicePtrMut;
16+
use cudarc::driver::LaunchConfig;
17+
use cudarc::driver::PushKernelArg;
18+
use cudarc::driver::sys::CUdevice_attribute;
1219
use futures::FutureExt;
1320
use futures::StreamExt;
1421
use futures::TryFutureExt;
1522
use futures::TryStreamExt;
1623
use futures::future::BoxFuture;
1724
use futures::stream;
25+
#[cfg(target_os = "linux")]
26+
use memmap2::Advice;
27+
use memmap2::Mmap;
28+
#[cfg(target_os = "linux")]
29+
use parking_lot::Mutex;
1830
use tokio::sync::Semaphore;
1931
use vortex::array::buffer::BufferHandle;
2032
use vortex::buffer::Alignment;
33+
use vortex::error::VortexExpect;
2134
use vortex::error::VortexResult;
2235
use vortex::error::vortex_ensure;
2336
use vortex::error::vortex_err;
@@ -29,6 +42,7 @@ use vortex::io::std_file::read_exact_at;
2942
#[cfg(target_os = "linux")]
3043
use self::direct::DirectFileReadBackend;
3144
use crate::CudaDeviceBuffer;
45+
use crate::CudaSession;
3246
use crate::pinned::PinnedByteBufferPool;
3347
use crate::pinned::PooledPinnedBuffer;
3448
use crate::stream::VortexCudaStream;
@@ -40,6 +54,15 @@ pub const DEFAULT_FILE_CONCURRENCY: usize = 32;
4054
// to limit blocking read sizes and start transfers before the whole segment is read.
4155
const FILE_READ_CHUNK_BYTES: usize = 8 << 20;
4256

57+
/// Upper bound on blocks per mapped copy launch: a few per SM saturates host-memory reads.
58+
const MAPPED_COPY_BLOCKS_MAX: u64 = 1024;
59+
const MAPPED_COPY_THREADS: u32 = 256;
60+
61+
/// Filling page tables from several threads contends on the process memory map and measured
62+
/// slower than a single thread, so fills are serialized across all mapped files.
63+
#[cfg(target_os = "linux")]
64+
static MAPPED_FILL: Mutex<()> = Mutex::new(());
65+
4366
/// Options controlling how [`PooledFileReadAt`] opens and reads a local file.
4467
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
4568
pub struct PooledFileReadAtOptions {
@@ -58,6 +81,10 @@ impl PooledFileReadAtOptions {
5881
self
5982
}
6083

84+
pub(crate) fn is_direct_io(self) -> bool {
85+
self.direct_io
86+
}
87+
6188
fn open(self, path: &Path) -> VortexResult<Arc<dyn FileReadBackend>> {
6289
#[cfg(target_os = "linux")]
6390
if self.direct_io {
@@ -103,6 +130,164 @@ impl FileReadBackend for File {
103130
}
104131
}
105132

133+
unsafe extern "C" {
134+
fn mincore(addr: *mut std::ffi::c_void, length: usize, vec: *mut u8) -> std::ffi::c_int;
135+
fn getpagesize() -> std::ffi::c_int;
136+
}
137+
138+
/// A buffered file whose page-cache pages the device copies directly.
139+
///
140+
/// A buffered `pread` copies cached pages into pinned staging on the CPU, the slowest step of a
141+
/// warm read. Devices that read pageable memory through the host page tables can instead copy
142+
/// from a mapping of those pages with a kernel. The CPU then only fills page tables, which is
143+
/// several times faster than copying the data.
144+
struct MappedFile {
145+
/// Present until drop; unmapped only after the copy stream has drained.
146+
mapping: Option<Mmap>,
147+
copy: CudaFunction,
148+
stream: VortexCudaStream,
149+
}
150+
151+
impl MappedFile {
152+
fn open(
153+
path: &Path,
154+
stream: &VortexCudaStream,
155+
session: &CudaSession,
156+
) -> VortexResult<Option<Self>> {
157+
if !cfg!(target_os = "linux") {
158+
return Ok(None);
159+
}
160+
let uses_host_page_tables = stream
161+
.context()
162+
.attribute(
163+
CUdevice_attribute::CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES,
164+
)
165+
.map_err(|error| vortex_err!("Failed to query pageable memory access: {error}"))?;
166+
if uses_host_page_tables != 1 {
167+
return Ok(None);
168+
}
169+
let file = File::open(path)?;
170+
// SAFETY: Vortex files are immutable while open. As with any file mapping, truncation
171+
// by another process would fault on access.
172+
let mapping = unsafe { Mmap::map(&file)? };
173+
let copy = session.load_function_with_suffixes("mapped_copy", &[])?;
174+
Ok(Some(Self {
175+
mapping: Some(mapping),
176+
copy,
177+
stream: stream.clone(),
178+
}))
179+
}
180+
181+
fn mapping(&self) -> &Mmap {
182+
self.mapping
183+
.as_ref()
184+
.vortex_expect("mapping is present until drop")
185+
}
186+
187+
/// Return the in-bounds file range if every overlapping page is in the page cache.
188+
///
189+
/// Filling page tables for uncached pages reads them from storage one fault at a time on the
190+
/// single fill thread, about 13x slower than parallel `pread`s on a cold Q6 scan. Uncached
191+
/// ranges therefore take the staged path, which also caches them for later reads.
192+
fn cached_range(&self, offset: u64, length: usize) -> Option<Range<usize>> {
193+
let mapping = self.mapping();
194+
let start = usize::try_from(offset).ok()?;
195+
let range = start..start.checked_add(length)?;
196+
let bytes = mapping.get(range.clone())?;
197+
if bytes.is_empty() {
198+
return Some(range);
199+
}
200+
// SAFETY: getpagesize has no preconditions.
201+
let page = usize::try_from(unsafe { getpagesize() }).ok()?;
202+
if !page.is_power_of_two() {
203+
return None;
204+
}
205+
let begin = (bytes.as_ptr() as usize) & !(page - 1);
206+
let end = bytes.as_ptr() as usize + bytes.len();
207+
let mut resident = vec![0u8; (end - begin).div_ceil(page)];
208+
// SAFETY: `begin..end` lies within the live mapping, rounded down to a page boundary, and
209+
// `resident` holds one byte per page of that range.
210+
let status = unsafe {
211+
mincore(
212+
begin as *mut std::ffi::c_void,
213+
end - begin,
214+
resident.as_mut_ptr(),
215+
)
216+
};
217+
(status == 0 && resident.iter().all(|page| page & 1 == 1)).then_some(range)
218+
}
219+
220+
/// Copy a range from [`Self::cached_range`] into its padded device allocation.
221+
///
222+
/// Works in chunks so filling the page tables of one chunk overlaps the device copy of the
223+
/// previous one. Returns `output` with the copies recorded as its pending write.
224+
fn copy_to_device(
225+
&self,
226+
range: Range<usize>,
227+
mut output: CudaSlice<u8>,
228+
) -> VortexResult<CudaSlice<u8>> {
229+
let mapping = self.mapping();
230+
let start = range.start;
231+
let length = range.len();
232+
let pad = start % 16;
233+
self.stream
234+
.context()
235+
.bind_to_thread()
236+
.map_err(|error| vortex_err!("Failed to bind CUDA context: {error}"))?;
237+
// The guard records the queued copies as the output's write when this scope ends,
238+
// including on error, so frees and readers on other streams wait for them.
239+
let (base, write) = output.device_ptr_mut(&self.stream);
240+
for chunk in (0..length).step_by(FILE_READ_CHUNK_BYTES) {
241+
let size = (length - chunk).min(FILE_READ_CHUNK_BYTES);
242+
#[cfg(target_os = "linux")]
243+
{
244+
let _fill = MAPPED_FILL.lock();
245+
mapping.advise_range(Advice::PopulateRead, start + chunk, size)?;
246+
}
247+
let source = mapping.as_ptr() as u64 + (start + chunk) as u64;
248+
let target = base + (pad + chunk) as u64;
249+
let len = size as u64;
250+
let blocks = len
251+
.div_ceil(u64::from(MAPPED_COPY_THREADS) * 16)
252+
.clamp(1, MAPPED_COPY_BLOCKS_MAX);
253+
let config = LaunchConfig {
254+
grid_dim: (u32::try_from(blocks)?, 1, 1),
255+
block_dim: (MAPPED_COPY_THREADS, 1, 1),
256+
shared_mem_bytes: 0,
257+
};
258+
let mut launch = self.stream.launch_builder(&self.copy);
259+
launch.arg(&source).arg(&target).arg(&len);
260+
// SAFETY: The source range is mapped and in bounds, the target covers `pad + length`
261+
// bytes, and both share the same offset within 16 bytes, as the kernel requires.
262+
unsafe { launch.launch(config) }
263+
.map_err(|error| vortex_err!("Failed to launch mapped copy: {error}"))?;
264+
}
265+
drop(write);
266+
Ok(output)
267+
}
268+
}
269+
270+
impl Drop for MappedFile {
271+
fn drop(&mut self) {
272+
let Some(mapping) = self.mapping.take() else {
273+
return;
274+
};
275+
// Queued copies may still read the mapping.
276+
if self.stream.synchronize().is_err() {
277+
// Without proof the copies finished, never unmap memory the device may read.
278+
let _leaked = ManuallyDrop::new(mapping);
279+
return;
280+
}
281+
// Unmapping filled page tables costs milliseconds per GiB; keep it off the caller. If
282+
// the thread cannot start, the returned closure unmaps here.
283+
drop(
284+
std::thread::Builder::new()
285+
.name("vortex-unmap".to_string())
286+
.spawn(move || drop(mapping)),
287+
);
288+
}
289+
}
290+
106291
/// File reader that uses CUDA pinned host memory for I/O buffers and transfers
107292
/// directly to the GPU.
108293
///
@@ -120,6 +305,8 @@ pub struct PooledFileReadAt {
120305
pool: Arc<PinnedByteBufferPool>,
121306
stream: VortexCudaStream,
122307
read_slots: Arc<Semaphore>,
308+
/// Set when buffered reads are copied by the device straight from the page cache.
309+
mapped: Option<Arc<MappedFile>>,
123310
}
124311

125312
impl PooledFileReadAt {
@@ -157,9 +344,48 @@ impl PooledFileReadAt {
157344
pool,
158345
stream,
159346
read_slots: Arc::new(Semaphore::new(DEFAULT_FILE_CONCURRENCY)),
347+
mapped: None,
160348
})
161349
}
162350

351+
/// Let the device copy buffered reads straight from the page cache when it reads pageable
352+
/// host memory through the host page tables. Otherwise the reader is unchanged.
353+
pub(crate) fn with_mapped_copy(
354+
mut self,
355+
path: &Path,
356+
session: &CudaSession,
357+
) -> VortexResult<Self> {
358+
self.mapped = MappedFile::open(path, &self.stream, session)?.map(Arc::new);
359+
Ok(self)
360+
}
361+
362+
/// Copy a cached range, or return `None` to use staged reads.
363+
///
364+
/// Residency checks and allocation stay on the submitting future; only page-table fills and
365+
/// copy launches run on a blocking worker.
366+
async fn read_mapped(&self, offset: u64, length: usize) -> VortexResult<Option<BufferHandle>> {
367+
let Some(mapped) = &self.mapped else {
368+
return Ok(None);
369+
};
370+
let Some(range) = mapped.cached_range(offset, length) else {
371+
return Ok(None);
372+
};
373+
// Start the target at the source's offset within 16 bytes so the copy moves aligned
374+
// words. The slice keeps the file's alignment, which the caller's alignment follows.
375+
let pad = range.start % 16;
376+
let output = self.stream.device_alloc::<u8>(pad + length)?;
377+
// The job owns the mapping and output until its copies are queued, even if cancelled.
378+
let mapped = Arc::clone(mapped);
379+
let output = self
380+
.handle
381+
.spawn_blocking(move || mapped.copy_to_device(range, output))
382+
.await?;
383+
Ok(Some(
384+
BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(output)))
385+
.slice(pad..pad + length),
386+
))
387+
}
388+
163389
/// Read into pinned memory under the file's shared concurrency limit.
164390
/// The returned `requested_range` excludes any backend alignment padding.
165391
async fn read_host(&self, offset: u64, length: usize) -> VortexResult<PooledHostRead> {
@@ -210,6 +436,9 @@ impl VortexReadAt for PooledFileReadAt {
210436
offset.checked_add(u64::try_from(length)?).is_some(),
211437
"file read range overflow: offset={offset}, length={length}"
212438
);
439+
if let Some(buffer) = reader.read_mapped(offset, length).await? {
440+
return Ok(buffer);
441+
}
213442
if length <= FILE_READ_CHUNK_BYTES {
214443
let read = reader.read_host(offset, length).await?;
215444
let cuda_buf = read.buffer.transfer_to_device(&reader.stream)?;

0 commit comments

Comments
 (0)