1use hrx::{
5 Result,
6 benchmark::percentile,
7 execution::{
8 Access, BindingContract, Buffer, ExecutableGraph, GpuKernel, KernelContract,
9 MemoryPlacement, Runtime,
10 },
11};
12use std::time::Instant;
13fn binding(bytes: usize, access: Access) -> BindingContract {
14 BindingContract {
15 bytes,
16 alignment: 4,
17 access,
18 layout: "contiguous row-major".into(),
19 }
20}
21fn gpu_kernel(
22 runtime: &Runtime,
23 source: &str,
24 symbol: &str,
25 config: &[(&str, String)],
26 elements: usize,
27 bindings: Vec<BindingContract>,
28) -> Result<GpuKernel> {
29 let compiler = hrx::loom::Compiler::for_target(None, runtime.gpu()?.target())?;
30 let module = compiler.module(source);
31 let mut spec = hrx::loom::Specialization::new(symbol);
32 for (key, value) in config {
33 spec.set_config(*key, value.clone());
34 }
35 let artifact = module.compile(&spec)?;
36 let stream = hrx::gpu::Stream::open()?;
37 let raw = unsafe { stream.load_artifact(&artifact) }?;
38 let mut constants = hrx::gpu::Constants::new();
39 match raw.info().constant_byte_length {
40 8 => constants.push(elements as u32)?,
41 12 => constants.push(elements as u64)?,
42 _ => return Err(hrx::Error::Message("unexpected scalar ABI".into())),
43 }
44 constants.push(1f32)?;
45 unsafe {
46 runtime.load_gpu_kernel(
47 artifact.path(),
48 symbol,
49 [elements.div_ceil(256) as u32, 1, 1],
50 [256, 1, 1],
51 KernelContract {
52 bindings,
53 constants: constants.as_bytes().to_vec(),
54 },
55 )
56 }
57}
58fn main() -> Result<()> {
59 let args: Vec<_> = std::env::args().collect();
60 if args.len() != 6 {
61 return Err(hrx::Error::Message(
62 "usage: gemm_pipeline <row-major-bf16.xclbin> <instructions.bin> M K N".into(),
63 ));
64 }
65 let parse = |i: usize| {
66 args[i]
67 .parse::<usize>()
68 .map_err(|e| hrx::Error::Message(e.to_string()))
69 };
70 let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
71 let extent = |a: usize, b: usize, width: usize| {
72 a.checked_mul(b)
73 .filter(|&count| count > 0 && count <= (1 << 30))
74 .and_then(|count| count.checked_mul(width))
75 .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
76 };
77 let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
78 let runtime = Runtime::new()?;
79 let program = unsafe { runtime.npu(0)?.load_program(&args[1]) }?;
80 let npu = unsafe {
81 program.kernel(
82 &std::fs::read(&args[2])?,
83 KernelContract {
84 bindings: vec![
85 binding(a_bytes, Access::Read),
86 binding(b_bytes, Access::Read),
87 binding(c_bytes, Access::Write),
88 ],
89 constants: vec![],
90 },
91 )
92 }?;
93 let prep = gpu_kernel(
94 &runtime,
95 include_str!("../tests/kernels/euler.loom"),
96 "krea2_euler",
97 &[
98 ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
99 ("krea2.euler.grid_y", "1".into()),
100 ],
101 m * k,
102 vec![
103 binding(a_bytes, Access::ReadWrite),
104 binding(a_bytes, Access::Read),
105 ],
106 )?;
107 let epilogue = gpu_kernel(
108 &runtime,
109 include_str!("../tests/kernels/add_f32.loom"),
110 "add_f32",
111 &[("add.grid", (m * n).div_ceil(256).to_string())],
112 m * n,
113 vec![binding(c_bytes, Access::ReadWrite)],
114 )?;
115 let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(program.clone()));
116 let weights = allocate(b_bytes)?;
117 let ones = allocate(a_bytes)?;
118 for buffer in [&weights, &ones] {
119 for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
120 word.copy_from_slice(&0x3f80u16.to_le_bytes());
121 }
122 }
123 let prepare = || -> Result<(ExecutableGraph, Buffer)> {
124 let a = allocate(a_bytes)?;
125 let c = allocate(c_bytes)?;
126 let mut graph = runtime.graph();
127 graph.fill(a.view(), 0)?;
128 graph.gpu(&prep, &[a.view(), ones.view()])?;
129 graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
130 graph.gpu(&epilogue, &[c.view()])?;
131 Ok((graph.prepare()?, c))
132 };
133 let (first, c0) = prepare()?;
134 let (second, c1) = prepare()?;
135 for _ in 0..5 {
136 first.submit()?.wait()?;
137 second.submit()?.wait()?;
138 }
139 let before = runtime.statistics();
140 let mut latency = Vec::with_capacity(50);
141 for _ in 0..50 {
142 let start = Instant::now();
143 first.submit()?.wait()?;
144 latency.push(start.elapsed().as_secs_f64());
145 }
146 let start = Instant::now();
147 for _ in 0..50 {
148 let a = first.submit()?;
149 let b = second.submit()?;
150 a.wait()?;
151 b.wait()?;
152 }
153 let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
154 for c in [&c0, &c1] {
155 if c.map_read()?
156 .as_chunks::<4>()
157 .0
158 .iter()
159 .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
160 {
161 return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
162 }
163 }
164 let after = runtime.statistics();
165 assert_eq!(before.allocations, after.allocations);
166 assert_eq!(before.imports, after.imports);
167 assert_eq!(after.copied_bytes, 0);
168 latency.sort_by(f64::total_cmp);
169 println!(
170 "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
171 percentile(&latency, 50) * 1e3,
172 percentile(&latency, 95) * 1e3
173 );
174 println!("{}", serde_json::to_string(&after)?);
175 Ok(())
176}