Skip to main content

mall_probe/
mall_probe.rs

1//! Sweep native XDNA copy working sets. Includes core copy and submission costs;
2//! throughput alone does not distinguish MALL caching from other bottlenecks.
3//! Usage: mall_probe [seconds-per-size]
4use hrx::{
5    Result,
6    fabric::{Device, Engine, XdnaBinding},
7    loom::{Compiler, Specialization},
8};
9use std::time::{Duration, Instant};
10fn main() -> Result<()> {
11    let seconds: f64 = std::env::args()
12        .nth(1)
13        .unwrap_or_else(|| "2".into())
14        .parse()
15        .map_err(|e: std::num::ParseFloatError| hrx::Error::Message(e.to_string()))?;
16    if !seconds.is_finite() || !(0.01..=3600.0).contains(&seconds) {
17        return Err(hrx::Error::Message("seconds must be in 0.01..=3600".into()));
18    }
19    let device = Device::open(Engine::Xdna, 0)?;
20    let compiler = Compiler::for_target(None, device.target())?;
21    println!("working_set_MiB GB/s dispatches/s");
22    for mib in [1, 4, 8, 16, 32, 64] {
23        let bytes = mib * 1024 * 1024;
24        let artifact = compiler
25            .module(include_str!("../tests/kernels/copy.xdna.loom"))
26            .compile(
27                &Specialization::new("copy")
28                    .with_config("copy.packets", (bytes / 4096).to_string()),
29            )?;
30        let input = device
31            .fabric()
32            .allocate(bytes, std::slice::from_ref(&device))?;
33        let output = device
34            .fabric()
35            .allocate(bytes, std::slice::from_ref(&device))?;
36        input.write(0, &vec![0x5a; bytes])?;
37        let program = unsafe {
38            device.prepare_xdna(
39                &artifact,
40                1,
41                &[
42                    XdnaBinding {
43                        buffer: &input,
44                        offset: 0,
45                        length: bytes,
46                    },
47                    XdnaBinding {
48                        buffer: &output,
49                        offset: 0,
50                        length: bytes,
51                    },
52                ],
53            )
54        }?;
55        unsafe { program.dispatch() }?.wait()?;
56        let mut check = vec![0; bytes];
57        output.read(0, &mut check)?;
58        assert!(check.iter().all(|&b| b == 0x5a));
59        let start = Instant::now();
60        let mut count = 0;
61        while start.elapsed() < Duration::from_secs_f64(seconds) {
62            unsafe { program.dispatch() }?.wait()?;
63            count += 1;
64        }
65        let elapsed = start.elapsed().as_secs_f64();
66        println!(
67            "{} {:.3} {:.1}",
68            2 * mib,
69            2.0 * bytes as f64 * count as f64 / elapsed / 1e9,
70            count as f64 / elapsed
71        );
72    }
73    Ok(())
74}