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

gpu_npu_parallel/
gpu_npu_parallel.rs
