Skip to main content

NpuDevice

Struct NpuDevice 

Source
pub struct NpuDevice { /* private fields */ }
Expand description

An explicitly selected NPU ordinal.

Implementations§

Source§

impl NpuDevice

Source

pub fn index(&self) -> i32

Physical NPU ordinal.

Source

pub fn target(&self) -> &Target

Exact XDNA deployment target for offline compilation.

Examples found in repository?
examples/shared_roundtrip.rs (line 28)
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}
More examples
Hide additional examples
examples/shared_bench.rs (line 27)
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}
examples/gemm_pipeline.rs (line 108)
60fn main() -> Result<()> {
61    let args: Vec<_> = std::env::args().collect();
62    let args = if args.len() == 1 {
63        vec![
64            args[0].clone(),
65            String::new(),
66            "gemm".into(),
67            "8".into(),
68            "8".into(),
69            "8".into(),
70        ]
71    } else {
72        args
73    };
74    if args.len() != 6 {
75        return Err(hrx::Error::Message(
76            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
77        ));
78    }
79    let parse = |i: usize| {
80        args[i]
81            .parse::<usize>()
82            .map_err(|e| hrx::Error::Message(e.to_string()))
83    };
84    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
85    let extent = |a: usize, b: usize, width: usize| {
86        a.checked_mul(b)
87            .filter(|&count| count > 0 && count <= (1 << 30))
88            .and_then(|count| count.checked_mul(width))
89            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
90    };
91    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
92    let runtime = Runtime::new()?;
93    let device = runtime.npu(0)?;
94    let (source, spec) = if args[1].is_empty() {
95        (
96            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
97            hrx::loom::Specialization::new("gemm"),
98        )
99    } else {
100        (
101            std::fs::read_to_string(&args[1])?,
102            hrx::loom::Specialization::new(&args[2])
103                .with_config("gemm.m", m.to_string())
104                .with_config("gemm.k", k.to_string())
105                .with_config("gemm.n", n.to_string()),
106        )
107    };
108    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
109        .module(&source)
110        .compile(&spec)?;
111    let npu = unsafe {
112        device.load_artifact(
113            &artifact,
114            1,
115            KernelContract {
116                bindings: vec![
117                    binding(a_bytes, Access::Read),
118                    binding(b_bytes, Access::Read),
119                    binding(c_bytes, Access::Write),
120                ],
121                constants: vec![],
122            },
123        )
124    }?;
125    let prep = gpu_kernel(
126        &runtime,
127        include_str!("../tests/kernels/euler.loom"),
128        "krea2_euler",
129        &[
130            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
131            ("krea2.euler.grid_y", "1".into()),
132        ],
133        m * k,
134        vec![
135            binding(a_bytes, Access::ReadWrite),
136            binding(a_bytes, Access::Read),
137        ],
138    )?;
139    let epilogue = gpu_kernel(
140        &runtime,
141        include_str!("../tests/kernels/add_f32.loom"),
142        "add_f32",
143        &[("add.grid", (m * n).div_ceil(256).to_string())],
144        m * n,
145        vec![binding(c_bytes, Access::ReadWrite)],
146    )?;
147    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
148    let weights = allocate(b_bytes)?;
149    let ones = allocate(a_bytes)?;
150    for buffer in [&weights, &ones] {
151        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
152            word.copy_from_slice(&0x3f80u16.to_le_bytes());
153        }
154    }
155    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
156        let a = allocate(a_bytes)?;
157        let c = allocate(c_bytes)?;
158        let mut graph = runtime.graph();
159        graph.fill(a.view(), 0)?;
160        graph.gpu(&prep, &[a.view(), ones.view()])?;
161        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
162        graph.gpu(&epilogue, &[c.view()])?;
163        Ok((graph.prepare()?, c))
164    };
165    let (first, c0) = prepare()?;
166    let (second, c1) = prepare()?;
167    for _ in 0..5 {
168        first.submit()?.wait()?;
169        second.submit()?.wait()?;
170    }
171    let before = runtime.statistics();
172    let mut latency = Vec::with_capacity(50);
173    for _ in 0..50 {
174        let start = Instant::now();
175        first.submit()?.wait()?;
176        latency.push(start.elapsed().as_secs_f64());
177    }
178    let start = Instant::now();
179    for _ in 0..50 {
180        let a = first.submit()?;
181        let b = second.submit()?;
182        a.wait()?;
183        b.wait()?;
184    }
185    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
186    for c in [&c0, &c1] {
187        if c.map_read()?
188            .as_chunks::<4>()
189            .0
190            .iter()
191            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
192        {
193            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
194        }
195    }
196    let after = runtime.statistics();
197    assert_eq!(before.allocations, after.allocations);
198    assert_eq!(before.imports, after.imports);
199    assert_eq!(after.copied_bytes, 0);
200    latency.sort_by(f64::total_cmp);
201    println!(
202        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
203        percentile(&latency, 50) * 1e3,
204        percentile(&latency, 95) * 1e3
205    );
206    println!("{}", serde_json::to_string(&after)?);
207    Ok(())
208}
examples/gpu_npu_parallel.rs (line 140)
104fn main() -> Result<()> {
105    let args: Vec<_> = std::env::args().skip(1).collect();
106    if args.len() > 4 || args.iter().any(|arg| arg == "--help") {
107        println!(
108            "usage: gpu_npu_parallel [gpu-elements=16777216] [npu-matrices=1024] [samples=21] [trace.json]"
109        );
110        return if args.len() <= 4 {
111            Ok(())
112        } else {
113            Err(Error::Message("too many arguments".into()))
114        };
115    }
116    let number = |index: usize, default: usize| -> Result<usize> {
117        args.get(index).map_or(Ok(default), |value| {
118            value
119                .parse()
120                .map_err(|_| Error::Message(format!("invalid integer: {value}")))
121        })
122    };
123    let elements = number(0, 1 << 24)?;
124    let matrices = number(1, 1024)?;
125    let samples = number(2, 21)?;
126    if !(256..=1 << 26).contains(&elements)
127        || !elements.is_multiple_of(256)
128        || !(1..=16384).contains(&matrices)
129        || !(3..=1000).contains(&samples)
130    {
131        return Err(Error::Message("gpu-elements must be a multiple of 256 in 256..=67108864; npu-matrices in 1..=16384; samples in 3..=1000".into()));
132    }
133    let trace_path = args.get(3).map(PathBuf::from);
134    let runtime = Runtime::new()?;
135    let gpu = runtime.gpu()?;
136    let npu = runtime.npu(0)?;
137    println!(
138        "GPU: {}; NPU: {}",
139        gpu.target().as_str(),
140        npu.target().as_str()
141    );
142    let gpu_compiler = Compiler::for_target(None, gpu.target())?;
143    let gpu_artifact = gpu_compiler
144        .import_cxx(CxxSource::new(
145            "affine.cpp",
146            format!(
147                r#"
148#include <hip/hip_runtime.h>
149__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
150void affine(const int* input, int* output) {{
151  unsigned i = blockIdx.x * 256u + threadIdx.x;
152  output[i] = input[i] * 3 + 7;
153}}
154"#,
155                elements / 256
156            ),
157        ))?
158        .compile(&Specialization::new("affine"))?;
159    // SAFETY: the fixed launch covers exactly elements i32 values in each
160    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
161    let gpu_kernel = unsafe {
162        runtime.load_gpu_artifact(
163            &gpu_artifact,
164            &[],
165            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
166        )
167    }?;
168    let npu_artifact = Compiler::for_target(None, npu.target())?
169        .module(include_str!("kernels/batched_gemm.xdna.loom"))
170        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
171    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
172    // produces the same number of 8x8 FP32 outputs on one NPU column.
173    let npu_kernel = unsafe {
174        npu.load_artifact(
175            &npu_artifact,
176            1,
177            contract(&[
178                (matrices * 128, Access::Read),
179                (matrices * 128, Access::Read),
180                (matrices * 256, Access::Write),
181            ]),
182        )
183    }?;
184    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
185    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
187    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let mut upload = runtime.graph();
189    upload.copy(input.view(), host_input.view())?;
190    let upload = upload.prepare()?;
191    let mut download = runtime.graph();
192    download.copy(host_output.view(), output.view())?;
193    let download = download.prepare()?;
194    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
195    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
197
198    let mut gpu_only = runtime.graph();
199    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
200    let gpu_only = gpu_only.prepare()?;
201    let mut npu_only = runtime.graph();
202    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
203    let npu_only = npu_only.prepare()?;
204
205    // Disjoint buffers imply no dependency edge. One submit admits both
206    // branches to separate GPU/NPU workers; one wait joins their completions.
207    let mut parallel = runtime.graph();
208    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
209    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
210    let parallel = parallel.prepare()?;
211    for pass in 0..3 {
212        initialize(&host_input, &a, &b, pass * 2)?;
213        upload.submit()?.wait()?;
214        gpu_only.submit()?.wait()?;
215        npu_only.submit()?.wait()?;
216        download.submit()?.wait()?;
217        validate(&host_output, &c, pass * 2)?;
218        // Change both inputs so validation cannot accept the sequential run's
219        // old output if either parallel branch fails to write its result.
220        initialize(&host_input, &a, &b, pass * 2 + 1)?;
221        upload.submit()?.wait()?;
222        parallel.submit()?.wait()?;
223        download.submit()?.wait()?;
224        validate(&host_output, &c, pass * 2 + 1)?;
225    }
226
227    let before = runtime.statistics();
228    let mut sequential_ms = Vec::with_capacity(samples);
229    let mut parallel_ms = Vec::with_capacity(samples);
230    for sample in 0..samples {
231        // Alternate order to reduce systematic clock/thermal ordering bias.
232        for concurrent in if sample % 2 == 0 {
233            [false, true]
234        } else {
235            [true, false]
236        } {
237            if concurrent {
238                parallel_ms.push(completed_ms(&parallel)?);
239            } else {
240                let start = Instant::now();
241                gpu_only.submit()?.wait()?;
242                npu_only.submit()?.wait()?;
243                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
244            }
245        }
246    }
247    let after = runtime.statistics();
248    if before.allocations != after.allocations || before.imports != after.imports {
249        return Err(Error::Message(
250            "replay unexpectedly allocated or imported storage".into(),
251        ));
252    }
253    download.submit()?.wait()?;
254    validate(&host_output, &c, 5)?;
255    let sequential = median(&mut sequential_ms);
256    let concurrent = median(&mut parallel_ms);
257    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
258    println!(
259        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
260        sequential / concurrent
261    );
262    println!(
263        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
264    );
265
266    // Capture separately so tracing does not affect the benchmark samples.
267    runtime.start_trace(32)?;
268    for _ in 0..4 {
269        parallel.submit()?.wait()?;
270    }
271    let trace = runtime
272        .finish_trace()
273        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
274    let overlap: f64 = trace
275        .events
276        .iter()
277        .filter(|e| e.lane == 1)
278        .flat_map(|gpu| {
279            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
280                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
281                    - gpu.start_us.max(npu.start_us))
282                .max(0.0)
283            })
284        })
285        .sum();
286    println!(
287        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
288    );
289    if let Some(path) = trace_path {
290        std::fs::write(&path, trace.to_json()?)?;
291        println!("Trace: {}", path.display());
292    }
293    Ok(())
294}
Source

