Skip to main content

Runtime

Struct Runtime 

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

Shared scheduler and allocation domain. Clones use the same dependency state.

Upload, compute, download, and NPU regions have independent scheduling lanes. Each lane runs one region at a time; conflicting memory accesses are ordered across all lanes and submissions. Actual device overlap depends on hardware.

Implementations§

Source§

impl Runtime

Source

pub unsafe fn with_gpu_access<T>( &self, stream: &mut Stream, accesses: &[GpuAccess], callback: impl FnOnce(&mut Stream, &[View<'_>]) -> Result<T>, ) -> Result<T>

Run existing Stream operations against coordinated allocations at a stage boundary. Reservations are allocation-wide; conflicting submitted work completes first. A live host mapping or another handoff returns Busy. The supplied stream is synchronized before releasing the reservations, including when the callback fails or panics.

§Safety

Declare every access to these allocations accurately. The callback must not retain pointers, buffers, or recorded uses, and must submit all such work on the supplied stream. No work may escape onto another stream. Kernel dispatch still carries the low-level GPU API’s safety contract.

Source§

impl Runtime

Source

pub fn start_trace(&self, event_capacity: usize) -> Result<()>

Begin an opt-in bounded trace of host-observed native-region latencies.

Examples found in repository?
examples/gpu_npu_parallel.rs (line 268)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn finish_trace(&self) -> Option<ExecutionTrace>

Stop capture without waiting for pending work. Wait for its completions first when the capture should include the entire request.

Examples found in repository?
examples/gpu_npu_parallel.rs (line 273)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn owns(&self, view: &BufferView) -> bool

Whether a view belongs to this runtime’s allocation and scheduling domain.

Source

pub fn same_domain(&self, other: &Self) -> bool

Whether two handles share the same allocation and scheduling domain.

Source

pub fn memory_budget(&self) -> Option<&MemoryBudget>

Shared allocation ceiling, if selected at construction. Native clients can attach it to their streams before loading weights or workspace.

Source

pub fn new() -> Result<Self>

Create a runtime. Native devices are opened lazily when requested.

Examples found in repository?
examples/graph_prepare.rs (line 5)
4fn main() -> hrx::Result<()> {
5    let runtime = Runtime::new()?;
6    let buffer = runtime.allocate(4096, MemoryPlacement::GpuLocal)?;
7    let stream = Stream::open()?;
8    let raw = stream.allocate(4096)?;
9    for count in [128, 512, 1024, 2048] {
10        let begin = Instant::now();
11        let mut graph = runtime.graph();
12        for _ in 0..count {
13            graph.fill(buffer.view(), 0)?;
14        }
15        let graph = graph.prepare()?;
16        let inferred = begin.elapsed();
17        let begin = Instant::now();
18        let mut native = stream.graph()?;
19        let mut last = None;
20        for _ in 0..count {
21            last = Some(native.fill(last.as_slice(), raw.binding(), 0)?);
22        }
23        let _native = native.finish()?;
24        println!(
25            "nodes={count} inferred_us={} explicit_us={}",
26            inferred.as_micros(),
27            begin.elapsed().as_micros()
28        );
29        graph.submit()?.wait()?;
30    }
31    Ok(())
32}
More examples
Hide additional examples
examples/shared_roundtrip.rs (line 26)
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}
examples/shared_bench.rs (line 25)
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 92)
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 134)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn with_options(options: RuntimeOptions) -> Result<Self>

Create a runtime with bounded pending submission and graph capacities.

Source

pub fn gpu(&self) -> Result<GpuDevice>

Open the selected GPU and return its actual target.

Examples found in repository?
examples/gemm_pipeline.rs (line 31)
23fn gpu_kernel(
24    runtime: &Runtime,
25    source: &str,
26    symbol: &str,
27    config: &[(&str, String)],
28    elements: usize,
29    bindings: Vec<BindingContract>,
30) -> Result<GpuKernel> {
31    let compiler = hrx::loom::Compiler::for_target(None, runtime.gpu()?.target())?;
32    let module = compiler.module(source);
33    let mut spec = hrx::loom::Specialization::new(symbol);
34    for (key, value) in config {
35        spec.set_config(*key, value.clone());
36    }
37    let artifact = module.compile(&spec)?;
38    let stream = hrx::gpu::Stream::open()?;
39    let raw = unsafe { stream.load_artifact(&artifact) }?;
40    let mut constants = hrx::gpu::Constants::new();
41    match raw.info().constant_byte_length {
42        8 => constants.push(elements as u32)?,
43        12 => constants.push(elements as u64)?,
44        _ => return Err(hrx::Error::Message("unexpected scalar ABI".into())),
45    }
46    constants.push(1f32)?;
47    unsafe {
48        runtime.load_gpu_kernel(
49            artifact.path(),
50            symbol,
51            [elements.div_ceil(256) as u32, 1, 1],
52            [256, 1, 1],
53            KernelContract {
54                bindings,
55                constants: constants.as_bytes().to_vec(),
56            },
57        )
58    }
59}
More examples
Hide additional examples
examples/gpu_npu_parallel.rs (line 135)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn npu(&self, index: i32) -> Result<NpuDevice>

Select an NPU. Opening its program validates actual device access.

Examples found in repository?
examples/shared_roundtrip.rs (line 27)
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 26)
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 93)
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 136)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn allocate( &self, bytes: usize, placement: MemoryPlacement, ) -> Result<Buffer>

Allocate initialized storage. Shared memory is exported and imported once.

