Matrix multiplication

1use std::{env, hint::black_box, time::Instant};
2
3use anyhow::{Context, bail, ensure};
4use oa::{Engine, Matrix, matrix};
5
6const CORRECTNESS_CASES: &[ShapeCase] = &[
7	ShapeCase::new("tiny-square", 8, 8, 8),
8	ShapeCase::new("square", 128, 128, 128),
9	ShapeCase::new("tall-skinny", 512, 64, 128),
10	ShapeCase::new("short-wide", 64, 512, 128),
11	ShapeCase::new("irregular", 100, 130, 77),
12	ShapeCase::new("gemv-decode", 1, 256, 256),
13	ShapeCase::new("mnist-hidden", 64, 128, 784),
14	ShapeCase::new("mnist-logits", 64, 10, 128),
15];
16
17const PERFORMANCE_CASES: &[ShapeCase] = &[
18	ShapeCase::new("square-512", 512, 512, 512),
19	ShapeCase::new("square-1024", 1024, 1024, 1024),
20	ShapeCase::new("square-2048", 2048, 2048, 2048),
21	ShapeCase::new("tall-skinny", 4096, 128, 1024),
22	ShapeCase::new("short-wide", 128, 4096, 1024),
23	ShapeCase::new("gemv-decode", 1, 4096, 4096),
24];
25
26#[derive(Clone, Copy)]
27struct ShapeCase {
28	name: &'static str,
29	m: usize,
30	n: usize,
31	k: usize,
32}
33
34impl ShapeCase {
35	const fn new(name: &'static str, m: usize, n: usize, k: usize) -> Self {
36		Self { name, m, n, k }
37	}
38}
39
40struct Options {
41	correctness: bool,
42	benchmark: bool,
43	warmup: usize,
44	samples: usize,
45	case: Option<String>,
46}
47
48struct PerfResult {
49	p50_ms: f64,
50	p95_ms: f64,
51	gflops: f64,
52}
53
54fn main() -> anyhow::Result<()> {
55	let options = parse_options()?;
56	let engine = Engine::new().context("create the OA Vulkan engine")?;
57
58	println!("OA Rust Core MatMul Intro — C = A @ Bᵀ");
59	println!("contract: FP32 [M,K] × [N,K] -> [M,N]");
60	demo_basic_syntax(&engine)?;
61
62	if options.correctness {
63		run_correctness_suite(&engine)?;
64	}
65	if options.benchmark {
66		run_performance_suite(&engine, &options)?;
67	}
68	Ok(())
69}
70
71fn parse_options() -> anyhow::Result<Options> {
72	let mut options = Options {
73		correctness: true,
74		benchmark: true,
75		warmup: 5,
76		samples: 20,
77		case: None,
78	};
79	let mut arguments = env::args().skip(1);
80	while let Some(argument) = arguments.next() {
81		match argument.as_str() {
82			"--correctness-only" => options.benchmark = false,
83			"--benchmark-only" => options.correctness = false,
84			"--warmup" => options.warmup = parse_count("--warmup", arguments.next())?,
85			"--samples" => options.samples = parse_count("--samples", arguments.next())?,
86			"--case" => {
87				options.case = Some(
88					arguments
89						.next()
90						.context("--case requires a performance-case name")?,
91				);
92			}
93			"--help" | "-h" => {
94				println!(
95					"usage: core_mat_mul_intro [--correctness-only|--benchmark-only] \
96					 [--case NAME] [--warmup N] [--samples N]"
97				);
98				std::process::exit(0);
99			}
100			_ => bail!("unknown argument: {argument}"),
101		}
102	}
103	ensure!(options.samples > 0, "--samples must be greater than zero");
104	Ok(options)
105}
106
107fn parse_count(flag: &str, value: Option<String>) -> anyhow::Result<usize> {
108	value
109		.with_context(|| format!("{flag} requires a non-negative integer"))?
110		.parse()
111		.with_context(|| format!("{flag} requires a non-negative integer"))
112}
113
114fn demo_basic_syntax(engine: &Engine) -> anyhow::Result<()> {
115	println!("\n1. Basic syntax");
116	let one = matrix::ones(engine, [2, 3])?;
117	let two = matrix::full(engine, [2, 3], 2.0)?;
118	let product = matrix::mat_mul_nt(&one, &two)?;
119	let values = product.read_f32()?;
120	ensure!(product.shape() == [2, 2], "unexpected output shape");
121	ensure!(
122		values == [6.0; 4],
123		"unexpected introductory result: {values:?}"
124	);
125	println!("   [2,3] × [2,3]ᵀ -> {:?}: {values:?}", product.shape());
126
127	// Operators, methods, associated calls and module functions share one body.
128	// Elementwise multiplication is distinct from A @ B-transpose above.
129	for elementwise in [
130		(&one * &two)?,
131		one.mul(&two)?,
132		Matrix::mul(&one, &two)?,
133		matrix::mul(&one, &two)?,
134	] {
135		ensure!(
136			elementwise.read_f32()? == [2.0; 6],
137			"elementwise API mismatch"
138		);
139	}
140	let added = (&one + &two)?;
141	ensure!(added.read_f32()? == [3.0; 6], "addition operator mismatch");
142	println!("   operator / method / associated / module calls agree");
143	Ok(())
144}
145
146fn run_correctness_suite(engine: &Engine) -> anyhow::Result<()> {
147	println!("\n2. Correctness vs independent CPU reference");
148	for (index, &case) in CORRECTNESS_CASES.iter().enumerate() {
149		let left_values = values(case.m * case.k, index * 2 + 3);
150		let right_values = values(case.n * case.k, index * 2 + 11);
151		let expected = cpu_mat_mul_nt(&left_values, &right_values, case.m, case.n, case.k);
152		let left = Matrix::from_f32(engine, [case.m, case.k], &left_values)?;
153		let right = Matrix::from_f32(engine, [case.n, case.k], &right_values)?;
154		let output = matrix::mat_mul_nt(&left, &right)?;
155		let actual = output.read_f32()?;
156		let error = normalized_error(&actual, &expected)?;
157		ensure!(error <= 1.0e-4, "{} exceeded the FP32 tolerance", case.name);
158		println!(
159			"   {:<16} [{:5},{:5},{:5}] norm_err={error:.2e} ok",
160			case.name, case.m, case.n, case.k
161		);
162	}
163	println!("   All correctness checks passed.");
164	Ok(())
165}
166
167fn run_performance_suite(engine: &Engine, options: &Options) -> anyhow::Result<()> {
168	println!("\n3. Exploratory eager end-to-end performance");
169	println!("   boundary: mat_mul_nt call through synchronized full output readback");
170	println!(
171		"   warmup={} samples={} (one process; not release benchmark evidence)",
172		options.warmup, options.samples
173	);
174
175	let mut matched = options.case.is_none();
176	for &case in PERFORMANCE_CASES {
177		if options
178			.case
179			.as_deref()
180			.is_some_and(|name| name != case.name)
181		{
182			continue;
183		}
184		matched = true;
185		let result = bench_one(engine, case, options.warmup, options.samples)?;
186		println!(
187			"   {:<16} [{:5},{:5},{:5}] p50={:9.3} ms p95={:9.3} ms {:9.2} GFLOP/s",
188			case.name, case.m, case.n, case.k, result.p50_ms, result.p95_ms, result.gflops
189		);
190		println!(
191			"BENCH name={} m={} n={} k={} boundary=eager_submit_readback p50_ms={:.6} p95_ms={:.6} gflops={:.6}",
192			case.name, case.m, case.n, case.k, result.p50_ms, result.p95_ms, result.gflops
193		);
194	}
195	ensure!(matched, "unknown performance case: {:?}", options.case);
196	Ok(())
197}
198
199fn bench_one(
200	engine: &Engine,
201	case: ShapeCase,
202	warmup: usize,
203	sample_count: usize,
204) -> anyhow::Result<PerfResult> {
205	let left = matrix::full(engine, [case.m, case.k], 0.25)?;
206	let right = matrix::full(engine, [case.n, case.k], 0.5)?;
207	for _ in 0..warmup {
208		let output = matrix::mat_mul_nt(&left, &right)?;
209		black_box(output.read_f32()?);
210	}
211
212	let mut samples = Vec::with_capacity(sample_count);
213	for _ in 0..sample_count {
214		let start = Instant::now();
215		let output = matrix::mat_mul_nt(&left, &right)?;
216		black_box(output.read_f32()?);
217		samples.push(start.elapsed().as_secs_f64() * 1_000.0);
218	}
219	samples.sort_by(f64::total_cmp);
220	let p50_ms = percentile(&samples, 0.50);
221	let p95_ms = percentile(&samples, 0.95);
222	let gflop = 2.0 * case.m as f64 * case.n as f64 * case.k as f64 / 1.0e9;
223	Ok(PerfResult {
224		p50_ms,
225		p95_ms,
226		gflops: gflop / (p50_ms / 1_000.0),
227	})
228}
229
230fn percentile(samples: &[f64], percentile: f64) -> f64 {
231	let index = ((percentile * samples.len() as f64).ceil() as usize)
232		.saturating_sub(1)
233		.min(samples.len() - 1);
234	samples[index]
235}
236
237fn values(length: usize, salt: usize) -> Vec<f32> {
238	(0..length)
239		.map(|index| (((index * 17 + salt) % 29) as f32 - 14.0) / 7.0)
240		.collect()
241}
242
243fn cpu_mat_mul_nt(left: &[f32], right: &[f32], m: usize, n: usize, k: usize) -> Vec<f32> {
244	let mut output = vec![0.0; m * n];
245	for row in 0..m {
246		for column in 0..n {
247			for inner in 0..k {
248				output[row * n + column] += left[row * k + inner] * right[column * k + inner];
249			}
250		}
251	}
252	output
253}
254
255fn normalized_error(actual: &[f32], expected: &[f32]) -> anyhow::Result<f64> {
256	ensure!(
257		actual.len() == expected.len(),
258		"output length differs from CPU oracle"
259	);
260	let mut max_error = 0.0_f64;
261	let mut max_reference = 1.0e-6_f64;
262	for (&actual, &expected) in actual.iter().zip(expected) {
263		max_error = max_error.max((f64::from(actual) - f64::from(expected)).abs());
264		max_reference = max_reference.max(f64::from(expected).abs());
265	}
266	Ok(max_error / max_reference)
267}
268