aprender_contracts_cli/commands/
roofline.rs1use std::path::Path;
2
3use provable_contracts::roofline::{self, Bottleneck, HardwareProfile};
4
5pub fn run(
7 contract_dir: &Path,
8 params: u64,
9 bits: u32,
10 hardware: &str,
11 format: &str,
12) -> Result<(), Box<dyn std::error::Error>> {
13 let hw = match hardware {
14 "apple-m" => HardwareProfile::apple_m_series(),
15 "a100" => HardwareProfile::nvidia_a100(),
16 _ => {
17 return Err(
18 format!("unknown hardware profile '{hardware}'. Use: apple-m, a100").into(),
19 );
20 }
21 };
22
23 if params == 0 {
24 return Err("--params must be > 0 (model has no parameters)".into());
25 }
26 if bits == 0 {
27 return Err("--bits must be > 0".into());
28 }
29
30 let desc = roofline::load_roofline_contract(contract_dir);
31 let r = roofline::compute_roofline(params, bits, &hw);
32
33 match format {
34 "json" => print_json(&r, &hw, desc.as_deref()),
35 _ => print_text(&r, &hw, hardware, params, bits, desc.as_deref()),
36 }
37
38 Ok(())
39}
40
41fn print_text(
42 r: &roofline::RooflineCeiling,
43 hw: &HardwareProfile,
44 hw_name: &str,
45 params: u64,
46 bits: u32,
47 desc: Option<&str>,
48) {
49 println!("Roofline Analysis ({})", r.contract_id);
50 if let Some(d) = desc {
51 println!(" {d}");
52 }
53 println!();
54 println!("Hardware: {hw_name}");
55 println!(" Bandwidth: {:.1} GB/s", hw.bandwidth_gb_s);
56 println!(" Compute: {:.1} GFLOPS", hw.compute_gflops);
57 println!();
58 println!("Model: {} Q{bits}", format_params(params));
59 println!(" Size: {:.2} GB", r.model_bytes / 1e9);
60 println!();
61 println!("Ceilings:");
62 println!(" BW ceiling: {:.1} tok/s", r.bw_ceiling);
63 println!(" Compute ceiling: {:.1} tok/s", r.compute_ceiling);
64 println!(" Effective: {:.1} tok/s", r.throughput_ceiling);
65 println!();
66 let marker = match r.bottleneck {
67 Bottleneck::Bandwidth => "MEMORY-BOUND (bw_ceiling < compute_ceiling)",
68 Bottleneck::Compute => "COMPUTE-BOUND (compute_ceiling < bw_ceiling)",
69 };
70 println!("Bottleneck: {marker}");
71}
72
73fn print_json(r: &roofline::RooflineCeiling, hw: &HardwareProfile, desc: Option<&str>) {
74 println!("{{");
75 println!(" \"contract_id\": \"{}\",", r.contract_id);
76 if let Some(d) = desc {
77 println!(" \"description\": \"{d}\",");
78 }
79 println!(" \"model_bytes\": {:.0},", r.model_bytes);
80 println!(" \"model_gb\": {:.4},", r.model_bytes / 1e9);
81 println!(" \"bw_ceiling_tok_s\": {:.2},", r.bw_ceiling);
82 println!(" \"compute_ceiling_tok_s\": {:.2},", r.compute_ceiling);
83 println!(
84 " \"throughput_ceiling_tok_s\": {:.2},",
85 r.throughput_ceiling
86 );
87 println!(
88 " \"bottleneck\": \"{}\",",
89 match r.bottleneck {
90 Bottleneck::Bandwidth => "bandwidth",
91 Bottleneck::Compute => "compute",
92 }
93 );
94 println!(" \"hardware\": {{");
95 println!(" \"bandwidth_gb_s\": {:.1},", hw.bandwidth_gb_s);
96 println!(" \"compute_gflops\": {:.1},", hw.compute_gflops);
97 println!(" \"ops_per_token\": {:.1}", hw.ops_per_token);
98 println!(" }}");
99 println!("}}");
100}
101
102fn format_params(n: u64) -> String {
103 #[allow(clippy::cast_precision_loss)]
104 let f = n as f64;
105 if f >= 1e12 {
106 format!("{:.1}T", f / 1e12)
107 } else if f >= 1e9 {
108 format!("{:.1}B", f / 1e9)
109 } else if f >= 1e6 {
110 format!("{:.0}M", f / 1e6)
111 } else {
112 format!("{n}")
113 }
114}