Skip to content

Commit cbbedfc

Browse files
committed
perf(cuda): copy cached file pages directly on the GPU
On devices that read pageable memory through host page tables, copy resident mapped-file ranges directly into device memory. Preserve staged reads for uncached ranges and direct I/O, source and destination alignment, cancellation safety, and mapping lifetime guards. Use the existing file-read chunk size so this feature is independent of chunk-size tuning. Original-commit: b546d93814 Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>
1 parent e083440 commit cbbedfc

6 files changed

Lines changed: 403 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
@@ -28,6 +28,7 @@ cudarc = { workspace = true, features = ["f16"] }
2828
futures = { workspace = true, features = ["executor"] }
2929
itertools = { workspace = true }
3030
kanal = { workspace = true }
31+
memmap2 = { workspace = true }
3132
num-traits = { workspace = true }
3233
object_store = { workspace = true, features = ["fs"] }
3334
parking_lot = { 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: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,31 @@
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+
use parking_lot::Mutex;
1829
use tokio::sync::Semaphore;
1930
use vortex::array::buffer::BufferHandle;
2031
use vortex::buffer::Alignment;
32+
use vortex::error::VortexExpect;
2133
use vortex::error::VortexResult;
2234
use vortex::error::vortex_ensure;
2335
use vortex::error::vortex_err;
@@ -29,6 +41,7 @@ use vortex::io::std_file::read_exact_at;
2941
#[cfg(target_os = "linux")]
3042
use self::direct::DirectFileReadBackend;
3143
use crate::CudaDeviceBuffer;
44+
use crate::CudaSession;
3245
use crate::pinned::PinnedByteBufferPool;
3346
use crate::pinned::PooledPinnedBuffer;
3447
use crate::stream::VortexCudaStream;
@@ -40,6 +53,15 @@ pub const DEFAULT_FILE_CONCURRENCY: usize = 32;
4053
// to limit blocking read sizes and start transfers before the whole segment is read.
4154
const FILE_READ_CHUNK_BYTES: usize = 4 << 20;
4255

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

83+
pub(crate) fn is_direct_io(self) -> bool {
84+
self.direct_io
85+
}
86+
6187
fn open(self, path: &Path) -> VortexResult<Arc<dyn FileReadBackend>> {
6288
#[cfg(target_os = "linux")]
6389
if self.direct_io {
@@ -103,6 +129,164 @@ impl FileReadBackend for File {
103129
}
104130
}
105131

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

125311
impl PooledFileReadAt {
@@ -157,9 +343,48 @@ impl PooledFileReadAt {
157343
pool,
158344
stream,
159345
read_slots: Arc::new(Semaphore::new(DEFAULT_FILE_CONCURRENCY)),
346+
mapped: None,
160347
})
161348
}
162349

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

0 commit comments

Comments
 (0)