1use hrx::{
7 Result,
8 benchmark::percentile,
9 execution::{
10 Access, BindingContract, Buffer, ExecutableGraph, GpuKernel, KernelContract,
11 MemoryPlacement, Runtime,
12 },
13};
14use std::time::Instant;
15fn binding(bytes: usize, access: Access) -> BindingContract {
16 BindingContract {
17 bytes,
18 alignment: 4,
19 access,
20 layout: "contiguous row-major".into(),
21 }
22}
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}
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}