55mod direct;
66
77use std:: fs:: File ;
8+ use std:: mem:: ManuallyDrop ;
89use std:: ops:: Range ;
910use std:: path:: Path ;
1011use 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 ;
1219use futures:: FutureExt ;
1320use futures:: StreamExt ;
1421use futures:: TryFutureExt ;
1522use futures:: TryStreamExt ;
1623use futures:: future:: BoxFuture ;
1724use 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 ;
1830use tokio:: sync:: Semaphore ;
1931use vortex:: array:: buffer:: BufferHandle ;
2032use vortex:: buffer:: Alignment ;
33+ use vortex:: error:: VortexExpect ;
2134use vortex:: error:: VortexResult ;
2235use vortex:: error:: vortex_ensure;
2336use vortex:: error:: vortex_err;
@@ -29,6 +42,7 @@ use vortex::io::std_file::read_exact_at;
2942#[ cfg( target_os = "linux" ) ]
3043use self :: direct:: DirectFileReadBackend ;
3144use crate :: CudaDeviceBuffer ;
45+ use crate :: CudaSession ;
3246use crate :: pinned:: PinnedByteBufferPool ;
3347use crate :: pinned:: PooledPinnedBuffer ;
3448use 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.
4155const 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 ) ]
4568pub 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
125312impl 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