1use 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}