gpu_npu_parallel/
gpu_npu_parallel.rs

1//! Independent GPU vector processing and batched NPU matrix multiplication.
2//! Usage: gpu_npu_parallel [gpu-elements] [npu-matrices] [samples] [trace.json]
3//! Timings cover warm, completed graphs, excluding compilation and host I/O.
4use hrx::{
5    Error, Result,
6    execution::{
7        Access, BindingContract, Buffer, ExecutableGraph, KernelContract, MemoryPlacement, Runtime,
8    },
9    loom::{Compiler, CxxSource, Specialization},
10};
11use std::{path::PathBuf, time::Instant};
12
13fn contract(bindings: &[(usize, Access)]) -> KernelContract {
14    KernelContract {
15        bindings: bindings
16            .iter()
17            .map(|&(bytes, access)| BindingContract {
18                bytes,
19                alignment: 64,
20                access,
21                layout: "contiguous row-major".into(),
22            })
23            .collect(),
24        constants: vec![],
25    }
26}
27
28fn gpu_value(index: usize, pass: usize) -> i32 {
29    ((index + pass * 17) % 2048) as i32 - 1024
30}
31
32fn matrix_value(matrix: usize, index: usize, operand: usize, pass: usize) -> f32 {
33    ((index * (operand + 3) + matrix + pass) % 7) as f32 - 3.0
34}
35
36fn initialize(gpu: &Buffer, a: &Buffer, b: &Buffer, pass: usize) -> Result<()> {
37    for (index, word) in gpu
38        .map_write()?
39        .as_chunks_mut::<4>()
40        .0
41        .iter_mut()
42        .enumerate()
43    {
44        word.copy_from_slice(&gpu_value(index, pass).to_le_bytes());
45    }
46    for (operand, buffer) in [a, b].into_iter().enumerate() {
47        for (index, word) in buffer
48            .map_write()?
49            .as_chunks_mut::<2>()
50            .0
51            .iter_mut()
52            .enumerate()
53        {
54            let value = matrix_value(index / 64, index % 64, operand, pass);
55            // Integers in -3..=3 are exactly representable in BF16 and BFP16.
56            word.copy_from_slice(&((value.to_bits() >> 16) as u16).to_le_bytes());
57        }
58    }
59    Ok(())
60}
61
62fn validate(gpu: &Buffer, matrices: &Buffer, pass: usize) -> Result<()> {
63    for (index, word) in gpu.map_read()?.as_chunks::<4>().0.iter().enumerate() {
64        let actual = i32::from_le_bytes(*word);
65        if actual != gpu_value(index, pass) * 3 + 7 {
66            return Err(Error::Message(format!("GPU mismatch at element {index}")));
67        }
68    }
69    for (index, word) in matrices.map_read()?.as_chunks::<4>().0.iter().enumerate() {
70        let matrix = index / 64;
71        let row = index % 64 / 8;
72        let column = index % 8;
73        let expected: f32 = (0..8)
74            .map(|k| {
75                matrix_value(matrix, row * 8 + k, 0, pass)
76                    * matrix_value(matrix, k * 8 + column, 1, pass)
77            })
78            .sum();
79        if f32::from_le_bytes(*word) != expected {
80            return Err(Error::Message(format!(
81                "NPU mismatch at matrix {matrix}, ({row}, {column})"
82            )));
83        }
84    }
85    Ok(())
86}
87
88fn completed_ms(graph: &ExecutableGraph) -> Result<f64> {
89    let start = Instant::now();
90    graph.submit()?.wait()?;
91    Ok(start.elapsed().as_secs_f64() * 1000.0)
92}
93
94fn median(values: &mut [f64]) -> f64 {
95    values.sort_by(f64::total_cmp);
96    let middle = values.len() / 2;
97    if values.len().is_multiple_of(2) {
98        (values[middle - 1] + values[middle]) / 2.0
99    } else {
100        values[middle]
101    }
102}
103
104fn main() -> Result<()> {
105    let args: Vec<_> = std::env::args().skip(1).collect();
106    if args.len() > 4 || args.iter().any(|arg| arg == "--help") {
107        println!(
108            "usage: gpu_npu_parallel [gpu-elements=16777216] [npu-matrices=1024] [samples=21] [trace.json]"
109        );
110        return if args.len() <= 4 {
111            Ok(())
112        } else {
113            Err(Error::Message("too many arguments".into()))
114        };
115    }
116    let number = |index: usize, default: usize| -> Result<usize> {
117        args.get(index).map_or(Ok(default), |value| {
118            value
119                .parse()
120                .map_err(|_| Error::Message(format!("invalid integer: {value}")))
121        })
122    };
123    let elements = number(0, 1 << 24)?;
124    let matrices = number(1, 1024)?;
125    let samples = number(2, 21)?;
126    if !(256..=1 << 26).contains(&elements)
127        || !elements.is_multiple_of(256)
128        || !(1..=16384).contains(&matrices)
129        || !(3..=1000).contains(&samples)
130    {
131        return Err(Error::Message("gpu-elements must be a multiple of 256 in 256..=67108864; npu-matrices in 1..=16384; samples in 3..=1000".into()));
132    }
133    let trace_path = args.get(3).map(PathBuf::from);
134    let runtime = Runtime::new()?;
135    let gpu = runtime.gpu()?;
136    let npu = runtime.npu(0)?;
137    println!(
138        "GPU: {}; NPU: {}",
139        gpu.target().as_str(),
140        npu.target().as_str()
141    );
142    let gpu_artifact = Compiler::for_target(None, gpu.target())?
143        .import_cxx(CxxSource::new(
144            "affine.cpp",
145            format!(
146                r#"
147#include <hip/hip_runtime.h>
148__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
149void affine(const int* input, int* output) {{
150  unsigned i = blockIdx.x * 256u + threadIdx.x;
151  output[i] = input[i] * 3 + 7;
152}}
153"#,
154                elements / 256
155            ),
156        ))?
157        .compile(&Specialization::new("affine"))?;
158    // SAFETY: the fixed launch covers exactly elements i32 values in each
159    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
160    let gpu_kernel = unsafe {
161        runtime.load_gpu_kernel(
162            gpu_artifact.path(),
163            "affine",
164            [(elements / 256) as u32, 1, 1],
165            [256, 1, 1],
166            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
167        )
168    }?;
169    let npu_artifact = Compiler::for_target(None, npu.target())?
170        .module(include_str!("kernels/batched_gemm.xdna.loom"))
171        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
172    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
173    // produces the same number of 8x8 FP32 outputs on one NPU column.
174    let npu_kernel = unsafe {
175        npu.load_artifact(
176            &npu_artifact,
177            1,
178            contract(&[
179                (matrices * 128, Access::Read),
180                (matrices * 128, Access::Read),
181                (matrices * 256, Access::Write),
182            ]),
183        )
184    }?;
185    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
186    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
187    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
188    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
189    let mut upload = runtime.graph();
190    upload.copy(input.view(), host_input.view())?;
191    let upload = upload.prepare()?;
192    let mut download = runtime.graph();
193    download.copy(host_output.view(), output.view())?;
194    let download = download.prepare()?;
195    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
196    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
197    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;
198
199    let mut gpu_only = runtime.graph();
200    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
201    let gpu_only = gpu_only.prepare()?;
202    let mut npu_only = runtime.graph();
203    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
204    let npu_only = npu_only.prepare()?;
205
206    // Disjoint buffers imply no dependency edge. One submit admits both
207    // branches to separate GPU/NPU workers; one wait joins their completions.
208    let mut parallel = runtime.graph();
209    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
210    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
211    let parallel = parallel.prepare()?;
212    for pass in 0..3 {
213        initialize(&host_input, &a, &b, pass * 2)?;
214        upload.submit()?.wait()?;
215        gpu_only.submit()?.wait()?;
216        npu_only.submit()?.wait()?;
217        download.submit()?.wait()?;
218        validate(&host_output, &c, pass * 2)?;
219        // Change both inputs so validation cannot accept the sequential run's
220        // old output if either parallel branch fails to write its result.
221        initialize(&host_input, &a, &b, pass * 2 + 1)?;
222        upload.submit()?.wait()?;
223        parallel.submit()?.wait()?;
224        download.submit()?.wait()?;
225        validate(&host_output, &c, pass * 2 + 1)?;
226    }
227
228    let before = runtime.statistics();
229    let mut sequential_ms = Vec::with_capacity(samples);
230    let mut parallel_ms = Vec::with_capacity(samples);
231    for sample in 0..samples {
232        // Alternate order to reduce systematic clock/thermal ordering bias.
233        for concurrent in if sample % 2 == 0 {
234            [false, true]
235        } else {
236            [true, false]
237        } {
238            if concurrent {
239                parallel_ms.push(completed_ms(&parallel)?);
240            } else {
241                let start = Instant::now();
242                gpu_only.submit()?.wait()?;
243                npu_only.submit()?.wait()?;
244                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
245            }
246        }
247    }
248    let after = runtime.statistics();
249    if before.allocations != after.allocations || before.imports != after.imports {
250        return Err(Error::Message(
251            "replay unexpectedly allocated or imported storage".into(),
252        ));
253    }
254    download.submit()?.wait()?;
255    validate(&host_output, &c, 5)?;
256    let sequential = median(&mut sequential_ms);
257    let concurrent = median(&mut parallel_ms);
258    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
259    println!(
260        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
261        sequential / concurrent
262    );
263    println!(
264        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
265    );
266
267    // Capture separately so tracing does not affect the benchmark samples.
268    runtime.start_trace(32)?;
269    for _ in 0..4 {
270        parallel.submit()?.wait()?;
271    }
272    let trace = runtime
273        .finish_trace()
274        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
275    let overlap: f64 = trace
276        .events
277        .iter()
278        .filter(|e| e.lane == 1)
279        .flat_map(|gpu| {
280            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
281                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
282                    - gpu.start_us.max(npu.start_us))
283                .max(0.0)
284            })
285        })
286        .sum();
287    println!(
288        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
289    );
290    if let Some(path) = trace_path {
291        std::fs::write(&path, trace.to_json()?)?;
292        println!("Trace: {}", path.display());
293    }
294    Ok(())
295}
gpu_npu_parallel/gpu_npu_parallel.rs

gpu_npu_parallel/
gpu_npu_parallel.rs