pub fn native(&self) -> &Device

Native device domain, independent of any loaded program.

Source

pub unsafe fn load_artifact( &self, artifact: &Artifact, columns: u16, contract: KernelContract, ) -> Result<NpuKernel>

Admit a canonical XDNA artifact and its external binding contract.

§Safety

The native code must be trusted and satisfy the supplied access contract.

Examples found in repository?
examples/shared_roundtrip.rs (line 44)
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}
More examples
Hide additional examples
examples/shared_bench.rs (lines 40-47)
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}
examples/gemm_pipeline.rs (lines 112-123)
60fn main() -> Result<()> {
61    let args: Vec<_> = std::env::args().collect();
62    let args = if args.len() == 1 {
63        vec![
64            args[0].clone(),
65            String::new(),
66            "gemm".into(),
67            "8".into(),
68            "8".into(),
69            "8".into(),
70        ]
71    } else {
72        args
73    };
74    if args.len() != 6 {
75        return Err(hrx::Error::Message(
76            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
77        ));
78    }
79    let parse = |i: usize| {
80        args[i]
81            .parse::<usize>()
82            .map_err(|e| hrx::Error::Message(e.to_string()))
83    };
84    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
85    let extent = |a: usize, b: usize, width: usize| {
86        a.checked_mul(b)
87            .filter(|&count| count > 0 && count <= (1 << 30))
88            .and_then(|count| count.checked_mul(width))
89            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
90    };
91    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
92    let runtime = Runtime::new()?;
93    let device = runtime.npu(0)?;
94    let (source, spec) = if args[1].is_empty() {
95        (
96            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
97            hrx::loom::Specialization::new("gemm"),
98        )
99    } else {
100        (
101            std::fs::read_to_string(&args[1])?,
102            hrx::loom::Specialization::new(&args[2])
103                .with_config("gemm.m", m.to_string())
104                .with_config("gemm.k", k.to_string())
105                .with_config("gemm.n", n.to_string()),
106        )
107    };
108    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
109        .module(&source)
110        .compile(&spec)?;
111    let npu = unsafe {
112        device.load_artifact(
113            &artifact,
114            1,
115            KernelContract {
116                bindings: vec![
117                    binding(a_bytes, Access::Read),
118                    binding(b_bytes, Access::Read),
119                    binding(c_bytes, Access::Write),
120                ],
121                constants: vec![],
122            },
123        )
124    }?;
125    let prep = gpu_kernel(
126        &runtime,
127        include_str!("../tests/kernels/euler.loom"),
128        "krea2_euler",
129        &[
130            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
131            ("krea2.euler.grid_y", "1".into()),
132        ],
133        m * k,
134        vec![
135            binding(a_bytes, Access::ReadWrite),
136            binding(a_bytes, Access::Read),
137        ],
138    )?;
139    let epilogue = gpu_kernel(
140        &runtime,
141        include_str!("../tests/kernels/add_f32.loom"),
142        "add_f32",
143        &[("add.grid", (m * n).div_ceil(256).to_string())],
144        m * n,
145        vec![binding(c_bytes, Access::ReadWrite)],
146    )?;
147    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
148    let weights = allocate(b_bytes)?;
149    let ones = allocate(a_bytes)?;
150    for buffer in [&weights, &ones] {
151        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
152            word.copy_from_slice(&0x3f80u16.to_le_bytes());
153        }
154    }
155    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
156        let a = allocate(a_bytes)?;
157        let c = allocate(c_bytes)?;
158        let mut graph = runtime.graph();
159        graph.fill(a.view(), 0)?;
160        graph.gpu(&prep, &[a.view(), ones.view()])?;
161        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
162        graph.gpu(&epilogue, &[c.view()])?;
163        Ok((graph.prepare()?, c))
164    };
165    let (first, c0) = prepare()?;
166    let (second, c1) = prepare()?;
167    for _ in 0..5 {
168        first.submit()?.wait()?;
169        second.submit()?.wait()?;
170    }
171    let before = runtime.statistics();
172    let mut latency = Vec::with_capacity(50);
173    for _ in 0..50 {
174        let start = Instant::now();
175        first.submit()?.wait()?;
176        latency.push(start.elapsed().as_secs_f64());
177    }
178    let start = Instant::now();
179    for _ in 0..50 {
180        let a = first.submit()?;
181        let b = second.submit()?;
182        a.wait()?;
183        b.wait()?;
184    }
185    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
186    for c in [&c0, &c1] {
187        if c.map_read()?
188            .as_chunks::<4>()
189            .0
190            .iter()
191            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
192        {
193            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
194        }
195    }
196    let after = runtime.statistics();
197    assert_eq!(before.allocations, after.allocations);
198    assert_eq!(before.imports, after.imports);
199    assert_eq!(after.copied_bytes, 0);
200    latency.sort_by(f64::total_cmp);
201    println!(
202        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
203        percentile(&latency, 50) * 1e3,
204        percentile(&latency, 95) * 1e3
205    );
206    println!("{}", serde_json::to_string(&after)?);
207    Ok(())
208}
examples/gpu_npu_parallel.rs (lines 174-182)
104fn main() -> Result<()> {
105    let args: Vec<_> = std::env::args().skip(1).collect();
106    if args.len() > 4 || args.iter().any(|arg| arg == "--help") {
107        println!(
108            "usage: gpu_npu_parallel [gpu-elements=16777216] [npu-matrices=1024] [samples=21] [trace.json]"
109        );
110        return if args.len() <= 4 {
111            Ok(())
112        } else {
113            Err(Error::Message("too many arguments".into()))
114        };
115    }
116    let number = |index: usize, default: usize| -> Result<usize> {
117        args.get(index).map_or(Ok(default), |value| {
118            value
119                .parse()
120                .map_err(|_| Error::Message(format!("invalid integer: {value}")))
121        })
122    };
123    let elements = number(0, 1 << 24)?;
124    let matrices = number(1, 1024)?;
125    let samples = number(2, 21)?;
126    if !(256..=1 << 26).contains(&elements)
127        || !elements.is_multiple_of(256)
128        || !(1..=16384).contains(&matrices)
129        || !(3..=1000).contains(&samples)
130    {
131        return Err(Error::Message("gpu-elements must be a multiple of 256 in 256..=67108864; npu-matrices in 1..=16384; samples in 3..=1000".into()));
132    }
133    let trace_path = args.get(3).map(PathBuf::from);
134    let runtime = Runtime::new()?;
135    let gpu = runtime.gpu()?;
136    let npu = runtime.npu(0)?;
137    println!(
138        "GPU: {}; NPU: {}",
139        gpu.target().as_str(),
140        npu.target().as_str()
141    );
142    let gpu_compiler = Compiler::for_target(None, gpu.target())?;
143    let gpu_artifact = gpu_compiler
144        .import_cxx(CxxSource::new(
145            "affine.cpp",
146            format!(
147                r#"
148#include <hip/hip_runtime.h>
149__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
150void affine(const int* input, int* output) {{
151  unsigned i = blockIdx.x * 256u + threadIdx.x;
152  output[i] = input[i] * 3 + 7;
153}}
154"#,
155                elements / 256
156            ),
157        ))?
158        .compile(&Specialization::new("affine"))?;
159    // SAFETY: the fixed launch covers exactly elements i32 values in each
160    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
161    let gpu_kernel = unsafe {
162        runtime.load_gpu_artifact(
163            &gpu_artifact,
164            &[],
165            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
166        )
167    }?;
168    let npu_artifact = Compiler::for_target(None, npu.target())?
169        .module(include_str!("kernels/batched_gemm.xdna.loom"))
170        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
171    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
172    // produces the same number of 8x8 FP32 outputs on one NPU column.
173    let npu_kernel = unsafe {
174        npu.load_artifact(
175            &npu_artifact,
176            1,
177            contract(&[
178                (matrices * 128, Access::Read),
179                (matrices * 128, Access::Read),
180                (matrices * 256, Access::Write),
181            ]),
182        )
183    }?;
184    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
185    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
187    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let mut upload = runtime.graph();
189    upload.copy(input.view(), host_input.view())?;
190    let upload = upload.prepare()?;
191    let mut download = runtime.graph();
192    download.copy(host_output.view(), output.view())?;
193    let download = download.prepare()?;
194    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
195    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
197
198    let mut gpu_only = runtime.graph();
199    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
200    let gpu_only = gpu_only.prepare()?;
201    let mut npu_only = runtime.graph();
202    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
203    let npu_only = npu_only.prepare()?;
204
205    // Disjoint buffers imply no dependency edge. One submit admits both
206    // branches to separate GPU/NPU workers; one wait joins their completions.
207    let mut parallel = runtime.graph();
208    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
209    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
210    let parallel = parallel.prepare()?;
211    for pass in 0..3 {
212        initialize(&host_input, &a, &b, pass * 2)?;
213        upload.submit()?.wait()?;
214        gpu_only.submit()?.wait()?;
215        npu_only.submit()?.wait()?;
216        download.submit()?.wait()?;
217        validate(&host_output, &c, pass * 2)?;
218        // Change both inputs so validation cannot accept the sequential run's
219        // old output if either parallel branch fails to write its result.
220        initialize(&host_input, &a, &b, pass * 2 + 1)?;
221        upload.submit()?.wait()?;
222        parallel.submit()?.wait()?;
223        download.submit()?.wait()?;
224        validate(&host_output, &c, pass * 2 + 1)?;
225    }
226
227    let before = runtime.statistics();
228    let mut sequential_ms = Vec::with_capacity(samples);
229    let mut parallel_ms = Vec::with_capacity(samples);
230    for sample in 0..samples {
231        // Alternate order to reduce systematic clock/thermal ordering bias.
232        for concurrent in if sample % 2 == 0 {
233            [false, true]
234        } else {
235            [true, false]
236        } {
237            if concurrent {
238                parallel_ms.push(completed_ms(&parallel)?);
239            } else {
240                let start = Instant::now();
241                gpu_only.submit()?.wait()?;
242                npu_only.submit()?.wait()?;
243                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
244            }
245        }
246    }
247    let after = runtime.statistics();
248    if before.allocations != after.allocations || before.imports != after.imports {
249        return Err(Error::Message(
250            "replay unexpectedly allocated or imported storage".into(),
251        ));
252    }
253    download.submit()?.wait()?;
254    validate(&host_output, &c, 5)?;
255    let sequential = median(&mut sequential_ms);
256    let concurrent = median(&mut parallel_ms);
257    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
258    println!(
259        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
260        sequential / concurrent
261    );
262    println!(
263        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
264    );
265
266    // Capture separately so tracing does not affect the benchmark samples.
267    runtime.start_trace(32)?;
268    for _ in 0..4 {
269        parallel.submit()?.wait()?;
270    }
271    let trace = runtime
272        .finish_trace()
273        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
274    let overlap: f64 = trace
275        .events
276        .iter()
277        .filter(|e| e.lane == 1)
278        .flat_map(|gpu| {
279            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
280                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
281                    - gpu.start_us.max(npu.start_us))
282                .max(0.0)
283            })
284        })
285        .sum();
286    println!(
287        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
288    );
289    if let Some(path) = trace_path {
290        std::fs::write(&path, trace.to_json()?)?;
291        println!("Trace: {}", path.display());
292    }
293    Ok(())
294}

