Skip to main content

gpu_load/
gpu_load.rs

1//! Sustained GPU memory-bandwidth load generator for heterogeneous contention probing.
2//!
3//! Strix Halo's iGPU and XDNA2 NPU share one LPDDR5X memory controller. This drives a
4//! bandwidth-bound stream of device-to-device copies (or fills) for a fixed wall-clock
5//! window and reports sustained GB/s, so the same load can be measured alone and while
6//! an NPU workload runs concurrently.
7//!
8//! Emits one JSON object on stdout; progress goes to stderr.
9//!
10//! Usage: gpu_load [--seconds F] [--mib N] [--mode copy|fill] [--batch N]
11
12use 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    /// Device-to-device copy: reads `size` and writes `size` per iteration.
25    Copy,
26    /// Fill: writes `size` per iteration, with no read traffic.
27    Fill,
28}
29
30impl Mode {
31    /// Bytes crossing the memory controller for one iteration over a `size`-byte buffer.
32    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        // Far past any cache: this must be a memory-controller measurement, not an L2 one.
51        mib: 512,
52        mode: Mode::Copy,
53        // Enqueue depth per synchronize, so submission overhead does not bound the rate.
54        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    // Touch both allocations so the steady-state loop never pays first-use faults.
101    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    // Warm up outside the measured window: the first batch pays queue and paging costs.
116    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    // Per-batch rates, so the harness can see whether contention was steady or bursty.
122    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}