qtt/
qtt.rs

1//! `qtt` — encode/decode raw little-endian f32 files with the QTT codec.
2//!
3//! ```text
4//! cargo run --release -p cram-codec --example qtt -- encode in.f32 out.qtt --dims 256x384x384 --psnr 60
5//! cargo run --release -p cram-codec --example qtt -- decode out.qtt back.f32
6//! cargo run --release -p cram-codec --example qtt -- bench in.f32 --dims 512x512 --max-rel 1e-3
7//! ```
8//!
9//! Bounds: `--psnr dB` (range-normalised), `--rmse x`, `--max-abs x`,
10//! `--max-rel r` (max-abs = r·range). Options: `--order morton|row`,
11//! `--pad mixed|pow2`, `--gd` (Gavish–Donoho rank cap), `--max-rank n`.
12//! `encode` and `bench` print one JSON line of stats (sizes, ranks, errors,
13//! single-thread wall times).
14
15use cram_codec::{decode, encode_f32, Config, Decoded, DigitOrder, ErrorBound, PadPolicy};
16use std::time::Instant;
17
18fn usage() -> ! {
19    eprintln!("usage: qtt encode IN.f32 OUT.qtt --dims AxBxC BOUND [opts]\n       qtt decode IN.qtt OUT.f32\n       qtt bench IN.f32 --dims AxBxC BOUND [opts]");
20    std::process::exit(2)
21}
22
23fn read_f32(path: &str) -> Vec<f32> {
24    let b = std::fs::read(path).unwrap_or_else(|e| panic!("read {path}: {e}"));
25    b.as_chunks::<4>()
26        .0
27        .iter()
28        .map(|&c| f32::from_le_bytes(c))
29        .collect()
30}
31
32fn parse(args: &[String], data: &[f32]) -> (Vec<usize>, Config) {
33    let mut dims = None;
34    let mut bound = None;
35    let mut order = DigitOrder::Morton;
36    let mut pad = PadPolicy::MixedRadix;
37    let mut gd = false;
38    let mut max_rank = 512;
39    let range = {
40        let (a, b) = data
41            .iter()
42            .fold((f32::INFINITY, f32::NEG_INFINITY), |(a, b), &v| {
43                (a.min(v), b.max(v))
44            });
45        (b - a) as f64
46    };
47    let mut i = 0;
48    while i < args.len() {
49        let v = || args.get(i + 1).cloned().unwrap_or_else(|| usage());
50        match args[i].as_str() {
51            "--dims" => {
52                dims = Some(
53                    v().split('x')
54                        .map(|s| s.parse().unwrap())
55                        .collect::<Vec<usize>>(),
56                )
57            }
58            "--psnr" => bound = Some(ErrorBound::Psnr(v().parse().unwrap())),
59            "--rmse" => bound = Some(ErrorBound::Rmse(v().parse().unwrap())),
60            "--max-abs" => bound = Some(ErrorBound::MaxAbs(v().parse().unwrap())),
61            "--max-rel" => bound = Some(ErrorBound::MaxAbs(v().parse::<f64>().unwrap() * range)),
62            "--order" => {
63                order = match v().as_str() {
64                    "morton" => DigitOrder::Morton,
65                    "row" => DigitOrder::RowMajor,
66                    _ => usage(),
67                }
68            }
69            "--pad" => {
70                pad = match v().as_str() {
71                    "mixed" => PadPolicy::MixedRadix,
72                    "pow2" => PadPolicy::Pow2,
73                    _ => usage(),
74                }
75            }
76            "--max-rank" => max_rank = v().parse().unwrap(),
77            "--gd" => {
78                gd = true;
79                i += 1;
80                continue;
81            }
82            _ => usage(),
83        }
84        i += 2;
85    }
86    let dims = dims.unwrap_or_else(|| vec![data.len()]);
87    assert_eq!(
88        dims.iter().product::<usize>(),
89        data.len(),
90        "--dims does not match file size"
91    );
92    let mut cfg = Config::new(bound.unwrap_or_else(|| usage()));
93    cfg.order = order;
94    cfg.pad = pad;
95    cfg.gavish_donoho = gd;
96    cfg.max_rank = max_rank;
97    (dims, cfg)
98}
99
100fn main() {
101    let args: Vec<String> = std::env::args().skip(1).collect();
102    match args.first().map(String::as_str) {
103        Some("encode") if args.len() >= 3 => {
104            let x = read_f32(&args[1]);
105            let (dims, cfg) = parse(&args[3..], &x);
106            let t = Instant::now();
107            let (bytes, st) = encode_f32(&x, &dims, &cfg);
108            let enc = t.elapsed().as_secs_f64();
109            std::fs::write(&args[2], &bytes).unwrap();
110            println!("{}", json(&st, enc, None));
111        }
112        Some("decode") if args.len() == 3 => {
113            let bytes = std::fs::read(&args[1]).unwrap();
114            let t = Instant::now();
115            let Decoded::F32 { data, .. } = decode(&bytes).expect("decode") else {
116                panic!("not an f32 stream")
117            };
118            let dec = t.elapsed().as_secs_f64();
119            let out: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
120            std::fs::write(&args[2], out).unwrap();
121            println!("{{\"decode_s\":{dec}}}");
122        }
123        Some("bench") if args.len() >= 2 => {
124            let x = read_f32(&args[1]);
125            let (dims, cfg) = parse(&args[2..], &x);
126            let t = Instant::now();
127            let (bytes, st) = encode_f32(&x, &dims, &cfg);
128            let enc = t.elapsed().as_secs_f64();
129            let t = Instant::now();
130            let Decoded::F32 { data, .. } = decode(&bytes).unwrap() else {
131                unreachable!()
132            };
133            let dec = t.elapsed().as_secs_f64();
134            let mx = x
135                .iter()
136                .zip(&data)
137                .map(|(&a, &b)| (a as f64 - b as f64).abs())
138                .fold(0.0, f64::max);
139            assert_eq!(
140                mx, st.max_abs,
141                "decode differs from the encoder's self-check"
142            );
143            println!("{}", json(&st, enc, Some(dec)));
144        }
145        _ => usage(),
146    }
147}
148
149fn json(st: &cram_codec::EncodeStats, enc: f64, dec: Option<f64>) -> String {
150    let ranks: Vec<String> = st.ranks.iter().map(|r| r.to_string()).collect();
151    let modes: Vec<String> = st.modes.iter().map(|r| r.to_string()).collect();
152    let psnr = 20.0 * (st.range / st.rmse).log10();
153    format!(
154        "{{\"bytes\":{},\"bpv\":{},\"header_bytes\":{},\"core_bytes\":{},\"residual_bytes\":{},\"n_residual\":{},\"rmse\":{},\"psnr\":{},\"max_abs\":{},\"range\":{},\"delta\":{},\"passes\":{},\"encode_s\":{},\"decode_s\":{},\"ranks\":[{}],\"modes\":[{}]}}",
155        st.bytes, st.bits_per_value, st.header_bytes, st.core_bytes, st.residual_bytes, st.n_residual,
156        st.rmse, if psnr.is_finite() { psnr } else { 999.0 }, st.max_abs, st.range, st.delta, st.passes, enc,
157        dec.map_or("null".into(), |d| d.to_string()), ranks.join(","), modes.join(",")
158    )
159}
qtt/qtt.rs

qtt/
qtt.rs