Trait Implementations§

Source§

impl Clone for NpuDevice

Source§

fn clone(&self) -> Self

Returns a duplicate of the value. Read more
1.0.0 (const: unstable) · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> CloneToUninit for T
where T: Clone,

Source§

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)
Performs copy-assignment from self to dest. Read more
Source§

impl<T> DropFlavorWrapper<T> for T

Source§

type Flavor = MayDrop

The DropFlavor that wraps T into Self
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T, W> HasTypeWitness<W> for T
where W: MakeTypeWitness<Arg = T>, T: ?Sized,

Source§

const WITNESS: W = W::MAKE

A constant of the type witness
Source§

impl<T> Identity for T
where T: ?Sized,

Source§

const TYPE_EQ: TypeEq<T, <T as Identity>::Type> = TypeEq::NEW

Proof that Self is the same type as Self::Type, provides methods for casting between Self and Self::Type.
Source§

type Type = T

The same type as Self, used to emulate type equality bounds (T == U) with associated type equality constraints (T: Identity<Type = U>).
Source§

impl<T> Instrument for T

Source§

fn instrument(self, span: Span) -> Instrumented<Self> ⓘ

Instruments this type with the provided Span, returning an Instrumented wrapper. Read more
Source§

fn in_current_span(self) -> Instrumented<Self> ⓘ

