@@ -9,15 +9,25 @@ use std::ops::Range;
99use std:: path:: Path ;
1010use 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 ;
1218use futures:: FutureExt ;
1319use futures:: StreamExt ;
1420use futures:: TryFutureExt ;
1521use futures:: TryStreamExt ;
1622use futures:: future:: BoxFuture ;
1723use futures:: stream;
24+ use memmap2:: Advice ;
25+ use memmap2:: Mmap ;
26+ use parking_lot:: Mutex ;
1827use tokio:: sync:: Semaphore ;
1928use vortex:: array:: buffer:: BufferHandle ;
2029use vortex:: buffer:: Alignment ;
30+ use vortex:: error:: VortexExpect ;
2131use vortex:: error:: VortexResult ;
2232use vortex:: error:: vortex_ensure;
2333use vortex:: error:: vortex_err;
@@ -29,6 +39,7 @@ use vortex::io::std_file::read_exact_at;
2939#[ cfg( target_os = "linux" ) ]
3040use self :: direct:: DirectFileReadBackend ;
3141use crate :: CudaDeviceBuffer ;
42+ use crate :: CudaSession ;
3243use crate :: pinned:: PinnedByteBufferPool ;
3344use crate :: pinned:: PooledPinnedBuffer ;
3445use 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.
4152const 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 ) ]
4564pub 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
125304impl 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