osobh
b2dce41532
CI / test (push) Failing after 3s
than h5py (5e) stable-worldmodel (arXiv 2605.21800, LeCun/Balestriero) supports HDF5 as one of three native formats and measures generic HDF5 at 1,416-1,474 samples/s for per-frame sample loading. This measures clawhdf5 against that shape, hardware-controlled: clawhdf5 and h5py reading the SAME file on the SAME machine. worldmodel_sampling example: mmap an (N,H,W,C) uint8 observation dataset, read each frame once per pass in shuffled (dataloader) order. The file is written by h5py (benchmarks/gen_worldmodel_frames.py) — clawhdf5 parsing an externally-produced HDF5 file is itself the interop result — and read by both clawhdf5 and the h5py counterpart (benchmarks/bench_worldmodel_h5py.py, opening exactly stable-worldmodel's HDF5Dataset: swmr + 256 MB cache). Results (tank, Ryzen 7 7800X3D, 20000x64x64x3 = 246 MB, in page cache, median of 3): clawhdf5 zero-copy view 593k samples/sec 8.1x clawhdf5 materialised copy 518k samples/sec 7.1x h5py (swmr, 256 MB cache) 73k samples/sec 1.0x The materialised-copy row is the fair equal-work comparison (to_vec per frame, matching h5py's numpy materialisation) and is still 7.1x faster; that the copy costs almost nothing shows the gap is h5py's per-frame call overhead, not data movement. Honest caveats in BENCHMARKS.md: absolute numbers are NOT comparable to the paper's (different hardware, smaller frames, no torch/transform), only the same-machine ratio is; this is an in-page-cache measurement isolating read-path overhead, not disk bandwidth. Adds only an example, two benchmark scripts, and a BENCHMARKS.md section — no library code. (Workspace clippy has pre-existing toolchain drift unrelated to this change; tracked separately.)
96 lines
3.4 KiB
Rust
96 lines
3.4 KiB
Rust
//! World-model sample-loading benchmark — clawhdf5 vs the h5py counterpart.
|
||
//!
|
||
//! Reproduces the access pattern of `stable-worldmodel`'s HDF5 dataloader
|
||
//! (arXiv 2605.21800): a dataset of `(N, H, W, C)` uint8 observation frames,
|
||
//! read one frame at a time in shuffled (dataloader) order. That paper
|
||
//! reports generic HDF5 at 1,416–1,474 samples/s (vs Lance 4,815); this
|
||
//! measures clawhdf5 and h5py on the **same machine and file**, so the
|
||
//! comparison is hardware-controlled. Absolute numbers are not comparable to
|
||
//! the paper's (different box, smaller frames, no torch/transform) — only
|
||
//! clawhdf5-vs-h5py *here* is.
|
||
//!
|
||
//! clawhdf5 mmaps the file once and takes a zero-copy `&[u8]` over the
|
||
//! contiguous observation dataset; frame `i` is a subslice, and the OS pages
|
||
//! it in on access. Two modes, because fairness demands both:
|
||
//! * default: sum the frame bytes through the zero-copy view — clawhdf5's
|
||
//! real advantage, no per-frame allocation;
|
||
//! * `--copy`: `to_vec()` each frame first, matching h5py's unavoidable
|
||
//! per-frame numpy materialization, so the two do equal work.
|
||
//!
|
||
//! Usage: `... --example worldmodel_sampling -- <file.h5> [passes] [--copy]`
|
||
|
||
use std::hint::black_box;
|
||
use std::time::Instant;
|
||
|
||
use clawhdf5::MmapFile;
|
||
|
||
fn main() {
|
||
let args: Vec<String> = std::env::args().collect();
|
||
let path = args
|
||
.get(1)
|
||
.expect("usage: worldmodel_sampling <file.h5> [passes] [--copy]");
|
||
let passes: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(5);
|
||
let copy = args.iter().any(|a| a == "--copy");
|
||
|
||
let file = MmapFile::open(path).expect("open");
|
||
let ds = file.dataset("observation").expect("observation dataset");
|
||
let shape = ds.shape().expect("shape");
|
||
let n = shape[0] as usize;
|
||
let frame_bytes: usize = shape[1..].iter().map(|&d| d as usize).product();
|
||
let raw = ds
|
||
.read_raw_slice()
|
||
.expect("read_raw_slice")
|
||
.expect("contiguous zero-copy slice");
|
||
assert_eq!(raw.len(), n * frame_bytes, "unexpected dataset size");
|
||
|
||
let order = shuffled(n);
|
||
|
||
let touch = |slice: &[u8]| -> u64 {
|
||
if copy {
|
||
let owned = slice.to_vec();
|
||
owned.iter().map(|&b| u64::from(b)).sum()
|
||
} else {
|
||
slice.iter().map(|&b| u64::from(b)).sum()
|
||
}
|
||
};
|
||
|
||
// Warm one pass (page-in), then time.
|
||
let mut sink = 0u64;
|
||
for &i in &order {
|
||
sink = sink.wrapping_add(touch(&raw[i * frame_bytes..(i + 1) * frame_bytes]));
|
||
}
|
||
black_box(sink);
|
||
|
||
let t0 = Instant::now();
|
||
let mut sink = 0u64;
|
||
for _ in 0..passes {
|
||
for &i in &order {
|
||
sink = sink.wrapping_add(touch(&raw[i * frame_bytes..(i + 1) * frame_bytes]));
|
||
}
|
||
}
|
||
black_box(sink);
|
||
let elapsed = t0.elapsed().as_secs_f64();
|
||
|
||
let total = (n * passes) as f64;
|
||
let mode = if copy {
|
||
"materialized copy"
|
||
} else {
|
||
"zero-copy view"
|
||
};
|
||
println!("clawhdf5 ({mode}): {n} frames x {passes} passes in {elapsed:.3}s");
|
||
println!("clawhdf5 ({mode}): {:.0} samples/sec", total / elapsed);
|
||
}
|
||
|
||
fn shuffled(n: usize) -> Vec<usize> {
|
||
let mut v: Vec<usize> = (0..n).collect();
|
||
let mut state: u64 = 0x9E37_79B9_7F4A_7C15;
|
||
for i in (1..n).rev() {
|
||
state = state
|
||
.wrapping_mul(6364136223846793005)
|
||
.wrapping_add(1442695040888963407);
|
||
let j = (state >> 33) as usize % (i + 1);
|
||
v.swap(i, j);
|
||
}
|
||
v
|
||
}
|