fn main() {
unsafe { std::env::set_var("CMF_GPU", "1") };
let cases = [
(384usize, 4096usize, 128usize),
(384, 128, 4096),
(384, 16384, 128),
(384, 128, 16384),
];
for (n, k, m) in cases {
let x = vec![0.01f32; n * k];
let w = vec![0.02f32; m * k];
let mut y = vec![0f32; n * m];
cortiq_engine::fcd_ops::gemm_nt(&x, &w, &mut y, n, k, m, None);
let t = std::time::Instant::now();
for _ in 0..10 {
cortiq_engine::fcd_ops::gemm_nt(&x, &w, &mut y, n, k, m, None);
}
let d = t.elapsed().as_secs_f64() / 10.0;
let g = 2.0 * n as f64 * k as f64 * m as f64 / d / 1e9;
let t2 = std::time::Instant::now();
for _ in 0..10 {
cortiq_engine::gpu::cpu_scope(|| {
cortiq_engine::fcd_ops::gemm_nt(&x, &w, &mut y, n, k, m, None)
});
}
let d2 = t2.elapsed().as_secs_f64() / 10.0;
let g2 = 2.0 * n as f64 * k as f64 * m as f64 / d2 / 1e9;
println!(
"{n}x{k}x{m}: default {:.2} ms ({g:.0} GF/s) cpu_scope {:.2} ms ({g2:.0} GF/s)",
d * 1e3,
d2 * 1e3
);
}
}