Skip to main content

queue_pool_bench/
queue_pool_bench.rs

1//! Queue-pool throughput and latency on a prepared GPU. Run with --release.
2//! Usage: queue_pool_bench QUEUES [SAMPLES], one fresh process per queue count.
3use hrx::{Buffer, Device, Result, Stream, benchmark::Distribution};
4use serde_json::json;
5use std::{sync::Barrier, time::Instant};
6
7const STREAMS: usize = 16;
8const WARMUP: usize = 3;
9
10struct Lane {
11    stream: Stream,
12    source: Buffer,
13    destination: Buffer,
14}
15impl Lane {
16    fn new(device: &Device, bytes: usize) -> Result<Self> {
17        let mut stream = device.stream()?;
18        let source = stream.allocate(bytes)?;
19        let destination = stream.allocate(bytes)?;
20        stream.fill(source.binding(), 0x3c)?;
21        stream.copy(destination.binding(), source.binding())?;
22        stream.synchronize()?;
23        Ok(Self {
24            stream,
25            source,
26            destination,
27        })
28    }
29    fn enqueue(&self) -> Result<()> {
30        self.stream
31            .copy(self.destination.binding(), self.source.binding())
32    }
33    fn check(&mut self) -> Result<()> {
34        let mut bytes = [0; 64];
35        self.stream
36            .read_blocking(self.destination.slice(0, 64), &mut bytes)?;
37        assert_eq!(bytes, [0x3c; 64]);
38        Ok(())
39    }
40}
41
42fn milliseconds(start: Instant) -> f64 {
43    start.elapsed().as_secs_f64() * 1_000.0
44}
45
46fn throughput(lanes: &mut [Lane], samples: usize, batch: usize) -> Result<Distribution> {
47    let mut timings = Vec::new();
48    for round in 0..samples + WARMUP {
49        let start = Instant::now();
50        for _ in 0..batch {
51            for lane in &*lanes {
52                lane.enqueue()?;
53            }
54        }
55        for lane in &mut *lanes {
56            lane.stream.synchronize()?;
57        }
58        if round >= WARMUP {
59            timings.push(milliseconds(start));
60        }
61    }
62    Distribution::from_samples(timings)
63}
64
65fn concurrent(lanes: &mut [Lane], samples: usize, batch: usize) -> Result<Distribution> {
66    let mut timings = Vec::new();
67    let barrier = Barrier::new(lanes.len());
68    for round in 0..samples + WARMUP {
69        // Creation and joining are excluded from each worker's elapsed time.
70        // The slowest worker determines the batch's throughput.
71        let elapsed = std::thread::scope(|scope| {
72            let workers: Vec<_> = lanes
73                .iter_mut()
74                .map(|lane| {
75                    let barrier = &barrier;
76                    scope.spawn(move || -> Result<f64> {
77                        barrier.wait();
78                        let start = Instant::now();
79                        for _ in 0..batch {
80                            lane.enqueue()?;
81                        }
82                        lane.stream.synchronize()?;
83                        Ok(milliseconds(start))
84                    })
85                })
86                .collect();
87            workers
88                .into_iter()
89                .map(|worker| worker.join().expect("benchmark worker panicked"))
90                .collect::<Result<Vec<_>>>()
91        })?;
92        if round >= WARMUP {
93            timings.push(elapsed.into_iter().fold(0.0, f64::max));
94        }
95    }
96    Distribution::from_samples(timings)
97}
98
99fn main() -> Result<()> {
100    let args: Vec<_> = std::env::args().skip(1).collect();
101    let parse = |value: &str| value.parse::<usize>().ok().filter(|&n| n > 0);
102    let queues = args.first().and_then(|arg| parse(arg));
103    let samples = args.get(1).map_or(Some(15), |arg| parse(arg));
104    let (queues, samples) = match (queues, samples, args.len()) {
105        (Some(queues), Some(samples), 1..=2) => (queues, samples),
106        _ => {
107            return Err(hrx::Error::Message(
108                "usage: queue_pool_bench QUEUES [SAMPLES] (positive integers)".into(),
109            ));
110        }
111    };
112    let device = Device::open(0)?;
113    device.set_stream_queue_count(queues)?;
114    let start = Instant::now();
115    let streams = (0..STREAMS)
116        .map(|_| device.stream())
117        .collect::<Result<Vec<_>>>()?;
118    let creation_ms = milliseconds(start);
119    // Keep these streams alive so every case uses the same already-created queues.
120    let mut small = (0..STREAMS)
121        .map(|_| Lane::new(&device, 4096))
122        .collect::<Result<Vec<_>>>()?;
123    let small_batch = 1024;
124    let small_times = throughput(&mut small, samples, small_batch)?;
125    let threaded_times = concurrent(&mut small, samples, small_batch)?;
126    let mut idle_latency = Vec::new();
127    for round in 0..samples + WARMUP {
128        for lane in &mut small {
129            let start = Instant::now();
130            lane.enqueue()?;
131            lane.stream.synchronize()?;
132            if round >= WARMUP {
133                idle_latency.push(milliseconds(start));
134            }
135        }
136    }
137    // A 1 GiB aggregate working set exceeds the on-device caches on gfx1151.
138    let bulk_bytes = 32 * 1024 * 1024;
139    let bulk_batch = 16;
140    let mut bulk = (0..STREAMS)
141        .map(|_| Lane::new(&device, bulk_bytes))
142        .collect::<Result<Vec<_>>>()?;
143    let bulk_times = throughput(&mut bulk, samples, bulk_batch)?;
144    let mut loaded_latency = Vec::new();
145    for round in 0..samples + WARMUP {
146        // Probe every queue position. Enqueue a bulk backlog on all streams,
147        // then measure one small copy's completion before draining the backlog.
148        for probe in &mut small {
149            for _ in 0..4 {
150                for lane in &bulk {
151                    lane.enqueue()?;
152                }
153            }
154            let start = Instant::now();
155            probe.enqueue()?;
156            probe.stream.synchronize()?;
157            if round >= WARMUP {
158                loaded_latency.push(milliseconds(start));
159            }
160            for lane in &mut bulk {
161                lane.stream.synchronize()?;
162            }
163        }
164    }
165    for lane in small.iter_mut().chain(bulk.iter_mut()) {
166        lane.check()?;
167    }
168    let idle = Distribution::from_samples(idle_latency)?;
169    let loaded = Distribution::from_samples(loaded_latency)?;
170    println!(
171        "{}",
172        json!({
173            "target": device.target().as_str(), "queues": queues, "streams": STREAMS,
174            "samples": samples, "warmup": WARMUP, "stream_creation_ms": creation_ms,
175            "small_commands_per_second": STREAMS as f64 * small_batch as f64 * 1_000.0 / small_times.median_ms,
176            "threaded_commands_per_second": STREAMS as f64 * small_batch as f64 * 1_000.0 / threaded_times.median_ms,
177            "bulk_copy_gb_s": 2.0 * bulk_bytes as f64 * STREAMS as f64 * bulk_batch as f64 / (bulk_times.median_ms * 1e6),
178            "small_batch_median_ms": small_times.median_ms, "small_batch_p95_ms": small_times.p95_ms,
179            "threaded_batch_median_ms": threaded_times.median_ms, "threaded_batch_p95_ms": threaded_times.p95_ms,
180            "bulk_batch_median_ms": bulk_times.median_ms, "bulk_batch_p95_ms": bulk_times.p95_ms,
181            "idle_latency_median_us": idle.median_ms * 1_000.0, "idle_latency_p95_us": idle.p95_ms * 1_000.0,
182            "loaded_latency_median_us": loaded.median_ms * 1_000.0, "loaded_latency_p95_us": loaded.p95_ms * 1_000.0,
183        })
184    );
185    drop(streams);
186    Ok(())
187}