pub fn pause_gpu() -> GpuPauseExpand description
Park the device for every thread until the returned guard drops.
Examples found in repository?
examples/zimage_stepcheck.rs (line 40)
32fn main() {
33 let a: Vec<String> = std::env::args().collect();
34 let model = Arc::new(cortiq_core::CmfModel::open(&a[1]).unwrap());
35 let (prompt, hh, ww) = (&a[2], a[3].parse::<usize>().unwrap(), a[4].parse::<usize>().unwrap());
36 let (steps, shift, i) = (a[5].parse::<usize>().unwrap(), a[6].parse::<f32>().unwrap(), a[7].parse::<usize>().unwrap());
37 let tok = Tokenizer::from_bytes(model.vocab.as_deref().unwrap()).unwrap();
38 let ids = cortiq_engine::zimagegen::prompt_ids(&tok, prompt, 512);
39 let cap = {
40 let _p = cortiq_engine::gpu::pause_gpu();
41 cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&ids)
42 };
43 let dit = ZImageDit::from_cmf(&model).unwrap();
44 let sig = zimage::sigmas_torch_f32(steps, shift);
45 let t = zimage::t_model(sig[i]);
46 let mods = dit.mods_for_steps(&[t]);
47 let fs = dit.final_scale_for_steps(&[t]);
48 let shape = ZShape::new(hh, ww, ids.len());
49 let prep = dit.prepare(&cap, shape, 1, None).unwrap();
50 println!("device prepared: {}", prep.device);
51 let (c, lh, lw) = (dit.cfg.in_channels, hh / 8, ww / 8);
52 let lat = |dir: &str| -> Vec<f32> {
53 if i == 0 {
54 read_f32(&std::env::var("CMF_INIT_LATENT").unwrap())
55 } else {
56 read_f32(&format!("{dir}/lat_{i}.f32"))
57 }
58 };
59 let xa = lat(&a[8]);
60 // `ZC_NEG=<negative prompt>`: the CFG pair as one batch-2 device
61 // forward against the two items stepped one by one.
62 if let Ok(neg) = std::env::var("ZC_NEG") {
63 let nids = cortiq_engine::zimagegen::prompt_ids(&tok, &neg, 512);
64 let ncap = {
65 let _p = cortiq_engine::gpu::pause_gpu();
66 cortiq_engine::qwen3te::Qwen3Encoder::from_cmf(&model).unwrap().encode(&nids)
67 };
68 let mut np = dit.prepare(&ncap, ZShape::new(hh, ww, nids.len()), 2, None).unwrap();
69 let tok_a = dit.tokens(&xa, &shape);
70 let vp = dit.step(&prep, i, &tok_a, &mods, &fs);
71 let vn = dit.step(&np, i, &tok_a, &mods, &fs);
72 np.device = false;
73 let vn_cpu = dit.step(&np, i, &tok_a, &mods, &fs);
74 let ok = dit.attach_device_pair(&prep, &np, 3, None);
75 println!("pair prepared: {ok} (neg L = {})", nids.len());
76 if let Some((pp, pn)) = dit.step_pair_device(3, shape.n_img, i, &tok_a, &mods, &fs) {
77 let nan = pp.iter().chain(&pn).filter(|v| !v.is_finite()).count();
78 println!("pair vs singles: pos {:.3e} neg {:.3e} non-finite {nan} single neg dev vs cpu {:.3e}",
79 rel(&pp, &vp), rel(&pn, &vn), rel(&vn, &vn_cpu));
80 }
81 return;
82 }
83 let tok_a = dit.tokens(&xa, &shape);
84 let va_dev = zimage::unpatchify(&dit.step(&prep, i, &tok_a, &mods, &fs), c, lh, lw);
85 let mut hp = zimage::ZPrepared { device: false, ..prep };
86 let va_cpu = zimage::unpatchify(&dit.step(&hp, i, &tok_a, &mods, &fs), c, lh, lw);
87 let va_tr = read_f32(&format!("{}/v_{i}.f32", a[8]));
88 println!("step {i} on A's lat: dev vs cpu {:.3e} cpu vs A's v {:.3e} dev vs A's v {:.3e}",
89 rel(&va_dev, &va_cpu), rel(&va_cpu, &va_tr), rel(&va_dev, &va_tr));
90 if let Some(b) = a.get(9) {
91 hp.device = true;
92 let xb = lat(b);
93 let vb_dev = zimage::unpatchify(&dit.step(&hp, i, &dit.tokens(&xb, &shape), &mods, &fs), c, lh, lw);
94 println!("inputs A vs B {:.3e} device outputs {:.3e} (B's own v {:.3e})",
95 rel(&xb, &xa), rel(&vb_dev, &va_dev), rel(&read_f32(&format!("{b}/v_{i}.f32")), &va_dev));
96 }
97}