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
impl Runtime
Sourcepub unsafe fn with_gpu_access<T>(
&self,
stream: &mut Stream,
accesses: &[GpuAccess],
callback: impl FnOnce(&mut Stream, &[View<'_>]) -> Result<T>,
) -> Result<T>
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
impl Runtime
Sourcepub fn start_trace(&self, event_capacity: usize) -> Result<()>
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?
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(¶llel)?);
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}Sourcepub fn finish_trace(&self) -> Option<ExecutionTrace>
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?
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(¶llel)?);
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}Sourcepub fn owns(&self, view: &BufferView) -> bool
pub fn owns(&self, view: &BufferView) -> bool
Whether a view belongs to this runtime’s allocation and scheduling domain.
Sourcepub fn same_domain(&self, other: &Self) -> bool
pub fn same_domain(&self, other: &Self) -> bool
Whether two handles share the same allocation and scheduling domain.
Sourcepub fn memory_budget(&self) -> Option<&MemoryBudget>
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.
Sourcepub fn new() -> Result<Self>
pub fn new() -> Result<Self>
Create a runtime. Native devices are opened lazily when requested.
Examples found in repository?
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
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}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}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}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(¶llel)?);
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}Sourcepub fn with_options(options: RuntimeOptions) -> Result<Self>
pub fn with_options(options: RuntimeOptions) -> Result<Self>
Create a runtime with bounded pending submission and graph capacities.
Sourcepub fn gpu(&self) -> Result<GpuDevice>
pub fn gpu(&self) -> Result<GpuDevice>
Open the selected GPU and return its actual target.
Examples found in repository?
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
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(¶llel)?);
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}Sourcepub fn npu(&self, index: i32) -> Result<NpuDevice>
pub fn npu(&self, index: i32) -> Result<NpuDevice>
Select an NPU. Opening its program validates actual device access.
Examples found in repository?
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
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}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}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(¶llel)?);
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}Sourcepub fn allocate(
&self,
bytes: usize,
placement: MemoryPlacement,
) -> Result<Buffer>
pub fn allocate( &self, bytes: usize, placement: MemoryPlacement, ) -> Result<Buffer>
Allocate initialized storage. Shared memory is exported and imported once.
Examples found in repository?
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
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}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}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}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(¶llel)?);
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}Sourcepub unsafe fn adopt_gpu_buffer(
&self,
buffer: Buffer,
placement: MemoryPlacement,
) -> Result<Buffer>
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.
Sourcepub fn statistics(&self) -> Statistics
pub fn statistics(&self) -> Statistics
Observe allocation, transfer, and execution counters.
Examples found in repository?
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
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(¶llel)?);
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}Sourcepub fn graph(&self) -> Graph
pub fn graph(&self) -> Graph
Start describing a reusable pipeline in this allocation domain.
Examples found in repository?
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
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}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}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}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(¶llel)?);
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}Sourcepub unsafe fn load_gpu_kernel(
&self,
path: impl AsRef<Path>,
symbol: &str,
grid: [u32; 3],
block: [u32; 3],
contract: KernelContract,
) -> Result<GpuKernel>
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?
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
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(¶llel)?);
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}Sourcepub unsafe fn adopt_gpu_kernel(
&self,
raw: Kernel,
grid: [u32; 3],
block: [u32; 3],
contract: KernelContract,
) -> Result<GpuKernel>
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§
Auto Trait Implementations§
impl Freeze for Runtime
impl RefUnwindSafe for Runtime
impl Send for Runtime
impl Sync for Runtime
impl Unpin for Runtime
impl UnsafeUnpin for Runtime
impl UnwindSafe for Runtime
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> DropFlavorWrapper<T> for T
impl<T> DropFlavorWrapper<T> for T
Source§impl<T, W> HasTypeWitness<W> for Twhere
W: MakeTypeWitness<Arg = T>,
T: ?Sized,
impl<T, W> HasTypeWitness<W> for Twhere
W: MakeTypeWitness<Arg = T>,
T: ?Sized,
Source§impl<T> Identity for Twhere
T: ?Sized,
impl<T> Identity for Twhere
T: ?Sized,
Source§impl<T> Instrument for T
impl<T> Instrument for T
Source§fn instrument(self, span: Span) -> Instrumented<Self> ⓘ
fn instrument(self, span: Span) -> Instrumented<Self> ⓘ
Source§fn in_current_span(self) -> Instrumented<Self> ⓘ
fn in_current_span(self) -> Instrumented<Self> ⓘ
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
fn into_either(self, into_left: bool) -> Either<Self, Self> ⓘ
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 moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self> ⓘ
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