use hrx::{
Buffer, Constants, GraphExec, Result, Stream,
loom::{Compiler, CompilerOptions, ProcessorMode, ReportMode, Specialization},
};
use std::{path::Path, time::Instant};
fn half(value: f32) -> u16 {
let bits = value.to_bits();
let sign = ((bits >> 16) & 0x8000) as u16;
let exponent = ((bits >> 23) & 255) as i32 - 127 + 15;
if exponent <= 0 {
return sign;
}
assert!(exponent < 31, "fixture value overflows fp16");
let mantissa = bits & 0x7fffff;
let rounded = mantissa + 0xfff + ((mantissa >> 13) & 1);
sign | (((exponent as u32) << 10) + (rounded >> 13)) as u16
}
fn float(value: u16) -> f64 {
let sign = if value & 0x8000 == 0 { 1.0 } else { -1.0 };
let exponent = (value >> 10) & 31;
let mantissa = f64::from(value & 1023);
match exponent {
0 => sign * mantissa * 2f64.powi(-24),
31 => f64::NAN,
_ => sign * (1.0 + mantissa / 1024.0) * 2f64.powi(i32::from(exponent) - 15),
}
}
fn bf16(value: f32) -> f32 {
let bits = value.to_bits();
f32::from_bits((bits + 0x7fff + ((bits >> 16) & 1)) & 0xffff0000)
}
fn halves(values: impl IntoIterator<Item = f32>) -> Vec<u8> {
values
.into_iter()
.flat_map(|value| half(value).to_le_bytes())
.collect()
}
fn floats(values: impl IntoIterator<Item = f32>) -> Vec<u8> {
values.into_iter().flat_map(f32::to_le_bytes).collect()
}
struct Case {
name: String,
source: &'static str,
spec: Specialization,
grid: [u32; 3],
block: [u32; 3],
indices: Vec<u32>,
data: Vec<Vec<u8>>,
output: usize,
expected: Vec<f64>,
tolerance: f64,
}
fn attention(tokens: usize, repack: bool) -> Case {
let capacity = tokens + 32;
let symbol = if repack {
"krea2_attention_query32"
} else {
"krea2_attention_gqa_lds_f16_wmma"
};
let prefix = if repack {
"krea2.attention_query32"
} else {
"krea2.attention_gqa_lds_f16_wmma"
};
let mut spec = Specialization::new(symbol).with_report(ReportMode::Summary);
for (key, value) in [
("q_stride", 512),
("kv_stride", 128),
("out_stride", 512),
("tokens", tokens),
("token_capacity", capacity),
] {
spec.set_config(format!("{prefix}.{key}"), value.to_string());
}
spec.set_config(format!("{prefix}.scale"), "0.08838834764831845");
let query = |row: usize, head: usize| ((row + head * 3) % 9) as f32 * 0.0625 - 0.25;
let key = |row: usize| (row % 7) as f32 * 0.0625 - 0.1875;
let value = |row: usize, channel: usize| {
(row % 7) as f32 * 0.125 - 0.375 + (channel % 8) as f32 * 0.0625
};
let q = halves((0..capacity * 512).map(|i| {
if i / 512 < tokens {
query(i / 512, (i % 512) / 128)
} else {
0.0
}
}));
let k = halves((0..capacity * 128).map(|i| if i / 128 < tokens { key(i / 128) } else { 0.0 }));
let v = halves((0..capacity * 128).map(|i| {
let (row, channel) = if repack {
(i % capacity, i / capacity)
} else {
(i / 128, i % 128)
};
if row < tokens {
value(row, channel)
} else {
0.0
}
}));
let mut expected = Vec::with_capacity(tokens * 512);
for row in 0..tokens {
for head in 0..4 {
let mut numerator = 0.0;
let mut denominator = 0.0;
for token in 0..tokens {
let score = (128.0 * f64::from(query(row, head)) * f64::from(key(token))
/ 128f64.sqrt())
.exp();
numerator += score * f64::from(value(token, 0));
denominator += score;
}
expected.extend(
(0..128).map(|channel| numerator / denominator + (channel % 8) as f64 * 0.0625),
);
}
}
Case {
name: format!("{symbol}-{tokens}"),
source: if repack {
include_str!("../native/qualification/attention_query32.loom")
} else {
include_str!("../native/qualification/attention_gqa_lds_f16_wmma.loom")
},
spec,
grid: [(tokens / if repack { 32 } else { 16 }) as u32, 1, 1],
block: [if repack { 256 } else { 128 }, 1, 1],
indices: vec![tokens as u32],
data: vec![q, k, v, vec![0; tokens * 512 * 2]],
output: 3,
expected,
tolerance: 0.001,
}
}
fn gemm(large: bool, integer: bool) -> Case {
let (m, n, k) = if large {
(256usize, 1024usize, 1024usize)
} else {
(256, 512, 256)
};
let symbol = if integer {
"krea2_gemm_i8_256"
} else {
"h3_gemm_f16_fast_256b"
};
let prefix = if integer {
"krea2.gemm_i8_256"
} else {
"h3.gemm_f16_fast_256b"
};
let mut spec = Specialization::new(symbol).with_report(ReportMode::Summary);
for (name, value) in [
("k_size", k),
("n_size", n),
("k_stride", k),
("m_group", 1),
] {
spec.set_config(format!("{prefix}.{name}"), value.to_string());
}
let a = |row: usize, col: usize| ((row * 3 + col * 5 + row * col) % 11) as i32 - 5;
let w = |row: usize, col: usize| ((row * 7 + col * 3 + row * col) % 13) as i32 - 6;
let bias = |col: usize| (col % 7) as f32 * 0.03125 - 0.09375;
let mut products = [[0f32; 13]; 11];
for (row, values) in products.iter_mut().enumerate() {
for (col, product) in values.iter_mut().enumerate() {
let sum: i32 = (0..k).map(|i| a(row, i) * w(col, i)).sum();
*product = sum as f32 / 64.0;
}
}
let expected = (0..m * n)
.map(|i| {
let dot = products[(i / n) % 11][(i % n) % 13];
float(half(if integer {
bf16(dot)
} else {
dot + bias(i % n)
}))
})
.collect();
let data = if integer {
vec![
(0..m * k).map(|i| a(i / k, i % k) as i8 as u8).collect(),
(0..n * k).map(|i| w(i / k, i % k) as i8 as u8).collect(),
floats(vec![0.125; n]),
floats(vec![0.125; m]),
vec![0; m * n * 2],
]
} else {
vec![
halves((0..m * k).map(|i| a(i / k, i % k) as f32 * 0.125)),
halves((0..n * k).map(|i| w(i / k, i % k) as f32 * 0.125)),
vec![0; m * n * 2],
floats((0..n).map(bias)),
]
};
Case {
name: format!("{symbol}-{m}x{n}x{k}"),
source: if integer {
include_str!("../native/qualification/gemm_i8_256.loom")
} else {
include_str!("../native/qualification/gemm_f16_fast_256b.loom")
},
spec,
grid: [
(n / if integer { 128 } else { 256 }) as u32,
(m / if integer { 256 } else { 128 }) as u32,
1,
],
block: [256, 1, 1],
indices: vec![m as u32],
data,
output: if integer { 4 } else { 2 },
expected,
tolerance: if integer { 0.0 } else { 2f64.powi(-14) },
}
}
struct Arm {
graph: GraphExec,
buffers: std::sync::Arc<Vec<Buffer>>,
identity: serde_json::Value,
}
fn prepare(
stream: &mut Stream,
compiler: &Compiler,
case: &Case,
shared: Option<std::sync::Arc<Vec<Buffer>>>,
) -> Result<Arm> {
let artifact = compiler.module(case.source).compile(&case.spec)?;
let kernel = unsafe { stream.load_artifact(&artifact) }?;
let constants = Constants::indices(&kernel, &case.indices)?;
let buffers = if let Some(buffers) = shared {
buffers
} else {
std::sync::Arc::new(
case.data
.iter()
.map(|data| {
let buffer = stream.allocate(data.len())?;
stream.upload_blocking(buffer.binding(), data)?;
Ok(buffer)
})
.collect::<Result<Vec<_>>>()?,
)
};
let bindings: Vec<_> = buffers.iter().map(Buffer::binding).collect();
let mut graph = stream.graph()?;
let mut prior = None;
for _ in 0..32 {
prior = Some(unsafe {
graph.dispatch(
prior.as_slice(),
&kernel,
case.grid,
case.block,
&constants,
&bindings,
)
}?);
}
let identity = serde_json::json!({"compiler":artifact.compiler_identity(), "artifact":hrx::bundle::digest(artifact.bytes()),
"report":artifact.report().map(|report| report.json())});
Ok(Arm {
graph: graph.finish()?,
buffers,
identity,
})
}
fn check(stream: &mut Stream, arm: &Arm, case: &Case) -> Result<f64> {
let mut bytes = vec![0; case.expected.len() * 2];
stream.read_blocking(arm.buffers[case.output].binding(), &mut bytes)?;
let mut maximum = 0f64;
for (index, (value, expected)) in bytes.chunks_exact(2).zip(&case.expected).enumerate() {
let value = float(u16::from_le_bytes(value.try_into().unwrap()));
let error = (value - expected).abs();
if !value.is_finite() || error > case.tolerance {
return Err(hrx::Error::Message(format!(
"{} output {index}: {value} vs {expected}, error {error}, limit {}",
case.name, case.tolerance
)));
}
maximum = maximum.max(error);
}
Ok(maximum)
}
fn main() -> Result<()> {
let args: Vec<_> = std::env::args().collect();
if args.len() < 3 || args.len() > 5 {
return Err(hrx::Error::Message(
"usage: patch_bench BASELINE_LIB CANDIDATE_LIB [pairs=30] [default|cu|wgp]".into(),
));
}
let pairs = args
.get(3)
.map(|s| s.parse::<usize>())
.transpose()
.map_err(|e| hrx::Error::Message(e.to_string()))?
.unwrap_or(30);
if pairs < 30 {
return Err(hrx::Error::Message("at least 30 pairs are required".into()));
}
let mode = match args.get(4).map(String::as_str).unwrap_or("default") {
"default" => ProcessorMode::Default,
"cu" => ProcessorMode::ComputeUnit,
"wgp" => ProcessorMode::WorkgroupProcessor,
_ => {
return Err(hrx::Error::Message(
"mode must be default, cu or wgp".into(),
));
}
};
let mut stream = Stream::open()?;
let mut options = CompilerOptions::for_target(stream.target());
options.processor_mode = mode;
let compilers = [
Compiler::with_options(Some(Path::new(&args[1])), options.clone())?,
Compiler::with_options(Some(Path::new(&args[2])), options)?,
];
for case in [
attention(128, false),
attention(512, false),
gemm(false, false),
gemm(true, false),
gemm(false, true),
gemm(true, true),
] {
let baseline = prepare(&mut stream, &compilers[0], &case, None)?;
let candidate = prepare(
&mut stream,
&compilers[1],
&case,
Some(baseline.buffers.clone()),
)?;
let mut arms = [baseline, candidate];
for arm in &mut arms {
for _ in 0..5 {
stream.launch(&mut arm.graph)?;
stream.synchronize()?;
}
check(&mut stream, arm, &case)?;
}
let errors = [
check(&mut stream, &arms[0], &case)?,
check(&mut stream, &arms[1], &case)?,
];
println!(
"{}",
serde_json::json!({"case":case.name,"mode":format!("{mode:?}"),"baseline":arms[0].identity,"candidate":arms[1].identity,"max_absolute_error":errors,"tolerance":case.tolerance,"replays":32})
);
for pair in 0..pairs {
let mut ns = [0.0; 2];
for index in if pair % 2 == 0 { [0, 1] } else { [1, 0] } {
let start = Instant::now();
stream.launch(&mut arms[index].graph)?;
stream.synchronize()?;
ns[index] = start.elapsed().as_secs_f64() * 1e9 / 32.0;
}
println!(
"{}",
serde_json::json!({"case":case.name,"pair":pair,"baseline_ns":ns[0],"candidate_ns":ns[1]})
);
}
for arm in &mut arms {
stream.launch(&mut arm.graph)?;
stream.synchronize()?;
check(&mut stream, arm, &case)?;
}
}
Ok(())
}