perf: copy contiguous hyperslab and point reads run by run
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]>
This commit is contained in:
osobh
@@ -0,0 +1,343 @@
|
||||
//! Copying a selection out of a row-major buffer one contiguous run at a time.
|
||||
//!
|
||||
//! A selection's elements, in output order, fall into runs that are adjacent
|
||||
//! in the source: a whole block along the last dimension, blocks that touch
|
||||
//! (`stride == block`), and whole rows when the inner dimensions are selected
|
||||
//! in full. Copying run by run turns a 256 x 256 hyperslab of a 1024-wide
|
||||
//! dataset into 256 `memcpy`s of 1 KiB, where the old extractor recursed and
|
||||
//! bounds-checked once per element.
|
||||
|
||||
#[cfg(not(feature = "std"))]
|
||||
use alloc::{vec, vec::Vec};
|
||||
|
||||
use crate::data_read::NativeElement;
|
||||
use crate::error::FormatError;
|
||||
use crate::selection::Selection;
|
||||
|
||||
/// Row-major element strides of `dims` (the last dimension has stride 1).
|
||||
fn strides(dims: &[u64]) -> Vec<u64> {
|
||||
let mut s = vec![1u64; dims.len()];
|
||||
for d in (0..dims.len().saturating_sub(1)).rev() {
|
||||
s[d] = s[d + 1].wrapping_mul(dims[d + 1]);
|
||||
}
|
||||
s
|
||||
}
|
||||
|
||||
/// Merges adjacent runs before handing them on.
|
||||
struct Coalesce<F: FnMut(u64, u64)> {
|
||||
start: u64,
|
||||
len: u64,
|
||||
emit: F,
|
||||
}
|
||||
|
||||
impl<F: FnMut(u64, u64)> Coalesce<F> {
|
||||
#[inline]
|
||||
fn push(&mut self, start: u64, len: u64) {
|
||||
if len == 0 {
|
||||
return;
|
||||
}
|
||||
if self.len > 0 && self.start.wrapping_add(self.len) == start {
|
||||
self.len += len;
|
||||
return;
|
||||
}
|
||||
self.flush();
|
||||
self.start = start;
|
||||
self.len = len;
|
||||
}
|
||||
|
||||
fn flush(&mut self) {
|
||||
if self.len > 0 {
|
||||
(self.emit)(self.start, self.len);
|
||||
self.len = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Call `emit(first_element, element_count)` for each run of a hyperslab's
|
||||
/// elements that is contiguous in a row-major dataset of shape `dims`, in
|
||||
/// the order the selection returns them. Adjacent runs are merged.
|
||||
///
|
||||
/// Coordinates at or past a dimension's extent are skipped, as the
|
||||
/// element-wise extractor always did; callers that want them to be an error
|
||||
/// validate the selection first. The four vectors must have `dims.len()`
|
||||
/// entries.
|
||||
pub(crate) fn hyperslab_runs(
|
||||
dims: &[u64],
|
||||
start: &[u64],
|
||||
stride: &[u64],
|
||||
count: &[u64],
|
||||
block: &[u64],
|
||||
emit: impl FnMut(u64, u64),
|
||||
) {
|
||||
let rank = dims.len();
|
||||
let mut out = Coalesce {
|
||||
start: 0,
|
||||
len: 0,
|
||||
emit,
|
||||
};
|
||||
if rank == 0 {
|
||||
out.push(0, 1);
|
||||
out.flush();
|
||||
return;
|
||||
}
|
||||
if (0..rank).any(|d| count[d] == 0 || block[d] == 0) {
|
||||
return;
|
||||
}
|
||||
let strides = strides(dims);
|
||||
let last = rank - 1;
|
||||
// Odometer over the outer dimensions: (block index, offset in block).
|
||||
let mut ci = vec![0u64; last];
|
||||
let mut bi = vec![0u64; last];
|
||||
'outer: loop {
|
||||
// Base offset of this row, or skip it if a coordinate is out of range.
|
||||
let mut base = 0u64;
|
||||
let mut in_range = true;
|
||||
for d in 0..last {
|
||||
let coord = start[d]
|
||||
.saturating_add(ci[d].saturating_mul(stride[d]))
|
||||
.saturating_add(bi[d]);
|
||||
if coord >= dims[d] {
|
||||
in_range = false;
|
||||
break;
|
||||
}
|
||||
base = base.wrapping_add(coord.wrapping_mul(strides[d]));
|
||||
}
|
||||
if in_range && (stride[last] == block[last] || count[last] == 1) {
|
||||
// Blocks that touch (the common unit-stride case: block 1,
|
||||
// stride 1) are one range; don't split it into per-element runs.
|
||||
let s = start[last];
|
||||
let e = s
|
||||
.saturating_add(count[last].saturating_mul(block[last]))
|
||||
.min(dims[last]);
|
||||
if s < e {
|
||||
out.push(base.wrapping_add(s), e - s);
|
||||
}
|
||||
} else if in_range {
|
||||
for c in 0..count[last] {
|
||||
let s = start[last].saturating_add(c.saturating_mul(stride[last]));
|
||||
if s >= dims[last] {
|
||||
continue;
|
||||
}
|
||||
let e = s.saturating_add(block[last]).min(dims[last]);
|
||||
out.push(base.wrapping_add(s), e - s);
|
||||
}
|
||||
}
|
||||
// Advance the odometer, last outer dimension fastest.
|
||||
let mut d = last;
|
||||
loop {
|
||||
if d == 0 {
|
||||
break 'outer;
|
||||
}
|
||||
d -= 1;
|
||||
bi[d] += 1;
|
||||
if bi[d] < block[d] {
|
||||
break;
|
||||
}
|
||||
bi[d] = 0;
|
||||
ci[d] += 1;
|
||||
if ci[d] < count[d] {
|
||||
break;
|
||||
}
|
||||
ci[d] = 0;
|
||||
}
|
||||
}
|
||||
out.flush();
|
||||
}
|
||||
|
||||
/// The selected elements of `src` — a row-major dataset of shape `dims` and
|
||||
/// `elem_size`-byte elements — copied into a fresh `Vec<T>`, one `memcpy` per
|
||||
/// contiguous run, with no zero-filling of the output first.
|
||||
///
|
||||
/// For `T` other than `u8`, `elem_size` must equal `size_of::<T>()`. The
|
||||
/// selection must be a validated hyperslab, point list or `None` (`All` is the
|
||||
/// caller's to handle); `src` must hold exactly the dataset. Anything that
|
||||
/// would read outside `src` is an error, never a partial result.
|
||||
pub(crate) fn gather<T: NativeElement>(
|
||||
src: &[u8],
|
||||
dims: &[u64],
|
||||
elem_size: usize,
|
||||
selection: &Selection,
|
||||
) -> Result<Vec<T>, FormatError> {
|
||||
let t_size = core::mem::size_of::<T>();
|
||||
if elem_size == 0 || (t_size != 1 && t_size != elem_size) {
|
||||
return Err(FormatError::DataSizeMismatch {
|
||||
expected: t_size,
|
||||
actual: elem_size,
|
||||
});
|
||||
}
|
||||
let n_elements = match selection {
|
||||
Selection::None => 0,
|
||||
Selection::Hyperslab { count, block, .. } => count
|
||||
.iter()
|
||||
.zip(block)
|
||||
.try_fold(1u64, |acc, (&c, &b)| acc.checked_mul(c.checked_mul(b)?))
|
||||
.ok_or_else(|| FormatError::Overflow("hyperslab count x block overflows".into()))?,
|
||||
Selection::Points(points) => points.len() as u64,
|
||||
Selection::All => {
|
||||
return Err(FormatError::SelectionOutOfBounds(
|
||||
"gather does not take Selection::All".into(),
|
||||
));
|
||||
}
|
||||
};
|
||||
let out_bytes = crate::chunked_read::checked_byte_len(n_elements, elem_size)?;
|
||||
let out_len = out_bytes / t_size;
|
||||
let mut out: Vec<T> = crate::bulk_alloc::vec_for_bulk(out_len);
|
||||
let dst = out.as_mut_ptr().cast::<u8>();
|
||||
let mut written = 0usize;
|
||||
let mut failed = false;
|
||||
let mut copy_run = |first: u64, n: u64| {
|
||||
if failed {
|
||||
return;
|
||||
}
|
||||
let range = usize::try_from(first)
|
||||
.ok()
|
||||
.and_then(|f| f.checked_mul(elem_size))
|
||||
.zip(
|
||||
usize::try_from(n)
|
||||
.ok()
|
||||
.and_then(|n| n.checked_mul(elem_size)),
|
||||
)
|
||||
.and_then(|(at, len)| Some((at, len, at.checked_add(len)?)));
|
||||
match range {
|
||||
Some((at, len, end)) if end <= src.len() && written + len <= out_bytes => {
|
||||
// SAFETY: `src[at..end]` is in bounds (checked above), and
|
||||
// `dst + written .. + len` lies within `out`'s capacity of
|
||||
// `out_bytes` bytes (checked above); `out` is a fresh
|
||||
// allocation, so the regions do not overlap.
|
||||
unsafe {
|
||||
core::ptr::copy_nonoverlapping(src.as_ptr().add(at), dst.add(written), len)
|
||||
};
|
||||
written += len;
|
||||
}
|
||||
_ => failed = true,
|
||||
}
|
||||
};
|
||||
let mut bad_point = false;
|
||||
match selection {
|
||||
Selection::Hyperslab {
|
||||
start,
|
||||
stride,
|
||||
count,
|
||||
block,
|
||||
} => {
|
||||
let rank = dims.len();
|
||||
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
|
||||
return Err(FormatError::SelectionOutOfBounds(
|
||||
"hyperslab rank does not match dataset rank".into(),
|
||||
));
|
||||
}
|
||||
hyperslab_runs(dims, start, stride, count, block, &mut copy_run);
|
||||
}
|
||||
Selection::Points(points) => {
|
||||
let strides = strides(dims);
|
||||
let mut runs = Coalesce {
|
||||
start: 0,
|
||||
len: 0,
|
||||
emit: &mut copy_run,
|
||||
};
|
||||
for p in points {
|
||||
if p.len() != dims.len() || p.iter().zip(dims).any(|(c, n)| c >= n) {
|
||||
bad_point = true;
|
||||
break;
|
||||
}
|
||||
let at = p
|
||||
.iter()
|
||||
.zip(&strides)
|
||||
.fold(0u64, |acc, (c, s)| acc.wrapping_add(c.wrapping_mul(*s)));
|
||||
runs.push(at, 1);
|
||||
}
|
||||
runs.flush();
|
||||
}
|
||||
Selection::None | Selection::All => {}
|
||||
}
|
||||
if failed || bad_point || written != out_bytes {
|
||||
return Err(FormatError::SelectionOutOfBounds(
|
||||
"selection addresses elements outside the dataset".into(),
|
||||
));
|
||||
}
|
||||
// SAFETY: all `out_bytes` bytes, i.e. `out_len` values of `T`, were
|
||||
// written above, and every bit pattern is a valid `T` (`NativeElement`).
|
||||
unsafe { out.set_len(out_len) };
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn runs(dims: &[u64], sel: [&[u64]; 4]) -> Vec<(u64, u64)> {
|
||||
let mut v = Vec::new();
|
||||
hyperslab_runs(dims, sel[0], sel[1], sel[2], sel[3], |s, n| v.push((s, n)));
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runs_merge_blocks_and_whole_rows() {
|
||||
// A box: one run per row.
|
||||
assert_eq!(
|
||||
runs(&[4, 10], [&[1, 2], &[1, 1], &[2, 3], &[1, 1]]),
|
||||
vec![(12, 3), (22, 3)]
|
||||
);
|
||||
// Whole rows: one run.
|
||||
assert_eq!(
|
||||
runs(&[4, 10], [&[1, 0], &[1, 1], &[3, 10], &[1, 1]]),
|
||||
vec![(10, 30)]
|
||||
);
|
||||
// stride == block: blocks merge.
|
||||
assert_eq!(
|
||||
runs(&[1, 10], [&[0, 1], &[1, 2], &[1, 4], &[1, 2]]),
|
||||
vec![(1, 8)]
|
||||
);
|
||||
// Strided with blocks along both dimensions.
|
||||
assert_eq!(
|
||||
runs(&[6, 10], [&[0, 1], &[3, 4], &[2, 2], &[2, 2]]),
|
||||
vec![
|
||||
(1, 2),
|
||||
(5, 2),
|
||||
(11, 2),
|
||||
(15, 2),
|
||||
(31, 2),
|
||||
(35, 2),
|
||||
(41, 2),
|
||||
(45, 2)
|
||||
]
|
||||
);
|
||||
// Empty.
|
||||
assert!(runs(&[4, 10], [&[0, 0], &[1, 1], &[0, 3], &[1, 1]]).is_empty());
|
||||
// Scalar.
|
||||
assert_eq!(runs(&[], [&[], &[], &[], &[]]), vec![(0, 1)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gather_matches_element_order_and_rejects_out_of_range() {
|
||||
let dims = [3u64, 4];
|
||||
let src: Vec<u8> = (0..12u16).flat_map(|v| v.to_le_bytes()).collect();
|
||||
let sel = Selection::Hyperslab {
|
||||
start: vec![0, 1],
|
||||
stride: vec![2, 2],
|
||||
count: vec![2, 2],
|
||||
block: vec![1, 1],
|
||||
};
|
||||
let got: Vec<u8> = gather(&src, &dims, 2, &sel).unwrap();
|
||||
let want: Vec<u8> = [1u16, 3, 9, 11]
|
||||
.iter()
|
||||
.flat_map(|v| v.to_le_bytes())
|
||||
.collect();
|
||||
assert_eq!(got, want);
|
||||
let pts = Selection::Points(vec![vec![2, 3], vec![0, 0], vec![0, 1]]);
|
||||
let got: Vec<u8> = gather(&src, &dims, 2, &pts).unwrap();
|
||||
let want: Vec<u8> = [11u16, 0, 1].iter().flat_map(|v| v.to_le_bytes()).collect();
|
||||
assert_eq!(got, want);
|
||||
// Past the extent, or a source shorter than the dataset: an error.
|
||||
let bad = Selection::Points(vec![vec![3, 0]]);
|
||||
assert!(gather::<u8>(&src, &dims, 2, &bad).is_err());
|
||||
let past = Selection::Hyperslab {
|
||||
start: vec![2, 0],
|
||||
stride: vec![1, 1],
|
||||
count: vec![2, 4],
|
||||
block: vec![1, 1],
|
||||
};
|
||||
assert!(gather::<u8>(&src, &dims, 2, &past).is_err());
|
||||
assert!(gather::<u8>(&src[..20], &dims, 2, &pts).is_err());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user