Examples found in repository?
examples/graph_prepare.rs (line 6)
4fn main() -> hrx::Result<()> {
5    let runtime = Runtime::new()?;
6    let buffer = runtime.allocate(4096, MemoryPlacement::GpuLocal)?;
7    let stream = Stream::open()?;
8    let raw = stream.allocate(4096)?;
9    for count in [128, 512, 1024, 2048] {
10        let begin = Instant::now();
11        let mut graph = runtime.graph();
12        for _ in 0..count {
13            graph.fill(buffer.view(), 0)?;
14        }
15        let graph = graph.prepare()?;
16        let inferred = begin.elapsed();
17        let begin = Instant::now();
18        let mut native = stream.graph()?;
19        let mut last = None;
20        for _ in 0..count {
21            last = Some(native.fill(last.as_slice(), raw.binding(), 0)?);
22        }
23        let _native = native.finish()?;
24        println!(
25            "nodes={count} inferred_us={} explicit_us={}",
26            inferred.as_micros(),
27            begin.elapsed().as_micros()
28        );
29        graph.submit()?.wait()?;
30    }
31    Ok(())
32}
More examples
Hide additional examples
examples/shared_roundtrip.rs (line 45)
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}
examples/shared_bench.rs (line 49)
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 147)
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 185)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub unsafe fn adopt_gpu_buffer( &self, buffer: Buffer, placement: MemoryPlacement, ) -> Result<Buffer>

Transfer an initialized allocation into this runtime, optionally importing it into an NPU program once. Only GpuLocal and Shared placement are accepted.

§Safety

All bytes must be initialized and prior uses complete. No old pointer, recorded graph or other external alias may access the allocation after adoption, except through this runtime’s scoped GPU handoff.

Source

pub fn statistics(&self) -> Statistics

Observe allocation, transfer, and execution counters.

Examples found in repository?
examples/gemm_pipeline.rs (line 171)
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}
More examples
Hide additional examples
examples/gpu_npu_parallel.rs (line 228)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub fn graph(&self) -> Graph

Start describing a reusable pipeline in this allocation domain.

Examples found in repository?
examples/graph_prepare.rs (line 11)
4fn main() -> hrx::Result<()> {
5    let runtime = Runtime::new()?;
6    let buffer = runtime.allocate(4096, MemoryPlacement::GpuLocal)?;
7    let stream = Stream::open()?;
8    let raw = stream.allocate(4096)?;
9    for count in [128, 512, 1024, 2048] {
10        let begin = Instant::now();
11        let mut graph = runtime.graph();
12        for _ in 0..count {
13            graph.fill(buffer.view(), 0)?;
14        }
15        let graph = graph.prepare()?;
16        let inferred = begin.elapsed();
17        let begin = Instant::now();
18        let mut native = stream.graph()?;
19        let mut last = None;
20        for _ in 0..count {
21            last = Some(native.fill(last.as_slice(), raw.binding(), 0)?);
22        }
23        let _native = native.finish()?;
24        println!(
25            "nodes={count} inferred_us={} explicit_us={}",
26            inferred.as_micros(),
27            begin.elapsed().as_micros()
28        );
29        graph.submit()?.wait()?;
30    }
31    Ok(())
32}
More examples
Hide additional examples
examples/shared_roundtrip.rs (line 48)
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}
examples/shared_bench.rs (line 51)
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 158)
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 189)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub unsafe fn load_gpu_kernel( &self, path: impl AsRef<Path>, symbol: &str, grid: [u32; 3], block: [u32; 3], contract: KernelContract, ) -> Result<GpuKernel>

Load a fixed GPU specialization with its trusted memory contract.

§Safety

The code must obey the contract for this grid/block and every accepted binding. Kernel memory access is not sandboxed by native binding lengths.

Examples found in repository?
examples/gemm_pipeline.rs (lines 48-57)
23fn gpu_kernel(
24    runtime: &Runtime,
25    source: &str,
26    symbol: &str,
27    config: &[(&str, String)],
28    elements: usize,
29    bindings: Vec<BindingContract>,
30) -> Result<GpuKernel> {
31    let compiler = hrx::loom::Compiler::for_target(None, runtime.gpu()?.target())?;
32    let module = compiler.module(source);
33    let mut spec = hrx::loom::Specialization::new(symbol);
34    for (key, value) in config {
35        spec.set_config(*key, value.clone());
36    }
37    let artifact = module.compile(&spec)?;
38    let stream = hrx::gpu::Stream::open()?;
39    let raw = unsafe { stream.load_artifact(&artifact) }?;
40    let mut constants = hrx::gpu::Constants::new();
41    match raw.info().constant_byte_length {
42        8 => constants.push(elements as u32)?,
43        12 => constants.push(elements as u64)?,
44        _ => return Err(hrx::Error::Message("unexpected scalar ABI".into())),
45    }
46    constants.push(1f32)?;
47    unsafe {
48        runtime.load_gpu_kernel(
49            artifact.path(),
50            symbol,
51            [elements.div_ceil(256) as u32, 1, 1],
52            [256, 1, 1],
53            KernelContract {
54                bindings,
55                constants: constants.as_bytes().to_vec(),
56            },
57        )
58    }
59}
More examples
Hide additional examples
examples/gpu_npu_parallel.rs (lines 161-167)
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_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
Source

pub unsafe fn adopt_gpu_kernel( &self, raw: Kernel, grid: [u32; 3], block: [u32; 3], contract: KernelContract, ) -> Result<GpuKernel>

Retain a loaded executable without loading or copying its code again.

§Safety

The kernel must obey the declared memory contract for this launch and every accepted binding, as for Self::load_gpu_kernel.

Trait Implementations§

Source§

impl Clone for Runtime

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