Skip to main content

shared_roundtrip/
shared_roundtrip.rs

1//! Execute GPU fill -> NPU passthrough -> GPU copy on one imported allocation.
2//! Usage: shared_roundtrip [bytes]
3//! Compiles the bundled Loom copy kernel for the active native XDNA device.
4use hrx::{
5    Result,
6    execution::{Access, BindingContract, KernelContract, MemoryPlacement, Runtime},
7};
8fn main() -> Result<()> {
9    let args: Vec<_> = std::env::args().collect();
10    if args.len() > 2 {
11        return Err(hrx::Error::Message(
12            "usage: shared_roundtrip [bytes]".into(),
13        ));
14    }
15    let bytes = args
16        .get(1)
17        .map(String::as_str)
18        .unwrap_or("1048576")
19        .parse::<usize>()
20        .map_err(|e| hrx::Error::Message(e.to_string()))?;
21    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
22        return Err(hrx::Error::Message(
23            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
24        ));
25    }
26    let runtime = Runtime::new()?;
27    let device = runtime.npu(0)?;
28    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
29        .module(include_str!("../tests/kernels/copy.xdna.loom"))
30        .compile(
31            &hrx::loom::Specialization::new("copy")
32                .with_config("copy.packets", (bytes / 4096).to_string()),
33        )?;
34    let binding = |bytes, access| BindingContract {
35        bytes,
36        access,
37        alignment: 4,
38        layout: "u32 passthrough".into(),
39    };
40    let contract = KernelContract {
41        bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
42        constants: vec![],
43    };
44    let kernel = unsafe { device.load_artifact(&artifact, 1, contract) }?;
45    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
46    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
47    let result = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
48    let mut graph = runtime.graph();
49    graph.fill(a.view(), 0x5a)?;
50    graph.npu(&kernel, &[a.view(), c.view()])?;
51    graph.copy(result.view(), c.view())?;
52    let graph = graph.prepare()?;
53    for _ in 0..10 {
54        graph.submit()?.wait()?;
55        let mapped = result.map_read()?;
56        if mapped.iter().any(|&b| b != 0x5a) {
57            return Err(hrx::Error::Message("GPU/NPU roundtrip mismatch".into()));
58        }
59    }
60    println!(
61        "GPU -> NPU -> GPU: all {bytes} bytes correct, 10 executions, {} prepared regions",
62        graph.regions()
63    );
64    Ok(())
65}