osobhandClaude Opus 5.5
2bc4cb46a6
A 256 x 256 hyperslab of a contiguous f32 dataset read at an eighth of h5py's speed: partial_read copied the bounding box out of the file, the extractor then walked it element by element (a recursive call and two bounds checks per element) into a second buffer, and read_f32_selection converted that into a third. Selections of contiguous data are now copied straight from the file, one memcpy per run of elements contiguous in the file (gather.rs: a block along the last dimension, touching blocks as one range, whole rows merged), with no zero-filled intermediate and no full copy for large selections. The typed selection readers copy into their Vec<T> directly when the dataset stores T natively (new data_read::read_selection_native and sealed NativeElement trait, which the read_as_* fast paths now share; read_as_u64 gains one) and convert as before otherwise. The general extractor used by the chunked paths runs on the same run walker, keeping its old handling of unvalidated selections. Checked against h5py (contiguous_read_interop.rs) for strided, blocked, adjacent-block and whole-row hyperslabs, points and empty selections of every 1-8-byte type in both byte orders, ranks 1-4. Also keeps the huge-page threshold constant out of no_std builds, where it was unused. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
80 lines
3.1 KiB
Rust
80 lines
3.1 KiB
Rust
//! Large output buffers backed by transparent huge pages where the OS offers
|
|
//! them.
|
|
//!
|
|
//! A fresh multi-megabyte `Vec` is mapped lazily by the kernel: the first
|
|
//! write to each 4 KiB page takes a page fault, and the kernel zeroes the page
|
|
//! before handing it over. For a 64 MiB read that is 16384 faults, and they
|
|
//! cost far more than the copy that fills the buffer — single-threaded
|
|
//! contiguous reads ran at about a quarter of h5py's speed because of them.
|
|
//! numpy (so h5py) avoids this by asking for transparent huge pages
|
|
//! (`madvise(MADV_HUGEPAGE)`) on every allocation of 4 MiB or more, which
|
|
//! turns 512 faults into one; this module does the same.
|
|
//!
|
|
//! The advice only changes how the pages are backed, never their contents, so
|
|
//! it is harmless when it cannot be honoured (THP disabled, not Linux, a
|
|
//! region that is part of the heap): the buffer is then exactly what it would
|
|
//! have been without it.
|
|
|
|
#[cfg(not(feature = "std"))]
|
|
use alloc::vec::Vec;
|
|
|
|
/// Buffers smaller than this are left alone (numpy uses the same threshold).
|
|
#[cfg(any(target_os = "linux", test))]
|
|
pub(crate) const HUGE_PAGE_THRESHOLD: usize = 4 << 20;
|
|
|
|
/// Advise the kernel to back `[ptr, ptr + len)` with transparent huge pages,
|
|
/// when `len` is large enough to benefit. Call it before the first write so
|
|
/// the faults happen at huge-page granularity.
|
|
#[inline]
|
|
pub(crate) fn advise_huge_pages(ptr: *const u8, len: usize) {
|
|
#[cfg(target_os = "linux")]
|
|
if len >= HUGE_PAGE_THRESHOLD {
|
|
const PAGE: usize = 4096;
|
|
let start = (ptr as usize).next_multiple_of(PAGE);
|
|
let end = (ptr as usize + len) & !(PAGE - 1);
|
|
if end > start {
|
|
// SAFETY: `[start, end)` lies inside an allocation of `len` bytes
|
|
// at `ptr` that the caller owns, and is page aligned as madvise
|
|
// requires. MADV_HUGEPAGE does not change the memory's contents or
|
|
// validity; on failure (EINVAL when THP is compiled out, etc.) the
|
|
// region is simply left as it was, so the result is ignored.
|
|
unsafe {
|
|
libc::madvise(start as *mut libc::c_void, end - start, libc::MADV_HUGEPAGE);
|
|
}
|
|
}
|
|
}
|
|
#[cfg(not(target_os = "linux"))]
|
|
let _ = (ptr, len);
|
|
}
|
|
|
|
/// `Vec::with_capacity(count)` for a buffer about to be filled in bulk, with
|
|
/// huge-page advice when it is large (see the module docs).
|
|
#[inline]
|
|
pub(crate) fn vec_for_bulk<T>(count: usize) -> Vec<T> {
|
|
let v: Vec<T> = Vec::with_capacity(count);
|
|
advise_huge_pages(
|
|
v.as_ptr().cast::<u8>(),
|
|
v.capacity().saturating_mul(core::mem::size_of::<T>()),
|
|
);
|
|
v
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn bulk_vec_is_an_ordinary_vec() {
|
|
for count in [0usize, 1, 1000, HUGE_PAGE_THRESHOLD / 4 + 3] {
|
|
let mut v: Vec<u32> = vec_for_bulk(count);
|
|
assert!(v.capacity() >= count);
|
|
v.extend((0..count as u32).map(|i| i.wrapping_mul(2654435761)));
|
|
assert!(
|
|
v.iter()
|
|
.enumerate()
|
|
.all(|(i, &x)| x == (i as u32).wrapping_mul(2654435761))
|
|
);
|
|
}
|
|
}
|
|
}
|