Instruments this type with the current Span, returning an Instrumented wrapper. Read more
Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> IntoEither for T

Source§

fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ

Converts self into a Left variant of Either<Self, Self> if into_left is true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
where F: FnOnce(&Self) -> bool,

Converts self into a Left variant of Either<Self, Self> if into_left(&self) returns true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more
Source§

impl<T> PolicyExt for T
where T: ?Sized,

Source§

fn and<P, B, E>(self, other: P) -> And<T, P>
where T: Sized + Policy<B, E>, P: Policy<B, E>,

Create a new Policy that returns Action::Follow only if self and other return Action::Follow. Read more
Source§

fn or<P, B, E>(self, other: P) -> Or<T, P>
where T: Sized + Policy<B, E>, P: Policy<B, E>,

Create a new Policy that returns Action::Follow if either self or other returns Action::Follow. Read more
Source§

impl<T> Same for T

Source§

type Output = T

Should always be Self
Source§

impl<T> ToOwned for T
where T: Clone,

Source§

type Owned = T

The resulting type after obtaining ownership.
Source§

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more
Source§

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = !

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, !>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<T> WithSubscriber for T

Source§

fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self> ⓘ
where S: Into<Dispatch>,

Attaches the provided Subscriber to this type, returning a WithDispatch wrapper. Read more
Source§

fn with_current_subscriber(self) -> WithDispatch<Self> ⓘ

Attaches the current default Subscriber to this type, returning a WithDispatch wrapper. Read more