Skip to main content

aprender_contracts_cli/commands/
roofline.rs

1use std::path::Path;
2
3use provable_contracts::roofline::{self, Bottleneck, HardwareProfile};
4
5/// Run the `pv roofline` command: compute performance ceilings from contract.
6pub 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}