Struct Ccx

Source

pub struct Ccx {
    pub t: HashMap<String, (Vec<usize>, Vec<f64>)>,
}

Expand description

A model exported by export_cc.py: named f64 arrays with shapes.

Fields§

§t: HashMap<String, (Vec<usize>, Vec<f64>)>

name → (shape, data)

Implementations§

Source §

impl Ccx

Source

pub fn parse(b: &[u8]) -> Result<Self, DecodeError>

Parse the CCX1 container.

Examples found in repository ?

examples/imgcodec.rs (line 11)

9fn main() {
10    let a: Vec<String> = std::env::args().collect();
11    let c = img::Ccx::parse(&std::fs::read(&a[1]).unwrap()).unwrap();
12    let lanes = a.get(2).map_or(32, |s| s.parse().unwrap());
13    // TILE as W or WxH (0 = whole grid)
14    let (tile_w, tile_h) = a.get(3).map_or((0, 0), |s| match s.split_once('x') {
15        Some((w, h)) => (w.parse().unwrap(), h.parse().unwrap()),
16        None => (s.parse().unwrap(), 0),
17    });
18    let t0 = Instant::now();
19    let (bytes, st) = img::encode(&c, lanes, tile_w, tile_h);
20    let te = t0.elapsed().as_secs_f64();
21    if let Some(p) = a.get(4) {
22        std::fs::write(p, &bytes).unwrap();
23    }
24    let t0 = Instant::now();
25    let (h, w, rgb) = img::decode(&bytes).unwrap();
26    let td = t0.elapsed().as_secs_f64();
27    let px = (h * w) as f64;
28    let cc = &c.t["bytes"].1;
29    let refrgb: Vec<u8> = c.t["ref_rgb"].1.iter().map(|&v| v as u8).collect();
30    let orig: Option<Vec<u8>> =
31        c.t.get("orig_rgb")
32            .map(|t| t.1.iter().map(|&v| v as u8).collect());
33    let lat = st.latent_bytes;
34    println!(
35        "{h}x{w} lanes {lanes} tile {tile_w}x{tile_h}: {} B = {:.4} bpp (model {} B, latents {} B, ideal {:.0} B) | coolchic {} B = {:.4} bpp (nn {} latent {})",
36        st.bytes,
37        st.bytes as f64 * 8.0 / px,
38        st.model_bytes,
39        lat,
40        st.ideal_bits / 8.0,
41        cc[0],
42        cc[0] * 8.0 / px,
43        cc[1],
44        cc[2]
45    );
46    if let Some(o) = &orig {
47        println!(
48            "PSNR ours {:.3} dB | coolchic ref {:.3} dB | ours vs ref {:.2} dB",
49            img::psnr(&rgb, o),
50            img::psnr(&refrgb, o),
51            img::psnr(&rgb, &refrgb)
52        );
53    } else {
54        println!("ours vs coolchic ref {:.2} dB", img::psnr(&rgb, &refrgb));
55    }
56    #[cfg(feature = "gpu")]
57    {
58        let gpu = cram_codec::imggpu::ImgGpu::new().unwrap();
59        let p = img::parse(&bytes).unwrap();
60        let pr = gpu.prepare(&p);
61        if std::env::var("TC_DUMP_WGSL").is_ok() {
62            std::fs::write("/tmp/tci.wgsl", &pr.src).unwrap();
63        }
64        let mut best = f64::MAX;
65        for _ in 0..5 {
66            best = best.min(gpu.run(&pr));
67        }
68        if std::env::var("TC_PROF").is_ok() {
69            let names = gpu.plan_names(&pr);
70            let mut prev = 0.0;
71            for k in 1..=names.len() {
72                let t = (0..3)
73                    .map(|_| gpu.run_upto(&pr, k))
74                    .fold(f64::MAX, f64::min);
75                println!(
76                    "  {:>8} cum {:7.2} ms (+{:.2})",
77                    names[k - 1],
78                    t * 1e3,
79                    (t - prev) * 1e3
80                );
81                prev = t;
82            }
83        }
84        let g = gpu.readback(&pr);
85        let same = g == rgb;
86        let diff = g.iter().zip(&rgb).filter(|(a, b)| a != b).count();
87        println!("GPU {}: decode {:.2} ms (best of 5), compile {:.2} s, byte-identical to CPU: {same} ({diff} bytes differ)", gpu.name, best * 1e3, pr.compile_s);
88    }
89    println!(
90        "encode {te:.2}s, CPU decode {td:.3}s; grid ideal bytes {:?}",
91        st.grid_bits
92            .iter()
93            .map(|b| (b / 8.0).round() as i64)
94            .collect::<Vec<_>>()
95    );
96}