Skip to main content

shared_bench/
shared_bench.rs

1//! Warm end-to-end latency for a real NPU passthrough pipeline.
2//! Usage: shared_bench [bytes]
3use hrx::{
4    Result,
5    benchmark::percentile,
6    execution::{Access, BindingContract, KernelContract, MemoryPlacement, Runtime},
7};
8use std::time::Instant;
9fn main() -> Result<()> {
10    let args: Vec<_> = std::env::args().collect();
11    if args.len() > 2 {
12        return Err(hrx::Error::Message("usage: shared_bench [bytes]".into()));
13    }
14    let bytes = args
15        .get(1)
16        .map(String::as_str)
17        .unwrap_or("1048576")
18        .parse::<usize>()
19        .map_err(|e| hrx::Error::Message(e.to_string()))?;
20    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
21        return Err(hrx::Error::Message(
22            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
23        ));
24    }
25    let runtime = Runtime::new()?;
26    let device = runtime.npu(0)?;
27    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
28        .module(include_str!("../tests/kernels/copy.xdna.loom"))
29        .compile(
30            &hrx::loom::Specialization::new("copy")
31                .with_config("copy.packets", (bytes / 4096).to_string()),
32        )?;
33    let binding = |bytes, access| BindingContract {
34        bytes,
35        access,
36        alignment: 4,
37        layout: "passthrough".into(),
38    };
39    let kernel = unsafe {
40        device.load_artifact(
41            &artifact,
42            1,
43            KernelContract {
44                bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
45                constants: vec![],
46            },
47        )
48    }?;
49    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
50    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
51    let mut graph = runtime.graph();
52    graph.fill(a.view(), 0x35)?;
53    graph.npu(&kernel, &[a.view(), c.view()])?;
54    let graph = graph.prepare()?;
55    for _ in 0..5 {
56        graph.submit()?.wait()?;
57    }
58    let mut samples = Vec::new();
59    for _ in 0..100 {
60        let start = Instant::now();
61        graph.submit()?.wait()?;
62        samples.push(start.elapsed().as_secs_f64());
63    }
64    if c.map_read()?.iter().any(|&b| b != 0x35) {
65        return Err(hrx::Error::Message("incorrect benchmark output".into()));
66    }
67    samples.sort_by(f64::total_cmp);
68    println!(
69        "bytes={bytes} p50_us={:.2} p95_us={:.2}",
70        percentile(&samples, 50) * 1e6,
71        percentile(&samples, 95) * 1e6
72    );
73    Ok(())
74}