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+ use parking_lot:: Mutex ;
1829use tokio:: sync:: Semaphore ;
1930use vortex:: array:: buffer:: BufferHandle ;
2031use vortex:: buffer:: Alignment ;
32+ use vortex:: error:: VortexExpect ;
2133use vortex:: error:: VortexResult ;
2234use vortex:: error:: vortex_ensure;
2335use vortex:: error:: vortex_err;
@@ -29,6 +41,7 @@ use vortex::io::std_file::read_exact_at;
2941#[ cfg( target_os = "linux" ) ]
3042use self :: direct:: DirectFileReadBackend ;
3143use crate :: CudaDeviceBuffer ;
44+ use crate :: CudaSession ;
3245use crate :: pinned:: PinnedByteBufferPool ;
3346use crate :: pinned:: PooledPinnedBuffer ;
3447use 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.
4154const 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 ) ]
4567pub 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
125311impl 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