metrale_storage/
bench.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Read micro-benchmarks for the storage backends, in MiB/s, from a file
4//! opened with `O_DIRECT` (`open_test_file`, which also advises the kernel to drop the
5//! file's cached pages) into device memory. `bench_cufile` reads with cuFile straight
6//! into the device buffer; `bench_io_uring` and `bench_posix` read into a pinned host
7//! buffer and copy each read to the device with `cuMemcpyHtoDAsync`. Sequential reads
8//! cycle over the first `SEQ_ITERS` reads of the file; random reads land on
9//! `RAND_IO_BYTES`-aligned offsets.
10//!
11//! Owner: storage.
12//! Invariants: none beyond the types.
13
14use anyhow::{Context, Result, bail};
15use std::ffi::c_void;
16use std::fs::OpenOptions;
17use std::os::fd::{AsRawFd, RawFd};
18use std::os::unix::fs::OpenOptionsExt;
19use std::path::Path;
20use std::time::Instant;
21
22use crate::cuda_min::{CudaCtx, DeviceBuffer, PinnedBuffer, copy_h_to_d_async, stream_sync};
23
24pub const SEQ_IO_BYTES: usize = 4 * 1024 * 1024;
25pub const RAND_IO_BYTES: usize = 64 * 1024;
26pub const SEQ_ITERS: usize = 64;
27pub const RAND_ITERS: usize = 1024;
28
29#[derive(Debug, Clone, Copy)]
30pub struct BenchResult {
31    pub mib_per_sec: f64,
32    pub iters: usize,
33    pub bytes_per_io: usize,
34}
35
36fn open_o_direct(path: &Path) -> Result<std::fs::File> {
37    OpenOptions::new()
38        .read(true)
39        .custom_flags(libc::O_DIRECT)
40        .open(path)
41        .with_context(|| format!("open O_DIRECT {}", path.display()))
42}
43
44fn drop_pagecache(fd: RawFd) {
45    unsafe {
46        libc::posix_fadvise(fd, 0, 0, libc::POSIX_FADV_DONTNEED);
47    }
48}
49
50fn rand_offset(file_bytes: u64, io_bytes: usize, i: usize) -> u64 {
51    let max = (file_bytes - io_bytes as u64) / RAND_IO_BYTES as u64;
52    let pseudo = (i as u64).wrapping_mul(2_654_435_761) % max.max(1);
53    pseudo * RAND_IO_BYTES as u64
54}
55
56pub fn bench_cufile(
57    cufile: &metrale_gpu_sys::cufile::CuFile,
58    fd: RawFd,
59    file_bytes: u64,
60    dev: &DeviceBuffer,
61    io_bytes: usize,
62    iters: usize,
63    sequential: bool,
64) -> Result<BenchResult> {
65    use metrale_gpu_sys::cufile::*;
66    let mut descr = CUfileDescr_t {
67        type_: CU_FILE_HANDLE_TYPE_OPAQUE_FD,
68        handle: CUfileDescrHandle::from_fd(fd),
69        fs_ops: std::ptr::null(),
70    };
71    let mut handle: CUfileHandle_t = std::ptr::null_mut();
72    let err = unsafe { (cufile.handle_register)(&mut handle, &mut descr) };
73    if err.err != CU_FILE_SUCCESS {
74        bail!(
75            "cuFileHandleRegister failed: {} ({})",
76            err.err,
77            err_to_str(err.err)
78        );
79    }
80    let _ = unsafe { (cufile.buf_register)(dev.ptr as *const c_void, dev.bytes, 0) };
81    let t = Instant::now();
82    for i in 0..iters {
83        let off = if sequential {
84            ((i % SEQ_ITERS) * io_bytes) as i64
85        } else {
86            rand_offset(file_bytes, io_bytes, i) as i64
87        };
88        let n = unsafe { (cufile.read)(handle, dev.ptr as *mut c_void, io_bytes, off as _, 0) };
89        if n != io_bytes as isize {
90            unsafe {
91                let _ = (cufile.buf_deregister)(dev.ptr as *const c_void);
92                (cufile.handle_deregister)(handle);
93            }
94            bail!("cuFileRead returned {n}, expected {io_bytes}");
95        }
96    }
97    let dt = t.elapsed().as_secs_f64();
98    unsafe {
99        let _ = (cufile.buf_deregister)(dev.ptr as *const c_void);
100        (cufile.handle_deregister)(handle);
101    }
102    let bytes = (io_bytes * iters) as f64;
103    Ok(BenchResult {
104        mib_per_sec: bytes / dt / (1024.0 * 1024.0),
105        iters,
106        bytes_per_io: io_bytes,
107    })
108}
109
110#[allow(clippy::too_many_arguments)]
111pub fn bench_io_uring(
112    ctx: &CudaCtx,
113    fd: RawFd,
114    file_bytes: u64,
115    pinned: &PinnedBuffer,
116    dev: &DeviceBuffer,
117    io_bytes: usize,
118    iters: usize,
119    sequential: bool,
120) -> Result<BenchResult> {
121    use io_uring::{IoUring, opcode, types};
122    let mut ring = IoUring::builder()
123        .setup_sqpoll(2_000)
124        .build(64)
125        .context("io_uring setup_sqpoll")?;
126    let t = Instant::now();
127    let chunk = io_bytes;
128    let host_ptr = pinned.ptr as *mut u8;
129    for i in 0..iters {
130        let off = if sequential {
131            ((i % SEQ_ITERS) * chunk) as u64
132        } else {
133            rand_offset(file_bytes, chunk, i)
134        };
135        let read_e = opcode::Read::new(types::Fd(fd), host_ptr, chunk as u32)
136            .offset(off)
137            .build()
138            .user_data(0);
139        unsafe {
140            ring.submission()
141                .push(&read_e)
142                .map_err(|_| anyhow::anyhow!("sq full"))?
143        };
144        ring.submit_and_wait(1)
145            .context("io_uring submit_and_wait")?;
146        let mut cq = ring.completion();
147        let entry: io_uring::cqueue::Entry = cq.next().expect("cqe");
148        if entry.result() != chunk as i32 {
149            bail!(
150                "io_uring read returned {}, expected {}",
151                entry.result(),
152                chunk
153            );
154        }
155        drop(cq);
156        copy_h_to_d_async(dev.ptr, host_ptr as *const c_void, chunk, ctx.stream)?;
157    }
158    stream_sync(ctx.stream)?;
159    let dt = t.elapsed().as_secs_f64();
160    let bytes = (io_bytes * iters) as f64;
161    Ok(BenchResult {
162        mib_per_sec: bytes / dt / (1024.0 * 1024.0),
163        iters,
164        bytes_per_io: io_bytes,
165    })
166}
167
168#[allow(clippy::too_many_arguments)]
169pub fn bench_posix(
170    ctx: &CudaCtx,
171    fd: RawFd,
172    file_bytes: u64,
173    pinned: &PinnedBuffer,
174    dev: &DeviceBuffer,
175    io_bytes: usize,
176    iters: usize,
177    sequential: bool,
178) -> Result<BenchResult> {
179    let host_ptr = pinned.ptr as *mut u8;
180    let t = Instant::now();
181    for i in 0..iters {
182        let off = if sequential {
183            ((i % SEQ_ITERS) * io_bytes) as i64
184        } else {
185            rand_offset(file_bytes, io_bytes, i) as i64
186        };
187        let n = unsafe { libc::pread(fd, host_ptr as *mut c_void, io_bytes, off) };
188        if n != io_bytes as isize {
189            bail!(
190                "pread returned {n}, expected {io_bytes} (errno {})",
191                std::io::Error::last_os_error()
192            );
193        }
194        copy_h_to_d_async(dev.ptr, host_ptr as *const c_void, io_bytes, ctx.stream)?;
195    }
196    stream_sync(ctx.stream)?;
197    let dt = t.elapsed().as_secs_f64();
198    let bytes = (io_bytes * iters) as f64;
199    Ok(BenchResult {
200        mib_per_sec: bytes / dt / (1024.0 * 1024.0),
201        iters,
202        bytes_per_io: io_bytes,
203    })
204}
205
206pub fn open_test_file(path: &Path) -> Result<(std::fs::File, u64)> {
207    let f = open_o_direct(path)?;
208    let len = f.metadata()?.len();
209    drop_pagecache(f.as_raw_fd());
210    Ok((f, len))
211}