nuvec_core/file.rs
1//! Arrays backed by a file's bytes, for loaders (safetensors, npy, ...).
2
3use std::path::Path;
4use std::sync::Arc;
5
6use crate::array::Data;
7use crate::buffer::Buffer;
8use crate::{Array, DType, Error, Result, shape};
9
10/// A file's bytes in array storage: on a Mac with a GPU, the file is memory-mapped and
11/// wrapped as GPU-visible memory without a copy, so arrays made from it are views whose
12/// pages are read when first touched (by the CPU or the GPU). Elsewhere the file is read
13/// into memory. The mapping lives as long as any array made from it.
14///
15/// The file must not be modified (written or truncated, by this or another process) while
16/// it is mapped: arrays made from it would change underneath their readers (undefined
17/// behaviour), and reads past a truncated end raise `SIGBUS`. Map only files nothing
18/// writes while they are in use, such as model weights.
19///
20/// ```
21/// use nuvec_core::{DType, MappedFile};
22///
23/// # fn main() -> nuvec_core::Result<()> {
24/// # if cfg!(miri) { return Ok(()); } // no file system under Miri's isolation
25/// let path = std::env::temp_dir().join("nuvec-mapped-file-doc.bin");
26/// let bytes: Vec<u8> = [1.0f32, 2.0, 3.0, 4.0].iter().flat_map(|v| v.to_le_bytes()).collect();
27/// std::fs::write(&path, &bytes).unwrap();
28/// let file = MappedFile::open(&path)?;
29/// let a = file.array(4, DType::F32, &[3])?; // bytes 4..16, no copy
30/// assert_eq!(a.to_vec::<f32>()?, vec![2.0, 3.0, 4.0]);
31/// # Ok(())
32/// # }
33/// ```
34#[derive(Clone)]
35pub struct MappedFile {
36 buffer: Arc<Buffer>,
37}
38
39impl std::fmt::Debug for MappedFile {
40 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 f.debug_struct("MappedFile")
42 .field("len", &self.len())
43 .finish()
44 }
45}
46
47impl MappedFile {
48 /// Maps (or reads) the file at `path`. The file must not be modified while it is
49 /// mapped (see [`MappedFile`]).
50 pub fn open(path: impl AsRef<Path>) -> Result<Self> {
51 Ok(Self {
52 buffer: Arc::new(Buffer::map_file(path.as_ref())?),
53 })
54 }
55
56 /// The bytes mapping a file of `len` bytes holds in the GPU allocator's `mapped`
57 /// statistic: the whole file, rounded up to pages.
58 pub fn mapping_size(len: u64) -> u64 {
59 nuvec_gpu::mapping_size(len as usize) as u64
60 }
61
62 /// The bytes storage of `len` bytes (as [`MappedFile::read_ranges`] allocates) holds in
63 /// the GPU allocator's `active` statistic: its size class.
64 pub fn copy_size(len: u64) -> u64 {
65 nuvec_gpu::allocation_size(len as usize).map_or(u64::MAX, |n| n as u64)
66 }
67
68 /// The bytes that may still be mapped or allocated before the GPU's memory limit
69 /// (`memory_limit − active − mapped`), or `None` without a GPU.
70 pub fn headroom() -> Option<u64> {
71 let m = nuvec_gpu::Device::shared()?.memory();
72 Some(m.memory_limit.saturating_sub(m.active_with_mapped) as u64)
73 }
74
75 /// The file's size in bytes.
76 pub fn len(&self) -> usize {
77 self.buffer.len()
78 }
79
80 /// Whether the file is empty.
81 pub fn is_empty(&self) -> bool {
82 self.len() == 0
83 }
84
85 /// The file's bytes.
86 pub fn bytes(&self) -> &[u8] {
87 self.buffer.bytes()
88 }
89
90 /// Reads parts of the file at `path` into new array storage of `len` bytes, without
91 /// mapping it: each `(file_offset, offset, n)` of `ranges` puts the file's bytes
92 /// `file_offset..file_offset + n` at `offset..offset + n`. For files whose arrays are
93 /// not aligned to their element sizes (safetensors written without header padding, as
94 /// MLX writes them): the ranges are placed so that every array is aligned, and arrays
95 /// made from the result are views of it. Bytes outside the ranges are unspecified.
96 ///
97 /// The ranges must be in increasing order of `offset` and must not overlap. They are
98 /// read with positioned reads in chunks of 16 MiB, on up to 8 threads (one per 64 MiB
99 /// read), straight into the storage: from the page cache this runs at memory
100 /// bandwidth, and nothing but the storage is allocated (the file is not mapped).
101 ///
102 /// ```
103 /// use nuvec_core::{DType, MappedFile};
104 ///
105 /// # fn main() -> nuvec_core::Result<()> {
106 /// # if cfg!(miri) { return Ok(()); }
107 /// let path = std::env::temp_dir().join("nuvec-read-ranges-doc.bin");
108 /// // One padding byte, then three f32s: misaligned in the file.
109 /// let mut bytes = vec![0u8];
110 /// bytes.extend([1.0f32, 2.0, 3.0].iter().flat_map(|v| v.to_le_bytes()));
111 /// std::fs::write(&path, &bytes).unwrap();
112 /// let file = MappedFile::read_ranges(&path, &[(1, 0, 12)], 12)?;
113 /// assert_eq!(file.array(0, DType::F32, &[3])?.to_vec::<f32>()?, vec![1.0, 2.0, 3.0]);
114 /// # Ok(())
115 /// # }
116 /// ```
117 pub fn read_ranges(
118 path: impl AsRef<Path>,
119 ranges: &[(u64, usize, usize)],
120 len: usize,
121 ) -> Result<Self> {
122 let path = path.as_ref();
123 let mut at = 0;
124 for &(_, offset, n) in ranges {
125 if offset < at || offset.checked_add(n).is_none_or(|end| end > len) {
126 return Err(Error::InvalidArgument(format!(
127 "{}: read ranges must be increasing, disjoint and within {len} bytes",
128 path.display()
129 )));
130 }
131 at = offset + n;
132 }
133 let io = |e: std::io::Error| Error::Io(format!("{}: {e}", path.display()));
134 let file = std::fs::File::open(path).map_err(io)?;
135 let mut buffer = Buffer::try_alloc(len)?;
136 read_parallel(&file, buffer.bytes_mut(), ranges).map_err(io)?;
137 Ok(Self {
138 buffer: Arc::new(buffer),
139 })
140 }
141
142 /// A row-major array of `dtype` and `shape` stored at byte `offset`: a view of the file
143 /// if `offset` is a multiple of the element size, otherwise a copy.
144 pub fn array(&self, offset: usize, dtype: DType, shape: &[usize]) -> Result<Array> {
145 let nbytes = shape::checked_nbytes(shape, dtype.size())?;
146 if offset
147 .checked_add(nbytes)
148 .is_none_or(|end| end > self.len())
149 {
150 return Err(Error::Io(format!(
151 "{nbytes} bytes at offset {offset} exceed the file's {} bytes",
152 self.len()
153 )));
154 }
155 if offset.is_multiple_of(dtype.size()) {
156 let data = Data {
157 buffer: self.buffer.clone(),
158 offset: (offset / dtype.size()) as isize,
159 };
160 return Array::leaf(Ok(data), dtype, shape.into());
161 }
162 let mut copy = Buffer::try_alloc(nbytes)?;
163 copy.bytes_mut()
164 .copy_from_slice(&self.bytes()[offset..offset + nbytes]);
165 Array::from_buffer(copy, dtype, shape)
166 }
167}
168
169/// Byte ranges of a file in array storage: each range memory-mapped on its own (without a
170/// copy, on a Mac with a GPU; read elsewhere), so the bytes between them are neither mapped,
171/// wired nor counted. For loads that keep some of a file's tensors. The file must not be
172/// modified while it is mapped (see [`MappedFile`]).
173///
174/// ```
175/// use nuvec_core::{DType, MappedRanges};
176///
177/// # fn main() -> nuvec_core::Result<()> {
178/// # if cfg!(miri) { return Ok(()); }
179/// let page = MappedRanges::page_size();
180/// let path = std::env::temp_dir().join("nuvec-mapped-ranges-doc.bin");
181/// let mut bytes = vec![0u8; 3 * page];
182/// bytes[2 * page..2 * page + 8].copy_from_slice(&[1.0f32, 2.0].map(f32::to_le_bytes).concat());
183/// std::fs::write(&path, &bytes).unwrap();
184/// let file = MappedRanges::open(&path, &[(2 * page as u64, page as u64)])?; // the last page only
185/// assert_eq!(file.array(2 * page, DType::F32, &[2])?.to_vec::<f32>()?, vec![1.0, 2.0]);
186/// assert!(file.array(0, DType::F32, &[2]).is_err()); // not mapped
187/// # Ok(())
188/// # }
189/// ```
190#[derive(Clone)]
191pub struct MappedRanges {
192 /// `(file offset, buffer)`, in file order.
193 runs: Vec<(usize, Arc<Buffer>)>,
194}
195
196impl std::fmt::Debug for MappedRanges {
197 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198 let runs: Vec<(usize, usize)> = self.runs.iter().map(|(o, b)| (*o, b.len())).collect();
199 f.debug_struct("MappedRanges").field("runs", &runs).finish()
200 }
201}
202
203impl MappedRanges {
204 /// The granularity of mappings: range offsets must be multiples of it.
205 pub fn page_size() -> usize {
206 nuvec_gpu::page_size()
207 }
208
209 /// Maps each `(offset, length)` range of the file at `path` (offsets multiples of
210 /// [`Self::page_size`], ranges disjoint and increasing).
211 pub fn open(path: impl AsRef<Path>, ranges: &[(u64, u64)]) -> Result<Self> {
212 let path = path.as_ref();
213 let mut end = 0;
214 let mut runs = Vec::with_capacity(ranges.len());
215 for &(offset, len) in ranges {
216 if offset < end {
217 return Err(Error::InvalidArgument(format!(
218 "{}: mapped ranges must be increasing and disjoint",
219 path.display()
220 )));
221 }
222 end = offset + len;
223 let (offset, len) = (offset as usize, len as usize);
224 runs.push((offset, Arc::new(Buffer::map_file_range(path, offset, len)?)));
225 }
226 Ok(Self { runs })
227 }
228
229 /// The bytes mapped, page-rounded (what the GPU allocator's `mapped` counts).
230 pub fn mapped_size(&self) -> u64 {
231 self.runs
232 .iter()
233 .map(|(_, b)| MappedFile::mapping_size(b.len() as u64))
234 .sum()
235 }
236
237 /// A row-major array of `dtype` and `shape` at byte `offset` of the file: a view of the
238 /// range that holds it (it must lie inside one range and be aligned to its element size).
239 pub fn array(&self, offset: usize, dtype: DType, shape: &[usize]) -> Result<Array> {
240 let nbytes = shape::checked_nbytes(shape, dtype.size())?;
241 let run = self.runs.iter().find(|(start, b)| {
242 *start <= offset
243 && offset
244 .checked_add(nbytes)
245 .is_some_and(|end| end <= start + b.len())
246 });
247 let Some((start, buffer)) = run else {
248 return Err(Error::Io(format!(
249 "{nbytes} bytes at offset {offset} are not inside a mapped range"
250 )));
251 };
252 let at = offset - start;
253 if !at.is_multiple_of(dtype.size()) {
254 return Err(Error::Io(format!(
255 "offset {offset} is not aligned to {} bytes",
256 dtype.size()
257 )));
258 }
259 let data = Data {
260 buffer: buffer.clone(),
261 offset: (at / dtype.size()) as isize,
262 };
263 Array::leaf(Ok(data), dtype, shape.into())
264 }
265}
266
267/// Reads chunk size: large enough that a read's cost is the copy, small enough to spread a
268/// file over the threads.
269const CHUNK: usize = 16 << 20;
270/// Bytes per reading thread (smaller reads use fewer threads).
271const PER_THREAD: usize = 64 << 20;
272/// At most this many reading threads: more do not read faster from the page cache.
273const MAX_THREADS: usize = 8;
274
275/// Fills `dest[offset..offset + n]` with the file's bytes `file_offset..file_offset + n` for
276/// each range (increasing and disjoint in `dest`, checked by the caller), in chunks of
277/// [`CHUNK`] spread over a few threads.
278fn read_parallel(
279 file: &std::fs::File,
280 dest: &mut [u8],
281 ranges: &[(u64, usize, usize)],
282) -> std::io::Result<()> {
283 // The chunks as disjoint mutable slices of `dest`, with their file offsets.
284 let mut chunks: Vec<(u64, &mut [u8])> = Vec::new();
285 let mut rest = dest;
286 let mut consumed = 0;
287 for &(file_offset, offset, n) in ranges {
288 let (_, tail) = std::mem::take(&mut rest).split_at_mut(offset - consumed);
289 let (mut range, tail) = tail.split_at_mut(n);
290 rest = tail;
291 consumed = offset + n;
292 let mut at = file_offset;
293 while !range.is_empty() {
294 let (chunk, tail) = range.split_at_mut(range.len().min(CHUNK));
295 let n = chunk.len() as u64;
296 chunks.push((at, chunk));
297 at += n;
298 range = tail;
299 }
300 }
301 let total: usize = ranges.iter().map(|r| r.2).sum();
302 let threads = std::thread::available_parallelism()
303 .map_or(1, usize::from)
304 .min(MAX_THREADS)
305 .min(total.div_ceil(PER_THREAD))
306 .max(1);
307 if threads == 1 {
308 return chunks
309 .into_iter()
310 .try_for_each(|(at, chunk)| read_exact_at(file, chunk, at));
311 }
312 let work = std::sync::Mutex::new(chunks.into_iter());
313 std::thread::scope(|s| {
314 let workers: Vec<_> = (0..threads)
315 .map(|_| {
316 s.spawn(|| {
317 loop {
318 let next = work
319 .lock()
320 .unwrap_or_else(std::sync::PoisonError::into_inner)
321 .next();
322 let Some((at, chunk)) = next else {
323 return Ok(());
324 };
325 read_exact_at(file, chunk, at)?;
326 }
327 })
328 })
329 .collect();
330 workers.into_iter().try_for_each(|w| {
331 w.join()
332 .unwrap_or_else(|_| Err(std::io::Error::other("a reading thread panicked")))
333 })
334 })
335}
336
337/// Fills `buf` with the file's bytes from `offset` (a positioned read: the file's cursor is
338/// neither used nor moved, so threads read one file concurrently).
339#[cfg(unix)]
340fn read_exact_at(file: &std::fs::File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
341 std::os::unix::fs::FileExt::read_exact_at(file, buf, offset)
342}
343
344#[cfg(windows)]
345fn read_exact_at(file: &std::fs::File, mut buf: &mut [u8], mut offset: u64) -> std::io::Result<()> {
346 use std::os::windows::fs::FileExt;
347 while !buf.is_empty() {
348 match file.seek_read(buf, offset) {
349 Ok(0) => return Err(std::io::ErrorKind::UnexpectedEof.into()),
350 Ok(n) => {
351 buf = &mut buf[n..];
352 offset += n as u64;
353 }
354 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
355 Err(e) => return Err(e),
356 }
357 }
358 Ok(())
359}
360
361#[cfg(test)]
362mod tests {
363 use super::*;
364
365 /// Ranges of several chunks, gaps and odd offsets land where they are asked to.
366 #[test]
367 #[cfg_attr(miri, ignore)] // no file system under Miri's isolation
368 fn read_ranges_places_every_byte() {
369 let path =
370 std::env::temp_dir().join(format!("nuvec-read-ranges-{}.bin", std::process::id()));
371 let bytes: Vec<u8> = (0..(CHUNK * 2 + 12_345))
372 .map(|i| (i * 31 % 251) as u8)
373 .collect();
374 std::fs::write(&path, &bytes).unwrap();
375 let big = CHUNK * 2 + 1000;
376 let ranges = [
377 (3u64, 0usize, 10usize),
378 (17, 16, big),
379 (bytes.len() as u64 - 5, big + 32, 5),
380 ];
381 let file = MappedFile::read_ranges(&path, &ranges, big + 37).unwrap();
382 for &(at, offset, n) in &ranges {
383 assert_eq!(
384 &file.bytes()[offset..offset + n],
385 &bytes[at as usize..at as usize + n]
386 );
387 }
388 // Past the end of the file, out of order, or past the storage: errors.
389 assert!(MappedFile::read_ranges(&path, &[(bytes.len() as u64 - 2, 0, 4)], 4).is_err());
390 assert!(MappedFile::read_ranges(&path, &[(0, 8, 4), (0, 0, 4)], 12).is_err());
391 assert!(MappedFile::read_ranges(&path, &[(0, 0, 8)], 4).is_err());
392 let empty = MappedFile::read_ranges(&path, &[], 0).unwrap();
393 assert!(empty.is_empty());
394 std::fs::remove_file(&path).unwrap();
395 }
396}