Skip to content

Commit d006295

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 a1fdb64 commit d006295

6 files changed

Lines changed: 395 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: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,25 @@ use std::ops::Range;
99
use std::path::Path;
1010
use std::sync::Arc;
1111

12+
use cudarc::driver::CudaFunction;
13+
use cudarc::driver::CudaSlice;
14+
use cudarc::driver::DevicePtrMut;
15+
use cudarc::driver::LaunchConfig;
16+
use cudarc::driver::PushKernelArg;
17+
use cudarc::driver::sys::CUdevice_attribute;
1218
use futures::FutureExt;
1319
use futures::StreamExt;
1420
use futures::TryFutureExt;
1521
use futures::TryStreamExt;
1622
use futures::future::BoxFuture;
1723
use futures::stream;
24+
use memmap2::Advice;
25+
use memmap2::Mmap;
26+
use parking_lot::Mutex;
1827
use tokio::sync::Semaphore;
1928
use vortex::array::buffer::BufferHandle;
2029
use vortex::buffer::Alignment;
30+
use vortex::error::VortexExpect;
2131
use vortex::error::VortexResult;
2232
use vortex::error::vortex_ensure;
2333
use vortex::error::vortex_err;
@@ -29,6 +39,7 @@ use vortex::io::std_file::read_exact_at;
2939
#[cfg(target_os = "linux")]
3040
use self::direct::DirectFileReadBackend;
3141
use crate::CudaDeviceBuffer;
42+
use crate::CudaSession;
3243
use crate::pinned::PinnedByteBufferPool;
3344
use crate::pinned::PooledPinnedBuffer;
3445
use crate::stream::VortexCudaStream;
@@ -40,6 +51,14 @@ pub const DEFAULT_FILE_CONCURRENCY: usize = 32;
4051
// to limit blocking read sizes and start transfers before the whole segment is read.
4152
const FILE_READ_CHUNK_BYTES: usize = 4 << 20;
4253

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

80+
pub(crate) fn is_direct_io(self) -> bool {
81+
self.direct_io
82+
}
83+
6184
fn open(self, path: &Path) -> VortexResult<Arc<dyn FileReadBackend>> {
6285
#[cfg(target_os = "linux")]
6386
if self.direct_io {
@@ -103,6 +126,160 @@ impl FileReadBackend for File {
103126
}
104127
}
105128

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

125304
impl PooledFileReadAt {
@@ -157,9 +336,48 @@ impl PooledFileReadAt {
157336
pool,
158337
stream,
159338
read_slots: Arc::new(Semaphore::new(DEFAULT_FILE_CONCURRENCY)),
339+
mapped: None,
160340
})
161341
}
162342

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

0 commit comments

Comments
 (0)