pub struct Net {
pub sizes: Vec<usize>,
pub wq: Vec<Vec<i32>>,
pub bq: Vec<Vec<i32>>,
pub apm: bool,
}Expand description
An integer network (Q8 weights, Q16 biases).
Fields§
§sizes: Vec<usize>Layer sizes, [NF, h…, NO].
wq: Vec<Vec<i32>>Q8 weights per layer (row-major out × in).
bq: Vec<Vec<i32>>Q16 biases per layer.
apm: boolWhether coding uses the adaptive probability map.
Implementations§
Source§impl Net
impl Net
Sourcepub fn fingerprint(&self) -> u64
pub fn fingerprint(&self) -> u64
64-bit FNV-1a fingerprint of the integer model.
Sourcepub fn num_params(&self) -> usize
pub fn num_params(&self) -> usize
Number of parameters.
Examples found in repository?
examples/pretrain3.rs (line 71)
17fn main() {
18 let args: Vec<String> = std::env::args().skip(1).collect();
19 let out = &args[0];
20 let hidden: Vec<usize> = opt(&args, "--hidden", "32,16")
21 .split(',')
22 .map(|v| v.parse().unwrap())
23 .collect();
24 let psnr: Vec<f64> = opt(&args, "--psnr", "40,50,60,70,80,90,100")
25 .split(',')
26 .map(|v| v.parse().unwrap())
27 .collect();
28 let lanes: usize = opt(&args, "--lanes", "8").parse().unwrap();
29 let block: usize = opt(&args, "--block", "32").parse().unwrap();
30 let per: usize = opt(&args, "--per", "100000").parse().unwrap();
31 let epochs: usize = opt(&args, "--epochs", "8").parse().unwrap();
32 let mut smp = Vec::new();
33 for a in args.iter().filter(|a| a.contains(".f32:")) {
34 let (path, d) = a.split_once(':').unwrap();
35 let dims: Vec<usize> = d.split('x').map(|s| s.parse().unwrap()).collect();
36 let b = std::fs::read(path).unwrap();
37 let x: Vec<f32> = b
38 .as_chunks::<4>()
39 .0
40 .iter()
41 .map(|&c| f32::from_le_bytes(c))
42 .collect();
43 for &p in &psnr {
44 let cfg = Gpu3Config {
45 block,
46 block2d: block * 4,
47 lanes,
48 ..Gpu3Config::new(ErrorBound::Psnr(p))
49 };
50 let s = samples3(&x, &dims, &cfg, 16, per);
51 eprintln!("{path} {p} dB: {} samples", s.len());
52 smp.extend(s);
53 }
54 }
55 let tc = TrainConfig {
56 hidden,
57 epochs,
58 per_epoch: 3_000_000,
59 batch: 256,
60 lr: 3e-3,
61 apm: false,
62 seed: 7,
63 };
64 let t = std::time::Instant::now();
65 let net = train(&smp, &tc);
66 eprintln!(
67 "trained on {} samples in {:.1}s: {:.4} bits/sample (float view), {} params",
68 smp.len(),
69 t.elapsed().as_secs_f64(),
70 mean_bits_float(&net, &smp),
71 net.num_params()
72 );
73 std::fs::write(out, net.to_bytes()).unwrap();
74}More examples
examples/pretrain.rs (line 82)
16fn main() {
17 let a: Vec<String> = std::env::args().skip(1).collect();
18 let out = &a[0];
19 let mut dims = vec![64, 128, 128];
20 let mut psnrs = vec![40.0, 60.0, 80.0, 100.0];
21 let mut cfg = TrainConfig::pretrain();
22 let mut per = 150_000usize;
23 let mut files = Vec::new();
24 let mut i = 1;
25 while i < a.len() {
26 match a[i].as_str() {
27 "--dims" => dims = a[i + 1].split('x').map(|s| s.parse().unwrap()).collect(),
28 "--psnr" => psnrs = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
29 "--hidden" => cfg.hidden = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
30 "--per" => per = a[i + 1].parse().unwrap(),
31 "--epochs" => cfg.epochs = a[i + 1].parse().unwrap(),
32 "--per-epoch" => cfg.per_epoch = a[i + 1].parse().unwrap(),
33 f => {
34 files.push(f.to_string());
35 i += 1;
36 continue;
37 }
38 }
39 i += 2;
40 }
41 let mut samples: Vec<Sample> = Vec::new();
42 let mut seed = 12345u64;
43 for f in &files {
44 let b = std::fs::read(f).unwrap();
45 let x: Vec<f32> = b
46 .as_chunks::<4>()
47 .0
48 .iter()
49 .map(|&c| f32::from_le_bytes(c))
50 .collect();
51 let xd: Vec<f64> = x.iter().map(|&v| v as f64).collect();
52 let co = Coeffs::new(&xd, &dims);
53 for &p in &psnrs {
54 let wc = WavConfig::new(ErrorBound::Psnr(p));
55 let (_, st) = encode_f32(&x, &dims, &wc);
56 let (q, _) = quantize(&co, st.delta, wc.rho);
57 let preds = fit_predictors(&co.p, &q);
58 let mut s = collect(&co.p, &q, &preds);
59 // seeded partial Fisher–Yates subsample
60 let k = per.min(s.len());
61 for j in 0..k {
62 seed = seed
63 .wrapping_mul(6364136223846793005)
64 .wrapping_add(1442695040888963407);
65 let r = j + (seed >> 11) as usize % (s.len() - j);
66 s.swap(j, r);
67 }
68 s.truncate(k);
69 eprintln!(
70 "{f} psnr {p}: delta {:.4e}, hand {:.4} bpv, {} samples",
71 st.delta,
72 st.bits_per_value,
73 s.len()
74 );
75 samples.extend(s);
76 }
77 }
78 let t = std::time::Instant::now();
79 let net = train(&samples, &cfg);
80 eprintln!(
81 "trained {} params on {} samples in {:.1}s: {:.4} bits/sample (float view)",
82 net.num_params(),
83 samples.len(),
84 t.elapsed().as_secs_f64(),
85 mean_bits_float(&net, &samples)
86 );
87 std::fs::write(out, Arc::new(net).to_bytes()).unwrap();
88 let _ = Net::from_bytes(&std::fs::read(out).unwrap()).unwrap();
89}Sourcepub fn to_bytes(&self) -> Vec<u8> ⓘ
pub fn to_bytes(&self) -> Vec<u8> ⓘ
Serialize (plain varints) — the pretrained-model file format.
Examples found in repository?
examples/pretrain3.rs (line 73)
17fn main() {
18 let args: Vec<String> = std::env::args().skip(1).collect();
19 let out = &args[0];
20 let hidden: Vec<usize> = opt(&args, "--hidden", "32,16")
21 .split(',')
22 .map(|v| v.parse().unwrap())
23 .collect();
24 let psnr: Vec<f64> = opt(&args, "--psnr", "40,50,60,70,80,90,100")
25 .split(',')
26 .map(|v| v.parse().unwrap())
27 .collect();
28 let lanes: usize = opt(&args, "--lanes", "8").parse().unwrap();
29 let block: usize = opt(&args, "--block", "32").parse().unwrap();
30 let per: usize = opt(&args, "--per", "100000").parse().unwrap();
31 let epochs: usize = opt(&args, "--epochs", "8").parse().unwrap();
32 let mut smp = Vec::new();
33 for a in args.iter().filter(|a| a.contains(".f32:")) {
34 let (path, d) = a.split_once(':').unwrap();
35 let dims: Vec<usize> = d.split('x').map(|s| s.parse().unwrap()).collect();
36 let b = std::fs::read(path).unwrap();
37 let x: Vec<f32> = b
38 .as_chunks::<4>()
39 .0
40 .iter()
41 .map(|&c| f32::from_le_bytes(c))
42 .collect();
43 for &p in &psnr {
44 let cfg = Gpu3Config {
45 block,
46 block2d: block * 4,
47 lanes,
48 ..Gpu3Config::new(ErrorBound::Psnr(p))
49 };
50 let s = samples3(&x, &dims, &cfg, 16, per);
51 eprintln!("{path} {p} dB: {} samples", s.len());
52 smp.extend(s);
53 }
54 }
55 let tc = TrainConfig {
56 hidden,
57 epochs,
58 per_epoch: 3_000_000,
59 batch: 256,
60 lr: 3e-3,
61 apm: false,
62 seed: 7,
63 };
64 let t = std::time::Instant::now();
65 let net = train(&smp, &tc);
66 eprintln!(
67 "trained on {} samples in {:.1}s: {:.4} bits/sample (float view), {} params",
68 smp.len(),
69 t.elapsed().as_secs_f64(),
70 mean_bits_float(&net, &smp),
71 net.num_params()
72 );
73 std::fs::write(out, net.to_bytes()).unwrap();
74}More examples
examples/pretrain.rs (line 87)
16fn main() {
17 let a: Vec<String> = std::env::args().skip(1).collect();
18 let out = &a[0];
19 let mut dims = vec![64, 128, 128];
20 let mut psnrs = vec![40.0, 60.0, 80.0, 100.0];
21 let mut cfg = TrainConfig::pretrain();
22 let mut per = 150_000usize;
23 let mut files = Vec::new();
24 let mut i = 1;
25 while i < a.len() {
26 match a[i].as_str() {
27 "--dims" => dims = a[i + 1].split('x').map(|s| s.parse().unwrap()).collect(),
28 "--psnr" => psnrs = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
29 "--hidden" => cfg.hidden = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
30 "--per" => per = a[i + 1].parse().unwrap(),
31 "--epochs" => cfg.epochs = a[i + 1].parse().unwrap(),
32 "--per-epoch" => cfg.per_epoch = a[i + 1].parse().unwrap(),
33 f => {
34 files.push(f.to_string());
35 i += 1;
36 continue;
37 }
38 }
39 i += 2;
40 }
41 let mut samples: Vec<Sample> = Vec::new();
42 let mut seed = 12345u64;
43 for f in &files {
44 let b = std::fs::read(f).unwrap();
45 let x: Vec<f32> = b
46 .as_chunks::<4>()
47 .0
48 .iter()
49 .map(|&c| f32::from_le_bytes(c))
50 .collect();
51 let xd: Vec<f64> = x.iter().map(|&v| v as f64).collect();
52 let co = Coeffs::new(&xd, &dims);
53 for &p in &psnrs {
54 let wc = WavConfig::new(ErrorBound::Psnr(p));
55 let (_, st) = encode_f32(&x, &dims, &wc);
56 let (q, _) = quantize(&co, st.delta, wc.rho);
57 let preds = fit_predictors(&co.p, &q);
58 let mut s = collect(&co.p, &q, &preds);
59 // seeded partial Fisher–Yates subsample
60 let k = per.min(s.len());
61 for j in 0..k {
62 seed = seed
63 .wrapping_mul(6364136223846793005)
64 .wrapping_add(1442695040888963407);
65 let r = j + (seed >> 11) as usize % (s.len() - j);
66 s.swap(j, r);
67 }
68 s.truncate(k);
69 eprintln!(
70 "{f} psnr {p}: delta {:.4e}, hand {:.4} bpv, {} samples",
71 st.delta,
72 st.bits_per_value,
73 s.len()
74 );
75 samples.extend(s);
76 }
77 }
78 let t = std::time::Instant::now();
79 let net = train(&samples, &cfg);
80 eprintln!(
81 "trained {} params on {} samples in {:.1}s: {:.4} bits/sample (float view)",
82 net.num_params(),
83 samples.len(),
84 t.elapsed().as_secs_f64(),
85 mean_bits_float(&net, &samples)
86 );
87 std::fs::write(out, Arc::new(net).to_bytes()).unwrap();
88 let _ = Net::from_bytes(&std::fs::read(out).unwrap()).unwrap();
89}Sourcepub fn from_bytes(b: &[u8]) -> Result<Self, DecodeError>
pub fn from_bytes(b: &[u8]) -> Result<Self, DecodeError>
Parse Net::to_bytes.
Examples found in repository?
examples/gpubench.rs (line 29)
14fn cfg3_of(bound: ErrorBound, bk: usize, args: &[String]) -> Gpu3Config {
15 let m = opt(args, "--model3", "static");
16 let model = if m == "static" {
17 Model3::Static
18 } else if let Some(h) = m.strip_prefix("overfit") {
19 let hidden = h.strip_prefix(':').map_or(vec![16], |h| {
20 h.split(',').map(|v| v.parse().unwrap()).collect()
21 });
22 Model3::Overfit {
23 hidden,
24 per_epoch: 200_000,
25 epochs: 4,
26 }
27 } else if let Some(p) = m.strip_prefix("pretrained:") {
28 let net = std::sync::Arc::new(
29 cram_codec::learned::Net::from_bytes(&std::fs::read(p).expect("model")).expect("parse"),
30 );
31 cram_codec::learned::register(net.clone());
32 Model3::Pretrained(net)
33 } else {
34 panic!("model3 {m}")
35 };
36 Gpu3Config {
37 block: bk,
38 block2d: bk,
39 resctx: !args.iter().any(|a| a == "--nores"),
40 lanes: opt(args, "--l3", "1").parse().unwrap(),
41 compact: !args.iter().any(|a| a == "--nocompact"),
42 axis: !args.iter().any(|a| a == "--noaxis"),
43 model,
44 ..Gpu3Config::new(bound)
45 }
46}
47use cram_codec::{wavcodec, Decoded};
48
49fn read_f32(path: &str) -> Vec<f32> {
50 let b = std::fs::read(path).unwrap_or_else(|e| panic!("read {path}: {e}"));
51 b.as_chunks::<4>()
52 .0
53 .iter()
54 .map(|&c| f32::from_le_bytes(c))
55 .collect()
56}
57
58fn err(x: &[f32], y: &[f32]) -> (f64, f64) {
59 let mut se = 0.0;
60 let mut mx = 0.0f64;
61 for (a, b) in x.iter().zip(y) {
62 let e = *a as f64 - *b as f64;
63 se += e * e;
64 mx = mx.max(e.abs());
65 }
66 ((se / x.len() as f64).sqrt(), mx)
67}
68
69fn opt(args: &[String], k: &str, d: &str) -> String {
70 args.iter()
71 .position(|a| a == k)
72 .map(|i| args[i + 1].clone())
73 .unwrap_or_else(|| d.to_string())
74}
75
76fn main() {
77 let args: Vec<String> = std::env::args().skip(1).collect();
78 let mode = args[0].as_str();
79 if mode == "cpu3" {
80 // gpubench cpu3 STREAM... [pretrained:PATH]: CPU decode vs saved hash (TC_CPU_IR=1: net via tang-expr)
81 for p in args.iter().filter_map(|a| a.strip_prefix("pretrained:")) {
82 cram_codec::learned::register(std::sync::Arc::new(
83 cram_codec::learned::Net::from_bytes(&std::fs::read(p).unwrap()).unwrap(),
84 ));
85 }
86 let threads: usize = opt(&args, "--threads", "16").parse().unwrap();
87 for f in args[1..]
88 .iter()
89 .filter(|a| a.ends_with(".tcg3") || a.ends_with(".tcs3"))
90 {
91 let b = std::fs::read(f).unwrap();
92 let want = std::fs::read_to_string(format!("{f}.hash")).unwrap();
93 let t0 = Instant::now();
94 let (_, y) = gpuwav3::decode3_f32(&b, threads).unwrap();
95 let s = t0.elapsed().as_secs_f64();
96 let h = y.iter().fold(0xcbf29ce484222325u64, |h, v| {
97 (h ^ v.to_bits() as u64).wrapping_mul(0x100000001b3)
98 });
99 println!(
100 "{{\"stream\":\"{f}\",\"identical\":{},\"dec_s\":{s},\"gbs\":{}}}",
101 format!("{h:016x}") == want.trim(),
102 y.len() as f64 * 4.0 / 1e9 / s
103 );
104 }
105 return;
106 }
107 if mode == "ship3" {
108 // gpubench ship3 STREAM.tcg3... [--mlp] [pretrained:PATH]: write STREAM.tcs3 carrying the
109 // built-in decoder (with the net's MLP as a graph when --mlp); report bytes and compile/cache time
110 for p in args.iter().filter_map(|a| a.strip_prefix("pretrained:")) {
111 cram_codec::learned::register(std::sync::Arc::new(
112 cram_codec::learned::Net::from_bytes(&std::fs::read(p).unwrap()).unwrap(),
113 ));
114 }
115 let mlp = args.iter().any(|a| a == "--mlp");
116 for f in args[1..].iter().filter(|a| a.ends_with(".tcg3")) {
117 let b = std::fs::read(f).unwrap();
118 let ps = gpuwav3::parse3(&b).unwrap();
119 let prog = cram_codec::shipdec::builtin_program(
120 if mlp { ps.net.as_ref() } else { None },
121 ps.net.is_some() && ps.ad.wa > 0,
122 );
123 let w = cram_codec::shipdec::wrap(&b, &prog);
124 let out = format!("{}.tcs3", f.trim_end_matches(".tcg3"));
125 std::fs::write(&out, &w).unwrap();
126 std::fs::copy(format!("{f}.hash"), format!("{out}.hash")).unwrap();
127 let t0 = Instant::now();
128 let d = cram_codec::shipdec::compile(&prog).unwrap();
129 let verify_lower_s = t0.elapsed().as_secs_f64();
130 let t1 = Instant::now();
131 let _ = cram_codec::shipdec::compile_cached(&prog).unwrap();
132 let first_s = t1.elapsed().as_secs_f64();
133 let t2 = Instant::now();
134 let _ = cram_codec::shipdec::compile_cached(&prog).unwrap();
135 let hit_s = t2.elapsed().as_secs_f64();
136 #[cfg(not(feature = "gpu"))]
137 let gpu = String::new();
138 #[cfg(feature = "gpu")]
139 let gpu = {
140 let ws = gpuwav3::parse3(&w).unwrap();
141 let g = cram_codec::gpudec::GpuDecoder::new(
142 ws.lanes,
143 ws.d3.iter().copied().max().unwrap(),
144 )
145 .unwrap();
146 let mut tt = cram_codec::gpudec::GpuTimes::default();
147 let pr = g.prepare3(&ws, &w, &mut tt);
148 let cold = g.warm3(&pr);
149 let hot = g.warm3(&pr);
150 format!(",\"gpu_pipeline_cold_s\":{cold},\"gpu_pipeline_hit_s\":{hot}")
151 };
152 println!(
153 "{{\"out\":\"{out}\",\"stream_bytes\":{},\"program_bytes\":{},\"added_bytes\":{},\"mlp\":{},\"verify_lower_s\":{verify_lower_s},\"cache_miss_s\":{first_s},\"cache_hit_s\":{hit_s},\"wgsl_chars\":{}{gpu}}}",
154 b.len(),
155 prog.len(),
156 w.len() - b.len(),
157 d.mlp.is_some(),
158 d.wgsl_res.len() + d.wgsl_mlp.as_ref().map_or(0, |m| m.len())
159 );
160 }
161 return;
162 }
163 #[cfg(feature = "gpu")]
164 if mode == "gpu3" {
165 // gpubench gpu3 STREAM... [--model3 pretrained:PATH]: GPU decode vs saved hash, kernel timing
166 for p in args.iter().filter_map(|a| a.strip_prefix("pretrained:")) {
167 cram_codec::learned::register(std::sync::Arc::new(
168 cram_codec::learned::Net::from_bytes(&std::fs::read(p).unwrap()).unwrap(),
169 ));
170 }
171 let reps: usize = opt(&args, "--reps", "3").parse().unwrap();
172 for f in args[1..]
173 .iter()
174 .filter(|a| a.ends_with(".tcg3") || a.ends_with(".tcs3"))
175 {
176 let b = std::fs::read(f).unwrap();
177 let want = std::fs::read_to_string(format!("{f}.hash")).unwrap();
178 let ps = gpuwav3::parse3(&b).unwrap();
179 let g =
180 cram_codec::gpudec::GpuDecoder::new(ps.lanes, ps.d3.iter().copied().max().unwrap())
181 .unwrap();
182 let (v, _) = g.decode3(&b).unwrap();
183 let h = v.iter().fold(0xcbf29ce484222325u64, |h, &x| {
184 (h ^ x as u64).wrapping_mul(0x100000001b3)
185 });
186 let mut t2 = cram_codec::gpudec::GpuTimes::default();
187 let pr = g.prepare3(&ps, &b, &mut t2);
188 let kern = (0..reps).map(|_| g.run3(&pr)).fold(f64::INFINITY, f64::min);
189 let mut tot = f64::INFINITY;
190 for _ in 0..reps {
191 let t0 = Instant::now();
192 let _ = g.decode3(&b).unwrap();
193 tot = tot.min(t0.elapsed().as_secs_f64());
194 }
195 let gb = v.len() as f64 * 4.0 / 1e9;
196 println!("{{\"stream\":\"{f}\",\"identical\":{},\"kernels_s\":{kern},\"kernels_gbs\":{},\"total_gbs\":{}}}", format!("{h:016x}") == want.trim(), gb / kern, gb / tot);
197 }
198 return;
199 }
200 let x = read_f32(&args[1]);
201 let dims: Vec<usize> = args[2].split('x').map(|s| s.parse().unwrap()).collect();
202 let name = std::path::Path::new(&args[1])
203 .file_stem()
204 .unwrap()
205 .to_string_lossy()
206 .to_string();
207 let bricks: Vec<usize> = opt(&args, "--bricks", "64")
208 .split(',')
209 .map(|s| s.parse().unwrap())
210 .collect();
211 let lanes: usize = opt(&args, "--lanes", "16").parse().unwrap();
212 let threads: usize = opt(&args, "--threads", "16").parse().unwrap();
213 let targets = opt(&args, "--targets", "p60,p80,p100,m1e-3,m1e-4");
214 let serial = !args.iter().any(|a| a == "--no-serial");
215 let (lo, hi) = x
216 .iter()
217 .fold((f32::INFINITY, f32::NEG_INFINITY), |(a, b), &v| {
218 (a.min(v), b.max(v))
219 });
220 let range = (hi - lo) as f64;
221 match mode {
222 "ratio" => {
223 for t in targets.split(',') {
224 let bound = if let Some(p) = t.strip_prefix('p') {
225 ErrorBound::Psnr(p.parse().unwrap())
226 } else {
227 ErrorBound::MaxAbs(t[1..].parse::<f64>().unwrap() * range)
228 };
229 if serial {
230 let t0 = Instant::now();
231 let (b, _) = wavcodec::encode_f32(&x, &dims, &wavcodec::WavConfig::new(bound));
232 let enc_s = t0.elapsed().as_secs_f64();
233 let t0 = Instant::now();
234 let Decoded::F32 { data, .. } = wavcodec::decode(&b).unwrap() else {
235 panic!()
236 };
237 let dec_s = t0.elapsed().as_secs_f64();
238 let (rmse, mx) = err(&x, &data);
239 println!(
240 "{{\"field\":\"{name}\",\"target\":\"{t}\",\"codec\":\"tc-wavelet\",\"brick\":0,\"bpv\":{},\"psnr\":{},\"max_abs_rel\":{},\"enc_s\":{enc_s},\"dec_s\":{dec_s}}}",
241 8.0 * b.len() as f64 / x.len() as f64,
242 20.0 * (range / rmse).log10(),
243 mx / range
244 );
245 }
246 let v3 = args.iter().any(|a| a == "--v3");
247 for &bk in &bricks {
248 let cfg = GpuConfig {
249 brick: bk,
250 lanes,
251 adaptive: args.iter().any(|a| a == "--adaptive"),
252 ..GpuConfig::new(bound)
253 };
254 let cfg3 = cfg3_of(bound, bk, &args);
255 let t0 = Instant::now();
256 let (b, st) = if v3 {
257 gpuwav3::encode3_f32(&x, &dims, &cfg3, threads)
258 } else {
259 gpuwav::encode_f32_mt(&x, &dims, &cfg, threads)
260 };
261 let enc_s = t0.elapsed().as_secs_f64();
262 let t0 = Instant::now();
263 let (_, y) = if v3 {
264 gpuwav3::decode3_f32(&b, 1).unwrap()
265 } else {
266 gpuwav::decode_f32(&b, 1).unwrap()
267 };
268 let dec_s = t0.elapsed().as_secs_f64();
269 let (rmse, mx) = err(&x, &y);
270 if let ErrorBound::MaxAbs(bb) = bound {
271 assert!(mx <= bb, "{name} {t}: {mx} > {bb}");
272 }
273 println!(
274 "{{\"field\":\"{name}\",\"target\":\"{t}\",\"codec\":\"{}\",\"brick\":{bk},\"lanes\":{lanes},\"bpv\":{},\"psnr\":{},\"max_abs_rel\":{},\"enc_s\":{enc_s},\"dec_s\":{dec_s},\"hdr\":{},\"idx\":{},\"outl\":{},\"esc\":{}}}",
275 if v3 { "tc-gpu3" } else { "tc-gpu" },
276 st.bits_per_value,
277 20.0 * (range / rmse).log10(),
278 mx / range,
279 st.header_bytes,
280 st.index_bytes,
281 st.outlier_bytes,
282 st.n_escapes
283 );
284 }
285 }
286 }
287 "speed" => {
288 // encode once, then time decoders on the same stream
289 let t = opt(&args, "--target", "p80");
290 let bound = if let Some(p) = t.strip_prefix('p') {
291 ErrorBound::Psnr(p.parse().unwrap())
292 } else {
293 ErrorBound::MaxAbs(t[1..].parse::<f64>().unwrap() * range)
294 };
295 let reps: usize = opt(&args, "--reps", "3").parse().unwrap();
296 for &bk in &bricks {
297 let cfg = GpuConfig {
298 brick: bk,
299 lanes,
300 adaptive: args.iter().any(|a| a == "--adaptive"),
301 ..GpuConfig::new(bound)
302 };
303 let t0 = Instant::now();
304 let (b, st) = gpuwav::encode_f32_mt(&x, &dims, &cfg, threads);
305 let enc_s = t0.elapsed().as_secs_f64();
306 let gb = x.len() as f64 * 4.0 / 1e9;
307 let best = |f: &mut dyn FnMut() -> f64| {
308 (0..reps).map(|_| f()).fold(f64::INFINITY, f64::min)
309 };
310 let ps = gpuwav::parse(&b).unwrap();
311 let mut cpu_bits = Vec::new();
312 let cpu1 = best(&mut || {
313 let t0 = Instant::now();
314 cpu_bits = gpuwav::decode_box_parsed(&ps, &b, [0; 3], ps.d3, 1).unwrap();
315 t0.elapsed().as_secs_f64()
316 });
317 let cpumt = best(&mut || {
318 let t0 = Instant::now();
319 let v = gpuwav::decode_box_parsed(&ps, &b, [0; 3], ps.d3, threads).unwrap();
320 assert!(v == cpu_bits);
321 t0.elapsed().as_secs_f64()
322 });
323 // random access: 64^3 (or 64^2) box at the field centre
324 let lo: [usize; 3] =
325 std::array::from_fn(|a| if ps.d3[a] > 64 { ps.d3[a] / 2 - 32 } else { 0 });
326 let hi: [usize; 3] = std::array::from_fn(|a| (lo[a] + 64).min(ps.d3[a]));
327 let mut sub = Vec::new();
328 let ra = best(&mut || {
329 let t0 = Instant::now();
330 sub = gpuwav::decode_box_bits(&b, lo, hi, 1).unwrap();
331 t0.elapsed().as_secs_f64()
332 });
333 let ra_mt = best(&mut || {
334 let t0 = Instant::now();
335 let _ = gpuwav::decode_box_bits(&b, lo, hi, threads).unwrap();
336 t0.elapsed().as_secs_f64()
337 });
338 let mut k = 0;
339 for z in lo[0]..hi[0] {
340 for y in lo[1]..hi[1] {
341 for xx in lo[2]..hi[2] {
342 assert_eq!(
343 sub[k],
344 cpu_bits[(z * ps.d3[1] + y) * ps.d3[2] + xx],
345 "random access mismatch"
346 );
347 k += 1;
348 }
349 }
350 }
351 #[allow(unused_mut, unused_assignments)]
352 let mut gpu_json = String::new();
353 #[cfg(feature = "gpu")]
354 {
355 let maxl = ps.shapes.iter().flat_map(|s| s.dims).max().unwrap();
356 let g = cram_codec::gpudec::GpuDecoder::new(lanes, maxl).expect("gpu");
357 // warm-up + check
358 let (v, _) = g.decode(&b).unwrap();
359 let diff = v.iter().zip(&cpu_bits).filter(|(a, b)| a != b).count();
360 assert_eq!(diff, 0, "GPU != CPU in {diff} cells");
361 let mut tt = cram_codec::gpudec::GpuTimes::default();
362 let mut tot = f64::INFINITY;
363 for _ in 0..reps {
364 let t0 = Instant::now();
365 let (v, t) = g.decode(&b).unwrap();
366 let w = t0.elapsed().as_secs_f64();
367 assert!(v == cpu_bits);
368 if w < tot {
369 tot = w;
370 tt = t;
371 }
372 }
373 let ps2 = gpuwav::parse(&b).unwrap();
374 let mut t2 = cram_codec::gpudec::GpuTimes::default();
375 let pr = g.prepare(&ps2, &b, &mut t2);
376 let kern = best(&mut || g.run(&pr));
377 gpu_json = format!(
378 ",\"gpu\":\"{}\",\"gpu_kernels_s\":{kern},\"gpu_kernels_gbs\":{},\"gpu_total_s\":{tot},\"gpu_total_gbs\":{},\"gpu_parse_s\":{},\"gpu_upload_s\":{},\"gpu_readback_s\":{},\"gpu_bit_identical\":true",
379 g.name,
380 gb / kern,
381 gb / tot,
382 tt.parse_s,
383 tt.upload_s,
384 tt.readback_s
385 );
386 }
387 println!(
388 "{{\"field\":\"{name}\",\"dims\":\"{}\",\"target\":\"{t}\",\"brick\":{bk},\"lanes\":{lanes},\"bpv\":{},\"enc_s\":{enc_s},\"cpu1_s\":{cpu1},\"cpu1_gbs\":{},\"cpumt_s\":{cpumt},\"cpumt_gbs\":{},\"threads\":{threads},\"ra_box\":\"{:?}-{:?}\",\"ra_s\":{ra},\"ra_mt_s\":{ra_mt}{gpu_json}}}",
389 args[2],
390 st.bits_per_value,
391 gb / cpu1,
392 gb / cpumt,
393 lo,
394 hi
395 );
396 }
397 }
398 #[cfg(feature = "gpu")]
399 "verify" => {
400 // CPU == GPU byte-for-byte, error bounds, and random-access crops
401 let g = cram_codec::gpudec::GpuDecoder::new(lanes, 1024).expect("gpu");
402 for t in targets.split(',') {
403 let bound = if let Some(p) = t.strip_prefix('p') {
404 ErrorBound::Psnr(p.parse().unwrap())
405 } else {
406 ErrorBound::MaxAbs(t[1..].parse::<f64>().unwrap() * range)
407 };
408 for &bk in &bricks {
409 let cfg = GpuConfig {
410 brick: bk,
411 lanes,
412 ..GpuConfig::new(bound)
413 };
414 let (b, st) = gpuwav::encode_f32_mt(&x, &dims, &cfg, threads);
415 let ps = gpuwav::parse(&b).unwrap();
416 let cpu = gpuwav::decode_box_parsed(&ps, &b, [0; 3], ps.d3, threads).unwrap();
417 let (gpu, _) = g.decode(&b).unwrap();
418 let diff = cpu.iter().zip(&gpu).filter(|(a, b)| a != b).count();
419 let y: Vec<f32> = cpu.iter().map(|&v| f32::from_bits(v)).collect();
420 let (rmse, mx) = err(&x, &y);
421 let ok_bound = match bound {
422 ErrorBound::MaxAbs(bb) => mx <= bb,
423 ErrorBound::Psnr(db) => 20.0 * (range / rmse).log10() >= db - 1e-9,
424 _ => true,
425 };
426 // random-access sub-box straddling bricks
427 let lo: [usize; 3] = std::array::from_fn(|a| ps.d3[a] / 3);
428 let hi: [usize; 3] =
429 std::array::from_fn(|a| (lo[a] + 64).min(ps.d3[a]).max(lo[a] + 1));
430 let sub = gpuwav::decode_box_bits(&b, lo, hi, 1).unwrap();
431 let mut ra_ok = true;
432 let mut k = 0;
433 for z in lo[0]..hi[0] {
434 for yy in lo[1]..hi[1] {
435 for xx in lo[2]..hi[2] {
436 ra_ok &= sub[k] == cpu[(z * ps.d3[1] + yy) * ps.d3[2] + xx];
437 k += 1;
438 }
439 }
440 }
441 println!(
442 "{{\"field\":\"{name}\",\"target\":\"{t}\",\"brick\":{bk},\"bpv\":{},\"cells\":{},\"gpu_diff_cells\":{diff},\"bound_ok\":{ok_bound},\"ra_ok\":{ra_ok},\"max_abs_rel\":{}}}",
443 st.bits_per_value,
444 x.len(),
445 mx / range
446 );
447 assert!(diff == 0 && ok_bound && ra_ok, "{name} {t} {bk}");
448 }
449 }
450 }
451 #[cfg(feature = "gpu")]
452 "verify3" | "speed3" => {
453 let speed = mode == "speed3";
454 let maxl = dims.iter().copied().max().unwrap();
455 let g = cram_codec::gpudec::GpuDecoder::new(lanes, maxl).expect("gpu");
456 let reps: usize = opt(&args, "--reps", "3").parse().unwrap();
457 let best =
458 |f: &mut dyn FnMut() -> f64| (0..reps).map(|_| f()).fold(f64::INFINITY, f64::min);
459 let gb = x.len() as f64 * 4.0 / 1e9;
460 for t in targets.split(',') {
461 let bound = if let Some(p) = t.strip_prefix('p') {
462 ErrorBound::Psnr(p.parse().unwrap())
463 } else {
464 ErrorBound::MaxAbs(t[1..].parse::<f64>().unwrap() * range)
465 };
466 for &bk in &bricks {
467 let cfg = cfg3_of(bound, bk, &args);
468 let t0 = Instant::now();
469 let (b, st) = gpuwav3::encode3_f32(&x, &dims, &cfg, threads);
470 let enc_s = t0.elapsed().as_secs_f64();
471 let ps = gpuwav3::parse3(&b).unwrap();
472 let full = |th: usize| {
473 gpuwav3::decode3_box_parsed(&ps, &b, [0; 3], ps.d3, th)
474 .unwrap()
475 .bits
476 };
477 let cpu = full(threads);
478 let (gpu, _) = g.decode3(&b).unwrap();
479 let diff = cpu.iter().zip(&gpu).filter(|(a, b)| a != b).count();
480 if diff > 0 {
481 let i = cpu.iter().zip(&gpu).position(|(a, b)| a != b).unwrap();
482 eprintln!(
483 "first diff at {i}: cpu {} gpu {}",
484 f32::from_bits(cpu[i]),
485 f32::from_bits(gpu[i])
486 );
487 }
488 let y: Vec<f32> = cpu.iter().map(|&v| f32::from_bits(v)).collect();
489 let (rmse, mx) = err(&x, &y);
490 let ok_bound = match bound {
491 ErrorBound::MaxAbs(bb) => mx <= bb,
492 ErrorBound::Psnr(db) => 20.0 * (range / rmse).log10() >= db - 1e-9,
493 _ => true,
494 };
495 // random access: 64^3 (2D: 64^2) box, off-centre so it straddles blocks
496 let lo: [usize; 3] =
497 std::array::from_fn(|a| if ps.d3[a] > 64 { ps.d3[a] / 3 } else { 0 });
498 let hi: [usize; 3] = std::array::from_fn(|a| (lo[a] + 64).min(ps.d3[a]));
499 let mut sub = None;
500 let ra = best(&mut || {
501 let t0 = Instant::now();
502 sub = Some(gpuwav3::decode3_box(&b, lo, hi, 1).unwrap());
503 t0.elapsed().as_secs_f64()
504 });
505 let ra_mt = best(&mut || {
506 let t0 = Instant::now();
507 let _ = gpuwav3::decode3_box(&b, lo, hi, threads).unwrap();
508 t0.elapsed().as_secs_f64()
509 });
510 let sub = sub.unwrap();
511 let mut ra_ok = true;
512 let mut k = 0;
513 for z in lo[0]..hi[0] {
514 for yy in lo[1]..hi[1] {
515 for xx in lo[2]..hi[2] {
516 ra_ok &= sub.bits[k] == cpu[(z * ps.d3[1] + yy) * ps.d3[2] + xx];
517 k += 1;
518 }
519 }
520 }
521 let mut sp = String::new();
522 if speed {
523 let cpu1 = best(&mut || {
524 let t0 = Instant::now();
525 let _ = full(1);
526 t0.elapsed().as_secs_f64()
527 });
528 let cpumt = best(&mut || {
529 let t0 = Instant::now();
530 let _ = full(threads);
531 t0.elapsed().as_secs_f64()
532 });
533 let mut tt = cram_codec::gpudec::GpuTimes::default();
534 let tot = best(&mut || {
535 let t0 = Instant::now();
536 let (v, t) = g.decode3(&b).unwrap();
537 assert!(v == cpu);
538 tt = t;
539 t0.elapsed().as_secs_f64()
540 });
541 let mut t2 = cram_codec::gpudec::GpuTimes::default();
542 let pr = g.prepare3(&ps, &b, &mut t2);
543 let kern = best(&mut || g.run3(&pr));
544 sp = format!(
545 ",\"enc_s\":{enc_s},\"cpu1_gbs\":{},\"cpumt_gbs\":{},\"gpu_kernels_gbs\":{},\"gpu_total_gbs\":{},\"gpu_readback_s\":{},\"ra_s\":{ra},\"ra_mt_s\":{ra_mt},\"ra_bytes\":{},\"ra_blocks\":{},\"total_blocks\":{}",
546 gb / cpu1,
547 gb / cpumt,
548 gb / kern,
549 gb / tot,
550 tt.readback_s,
551 sub.bytes_read,
552 sub.blocks,
553 ps.blocks.len()
554 );
555 }
556 println!(
557 "{{\"field\":\"{name}\",\"target\":\"{t}\",\"block\":{bk},\"bpv\":{},\"bytes\":{},\"cells\":{},\"gpu_diff_cells\":{diff},\"bound_ok\":{ok_bound},\"ra_ok\":{ra_ok},\"ra_bytes\":{},\"ra_s\":{ra}{sp}}}",
558 st.bits_per_value,
559 b.len(),
560 x.len(),
561 sub.bytes_read
562 );
563 assert!(diff == 0 && ok_bound && ra_ok, "{name} {t} {bk}");
564 }
565 }
566 }
567 "encode3" => {
568 // write a v3 stream for later GPU-only timing: gpubench encode3 FIELD dims OUT --bricks B --targets T ...
569 let out = opt(&args, "--out", "out.tcg3");
570 let t = targets.split(',').next().unwrap();
571 let bound = if let Some(p) = t.strip_prefix('p') {
572 ErrorBound::Psnr(p.parse().unwrap())
573 } else {
574 ErrorBound::MaxAbs(t[1..].parse::<f64>().unwrap() * range)
575 };
576 let (b, st) =
577 gpuwav3::encode3_f32(&x, &dims, &cfg3_of(bound, bricks[0], &args), threads);
578 std::fs::write(&out, &b).unwrap();
579 let (_, y) = gpuwav3::decode3_f32(&b, threads).unwrap();
580 let h = y.iter().fold(0xcbf29ce484222325u64, |h, v| {
581 (h ^ v.to_bits() as u64).wrapping_mul(0x100000001b3)
582 });
583 std::fs::write(format!("{out}.hash"), format!("{h:016x}")).unwrap();
584 println!(
585 "{{\"out\":\"{out}\",\"bpv\":{},\"hash\":\"{h:016x}\"}}",
586 st.bits_per_value
587 );
588 }
589 _ => panic!("mode"),
590 }
591}More examples
examples/pretrain.rs (line 88)
16fn main() {
17 let a: Vec<String> = std::env::args().skip(1).collect();
18 let out = &a[0];
19 let mut dims = vec![64, 128, 128];
20 let mut psnrs = vec![40.0, 60.0, 80.0, 100.0];
21 let mut cfg = TrainConfig::pretrain();
22 let mut per = 150_000usize;
23 let mut files = Vec::new();
24 let mut i = 1;
25 while i < a.len() {
26 match a[i].as_str() {
27 "--dims" => dims = a[i + 1].split('x').map(|s| s.parse().unwrap()).collect(),
28 "--psnr" => psnrs = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
29 "--hidden" => cfg.hidden = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
30 "--per" => per = a[i + 1].parse().unwrap(),
31 "--epochs" => cfg.epochs = a[i + 1].parse().unwrap(),
32 "--per-epoch" => cfg.per_epoch = a[i + 1].parse().unwrap(),
33 f => {
34 files.push(f.to_string());
35 i += 1;
36 continue;
37 }
38 }
39 i += 2;
40 }
41 let mut samples: Vec<Sample> = Vec::new();
42 let mut seed = 12345u64;
43 for f in &files {
44 let b = std::fs::read(f).unwrap();
45 let x: Vec<f32> = b
46 .as_chunks::<4>()
47 .0
48 .iter()
49 .map(|&c| f32::from_le_bytes(c))
50 .collect();
51 let xd: Vec<f64> = x.iter().map(|&v| v as f64).collect();
52 let co = Coeffs::new(&xd, &dims);
53 for &p in &psnrs {
54 let wc = WavConfig::new(ErrorBound::Psnr(p));
55 let (_, st) = encode_f32(&x, &dims, &wc);
56 let (q, _) = quantize(&co, st.delta, wc.rho);
57 let preds = fit_predictors(&co.p, &q);
58 let mut s = collect(&co.p, &q, &preds);
59 // seeded partial Fisher–Yates subsample
60 let k = per.min(s.len());
61 for j in 0..k {
62 seed = seed
63 .wrapping_mul(6364136223846793005)
64 .wrapping_add(1442695040888963407);
65 let r = j + (seed >> 11) as usize % (s.len() - j);
66 s.swap(j, r);
67 }
68 s.truncate(k);
69 eprintln!(
70 "{f} psnr {p}: delta {:.4e}, hand {:.4} bpv, {} samples",
71 st.delta,
72 st.bits_per_value,
73 s.len()
74 );
75 samples.extend(s);
76 }
77 }
78 let t = std::time::Instant::now();
79 let net = train(&samples, &cfg);
80 eprintln!(
81 "trained {} params on {} samples in {:.1}s: {:.4} bits/sample (float view)",
82 net.num_params(),
83 samples.len(),
84 t.elapsed().as_secs_f64(),
85 mean_bits_float(&net, &samples)
86 );
87 std::fs::write(out, Arc::new(net).to_bytes()).unwrap();
88 let _ = Net::from_bytes(&std::fs::read(out).unwrap()).unwrap();
89}examples/sweep3.rs (line 36)
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}Trait Implementations§
Auto Trait Implementations§
impl Freeze for Net
impl RefUnwindSafe for Net
impl Send for Net
impl Sync for Net
impl Unpin for Net
impl UnsafeUnpin for Net
impl UnwindSafe for Net
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more