quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 1 Wiki Activity

perf: copy contiguous hyperslab and point reads run by run

Browse Source 
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
2026-09-26 08:21:47 -05:00
co-authored by Claude Opus 5.5
parent 78c769f179
commit 2bc4cb46a6
8 changed files with 586 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-format/src/data_read.rs
+151 -128
View File
@@ -531,6 +531,11 @@ pub fn extract_selection_from_buffer(
block,
} => {
let rank = dims.len();
if [start.len(), stride.len(), count.len(), block.len()] != [rank; 4] {
return Err(FormatError::SelectionOutOfBounds(format!(
"hyperslab rank does not match dataset rank {rank}"
)));
}
let output_elements = count
.iter()
.zip(block.iter())
@@ -540,96 +545,40 @@ pub fn extract_selection_from_buffer(
crate::chunked_read::checked_byte_len(output_elements, elem_size)?,
)?;
// Compute dataset strides (row-major)
let mut ds_strides = vec![1usize; rank];
for i in (0..rank.saturating_sub(1)).rev() {
ds_strides[i] = ds_strides[i + 1] * dims[i + 1] as usize;
}
// Compute output shape and strides
let output_dims: Vec<usize> = count
.iter()
.zip(block.iter())
.map(|(&c, &b)| (c * b) as usize)
.collect();
let mut out_strides = vec![1usize; rank];
for i in (0..rank.saturating_sub(1)).rev() {
out_strides[i] = out_strides[i + 1] * output_dims[i + 1];
}
// Iterate over all selected elements
// For each block in the hyperslab, copy the elements
let mut out_linear = 0usize;
let _block_coords = vec![0u64; rank];
#[allow(clippy::too_many_arguments)]
fn iterate_hyperslab(
d: usize,
rank: usize,
start: &[u64],
stride: &[u64],
count: &[u64],
block: &[u64],
dims: &[u64],
ds_strides: &[usize],
elem_size: usize,
full_data: &[u8],
output: &mut [u8],
out_linear: &mut usize,
current_ds_offset: usize,
) {
if d == rank {
// Copy one element
let src = current_ds_offset * elem_size;
let dst = *out_linear * elem_size;
if src + elem_size <= full_data.len() && dst + elem_size <= output.len() {
output[dst..dst + elem_size]
.copy_from_slice(&full_data[src..src + elem_size]);
}
*out_linear += 1;
return;
}
for bi in 0..count[d] {
let block_start = start[d] + bi * stride[d];
for bj in 0..block[d] {
let coord = block_start + bj;
if coord < dims[d] {
iterate_hyperslab(
d + 1,
rank,
start,
stride,
count,
block,
dims,
ds_strides,
elem_size,
full_data,
output,
out_linear,
current_ds_offset + coord as usize * ds_strides[d],
);
// One copy per run of elements contiguous in `full_data`
// (`gather`'s runs). Coordinates past the extent are skipped and
// runs past the end of `full_data` left as zeros, element by
// element, as this extractor always did; validated selections
// never hit either.
let mut out_at = 0usize;
crate::gather::hyperslab_runs(dims, start, stride, count, block, |first, n| {
let big = |v: u64| usize::try_from(v).unwrap_or(usize::MAX);
let (first, n) = (big(first), big(n));
let len = n.saturating_mul(elem_size);
let src = first.saturating_mul(elem_size);
let out_end = out_at.saturating_add(len);
if let (Some(from), Some(to)) = (
full_data.get(src..src.saturating_add(len)),
output.get_mut(out_at..out_end),
) {
to.copy_from_slice(from);
} else {
for k in 0..n {
let s = first.saturating_add(k).saturating_mul(elem_size);
let o = out_at.saturating_add(k.saturating_mul(elem_size));
if o >= output.len() {
break;
}
if let (Some(from), Some(to)) = (
full_data.get(s..s.saturating_add(elem_size)),
output.get_mut(o..o.saturating_add(elem_size)),
) {
to.copy_from_slice(from);
}
}
}
}
iterate_hyperslab(
0,
rank,
start,
stride,
count,
block,
dims,
&ds_strides,
elem_size,
full_data,
&mut output,
&mut out_linear,
0,
);
out_at = out_end;
});
Ok(output)
}
@@ -757,22 +706,76 @@ fn get_size(dt: &Datatype) -> usize {
dt.type_size() as usize
}
/// Reinterpret little-endian bytes as `count` native values of `T` on a
/// little-endian target, in one copy.
mod sealed {
pub trait Sealed {}
}
/// A numeric type whose values can be copied straight out of a dataset's
/// bytes when the dataset stores exactly that type in the target's byte
/// order: `u8`, `i32`, `i64`, `u64`, `f32` and `f64`.
///
/// # Safety
///
/// Implementors have no padding and no invalid bit patterns, so a buffer of
/// them may be filled by copying bytes. The trait is sealed.
pub unsafe trait NativeElement: sealed::Sealed + Copy + 'static {
/// Whether `datatype`'s stored bytes are this type's native in-memory
/// representation (same size, byte order, signedness, full precision,
/// IEEE layout), so reading needs a copy and no conversion.
fn is_native(datatype: &Datatype) -> bool;
}
/// A full-width fixed-point type of `size` bytes and the given signedness in
/// the target's byte order.
fn is_native_int(datatype: &Datatype, size: u32, want_signed: bool) -> bool {
let order = if cfg!(target_endian = "little") {
DatatypeByteOrder::LittleEndian
} else {
DatatypeByteOrder::BigEndian
};
matches!(
datatype,
Datatype::FixedPoint { size: s, signed, byte_order, .. }
if *s == size && *signed == want_signed && (size == 1 || *byte_order == order)
) && is_full_width(datatype)
}
macro_rules! native_element {
($($t:ty => |$dt:ident| $check:expr;)*) => {$(
impl sealed::Sealed for $t {}
// SAFETY: a primitive integer or float: no padding, and every bit
// pattern is a valid value.
unsafe impl NativeElement for $t {
fn is_native($dt: &Datatype) -> bool {
$check
}
}
)*};
}
native_element! {
u8 => |dt| is_native_int(dt, 1, false);
i32 => |dt| is_native_int(dt, 4, true);
i64 => |dt| is_native_int(dt, 8, true);
u64 => |dt| is_native_int(dt, 8, false);
f32 => |dt| cfg!(target_endian = "little") && is_native_le_float(dt, FloatFormat::Single);
f64 => |dt| cfg!(target_endian = "little") && is_native_le_float(dt, FloatFormat::Double);
}
/// Copy `count` values of `T` out of `raw`, which holds them in `T`'s native
/// representation (see [`NativeElement::is_native`]), in one copy.
///
/// The buffer is allocated uninitialised and filled by the copy. It used to be
/// `vec![0; count]` first, which for a large dataset meant writing every page
/// twice (zero it, then overwrite it) — about as expensive as the copy itself.
#[cfg(target_endian = "little")]
fn native_le_to_vec<T: Copy>(raw: &[u8], count: usize) -> Vec<T> {
fn native_to_vec<T: NativeElement>(raw: &[u8], count: usize) -> Vec<T> {
let bytes = count * core::mem::size_of::<T>();
debug_assert!(bytes <= raw.len());
assert!(bytes <= raw.len(), "native_to_vec: source too short");
let mut result: Vec<T> = crate::bulk_alloc::vec_for_bulk(count);
// SAFETY: `result` has capacity for `count` values of `T`, i.e. `bytes`
// bytes; `raw` holds at least `bytes` bytes (callers derive `count` from
// `raw.len() / size_of::<T>()`); the regions cannot overlap because
// `result` was just allocated. Every `T` used here (f32/f64/i32/i64) is
// valid for any bit pattern, so after the copy all `count` values are
// bytes; `raw` holds at least `bytes` bytes (asserted); the regions
// cannot overlap because `result` was just allocated. `T: NativeElement`
// is valid for any bit pattern, so after the copy all `count` values are
// initialised and `set_len` is sound.
unsafe {
core::ptr::copy_nonoverlapping(raw.as_ptr(), result.as_mut_ptr().cast::<u8>(), bytes);
@@ -781,6 +784,44 @@ fn native_le_to_vec<T: Copy>(raw: &[u8], count: usize) -> Vec<T> {
result
}
/// Read `selection` of a dataset whose raw bytes (all of them, row-major, of
/// shape `dims`) are `raw` — typically a contiguous dataset's bytes borrowed
/// from the file — straight into a `Vec<T>`, copying each contiguous run of
/// selected elements once.
///
/// Returns `Ok(None)` when `datatype` is not `T`'s native representation
/// ([`NativeElement::is_native`]); the caller then converts through
/// [`read_raw_data_selection`] and the `read_as_*` functions. The selection is
/// validated like every selection read: out-of-range coordinates are
/// [`FormatError::SelectionOutOfBounds`].
pub fn read_selection_native<T: NativeElement>(
raw: &[u8],
dims: &[u64],
datatype: &Datatype,
selection: &crate::selection::Selection,
) -> Result<Option<Vec<T>>, FormatError> {
if !T::is_native(datatype) {
return Ok(None);
}
let elem_size = core::mem::size_of::<T>();
let total = dims
.iter()
.try_fold(1u64, |acc, &d| acc.checked_mul(d))
.ok_or_else(|| FormatError::Overflow("dataset shape overflows".into()))?;
let expected = crate::chunked_read::checked_byte_len(total, elem_size)?;
if raw.len() != expected {
return Err(FormatError::DataSizeMismatch {
expected,
actual: raw.len(),
});
}
if let crate::selection::Selection::All = selection {
return Ok(Some(native_to_vec(raw, expected / elem_size)));
}
crate::partial_read::validate(selection, dims)?;
crate::gather::gather::<T>(raw, dims, elem_size, selection).map(Some)
}
/// Convert raw bytes to `f64` values.
pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatError> {
// Array datatypes read as a flat sequence of their base elements, and
@@ -799,9 +840,8 @@ pub fn read_as_f64(raw: &[u8], datatype: &Datatype) -> Result<Vec<f64>, FormatEr
let count = raw.len() / elem_size;
// Fast path: native-endian f64 — single bulk memcpy
#[cfg(target_endian = "little")]
if is_native_le_float(datatype, FloatFormat::Double) {
return Ok(native_le_to_vec::<f64>(raw, count));
if f64::is_native(datatype) {
return Ok(native_to_vec::<f64>(raw, count));
}
let order = get_byte_order(datatype);
@@ -941,19 +981,8 @@ pub fn read_as_i64(raw: &[u8], datatype: &Datatype) -> Result<Vec<i64>, FormatEr
let count = raw.len() / elem_size;
// Fast path: native LE i64 — single bulk memcpy
#[cfg(target_endian = "little")]
if elem_size == 8
&& is_full_width(datatype)
&& matches!(
datatype,
Datatype::FixedPoint {
byte_order: DatatypeByteOrder::LittleEndian,
signed: true,
..
}
)
{
return Ok(native_le_to_vec::<i64>(raw, count));
if i64::is_native(datatype) {
return Ok(native_to_vec::<i64>(raw, count));
}
let order = get_byte_order(datatype);
@@ -986,6 +1015,12 @@ pub fn read_as_u64(raw: &[u8], datatype: &Datatype) -> Result<Vec<u64>, FormatEr
});
}
let count = raw.len() / elem_size;
// Fast path: native u64 — single bulk memcpy
if u64::is_native(datatype) {
return Ok(native_to_vec::<u64>(raw, count));
}
let order = get_byte_order(datatype);
let mut result = crate::bulk_alloc::vec_for_bulk(count);
for i in 0..count {
@@ -1013,9 +1048,8 @@ pub fn read_as_f32(raw: &[u8], datatype: &Datatype) -> Result<Vec<f32>, FormatEr
let count = raw.len() / elem_size;
// Fast path: native-endian f32 — single bulk memcpy
#[cfg(target_endian = "little")]
if is_native_le_float(datatype, FloatFormat::Single) {
return Ok(native_le_to_vec::<f32>(raw, count));
if f32::is_native(datatype) {
return Ok(native_to_vec::<f32>(raw, count));
}
// Little-endian IEEE half precision (numpy float16): widen directly.
if is_native_le_float(datatype, FloatFormat::Half) {
@@ -1103,19 +1137,8 @@ pub fn read_as_i32(raw: &[u8], datatype: &Datatype) -> Result<Vec<i32>, FormatEr
let count = raw.len() / elem_size;
// Fast path: native LE i32 — single bulk memcpy
#[cfg(target_endian = "little")]
if elem_size == 4
&& is_full_width(datatype)
&& matches!(
datatype,
Datatype::FixedPoint {
byte_order: DatatypeByteOrder::LittleEndian,
signed: true,
..
}
)
{
return Ok(native_le_to_vec::<i32>(raw, count));
if i32::is_native(datatype) {
return Ok(native_to_vec::<i32>(raw, count));
}
let order = get_byte_order(datatype);