sweep3/
sweep3.rs

1//! `sweep3` — CPU-only TCG3 ratio sweep over the 19 bench fields (no GPU decode).
2//!
3//! ```text
4//! sweep3 DATA_ROOT [--model3 pretrained:M.tcm] [--l3 8] [--targets p80] [--threads 16]
5//! ```
6//! Prints one JSON line per (field, target) and the geomean bpv ratio vs
7//! serial tc-wavelet (from `DATA_ROOT/out/gpu_ratio.jsonl`) on all / orig / held.
8use cram_codec::codec::ErrorBound;
9use cram_codec::gpuwav3::{self, Gpu3Config, Model3};
10
11fn opt(args: &[String], k: &str, d: &str) -> String {
12    args.iter()
13        .position(|a| a == k)
14        .map(|i| args[i + 1].clone())
15        .unwrap_or_else(|| d.to_string())
16}
17
18const ORIG: [&str; 5] = [
19    "miranda-density",
20    "miranda-velocityx",
21    "hurricane-TC",
22    "hurricane-U",
23    "nyx-temperature",
24];
25
26fn main() {
27    let args: Vec<String> = std::env::args().skip(1).collect();
28    let root = std::path::PathBuf::from(&args[0]);
29    let threads: usize = opt(&args, "--threads", "16").parse().unwrap();
30    let lanes: usize = opt(&args, "--l3", "8").parse().unwrap();
31    let m = opt(&args, "--model3", "static");
32    let model = if m == "static" {
33        Model3::Static
34    } else if let Some(p) = m.strip_prefix("pretrained:") {
35        let net = std::sync::Arc::new(
36            cram_codec::learned::Net::from_bytes(&std::fs::read(p).expect("model")).expect("parse"),
37        );
38        cram_codec::learned::register(net.clone());
39        Model3::Pretrained(net)
40    } else {
41        panic!("model3 {m}")
42    };
43    let mut fields: Vec<(String, std::path::PathBuf, Vec<usize>)> = ORIG
44        .iter()
45        .map(|n| {
46            (
47                n.to_string(),
48                root.join(format!("out/{n}-64x128x128.f32")),
49                vec![64, 128, 128],
50            )
51        })
52        .collect();
53    let man = std::fs::read_to_string(root.join("data/heldout/manifest.json")).unwrap();
54    // minimal manifest parse: [{"name": "...", "shape": [..]}, ...]
55    for obj in man.split('{').skip(1) {
56        let name = obj
57            .split("\"name\"")
58            .nth(1)
59            .unwrap()
60            .split('"')
61            .nth(1)
62            .unwrap()
63            .to_string();
64        let sh = obj
65            .split("\"shape\"")
66            .nth(1)
67            .unwrap()
68            .split('[')
69            .nth(1)
70            .unwrap()
71            .split(']')
72            .next()
73            .unwrap();
74        let shape: Vec<usize> = sh.split(',').map(|v| v.trim().parse().unwrap()).collect();
75        fields.push((
76            name.clone(),
77            root.join(format!("data/heldout/{name}.f32")),
78            shape,
79        ));
80    }
81    let only = opt(&args, "--only", "");
82    let ser: std::collections::HashMap<(String, String), f64> =
83        std::fs::read_to_string(root.join("out/gpu_ratio.jsonl"))
84            .unwrap()
85            .lines()
86            .filter(|l| l.contains("\"tc-wavelet\""))
87            .map(|l| {
88                let g = |k: &str| {
89                    l.split(&format!("\"{k}\":"))
90                        .nth(1)
91                        .unwrap()
92                        .split([',', '}'])
93                        .next()
94                        .unwrap()
95                        .trim()
96                        .trim_matches('"')
97                        .to_string()
98                };
99                (
100                    (g("field").replace("-64x128x128", ""), g("target")),
101                    g("bpv").parse().unwrap(),
102                )
103            })
104            .collect();
105    let mut acc: std::collections::HashMap<(String, String), Vec<f64>> = Default::default();
106    for t in opt(&args, "--targets", "p80").split(',') {
107        for (name, path, dims) in &fields {
108            if !only.is_empty() && !only.split(',').any(|o| name.contains(o)) {
109                continue;
110            }
111            let b = std::fs::read(path).unwrap();
112            let x: Vec<f32> = b
113                .as_chunks::<4>()
114                .0
115                .iter()
116                .map(|&c| f32::from_le_bytes(c))
117                .collect();
118            let range = x.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)) as f64
119                - x.iter().fold(f32::INFINITY, |a, &b| a.min(b)) as f64;
120            let bound = if let Some(p) = t.strip_prefix('p') {
121                ErrorBound::Psnr(p.parse().unwrap())
122            } else {
123                ErrorBound::MaxAbs(t[1..].parse::<f64>().unwrap() * range)
124            };
125            let cfg = Gpu3Config {
126                lanes,
127                compact: !args.iter().any(|a| a == "--nocompact"),
128                axis: !args.iter().any(|a| a == "--noaxis"),
129                model: model.clone(),
130                ..Gpu3Config::new(bound)
131            };
132            let (s, _) = gpuwav3::encode3_f32(&x, dims, &cfg, threads);
133            let bpv = s.len() as f64 * 8.0 / x.len() as f64;
134            let r = ser.get(&(name.clone(), t.to_string())).map(|sb| bpv / sb);
135            println!(
136                "{{\"field\": \"{name}\", \"target\": \"{t}\", \"bpv\": {bpv}, \"vs_serial\": {}}}",
137                r.map_or("null".into(), |v| format!("{v:.4}"))
138            );
139            if let Some(r) = r {
140                let g = if ORIG.contains(&name.as_str()) {
141                    "orig"
142                } else {
143                    "held"
144                };
145                acc.entry((t.to_string(), "all".into()))
146                    .or_default()
147                    .push(r);
148                acc.entry((t.to_string(), g.into())).or_default().push(r);
149            }
150        }
151    }
152    let mut keys: Vec<_> = acc.keys().cloned().collect();
153    keys.sort();
154    for k in keys {
155        let v = &acc[&k];
156        let gm = (v.iter().map(|x| x.ln()).sum::<f64>() / v.len() as f64).exp();
157        println!("GEOMEAN {} {} n={} {:.4}", k.0, k.1, v.len(), gm);
158    }
159}
sweep3/sweep3.rs

sweep3/
sweep3.rs
