spark_storage/
bench.rs

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

spark_storage/
bench.rs
