pub fn register(net: Arc<Net>)Expand description
Make a pretrained network available to decoders in this process.
Examples found in repository?
examples/gpubench.rs (line 31)
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/sweep3.rs (line 38)
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}