1use hrx::{Result, Stream};
13use std::time::{Duration, Instant};
14
15struct Options {
16 seconds: f64,
17 mib: usize,
18 mode: Mode,
19 batch: usize,
20}
21
22#[derive(Clone, Copy, PartialEq)]
23enum Mode {
24 Copy,
26 Fill,
28}
29
30impl Mode {
31 fn traffic(self, size: usize) -> usize {
33 match self {
34 Mode::Copy => 2 * size,
35 Mode::Fill => size,
36 }
37 }
38
39 fn name(self) -> &'static str {
40 match self {
41 Mode::Copy => "copy",
42 Mode::Fill => "fill",
43 }
44 }
45}
46
47fn parse() -> std::result::Result<Options, String> {
48 let mut options = Options {
49 seconds: 10.0,
50 mib: 512,
52 mode: Mode::Copy,
53 batch: 16,
55 };
56 let args: Vec<String> = std::env::args().skip(1).collect();
57 let mut index = 0;
58 while index < args.len() {
59 let value = || -> std::result::Result<String, String> {
60 args.get(index + 1)
61 .cloned()
62 .ok_or_else(|| format!("{} needs a value", args[index]))
63 };
64 match args[index].as_str() {
65 "--seconds" => options.seconds = value()?.parse().map_err(|e| format!("{e}"))?,
66 "--mib" => options.mib = value()?.parse().map_err(|e| format!("{e}"))?,
67 "--batch" => options.batch = value()?.parse().map_err(|e| format!("{e}"))?,
68 "--mode" => {
69 options.mode = match value()?.as_str() {
70 "copy" => Mode::Copy,
71 "fill" => Mode::Fill,
72 other => return Err(format!("unknown mode {other}")),
73 }
74 }
75 other => return Err(format!("unknown argument {other}")),
76 }
77 index += 2;
78 }
79 if options.mib == 0
80 || options.batch == 0
81 || (!options.seconds.is_finite() || options.seconds <= 0.0)
82 {
83 return Err("seconds, mib and batch must all be positive".into());
84 }
85 Ok(options)
86}
87
88fn main() -> Result<()> {
89 let options = match parse() {
90 Ok(options) => options,
91 Err(message) => {
92 eprintln!("gpu_load: {message}");
93 std::process::exit(2);
94 }
95 };
96 let size = options.mib * 1024 * 1024;
97 let mut stream = Stream::open()?;
98 let source = stream.allocate(size)?;
99 let destination = stream.allocate(size)?;
100 stream.fill(source.binding(), 0x3c)?;
102 stream.fill(destination.binding(), 0)?;
103 stream.synchronize()?;
104
105 let run_batch = |stream: &mut Stream| -> Result<()> {
106 for _ in 0..options.batch {
107 match options.mode {
108 Mode::Copy => stream.copy(destination.binding(), source.binding())?,
109 Mode::Fill => stream.fill(destination.binding(), 0x5a)?,
110 }
111 }
112 stream.synchronize()
113 };
114
115 run_batch(&mut stream)?;
117
118 let window = Duration::from_secs_f64(options.seconds);
119 let start = Instant::now();
120 let mut iterations = 0usize;
121 let mut batch_rates = Vec::new();
123 while start.elapsed() < window {
124 let batch_start = Instant::now();
125 run_batch(&mut stream)?;
126 let seconds = batch_start.elapsed().as_secs_f64();
127 iterations += options.batch;
128 let bytes = options.mode.traffic(size) * options.batch;
129 batch_rates.push(bytes as f64 / seconds / 1e9);
130 }
131 let elapsed = start.elapsed().as_secs_f64();
132 let total = options.mode.traffic(size) as f64 * iterations as f64;
133 let mean = total / elapsed / 1e9;
134 batch_rates.sort_by(f64::total_cmp);
135 let median = batch_rates[batch_rates.len() / 2];
136 let low = batch_rates[batch_rates.len() / 20];
137
138 println!(
139 concat!(
140 r#"{{"device":"gpu","mode":"{}","buffer_mib":{},"batch":{},"#,
141 r#""seconds":{:.4},"iterations":{},"gb_s_mean":{:.3},"#,
142 r#""gb_s_median":{:.3},"gb_s_p5":{:.3}}}"#
143 ),
144 options.mode.name(),
145 options.mib,
146 options.batch,
147 elapsed,
148 iterations,
149 mean,
150 median,
151 low
152 );
153 eprintln!(
154 "gpu_load: {} {:.1} GB/s mean over {:.1}s ({} iterations)",
155 options.mode.name(),
156 mean,
157 elapsed,
158 iterations
159 );
160 Ok(())
161}