From f625c74c7a87513bcc0cf9c981eef940b79220d6 Mon Sep 17 00:00:00 2001 From: Andrew Frahn Date: Wed, 14 Oct 2026 19:43:44 +1100 Subject: [PATCH] Faces: whole networks on the GPU, batched faces, CPU-only switch, paced GPU use Activations stay on the device between layers (net/program.rs lowers the graph; every operator is a WGSL kernel, one submission and read-back per run); faces of a photo go in batches of up to 16; Settings > Faces shows each model's device and a CPU-only switch (faces.gpu); the scan's pace limits GPU use; the GPU check is an activity row. YuNet 58 -> 21 ms, SFace 40 -> 9.6 ms, AuraFace 295 -> 71 ms (heavily loaded machine). Sources: own design. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/engine/src/cmd/face_detect.rs | 140 ++++- crates/engine/src/cmd/face_models.rs | 12 + crates/engine/src/cmd/face_recognize.rs | 5 +- crates/engine/src/faces_worker.rs | 28 +- crates/engine/src/tests_face_models.rs | 14 + crates/faces/src/net.rs | 241 ++++++++- crates/faces/src/net/kernels.rs | 8 +- crates/faces/src/net/program.rs | 315 +++++++++++ crates/faces/src/net/shape.rs | 10 +- crates/faces/src/runtime.rs | 76 ++- crates/faces/src/synthetic.rs | 60 +++ crates/gpu/src/faces.rs | 671 +++++++++++++++++++----- crates/gpu/src/wgsl/faces_ops.wgsl | 188 +++++++ crates/ui-egui/locales/de.json | 10 +- crates/ui-egui/locales/es.json | 10 +- crates/ui-egui/locales/fr.json | 10 +- crates/ui-egui/locales/ja.json | 10 +- crates/ui-egui/locales/pt-br.json | 10 +- crates/ui-egui/locales/ru.json | 10 +- crates/ui-egui/locales/uk.json | 10 +- crates/ui-egui/locales/zh-hans.json | 10 +- crates/ui-egui/locales/zh-hant.json | 10 +- crates/ui-egui/src/panels/faces.rs | 37 ++ docs/faces.md | 49 +- 24 files changed, 1753 insertions(+), 191 deletions(-) create mode 100644 crates/faces/src/net/program.rs create mode 100644 crates/gpu/src/wgsl/faces_ops.wgsl diff --git a/crates/engine/src/cmd/face_detect.rs b/crates/engine/src/cmd/face_detect.rs index 8e706c367..718420ffc 100644 --- a/crates/engine/src/cmd/face_detect.rs +++ b/crates/engine/src/cmd/face_detect.rs @@ -17,7 +17,7 @@ use lightcraft_faces::yunet::{Detector, Face, Options}; use lightcraft_meta::{Region, RegionKind}; use serde_json::{Value, json}; -use super::{CommandSpec, bad, cmd, has_selection}; +use super::{CommandSpec, always, bad, cmd, has_selection}; use crate::{Result, Session}; const C: &str = "faces.detect"; @@ -34,10 +34,12 @@ const EXISTING_IOU: f64 = 0.3; /// What to tell someone who has not downloaded the detector yet. const NOT_INSTALLED: &str = "finding faces needs the YuNet face detector (232 KB, MIT): download it in Settings > Faces"; +/// The YuNet loaded last, and from where. +static LOADED: Mutex)>> = Mutex::new(None); + /// The installed YuNet, loaded once and kept (a new one is loaded when the model folder changes). /// (also used by the background scan, which aligns faces by the detector's landmarks and finds faces in photos that have none) pub(crate) fn detector(s: &Session) -> std::result::Result, String> { - static LOADED: Mutex)>> = Mutex::new(None); let home = s.face_models_dir.as_deref().ok_or("this build has nowhere to keep face models (the desktop app does)")?.join(YUNET_ID); let path = home.join("model.onnx"); // installing writes the manifest last, so a folder without one is not (yet) a model @@ -59,35 +61,130 @@ pub(crate) fn detector(s: &Session) -> std::result::Result, String #[cfg(not(target_arch = "wasm32"))] { let net = d.clone(); - offer_gpu("detection", move |gpu| net.accelerate(gpu, true)); + offer_gpu(s, "detection", move |gpu| net.accelerate(gpu, true)); } *loaded = Some((path, d.clone())); Ok(d) } +/// Where a face model runs, for Settings ▸ Faces and `faces.gpu`. +#[derive(Clone, Debug, Default, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct GpuState { + /// The background check is still running. + pub checking: bool, + /// The GPU it runs on; `None` when it runs on the CPU. + pub device: Option, + /// Why it runs on the CPU. + pub why: Option, +} + +/// Each model's [`GpuState`], by what it does ("detection", "recognition"). +static GPU_STATE: Mutex> = Mutex::new(Vec::new()); + +pub(crate) fn gpu_state(what: &str) -> Option { + GPU_STATE.lock().unwrap_or_else(PoisonError::into_inner).iter().find(|(w, _)| *w == what).map(|(_, st)| st.clone()) +} + +fn set_gpu_state(what: &'static str, st: GpuState) { + let mut all = GPU_STATE.lock().unwrap_or_else(PoisonError::into_inner); + match all.iter_mut().find(|(w, _)| *w == what) { + Some((_, old)) => *old = st, + None => all.push((what, st)), + } +} + +/// The user chose the CPU for face models (Settings ▸ Faces). +pub(crate) fn cpu_only(s: &Session) -> bool { + s.face_models_dir.as_deref().is_some_and(|d| super::face_models::read_settings(d).cpu_only) +} + /// Check in the background whether the GPU runs a face model faster than the CPU and gives its results; the model -/// runs on the CPU meanwhile, and stays there when the GPU is off, missing, wrong or slower. +/// runs on the CPU meanwhile, and stays there when the GPU is off, missing, wrong or slower, or the user chose the CPU. +/// The check shows in the activity stack while it runs. #[cfg(not(target_arch = "wasm32"))] pub(crate) fn offer_gpu( + s: &Session, what: &'static str, attach: impl FnOnce(Arc) -> std::result::Result<(), String> + Send + 'static, ) { + if cpu_only(s) { + set_gpu_state(what, GpuState { why: Some("the CPU is chosen in Settings".into()), ..GpuState::default() }); + return; + } + let task = s.activity.start("facesGpu", "Checking the GPU for faces", crate::activity::Cancel::No); + task.detail(what); + set_gpu_state(what, GpuState { checking: true, ..GpuState::default() }); let spawned = std::thread::Builder::new().name(format!("faces-gpu-{what}")).spawn(move || { + let _task = task; let gpu = match lightcraft_gpu::faces::accelerator() { Ok(gpu) => gpu, - Err(why) => return log::info!("faces: {what} runs on the CPU ({why})"), + Err(why) => { + log::info!("faces: {what} runs on the CPU ({why})"); + return set_gpu_state(what, GpuState { why: Some(why), ..GpuState::default() }); + } }; let device = lightcraft_faces::net::Accelerator::device(&*gpu); match attach(gpu) { - Ok(()) => log::info!("faces: {what} runs on {device}"), - Err(why) => log::info!("faces: {what} stays on the CPU: {why}"), + Ok(()) => { + log::info!("faces: {what} runs on {device}"); + set_gpu_state(what, GpuState { device: Some(device), ..GpuState::default() }); + } + Err(why) => { + log::info!("faces: {what} stays on the CPU: {why}"); + set_gpu_state(what, GpuState { why: Some(why), ..GpuState::default() }); + } } }); if let Err(e) = spawned { log::warn!("faces: could not check the GPU for {what}: {e}"); + set_gpu_state(what, GpuState { why: Some(format!("could not check the GPU: {e}")), ..GpuState::default() }); } } +/// `faces.gpu`: where the face models run, and the switch that keeps them on the CPU. +fn gpu(s: &mut Session, p: &Value) -> Result { + const C: &str = "faces.gpu"; + if let Some(cpu) = p.get("cpuOnly") { + let cpu = cpu.as_bool().ok_or_else(|| bad(C, "`cpuOnly` is true or false"))?; + let dir = s.face_models_dir.clone().ok_or_else(|| bad(C, "this build has nowhere to keep face models (the desktop app does)"))?; + super::face_models::set_cpu_only(&dir, cpu)?; + let detector = LOADED.lock().unwrap_or_else(PoisonError::into_inner).as_ref().map(|(_, d)| d.clone()); + #[cfg(not(target_arch = "wasm32"))] + let embedder = s.faces.embedder.as_ref().map(|(_, e)| e.clone()); + #[cfg(target_arch = "wasm32")] + let embedder: Option> = None; + if cpu { + for (what, net) in [("detection", detector.as_ref().map(|d| d.net())), ("recognition", embedder.as_ref().map(|e| e.net()))] { + if let Some(net) = net { + net.detach(); + } + set_gpu_state(what, GpuState { why: Some("the CPU is chosen in Settings".into()), ..GpuState::default() }); + } + } else { + #[cfg(not(target_arch = "wasm32"))] + { + if let Some(d) = detector { + offer_gpu(s, "detection", move |gpu| d.accelerate(gpu, true)); + } + if let Some(e) = embedder { + offer_gpu(s, "recognition", move |gpu| e.accelerate(gpu, true).map_err(|e| e.to_string())); + } + } + } + } + Ok(gpu_json(s)) +} + +/// What `faces.gpu` and `faces.models.list` say about the GPU. +pub(crate) fn gpu_json(s: &Session) -> Value { + json!({ + "cpuOnly": cpu_only(s), + "detection": gpu_state("detection"), + "recognition": gpu_state("recognition"), + }) +} + fn face_json(f: &Face) -> Value { json!({ "rect": {"x0": f.x0, "y0": f.y0, "x1": f.x1, "y1": f.y1}, @@ -175,15 +272,26 @@ fn detect(s: &mut Session, p: &Value) -> Result { } pub fn specs() -> Vec { - vec![cmd!( - "faces.detect", - "Detect Faces", - ["Photo"], - None, - "{ids?, apply?: bool (default true), score?: 0.6, nmsIou?: 0.3, maxFaces?} → {photos: [{id, faces: [{rect, score, landmarks}]}]}. Finds faces with the YuNet detector (download it first: `faces.models.download {id: \"yunet-2023mar\"}`) and, unless `apply` is false, adds them as unnamed face regions (one undo step), replacing earlier detections but never regions from XMP or ones you named; the sidecar is not touched", - has_selection, - detect - )] + vec![ + cmd!( + "faces.detect", + "Detect Faces", + ["Photo"], + None, + "{ids?, apply?: bool (default true), score?: 0.6, nmsIou?: 0.3, maxFaces?} → {photos: [{id, faces: [{rect, score, landmarks}]}]}. Finds faces with the YuNet detector (download it first: `faces.models.download {id: \"yunet-2023mar\"}`) and, unless `apply` is false, adds them as unnamed face regions (one undo step), replacing earlier detections but never regions from XMP or ones you named; the sidecar is not touched", + has_selection, + detect + ), + cmd!( + query "faces.gpu", + "Face Models on the GPU", + [], + None, + "{cpuOnly?: bool} → {cpuOnly, detection, recognition}: where the face models run, each as {checking, device, why} (null until it is loaded). The GPU is used only after a background check shows it gives the CPU's results faster; `cpuOnly: true` keeps every face model on the CPU (Settings ▸ Faces), `false` checks the GPU again", + always, + gpu + ), + ] } #[cfg(test)] diff --git a/crates/engine/src/cmd/face_models.rs b/crates/engine/src/cmd/face_models.rs index 6033c4dcc..1ea39d04c 100644 --- a/crates/engine/src/cmd/face_models.rs +++ b/crates/engine/src/cmd/face_models.rs @@ -27,6 +27,9 @@ pub(super) struct FaceSettings { pub(super) enabled: bool, /// The installed recogniser in use. pub(super) embedder: Option, + /// Keep every face model on the CPU, even when the GPU would be faster. + #[serde(skip_serializing_if = "std::ops::Not::not")] + pub(super) cpu_only: bool, } fn fail(what: &str, e: impl std::fmt::Display) -> EngineError { @@ -57,6 +60,13 @@ fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { }) } +/// Keep the face models on the CPU (or let them use the GPU again). +pub(super) fn set_cpu_only(dir: &Path, on: bool) -> Result<()> { + let mut st = read_settings(dir); + st.cpu_only = on; + write_settings(dir, &st) +} + fn write_settings(dir: &Path, st: &FaceSettings) -> Result<()> { std::fs::create_dir_all(dir).map_err(|e| fail("could not create the face models folder", e))?; write_atomic(&dir.join("settings.json"), &serde_json::to_vec_pretty(st).map_err(|e| fail("settings", e))?) @@ -195,6 +205,8 @@ fn list(s: &mut Session, _: &Value) -> Result { "dir": dir.as_ref().map(|d| d.display().to_string()), "enabled": settings.enabled, "embedder": settings.embedder, + // where the models run: {cpuOnly, detection, recognition} as `faces.gpu` gives it + "gpu": super::face_detect::gpu_json(s), // Native builds always have the pure-Rust runner; model files/background scanning remain native-only. "runtime": cfg!(not(target_arch = "wasm32")), "models": models, diff --git a/crates/engine/src/cmd/face_recognize.rs b/crates/engine/src/cmd/face_recognize.rs index 094735c6d..509a5bc70 100644 --- a/crates/engine/src/cmd/face_recognize.rs +++ b/crates/engine/src/cmd/face_recognize.rs @@ -235,7 +235,7 @@ mod imp { } let embedder = Arc::new(Embedder::load(&dir.join(&id).join("model.onnx"), &installed.manifest).map_err(|e| bad(cmd, e.to_string()))?); let net = embedder.clone(); - super::super::face_detect::offer_gpu("recognition", move |gpu| net.accelerate(gpu, true).map_err(|e| e.to_string())); + super::super::face_detect::offer_gpu(s, "recognition", move |gpu| net.accelerate(gpu, true).map_err(|e| e.to_string())); let dim = match installed.manifest.output { lightcraft_faces::OutputSpec::Embedding { dim } => dim as usize, _ => return Err(bad(cmd, "the chosen model is not a recognition model")), @@ -584,6 +584,9 @@ mod imp { if let Some(w) = &s.faces.worker { w.allow(target, pace); } + // the GPU is held back the same way: a light pace leaves it mostly free for the window's own renders + let (runs, share) = crate::faces_worker::gpu_pace(pace); + lightcraft_gpu::faces::set_pace(runs, share); while target > 0 && s.faces.in_flight.len() < target + PREFETCH { let Some(id) = s.faces.queue.pop() else { break }; let Some(job) = prepare(s, id, true) else { continue }; diff --git a/crates/engine/src/faces_worker.rs b/crates/engine/src/faces_worker.rs index 81ff6adae..e44ad77c9 100644 --- a/crates/engine/src/faces_worker.rs +++ b/crates/engine/src/faces_worker.rs @@ -127,6 +127,17 @@ fn pool_threads_for(cores: usize, pace: Pace, cap: Option) -> usize { cap.map_or(n, |c| n.min(c.max(1))) } +/// How much of the GPU face models may take at `pace`: face model runs at once, and the share of the time it may be +/// busy. Paused or light: one run at a time, busy a quarter of the time (the user is working, and the window renders +/// on the same GPU); normal: two at once, busy most of the time; full: as fast as it goes. +pub(crate) fn gpu_pace(pace: Pace) -> (usize, f32) { + match pace { + Pace::Pause | Pace::Light => (1, 0.25), + Pace::Normal => (2, 0.75), + Pace::Full => (4, 1.0), + } +} + /// One photo's faces to embed (and, if `detect`, to find first). pub(crate) struct Prepared { /// `@` the embeddings are for. @@ -197,7 +208,8 @@ pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done { let img = Rgb { data: &rgb, width: w, height: h }; let (iw, ih) = embedder.input_size(); let scaled = |f: &Face| f.landmarks.map(|(x, y)| (x * w as f32, y * h as f32)); - let mut results: Vec<(Rect, Option>, Rect)> = todo + // every face's crop first, then all of them through the model at once (one GPU dispatch for a group photo) + let mut faces: Vec<(Rect, Option>, Rect)> = todo .into_iter() .map(|rect| { let best = found.iter().map(|f| (overlap(&rect, f), f)).filter(|(o, _)| *o >= MATCH_IOU).max_by(|a, b| a.0.total_cmp(&b.0)); @@ -214,20 +226,18 @@ pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done { ih, ), }; - let v = crop.and_then(|c| embedder.embed(&c).ok()); // shown by the detector's box where it found this face: every face then has the same tightness let view = best.and_then(|(_, f)| face_rect(f)).unwrap_or(rect); - (rect, v, view) + (rect, crop, view) }) .collect(); // a photo with no face regions: what the detector is sure of becomes its faces, embedded in the same pass let new_faces = if detect { - detected.map(|faces| { + detected.map(|found| { let mut rects = Vec::new(); - for f in faces.iter().filter(|f| f.score >= REGION_SCORE) { + for f in found.iter().filter(|f| f.score >= REGION_SCORE) { let Some(rect) = face_rect(f) else { continue }; - let v = align_to_template(&img, &scaled(f), iw, ih).and_then(|c| embedder.embed(&c).ok()); - results.push((rect, v, rect)); + faces.push((rect, align_to_template(&img, &scaled(f), iw, ih), rect)); rects.push(rect); } rects @@ -235,6 +245,10 @@ pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done { } else { None }; + let crops: Vec<&[u8]> = faces.iter().filter_map(|(_, c, _)| c.as_deref()).collect(); + let mut embedded = embedder.embed_batch(&crops).into_iter(); + let results: Vec<(Rect, Option>, Rect)> = + faces.into_iter().map(|(rect, crop, view)| (rect, crop.and_then(|_| embedded.next().and_then(Result::ok)), view)).collect(); (results, new_faces) })); let (results, found) = run.unwrap_or_else(|_| (fallback(rects), None)); diff --git a/crates/engine/src/tests_face_models.rs b/crates/engine/src/tests_face_models.rs index f1c27c849..00d6dc0cc 100644 --- a/crates/engine/src/tests_face_models.rs +++ b/crates/engine/src/tests_face_models.rs @@ -613,3 +613,17 @@ fn recognition_cannot_be_turned_on_without_a_chosen_model() { assert!(s.execute("faces.enable", &json!({"enabled": true})).is_err()); let _ = std::fs::remove_dir_all(&d); } + +#[test] +fn the_cpu_only_switch_is_kept_and_bad_values_are_refused() { + let d = temp("cpu-only"); + let mut s = session(&d); + assert_eq!(s.execute("faces.gpu", &json!({})).unwrap()["cpuOnly"], false); + assert_eq!(s.execute("faces.gpu", &json!({"cpuOnly": true})).unwrap()["cpuOnly"], true); + // kept with the models, so a new session starts on the CPU too + let mut again = session(&d); + assert_eq!(again.execute("faces.models.list", &json!({})).unwrap()["gpu"]["cpuOnly"], true); + assert!(again.execute("faces.gpu", &json!({"cpuOnly": "yes"})).is_err()); + assert_eq!(again.execute("faces.gpu", &json!({"cpuOnly": false})).unwrap()["cpuOnly"], false); + let _ = std::fs::remove_dir_all(d); +} diff --git a/crates/faces/src/net.rs b/crates/faces/src/net.rs index 3983404c5..430335840 100644 --- a/crates/faces/src/net.rs +++ b/crates/faces/src/net.rs @@ -12,6 +12,7 @@ use std::rc::Rc; use std::sync::{Arc, PoisonError, RwLock}; mod kernels; +pub mod program; mod shape; use crate::graph::{Data, Graph, Node}; @@ -111,6 +112,12 @@ pub trait Accelerator: Send + Sync { fn device(&self) -> String; /// `[cout][oh][ow]`, or `Ok(None)` to leave this one to the CPU (too small to be worth a dispatch, say). fn conv(&self, job: &ConvJob) -> std::result::Result>, String>; + /// The whole network on the device, keeping every value there between steps, for each of `inputs` (all of + /// `program.lens[INPUT]` values) in one go: per input, each of `program.outputs` in order. `Ok(None)` leaves the + /// run to the CPU path (which may still use [`Accelerator::conv`]). + fn run_program(&self, _program: &program::Program, _inputs: &[&[f32]]) -> std::result::Result>>>, String> { + Ok(None) + } } /// A loaded network, ready to run. @@ -190,7 +197,7 @@ impl Net { /// A convolution on the accelerator, if one is attached and takes it. A failing accelerator is dropped. fn accelerated_conv(&self, layer: usize, node: &Node, x: &Tensor) -> Result> { let Some(accelerator) = self.accelerator() else { return Ok(None) }; - let job = kernels::conv_job(layer, node, x, &self.graph)?; + let job = kernels::conv_job(layer, node, &x.shape, &x.data, &self.graph)?; match accelerator.conv(&job) { Ok(Some(out)) if out.len() == job.cout * job.oh * job.ow => Ok(Some(Tensor::new(vec![1, job.cout, job.oh, job.ow], out)?)), Ok(None) => Ok(None), @@ -214,7 +221,67 @@ impl Net { /// Run the network on `input` (its single input); returns each output by name. pub fn run(&self, input: Tensor) -> Result> { - shape::tensor(&input.shape, input.data.len(), &input.data)?; + self.run_batch(vec![input])?.pop().ok_or_else(|| NetError::Shape("the run gave nothing".into())) + } + + /// Run the network on each of `inputs`; an accelerator that keeps the values on its device runs inputs of one shape + /// together (one upload and one read-back for all of them). + pub fn run_batch(&self, inputs: Vec) -> Result>> { + for input in &inputs { + shape::tensor(&input.shape, input.data.len(), &input.data)?; + } + if let Some(out) = self.run_on_device(&inputs)? { + return Ok(out); + } + inputs.into_iter().map(|input| self.run_one(input)).collect() + } + + /// The whole batch on the attached accelerator, or `None` to run it on the CPU path. A failing or misbehaving + /// accelerator is dropped. + fn run_on_device(&self, inputs: &[Tensor]) -> Result>>> { + let (Some(accelerator), Some(first)) = (self.accelerator(), inputs.first()) else { return Ok(None) }; + if inputs.iter().any(|t| t.shape != first.shape) { + return Ok(None); + } + let Some(program) = self.program(&first.shape)? else { return Ok(None) }; + let slices: Vec<&[f32]> = inputs.iter().map(|t| t.data.as_slice()).collect(); + let started = web_time::Instant::now(); + let outs = match accelerator.run_program(&program, &slices) { + Ok(Some(outs)) => outs, + Ok(None) => return Ok(None), + Err(_) => { + self.detach(); + return Ok(None); + } + }; + let gathered: Option>> = (outs.len() == inputs.len()) + .then(|| { + outs.into_iter() + .map(|values| { + (values.len() == program.outputs.len()) + .then(|| { + values + .into_iter() + .zip(&program.outputs) + .map(|(data, (name, _, shape))| Tensor::new(shape.clone(), data).ok().map(|t| (name.clone(), t))) + .collect::>>() + }) + .flatten() + }) + .collect() + }) + .flatten(); + if std::env::var_os("LIGHTCRAFT_PROFILE").is_some() { + eprintln!("[faces-gpu] {} steps x {}: {:.3} ms", program.steps.len(), inputs.len(), started.elapsed().as_secs_f64() * 1000.0); + } + if gathered.is_none() { + // wrong sizes or values that are not finite: the device is not trusted any more + self.detach(); + } + Ok(gathered) + } + + fn run_one(&self, input: Tensor) -> Result> { self.validate_input(&input.shape)?; let mut scratch = kernels::Scratch::default(); let profile = std::env::var_os("LIGHTCRAFT_PROFILE").is_some(); @@ -745,6 +812,176 @@ mod tests { } } + /// An accelerator that runs whole programs with plain loops (the definition of each step), counting its calls; + /// `poison` makes it give a value that is not finite. + struct ProgramFake { + calls: std::sync::atomic::AtomicUsize, + poison: bool, + } + + fn run_steps(p: &program::Program, input: &[f32]) -> Vec> { + use program::{Arg, Binary, Op, Unary}; + let mut values: Vec> = vec![Vec::new(); p.lens.len()]; + values[program::INPUT] = input.to_vec(); + for step in &p.steps { + let arg = |i: usize| -> Vec { + match &step.args[i] { + Arg::Value(v) => values[*v].clone(), + Arg::Const { data, .. } => data.to_vec(), + } + }; + let len = p.lens[step.out]; + let x = arg(0); + let out: Vec = match &step.op { + Op::Conv(j) => { + let (m, cg) = (j.cout / j.groups, j.cin / j.groups); + let mut out = vec![0f32; len]; + for co in 0..j.cout { + for oy in 0..j.oh { + for ox in 0..j.ow { + let mut v = j.bias.map_or(0.0, |b| b[co]); + for ci in 0..cg { + for ky in 0..j.kh { + for kx in 0..j.kw { + let (iy, ix) = ((oy * j.sh + ky) as isize - j.pt as isize, (ox * j.sw + kx) as isize - j.pl as isize); + if iy >= 0 && ix >= 0 && (iy as usize) < j.h && (ix as usize) < j.w { + v += j.weights[((co * cg + ci) * j.kh + ky) * j.kw + kx] + * x[((co / m * cg + ci) * j.h + iy as usize) * j.w + ix as usize]; + } + } + } + } + out[(co * j.oh + oy) * j.ow + ox] = v; + } + } + } + out + } + Op::Unary(Unary::Relu) => x.iter().map(|v| v.max(0.0)).collect(), + Op::Unary(Unary::Sigmoid) => x.iter().map(|v| 1.0 / (1.0 + (-v).exp())).collect(), + Op::ChannelAffine { channels, plane } => { + let (s, o) = (arg(1), arg(2)); + x.iter().enumerate().map(|(i, v)| v * s[(i / plane) % channels] + o[(i / plane) % channels]).collect() + } + Op::ChannelPrelu { channels, plane } => { + let a = arg(1); + x.iter().enumerate().map(|(i, v)| if *v < 0.0 { v * a[(i / plane) % channels] } else { *v }).collect() + } + Op::Binary { op, dims, a, b } => { + let y = arg(1); + (0..len) + .map(|i| { + let (mut rest, mut ia, mut ib) = (i, 0, 0); + for d in (0..4).rev() { + let at = rest % dims[d]; + rest /= dims[d]; + ia += at * a[d]; + ib += at * b[d]; + } + match op { + Binary::Add => x[ia] + y[ib], + Binary::Sub => x[ia] - y[ib], + Binary::Mul => x[ia] * y[ib], + } + }) + .collect() + } + Op::Transpose { dims, strides } => (0..len) + .map(|i| { + let (mut rest, mut from) = (i, 0); + for d in (0..4).rev() { + from += (rest % dims[d]) * strides[d]; + rest /= dims[d]; + } + x[from] + }) + .collect(), + Op::MaxPool { h, w, kh, kw, sh, sw, oh, ow, .. } => (0..len) + .map(|i| { + let (ox, oy, c) = (i % ow, (i / ow) % oh, i / (ow * oh)); + let mut m = f32::NEG_INFINITY; + for ky in 0..*kh { + for kx in 0..*kw { + m = m.max(x[(c * h + oy * sh + ky) * w + ox * sw + kx]); + } + } + m + }) + .collect(), + Op::Resize { h, w, sh, sw, .. } => (0..len) + .map(|i| { + let (ow, oh) = (w * sw, h * sh); + let (ox, oy, c) = (i % ow, (i / ow) % oh, i / (ow * oh)); + x[(c * h + oy / sh) * w + ox / sw] + }) + .collect(), + Op::GlobalAverage { plane, .. } => x.chunks(*plane).map(|c| c.iter().sum::() / *plane as f32).collect(), + Op::Gemv { k, n, trans_b, alpha } => { + let (wt, bias) = (arg(1), arg(2)); + (0..*n).map(|o| alpha * (0..*k).map(|j| x[j] * wt[if *trans_b { o * k + j } else { j * n + o }]).sum::() + bias[o]).collect() + } + }; + assert_eq!(out.len(), len, "{:?}", step.op); + values[step.out] = out; + for f in &step.free { + values[*f] = Vec::new(); + } + } + p.outputs.iter().map(|(_, v, _)| values[*v].clone()).collect() + } + + impl Accelerator for ProgramFake { + fn device(&self) -> String { + "program test".into() + } + fn conv(&self, _: &ConvJob) -> std::result::Result>, String> { + Ok(None) + } + fn run_program(&self, p: &program::Program, inputs: &[&[f32]]) -> std::result::Result>>>, String> { + self.calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let mut out: Vec>> = inputs.iter().map(|x| run_steps(p, x)).collect(); + if self.poison { + out[0][0][0] = f32::NAN; + } + Ok(Some(out)) + } + } + + fn every_operator_input(seed: usize) -> Tensor { + input(4, 8, 8, |i| (((i * 31 + seed * 7) % 37) as f32 - 18.0) * 0.05) + } + + #[test] + fn a_program_of_every_operator_gives_what_the_interpreter_gives() { + let net = Net::new(crate::synthetic::every_operator_graph()).unwrap(); + let x = every_operator_input(0); + let cpu = net.run(x.clone()).unwrap(); + let fake = Arc::new(ProgramFake { calls: 0.into(), poison: false }); + net.attach_checked(fake.clone(), &x, 1e-5, false).unwrap(); + let calls = fake.calls.load(std::sync::atomic::Ordering::Relaxed); + assert!(calls > 0, "the whole network ran as one program"); + // several inputs go to the device together, and each gets its own outputs + let inputs: Vec = (0..3).map(every_operator_input).collect(); + let together = net.run_batch(inputs.clone()).unwrap(); + assert_eq!(fake.calls.load(std::sync::atomic::Ordering::Relaxed), calls + 1, "one dispatch for the batch"); + net.detach(); + for (x, got) in inputs.into_iter().zip(&together) { + let want = net.run(x).unwrap(); + agree(&want, got, 1e-5).unwrap(); + } + assert_eq!(cpu.len(), 2); + } + + #[test] + fn a_device_that_gives_values_that_are_not_finite_is_dropped() { + let net = Net::new(crate::synthetic::every_operator_graph()).unwrap(); + let x = every_operator_input(1); + let cpu = net.run(x.clone()).unwrap(); + assert!(net.attach_checked(Arc::new(ProgramFake { calls: 0.into(), poison: true }), &x, 1e-5, false).is_err()); + assert!(net.accelerator().is_none()); + assert_eq!(net.run(x).unwrap(), cpu); + } + #[test] fn unsupported_and_inconsistent_models_are_errors() { let x = input(1, 3, 3, |i| i as f32); diff --git a/crates/faces/src/net/kernels.rs b/crates/faces/src/net/kernels.rs index 6d7c1dbb9..71ae95874 100644 --- a/crates/faces/src/net/kernels.rs +++ b/crates/faces/src/net/kernels.rs @@ -60,8 +60,8 @@ fn shifted_plane(plane: &[f32], width: usize, ky: usize, kx: usize, pt: usize, p Ok(()) } /// The checked geometry of a convolution node, for an accelerator. -pub(super) fn conv_job<'a>(layer: usize, n: &Node, x: &'a Tensor, g: &'a Graph) -> Result> { - let (cin, h, w) = chw(x)?; +pub(super) fn conv_job<'a>(layer: usize, n: &Node, shape: &[usize], input: &'a [f32], g: &'a Graph) -> Result> { + let (cin, h, w) = program::chw_of(shape)?; let wt = shape::weight(n, g, 1)?; let weights = shape::floats(n, g, 1)?; let [cout, cg, kh, kw] = wt.dims.as_slice() else { @@ -84,7 +84,7 @@ pub(super) fn conv_job<'a>(layer: usize, n: &Node, x: &'a Tensor, g: &'a Graph) return shape_err("the convolution bias does not match"); } let (oh, ow) = ((ph - kh) / sh + 1, (pw - kw) / sw + 1); - Ok(ConvJob { layer, cin, h, w, cout, groups, kh, kw, sh, sw, pt, pl, oh, ow, weights, bias, input: &x.data }) + Ok(ConvJob { layer, cin, h, w, cout, groups, kh, kw, sh, sw, pt, pl, oh, ow, weights, bias, input }) } pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Result { let (cin, h, w) = chw(x)?; @@ -208,7 +208,7 @@ pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Re Ok(out) } /// Flat offset under ONNX's trailing-dimension broadcasting; shapes are checked before execution. -fn broadcast_offset(mut at: usize, out: &[usize], input: &[usize]) -> usize { +pub(super) fn broadcast_offset(mut at: usize, out: &[usize], input: &[usize]) -> usize { let mut offset = 0; let mut stride = 1; for (axis, &size) in out.iter().enumerate().rev() { diff --git a/crates/faces/src/net/program.rs b/crates/faces/src/net/program.rs new file mode 100644 index 000000000..e00e51da1 --- /dev/null +++ b/crates/faces/src/net/program.rs @@ -0,0 +1,315 @@ +//! A network lowered to a flat list of steps over numbered values, for an [`Accelerator`] that keeps every activation +//! on its device between layers (one upload, one read-back per run instead of one per convolution). +//! +//! Shapes are resolved for one input shape; every step names the values it reads and the one it writes, and which +//! values nothing reads after it. Shape-only operators (`Reshape`, `Flatten`, `Identity`, `Dropout`) make no step: their +//! output is the same value under another name. Anything this cannot express (a broadcast past four dimensions, a +//! transposed or batched `Gemm`, a slope that is not per channel) makes the whole lowering `None`, and the net runs +//! the CPU path instead. + +use std::borrow::Cow; + +use super::*; + +/// Value 0 is the network's input. +pub const INPUT: usize = 0; + +/// An operand: a value computed on the device, or a constant (uploaded once and kept, by `id`). +#[derive(Clone, Debug)] +pub enum Arg<'a> { + Value(usize), + /// `id` is the same for every run of the same net, so the device can keep the data. + Const { + id: usize, + data: Cow<'a, [f32]>, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Unary { + Relu, + Sigmoid, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Binary { + Add, + Sub, + Mul, +} + +/// What a step computes. Element counts come from the program's `lens` of the step's output. +#[derive(Clone, Debug)] +pub enum Op<'a> { + /// A convolution; its `input` is empty (it reads the step's first argument). + Conv(ConvJob<'a>), + /// `out[i] = f(x[i])`. Args: `[x]`. + Unary(Unary), + /// `out[i] = x[i] * scale[c] + offset[c]`, `c = (i / plane) % channels`. Args: `[x, scale, offset]`. + ChannelAffine { channels: usize, plane: usize }, + /// `out[i] = x[i] < 0 ? x[i] * slope[c] : x[i]`. Args: `[x, slope]` (`channels` values). + ChannelPrelu { channels: usize, plane: usize }, + /// Broadcasting over four dimensions: `dims` of the output, element strides of each operand (0 where it repeats). + /// Args: `[a, b]`. + Binary { op: Binary, dims: [usize; 4], a: [usize; 4], b: [usize; 4] }, + /// `out[n] = alpha · Σk x[k] · B(k, n) + bias[n]`, `B(k, n) = weights[n·k_len + k]` when `trans_b`, else + /// `weights[k·n_len + n]`. Args: `[x, weights, bias]` (bias already times beta; zeros when the node has none). + Gemv { k: usize, n: usize, trans_b: bool, alpha: f32 }, + /// The mean of each `plane`-sized run. Args: `[x]`. + GlobalAverage { channels: usize, plane: usize }, + /// Unpadded max pooling. Args: `[x]`. + MaxPool { c: usize, h: usize, w: usize, kh: usize, kw: usize, sh: usize, sw: usize, oh: usize, ow: usize }, + /// Nearest-neighbour upscale by whole numbers. Args: `[x]`. + Resize { c: usize, h: usize, w: usize, sh: usize, sw: usize }, + /// `out` in its own order (`dims`), read from the input at `strides` (the input's strides, permuted). Args: `[x]`. + Transpose { dims: [usize; 4], strides: [usize; 4] }, +} + +#[derive(Clone, Debug)] +pub struct Step<'a> { + pub op: Op<'a>, + pub args: Vec>, + pub out: usize, + /// Values nothing reads after this step (never an output): their memory can be reused. + pub free: Vec, +} + +/// A whole network for one input shape. +#[derive(Clone, Debug)] +pub struct Program<'a> { + /// Elements of each value; `lens[INPUT]` is the input's. + pub lens: Vec, + pub steps: Vec>, + /// Each graph output: its name, the value that holds it and its shape. + pub outputs: Vec<(String, usize, Vec)>, +} + +/// Leading ones up to rank four, or `None` past four. +fn rank4(s: &[usize]) -> Option<[usize; 4]> { + let pad = 4usize.checked_sub(s.len())?; + let mut out = [1; 4]; + for (i, d) in s.iter().enumerate() { + *out.get_mut(pad + i)? = *d; + } + Some(out) +} + +/// Element strides of `input` read as `out` under trailing-dimension broadcasting (0 along repeated axes). +fn broadcast_strides(out: [usize; 4], input: &[usize]) -> Option<[usize; 4]> { + let dims = rank4(input)?; + let mut strides = [0; 4]; + let mut step = 1; + for axis in (0..4).rev() { + let (d, o) = (*dims.get(axis)?, *out.get(axis)?); + if d == o && d != 1 { + *strides.get_mut(axis)? = step; + } else if d != 1 { + return None; + } + step *= d; + } + Some(strides) +} + +impl Net { + /// This network lowered for `input` (its shape), or `None` when a step has no device form. + pub fn program(&self, input: &[usize]) -> Result>> { + let g = &self.graph; + let (shapes, _) = shape::infer_all(g, input)?; + let Some((input_name, _)) = g.inputs.first() else { return shape_err("the model has no input") }; + let mut lens = vec![shape::elements(input)?]; + let mut names: HashMap<&str, (usize, Vec)> = HashMap::new(); + names.insert(input_name.as_str(), (INPUT, input.to_vec())); + let mut steps: Vec = Vec::new(); + for (i, node) in g.nodes.iter().enumerate() { + let Some(out_shape) = shapes.get(i) else { return shape_err("a node has no shape") }; + let Some(out_name) = node.outputs.first() else { return shape_err("a node has no output") }; + let value = |idx: usize| -> Option<&(usize, Vec)> { node.inputs.get(idx).and_then(|n| names.get(n.as_str())) }; + let constant = |idx: usize| -> Option<(Cow<'_, [f32]>, &[usize])> { + let w = node.inputs.get(idx).and_then(|n| g.weights.get(n))?; + match &w.data { + Data::F32(v) => Some((Cow::Borrowed(v.as_slice()), w.dims.as_slice())), + _ => None, + } + }; + let id = |k: usize| i * 4 + k; + let Some((x, x_shape)) = value(0).cloned() else { return Ok(None) }; + let (op, args) = match node.op.as_str() { + "Reshape" | "Flatten" | "Identity" | "Dropout" => { + names.insert(out_name.as_str(), (x, out_shape.clone())); + continue; + } + "Conv" => (Op::Conv(kernels::conv_job(i, node, &x_shape, &[], g)?), vec![Arg::Value(x)]), + "Relu" => (Op::Unary(Unary::Relu), vec![Arg::Value(x)]), + "Sigmoid" => (Op::Unary(Unary::Sigmoid), vec![Arg::Value(x)]), + "BatchNormalization" => { + let channels = *x_shape.get(1).ok_or_else(|| NetError::Shape("BatchNormalization rank".into()))?; + let plane = shape::elements(x_shape.get(2..).unwrap_or(&[]))?; + let (s, b, m, v) = + (shape::floats(node, g, 1)?, shape::floats(node, g, 2)?, shape::floats(node, g, 3)?, shape::floats(node, g, 4)?); + let eps = shape::float(node, "epsilon", 1e-5); + let mut scale = Vec::with_capacity(channels); + let mut offset = Vec::with_capacity(channels); + for c in 0..channels { + let (Some(s), Some(b), Some(m), Some(v)) = (s.get(c), b.get(c), m.get(c), v.get(c)) else { + return shape_err("short BatchNormalization parameters"); + }; + let k = s / (v + eps).sqrt(); + scale.push(k); + offset.push(b - m * k); + } + ( + Op::ChannelAffine { channels, plane }, + vec![Arg::Value(x), Arg::Const { id: id(1), data: scale.into() }, Arg::Const { id: id(2), data: offset.into() }], + ) + } + "PRelu" => { + let Some((slope, dims)) = constant(1) else { return Ok(None) }; + let channels = x_shape.get(1).copied().unwrap_or(1); + let plane = shape::elements(x_shape.get(2..).unwrap_or(&[]))?; + let mut expanded = vec![1; x_shape.len().saturating_sub(dims.len())]; + expanded.extend_from_slice(dims); + let per_channel = expanded.iter().enumerate().all(|(i, &d)| d == 1 || i == 1) && slope.len() == channels; + let slope: Cow<[f32]> = if per_channel { + slope + } else if let [one] = slope.as_ref() { + vec![*one; channels].into() + } else { + return Ok(None); + }; + (Op::ChannelPrelu { channels, plane }, vec![Arg::Value(x), Arg::Const { id: id(1), data: slope }]) + } + "Add" | "Sub" | "Mul" => { + let op = match node.op.as_str() { + "Add" => Binary::Add, + "Sub" => Binary::Sub, + _ => Binary::Mul, + }; + let (b, b_shape) = match (value(1), constant(1)) { + (Some((v, s)), _) => (Arg::Value(*v), s.clone()), + (None, Some((data, dims))) => (Arg::Const { id: id(1), data }, dims.to_vec()), + _ => return Ok(None), + }; + let Some(dims) = rank4(out_shape) else { return Ok(None) }; + let (Some(sa), Some(sb)) = (broadcast_strides(dims, &x_shape), broadcast_strides(dims, &b_shape)) else { return Ok(None) }; + (Op::Binary { op, dims, a: sa, b: sb }, vec![Arg::Value(x), b]) + } + "Gemm" => { + let Some((weights, wdims)) = constant(1) else { return Ok(None) }; + let ([1, k], [n_out]) = (x_shape.as_slice(), out_shape.get(1..).unwrap_or(&[])) else { return Ok(None) }; + if node.int("transA").unwrap_or(0) != 0 || out_shape.first() != Some(&1) { + return Ok(None); + } + let trans_b = node.int("transB").unwrap_or(0) != 0; + if wdims.len() != 2 || weights.len() != k * n_out { + return shape_err("Gemm weights do not match"); + } + let mut bias = vec![0.0f32; *n_out]; + if node.inputs.get(2).is_some_and(|s| !s.is_empty()) { + let wt = shape::weight(node, g, 2)?; + let b = shape::floats(node, g, 2)?; + let beta = shape::float(node, "beta", 1.0); + for (j, to) in bias.iter_mut().enumerate() { + *to = beta + * b.get(kernels::broadcast_offset(j, out_shape, &wt.dims)) + .copied() + .ok_or_else(|| NetError::Shape("short matrix bias".into()))?; + } + } + ( + Op::Gemv { k: *k, n: *n_out, trans_b, alpha: shape::float(node, "alpha", 1.0) }, + vec![Arg::Value(x), Arg::Const { id: id(1), data: weights }, Arg::Const { id: id(2), data: bias.into() }], + ) + } + "GlobalAveragePool" => { + let (c, h, w) = chw_of(&x_shape)?; + (Op::GlobalAverage { channels: c, plane: h * w }, vec![Arg::Value(x)]) + } + "MaxPool" => { + let (c, h, w) = chw_of(&x_shape)?; + let (kh, kw) = pair(node, "kernel_shape", 1)?; + let (sh, sw) = pair(node, "strides", 1)?; + if pads(node)? != (0, 0, 0, 0) { + return Ok(None); + } + let (_, oh, ow) = chw_of(out_shape)?; + (Op::MaxPool { c, h, w, kh, kw, sh, sw, oh, ow }, vec![Arg::Value(x)]) + } + "Resize" => { + let (c, h, w) = chw_of(&x_shape)?; + let (_, oh, ow) = chw_of(out_shape)?; + if oh % h != 0 || ow % w != 0 { + return shape_err("a Resize that is not by whole numbers"); + } + (Op::Resize { c, h, w, sh: oh / h, sw: ow / w }, vec![Arg::Value(x)]) + } + "Transpose" => { + let rank = x_shape.len(); + let default_perm: Vec = (0..rank as i64).rev().collect(); + let perm = node.ints("perm").unwrap_or(&default_perm); + if rank > 4 || perm.len() != rank { + return Ok(None); + } + let mut in_strides = vec![1usize; rank]; + for a in (0..rank.saturating_sub(1)).rev() { + let next = in_strides.get(a + 1).copied().unwrap_or(1) * x_shape.get(a + 1).copied().unwrap_or(1); + if let Some(s) = in_strides.get_mut(a) { + *s = next; + } + } + let Some(dims) = rank4(out_shape) else { return Ok(None) }; + let mut strides = [0usize; 4]; + for (j, p) in perm.iter().enumerate() { + let p = usize::try_from(*p).map_err(|_| NetError::Shape("a negative permutation".into()))?; + let s = in_strides.get(p).copied().ok_or_else(|| NetError::Shape("a bad permutation".into()))?; + *strides.get_mut(4 - rank + j).ok_or_else(|| NetError::Shape("a bad permutation".into()))? = s; + } + (Op::Transpose { dims, strides }, vec![Arg::Value(x)]) + } + _ => return Ok(None), + }; + let out = lens.len(); + lens.push(shape::elements(out_shape)?); + names.insert(out_name.as_str(), (out, out_shape.clone())); + steps.push(Step { op, args, out, free: Vec::new() }); + } + let outputs = g + .outputs + .iter() + .map(|name| { + let (v, s) = names.get(name.as_str()).ok_or_else(|| NetError::Shape(format!("output `{name}` was never computed")))?; + Ok((name.clone(), *v, s.clone())) + }) + .collect::>>()?; + // the last step that reads each value; a value nothing reads is free right after it is made + let mut last: Vec> = vec![None; lens.len()]; + for (at, step) in steps.iter().enumerate() { + if let Some(l) = last.get_mut(step.out) { + *l = Some(at); + } + for a in &step.args { + if let Arg::Value(v) = a + && let Some(l) = last.get_mut(*v) + { + *l = Some(at); + } + } + } + for (v, at) in last.iter().enumerate() { + if outputs.iter().any(|(_, o, _)| *o == v) { + continue; + } + if let Some(step) = at.and_then(|at| steps.get_mut(at)) { + step.free.push(v); + } + } + Ok(Some(Program { lens, steps, outputs })) + } +} + +pub(super) fn chw_of(s: &[usize]) -> Result<(usize, usize, usize)> { + match s { + [1, c, h, w] => Ok((*c, *h, *w)), + s => unsupported(format!("a tensor of shape {s:?} (only [1, channels, height, width])")), + } +} diff --git a/crates/faces/src/net/shape.rs b/crates/faces/src/net/shape.rs index d76271bc7..5b9fc68b1 100644 --- a/crates/faces/src/net/shape.rs +++ b/crates/faces/src/net/shape.rs @@ -261,7 +261,12 @@ fn image(s: &[usize]) -> Result<(usize, usize, usize)> { } } pub(super) fn infer(g: &Graph, input: &[usize]) -> Result>> { + Ok(infer_all(g, input)?.1) +} +/// Each node's output shape, and the graph's outputs' shapes. +pub(super) fn infer_all(g: &Graph, input: &[usize]) -> Result<(Vec>, Vec>)> { elements(input)?; + let mut per_node = Vec::with_capacity(g.nodes.len()); let mut shapes: HashMap<&str, Vec> = g.weights.iter().map(|(n, w)| (n.as_str(), w.dims.clone())).collect(); shapes.insert(g.inputs.first().ok_or_else(|| NetError::Shape("no input".into()))?.0.as_str(), input.to_vec()); let mut live = elements(input)?; @@ -484,6 +489,7 @@ pub(super) fn infer(g: &Graph, input: &[usize]) -> Result>> { .checked_add(out_count) .filter(|&n| n <= 128 * 1024 * 1024) .ok_or_else(|| NetError::Shape("live tensors exceed the 512 MiB working limit".into()))?; + per_node.push(out.clone()); shapes.insert(n.outputs.first().ok_or_else(|| NetError::Shape("no output".into()))?, out); let mut removed = std::collections::HashSet::new(); for name in &n.inputs { @@ -492,5 +498,7 @@ pub(super) fn infer(g: &Graph, input: &[usize]) -> Result>> { } } } - g.outputs.iter().map(|s| shapes.get(s.as_str()).cloned().ok_or_else(|| NetError::Shape("missing output".into()))).collect() + let outputs = + g.outputs.iter().map(|s| shapes.get(s.as_str()).cloned().ok_or_else(|| NetError::Shape("missing output".into()))).collect::>()?; + Ok((per_node, outputs)) } diff --git a/crates/faces/src/runtime.rs b/crates/faces/src/runtime.rs index 7a8ab948a..92f9e8134 100644 --- a/crates/faces/src/runtime.rs +++ b/crates/faces/src/runtime.rs @@ -40,6 +40,9 @@ pub struct SelfTest { type Plan = Net; +/// Most faces [`Embedder::embed_batch`] sends to a GPU in one go. +pub const BATCH: usize = 16; + /// A loaded recognition model. #[derive(Clone)] pub struct Embedder { @@ -135,6 +138,66 @@ impl Embedder { /// The face as a unit-length vector. `rgb` is an aligned face at exactly [`Self::input_size`], 8-bit RGB. pub fn embed(&self, rgb: &[u8]) -> Result, RuntimeError> { + let input = self.input(rgb)?; + let out = guarded("running it", || { + let mut result = self.plan.run(input).map_err(|e| RuntimeError::Run(e.to_string()))?; + if result.len() != 1 { + return Err(RuntimeError::Run("the model must give one embedding".into())); + } + Ok(result.pop().ok_or_else(|| RuntimeError::Run("no output".into()))?.1.data) + })?; + self.unit(out) + } + + /// Several faces at once, each as [`Self::embed`] gives it. On a GPU they run together, up to [`BATCH`] per + /// dispatch (one upload and one read-back for all of them); on the CPU one after another. + pub fn embed_batch(&self, faces: &[&[u8]]) -> Vec, RuntimeError>> { + if self.plan.accelerator().is_none() || faces.len() < 2 { + return faces.iter().map(|f| self.embed(f)).collect(); + } + let mut results: Vec, RuntimeError>>> = faces.iter().map(|_| None).collect(); + let mut ready: Vec<(usize, Tensor)> = Vec::new(); + for (i, face) in faces.iter().enumerate() { + match self.input(face) { + Ok(t) => ready.push((i, t)), + Err(e) => { + if let Some(r) = results.get_mut(i) { + *r = Some(Err(e)); + } + } + } + } + while !ready.is_empty() { + let chunk: Vec<(usize, Tensor)> = ready.drain(..ready.len().min(BATCH)).collect(); + let (at, inputs): (Vec, Vec) = chunk.into_iter().unzip(); + let ran = guarded("running it", || self.plan.run_batch(inputs).map_err(|e| RuntimeError::Run(e.to_string()))); + match ran { + Ok(outs) if outs.len() == at.len() => { + for (i, mut out) in at.into_iter().zip(outs) { + let r = match (out.len(), out.pop()) { + (1, Some((_, t))) => self.unit(t.data), + _ => Err(RuntimeError::Run("the model must give one embedding".into())), + }; + if let Some(slot) = results.get_mut(i) { + *slot = Some(r); + } + } + } + // one bad face spoils a batch: run them one by one so the others still get theirs + _ => { + for i in at { + if let (Some(slot), Some(face)) = (results.get_mut(i), faces.get(i)) { + *slot = Some(self.embed(face)); + } + } + } + } + } + results.into_iter().map(|r| r.unwrap_or_else(|| Err(RuntimeError::Run("the face was not embedded".into())))).collect() + } + + /// The model's input for an aligned face. + fn input(&self, rgb: &[u8]) -> Result { if rgb.len() != self.width * self.height * 3 { return Err(RuntimeError::Input("the face picture is not the size the model wants".into())); } @@ -149,14 +212,11 @@ impl Embedder { } } } - let out = guarded("running it", || { - let input = Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))?; - let mut result = self.plan.run(input).map_err(|e| RuntimeError::Run(e.to_string()))?; - if result.len() != 1 { - return Err(RuntimeError::Run("the model must give one embedding".into())); - } - Ok(result.pop().ok_or_else(|| RuntimeError::Run("no output".into()))?.1.data) - })?; + Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string())) + } + + /// The model's output, checked and scaled to unit length. + fn unit(&self, out: Vec) -> Result, RuntimeError> { if out.len() != self.dim { return Err(RuntimeError::Run(format!("the model gave {} numbers per face, not the {} its description says", out.len(), self.dim))); } diff --git a/crates/faces/src/synthetic.rs b/crates/faces/src/synthetic.rs index 61010a243..6bb792a1f 100644 --- a/crates/faces/src/synthetic.rs +++ b/crates/faces/src/synthetic.rs @@ -248,3 +248,63 @@ pub fn tiny_detector_manifest() -> crate::manifest::ModelManifest { m.size_bytes = Some(bytes.len() as u64); m } + +/// A small network that uses every operator the interpreter runs (both kinds of convolution, batch norm, PRelu, +/// sigmoid, broadcasting add/sub/mul with values and constants, max pooling, nearest upscaling, transposes, reshape, +/// global pooling, flatten and a fully connected layer), with two outputs, for checking an accelerator's program +/// against the CPU. Its input is `[1, 4, 8, 8]`. +pub fn every_operator_graph() -> crate::graph::Graph { + use crate::graph::{Attr, Data, Graph, Node, Weight}; + let node = |op: &str, inputs: &[&str], output: &str, attrs: &[(&str, Attr)]| Node { + op: op.into(), + inputs: inputs.iter().map(|s| s.to_string()).collect(), + outputs: vec![output.into()], + attrs: attrs.iter().map(|(k, v)| (k.to_string(), v.clone())).collect(), + }; + let wave = |n: usize, k: usize| -> Vec { (0..n).map(|i| (((i * k) % 23) as f32 - 11.0) * 0.04).collect() }; + let w = |dims: &[usize], v: Vec| Weight { dims: dims.to_vec(), data: Data::F32(v) }; + let weights = [ + ("w1", w(&[8, 4, 3, 3], wave(288, 7))), + ("b1", w(&[8], wave(8, 5))), + ("bn_s", w(&[8], (0..8).map(|i| 0.5 + i as f32 * 0.1).collect())), + ("bn_b", w(&[8], wave(8, 3))), + ("bn_m", w(&[8], wave(8, 11))), + ("bn_v", w(&[8], (0..8).map(|i| 0.2 + i as f32 * 0.05).collect())), + ("slope", w(&[8, 1, 1], (0..8).map(|i| 0.05 * i as f32).collect())), + ("w2", w(&[8, 1, 3, 3], wave(72, 13))), + ("scales", w(&[4], vec![1.0, 1.0, 2.0, 2.0])), + ("half", w(&[1], vec![0.5])), + ("per_channel", w(&[1, 8, 1, 1], wave(8, 17))), + ("wg", w(&[5, 8], wave(40, 19))), + ("bg", w(&[5], wave(5, 2))), + ("flat", Weight { dims: vec![2], data: Data::I64(vec![1, -1]) }), + ]; + let pads = |p: i64| ("pads", Attr::Ints(vec![p; 4])); + let nodes = vec![ + node("Conv", &["x", "w1", "b1"], "c1", &[pads(1)]), + node("BatchNormalization", &["c1", "bn_s", "bn_b", "bn_m", "bn_v"], "bn", &[]), + node("PRelu", &["bn", "slope"], "pr", &[]), + node("Conv", &["pr", "w2"], "c2", &[("group", Attr::Int(8)), ("strides", Attr::Ints(vec![2, 2])), pads(1)]), + node("Sigmoid", &["c2"], "sg", &[]), + node("Mul", &["c2", "sg"], "swish", &[]), + node("Resize", &["swish", "", "scales"], "up", &[]), + node("Add", &["up", "pr"], "sum", &[]), + node("Sub", &["sum", "half"], "less", &[]), + node("MaxPool", &["less"], "pool", &[("kernel_shape", Attr::Ints(vec![2, 2])), ("strides", Attr::Ints(vec![2, 2]))]), + node("Relu", &["pool"], "relu", &[]), + node("Transpose", &["relu"], "nhwc", &[("perm", Attr::Ints(vec![0, 2, 3, 1]))]), + node("Transpose", &["nhwc"], "nchw", &[("perm", Attr::Ints(vec![0, 3, 1, 2]))]), + node("Mul", &["nchw", "per_channel"], "scaled", &[]), + node("Identity", &["scaled"], "same", &[]), + node("GlobalAveragePool", &["same"], "avg", &[]), + node("Flatten", &["avg"], "row", &[]), + node("Reshape", &["row", "flat"], "row2", &[]), + node("Gemm", &["row2", "wg", "bg"], "y", &[("transB", Attr::Int(1)), ("alpha", Attr::Float(0.5))]), + ]; + Graph { + nodes, + weights: weights.into_iter().map(|(k, v)| (k.to_string(), v)).collect(), + inputs: vec![("x".into(), vec![1, 4, 8, 8])], + outputs: vec!["y".into(), "nhwc".into()], + } +} diff --git a/crates/gpu/src/faces.rs b/crates/gpu/src/faces.rs index ae982d30b..bee9929ef 100644 --- a/crates/gpu/src/faces.rs +++ b/crates/gpu/src/faces.rs @@ -14,6 +14,7 @@ use std::sync::{Arc, Mutex, OnceLock, PoisonError}; use std::time::Duration; use bytemuck::{Pod, Zeroable}; +use lightcraft_faces::net::program::{Arg, Binary, INPUT, Op, Program, Unary}; use lightcraft_faces::net::{Accelerator, ConvJob}; use wgpu::util::DeviceExt; @@ -35,6 +36,22 @@ struct Dev { tiled: wgpu::ComputePipeline, direct: wgpu::ComputePipeline, reduce: wgpu::ComputePipeline, + ops: Ops, + /// Bound where a kernel has fewer operands than bindings. + dummy: wgpu::Buffer, +} + +/// The kernels of `wgsl/faces_ops.wgsl`. +struct Ops { + unary: wgpu::ComputePipeline, + affine: wgpu::ComputePipeline, + prelu: wgpu::ComputePipeline, + binary: wgpu::ComputePipeline, + transpose: wgpu::ComputePipeline, + maxpool: wgpu::ComputePipeline, + resize: wgpu::ComputePipeline, + gavg: wgpu::ComputePipeline, + gemv: wgpu::ComputePipeline, } static DEV: OnceLock> = OnceLock::new(); @@ -146,20 +163,43 @@ fn make_device(backends: wgpu::Backends) -> Result { source: wgpu::ShaderSource::Wgsl(include_str!("wgsl/faces_conv.wgsl").into()), }); let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: None, bind_group_layouts: &[Some(&layout)], immediate_size: 0 }); - let pipe = |entry: &str| { + let ops_module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("faces ops"), + source: wgpu::ShaderSource::Wgsl(include_str!("wgsl/faces_ops.wgsl").into()), + }); + let pipe_in = |module: &wgpu::ShaderModule, entry: &str| { device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: Some(entry), layout: Some(&pl), - module: &module, + module, entry_point: Some(entry), // both kernels write all of their workgroup memory before reading it compilation_options: wgpu::PipelineCompilationOptions { zero_initialize_workgroup_memory: false, ..Default::default() }, cache: None, }) }; + let pipe = |entry: &str| pipe_in(&module, entry); + let op = |entry: &str| pipe_in(&ops_module, entry); let (tiled, direct, reduce) = (pipe("tiled"), pipe("direct"), pipe("reduce")); + let ops = Ops { + unary: op("unary"), + affine: op("affine"), + prelu: op("prelu"), + binary: op("binary"), + transpose: op("transpose"), + maxpool: op("maxpool"), + resize: op("resize"), + gavg: op("gavg"), + gemv: op("gemv"), + }; + let dummy = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("faces unused"), + size: 16, + usage: wgpu::BufferUsages::STORAGE, + mapped_at_creation: false, + }); log::info!("gpu (faces): {} ({:?})", info.name, info.backend); - Ok(Dev { device, queue, info, limits, layout, tiled, direct, reduce }) + Ok(Dev { device, queue, info, limits, layout, tiled, direct, reduce, ops, dummy }) } /// The adapter face scans would use ("AMD Radeon RX 7900 XTX (Vulkan)"), or why they cannot use one. Creates the @@ -203,53 +243,151 @@ struct Resident { pub struct FaceGpu { dev: &'static Dev, layers: Mutex>>, + consts: Mutex>, resident: AtomicU64, } /// A GPU accelerator for one face model, or why there is none (GPU off, missing or broken). pub fn accelerator() -> Result, String> { - Ok(Arc::new(FaceGpu { dev: dev()?, layers: Mutex::new(HashMap::new()), resident: AtomicU64::new(0) })) + Ok(Arc::new(FaceGpu { dev: dev()?, layers: Mutex::new(HashMap::new()), consts: Mutex::new(HashMap::new()), resident: AtomicU64::new(0) })) } fn u32_of(v: usize) -> Result { u32::try_from(v).ok().filter(|v| *v < (1 << 31)).ok_or_else(|| "a layer too large for the GPU kernel".to_string()) } +/// How much of the GPU face work may take: at most `concurrent` runs at once, and busy for at most `share` of the time +/// (after a run of `t`, the next one waits `t · (1 − share) / share`). Set by the background scan's pace. +struct Pace { + concurrent: usize, + share: f32, + running: usize, +} + +static PACE: Mutex = Mutex::new(Pace { concurrent: usize::MAX, share: 1.0, running: 0 }); +static PACE_FREED: std::sync::Condvar = std::sync::Condvar::new(); + +/// Limit the GPU work face models do: `concurrent` runs at once (at least one) and `share` of the time (0.05 to 1). +/// The background face scan sets it from its pace; anything else that runs a face model shares the same limit. +pub fn set_pace(concurrent: usize, share: f32) { + let mut p = PACE.lock().unwrap_or_else(PoisonError::into_inner); + p.concurrent = concurrent.max(1); + p.share = if share.is_finite() { share.clamp(0.05, 1.0) } else { 1.0 }; + PACE_FREED.notify_all(); +} + +/// A turn on the GPU under the pace; waits for one. Dropping it gives it back. +struct Turn { + share: f32, + started: std::time::Instant, +} + +impl Turn { + fn take() -> Turn { + let mut p = PACE.lock().unwrap_or_else(PoisonError::into_inner); + while p.running >= p.concurrent { + p = PACE_FREED.wait_timeout(p, Duration::from_millis(200)).map(|(g, _)| g).unwrap_or_else(|e| e.into_inner().0); + } + p.running += 1; + Turn { share: p.share, started: std::time::Instant::now() } + } +} + +impl Drop for Turn { + fn drop(&mut self) { + // rest while still holding the turn, so the GPU stays idle for the rest of the period + if self.share < 1.0 { + let rest = self.started.elapsed().mul_f32((1.0 - self.share) / self.share).min(Duration::from_secs(2)); + std::thread::sleep(rest); + } + let mut p = PACE.lock().unwrap_or_else(PoisonError::into_inner); + p.running = p.running.saturating_sub(1); + PACE_FREED.notify_one(); + } +} + +/// One dispatch to record: a pipeline, its parameters, its four buffers (three read, the last written) and its +/// workgroup counts. +struct Dispatch { + pipeline: &'static wgpu::ComputePipeline, + params: Vec, + buffers: [wgpu::Buffer; 4], + groups: [u32; 3], +} + +/// Workgroups for `count` items of 256 (or `count` workgroups when `per_group`): x, y and the row length the kernel +/// reads back as `q(1)`. +fn flat(count: usize, limit: usize) -> Option<[u32; 3]> { + let blocks = count.max(1); + let row = blocks.clamp(1, limit.min(65_535)); + let rows = blocks.div_ceil(row); + (rows <= limit).then_some([u32::try_from(row).ok()?, u32::try_from(rows).ok()?, u32::try_from(row).ok()?]) +} + +fn params24(values: &[u32]) -> Vec { + let mut all = [0u32; 24]; + for (to, v) in all.iter_mut().zip(values) { + *to = *v; + } + bytemuck::cast_slice(&all).to_vec() +} + impl FaceGpu { /// Bytes of weights this model keeps on the device. pub fn resident_bytes(&self) -> u64 { self.resident.load(Ordering::Relaxed) } - fn resident(&self, job: &ConvJob) -> Result>, String> { - let mut layers = self.layers.lock().unwrap_or_else(PoisonError::into_inner); - if let Some(r) = layers.get(&job.layer) { - return Ok(Some(r.clone())); - } - let bytes = (job.weights.len() as u64 + job.cout as u64) * 4; - if self.resident_bytes().saturating_add(bytes) > MAX_RESIDENT { - return Ok(None); - } - let d = self.dev; - let usage = wgpu::BufferUsages::STORAGE; - let weights = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("faces weights"), - contents: bytemuck::cast_slice(job.weights), - usage, - }); - let zero = [0f32; 1]; - let bias = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("faces bias"), - contents: bytemuck::cast_slice(job.bias.unwrap_or(&zero)), - usage, - }); - let r = Arc::new(Resident { weights, bias }); - self.resident.fetch_add(bytes, Ordering::Relaxed); - layers.insert(job.layer, r.clone()); - Ok(Some(r)) + fn upload(&self, label: &str, data: &[f32]) -> wgpu::Buffer { + let zero = [0f32; 4]; + self.dev.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some(label), + contents: bytemuck::cast_slice(if data.is_empty() { &zero } else { data }), + usage: wgpu::BufferUsages::STORAGE, + }) } - fn run(&self, job: &ConvJob) -> Result>, String> { + /// Whether `bytes` more may stay on the device; counts them if so. + fn keep(&self, bytes: u64) -> bool { + if self.resident_bytes().saturating_add(bytes) > MAX_RESIDENT { + return false; + } + self.resident.fetch_add(bytes, Ordering::Relaxed); + true + } + + /// A convolution's weights and bias: kept on the device while the budget lasts, uploaded for this run past it. + fn resident(&self, job: &ConvJob) -> Arc { + let mut layers = self.layers.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(r) = layers.get(&job.layer) { + return r.clone(); + } + let r = Arc::new(Resident { weights: self.upload("faces weights", job.weights), bias: self.upload("faces bias", job.bias.unwrap_or(&[])) }); + if self.keep((job.weights.len() as u64 + job.cout as u64) * 4) { + layers.insert(job.layer, r.clone()); + } + r + } + + /// A program's constant, kept on the device like the weights. + fn constant(&self, id: usize, data: &[f32]) -> wgpu::Buffer { + let mut consts = self.consts.lock().unwrap_or_else(PoisonError::into_inner); + if let Some(b) = consts.get(&id) { + return b.clone(); + } + let b = self.upload("faces constant", data); + if self.keep(data.len() as u64 * 4) { + consts.insert(id, b.clone()); + } + b + } + + fn largest_binding(&self) -> u64 { + self.dev.limits.max_storage_buffer_binding_size.min(self.dev.limits.max_buffer_size) + } + + /// The dispatches of one convolution reading `input` and writing `output`, or `None` when it exceeds the device. + fn conv_dispatches(&self, job: &ConvJob, input: &wgpu::Buffer, output: &wgpu::Buffer) -> Result>, String> { let d = self.dev; let groups = job.groups.max(1); let (m, cg) = (job.cout / groups, job.cin / groups); @@ -264,22 +402,15 @@ impl FaceGpu { let kchunk = k.div_ceil(splits).div_ceil(16) * 16; let splits = k.div_ceil(kchunk.max(1)).max(1); let partial = total.checked_mul(splits).ok_or("a layer too large")?; - let largest = [job.input.len(), job.weights.len(), partial].into_iter().max().unwrap_or(0) as u64 * 4; - if largest > d.limits.max_storage_buffer_binding_size.min(d.limits.max_buffer_size) { - return Ok(None); - } - let flat = |count: usize| { - let blocks = count.div_ceil(256); - let row = blocks.clamp(1, 65_535); - (row, blocks.div_ceil(row), 1, row) - }; - let (wx, wy, wz, row) = if direct { flat(total) } else { (n.div_ceil(64), m.div_ceil(64), groups * splits, flat(total).3) }; - let (rx, ry, _, _) = flat(total); + let cin_len = job.cin.checked_mul(job.h).and_then(|v| v.checked_mul(job.w)).ok_or("a layer too large")?; + let largest = [cin_len, job.weights.len(), partial].into_iter().max().unwrap_or(0) as u64 * 4; let limit = d.limits.max_compute_workgroups_per_dimension as usize; - if [wx, wy, wz, rx, ry].into_iter().any(|v| v > limit) { + let Some(sum) = flat(total.div_ceil(256), limit) else { return Ok(None) }; + let (wx, wy, wz) = if direct { (sum[0] as usize, sum[1] as usize, 1) } else { (n.div_ceil(64), m.div_ceil(64), groups * splits) }; + if largest > self.largest_binding() || [wx, wy, wz].into_iter().any(|v| v > limit) { return Ok(None); } - let Some(res) = self.resident(job)? else { return Ok(None) }; + let res = self.resident(job); let params = Params { m: u32_of(m)?, n: u32_of(n)?, @@ -296,82 +427,93 @@ impl FaceGpu { pl: u32_of(job.pl)?, bias: u32::from(job.bias.is_some()), total: u32_of(total)?, - row: u32_of(row)?, + row: sum[2], splits: u32_of(splits)?, kchunk: u32_of(kchunk)?, _a: 0, _b: 0, }; - let Some(_work) = crate::exit::enter() else { return Err("the app is closing".into()) }; - let scopes = [d.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), d.device.push_error_scope(wgpu::ErrorFilter::Validation)]; + let params = bytemuck::bytes_of(¶ms).to_vec(); + let main = |out: &wgpu::Buffer| Dispatch { + pipeline: if direct { &d.direct } else { &d.tiled }, + params: params.clone(), + buffers: [input.clone(), res.weights.clone(), res.bias.clone(), out.clone()], + groups: [wx as u32, wy as u32, wz as u32], + }; + if splits == 1 { + return Ok(Some(vec![main(output)])); + } + let partials = d.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("faces partial sums"), + size: partial as u64 * 4, + usage: wgpu::BufferUsages::STORAGE, + mapped_at_creation: false, + }); + Ok(Some(vec![ + main(&partials), + Dispatch { + pipeline: &d.reduce, + params, + buffers: [partials, res.weights.clone(), res.bias.clone(), output.clone()], + groups: [sum[0], sum[1], 1], + }, + ])) + } + + /// Record `list` into `enc` as one compute pass, with all their parameters in one uniform buffer. + fn record(&self, enc: &mut wgpu::CommandEncoder, list: &[Dispatch]) -> Result<(), String> { + let d = self.dev; + if list.is_empty() { + return Ok(()); + } + let align = (d.limits.min_uniform_buffer_offset_alignment as usize).max(256); + let mut blob = vec![0u8; list.len() * align]; + for (i, x) in list.iter().enumerate() { + blob.get_mut(i * align..i * align + x.params.len()).ok_or("parameters too large")?.copy_from_slice(&x.params); + } let uniform = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: Some("faces params"), - contents: bytemuck::bytes_of(¶ms), + contents: &blob, usage: wgpu::BufferUsages::UNIFORM, }); - let input = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("faces input"), - contents: bytemuck::cast_slice(job.input), - usage: wgpu::BufferUsages::STORAGE, - }); - let bytes = total as u64 * 4; - let output = d.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("faces output"), - size: bytes, - usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, - mapped_at_creation: false, - }); - let partials = (splits > 1).then(|| { - d.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("faces partial sums"), - size: partial as u64 * 4, - usage: wgpu::BufferUsages::STORAGE, - mapped_at_creation: false, + let binds = list + .iter() + .enumerate() + .map(|(i, x)| { + let [a, b, c, out] = &x.buffers; + d.device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &d.layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding { + buffer: &uniform, + offset: (i * align) as u64, + size: std::num::NonZeroU64::new(x.params.len() as u64), + }), + }, + wgpu::BindGroupEntry { binding: 1, resource: a.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 2, resource: b.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 3, resource: c.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 4, resource: out.as_entire_binding() }, + ], + }) }) - }); - let staging = d.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("faces read-back"), - size: bytes, - usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, - mapped_at_creation: false, - }); - let bind = d.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: None, - layout: &d.layout, - entries: &[ - wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 1, resource: input.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 2, resource: res.weights.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 3, resource: res.bias.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 4, resource: partials.as_ref().unwrap_or(&output).as_entire_binding() }, - ], - }); - let summing = partials.as_ref().map(|partials| { - d.device.create_bind_group(&wgpu::BindGroupDescriptor { - label: None, - layout: &d.layout, - entries: &[ - wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 1, resource: partials.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 2, resource: res.weights.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 3, resource: res.bias.as_entire_binding() }, - wgpu::BindGroupEntry { binding: 4, resource: output.as_entire_binding() }, - ], - }) - }); - let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); - { - let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("faces conv"), timestamp_writes: None }); - pass.set_pipeline(if direct { &d.direct } else { &d.tiled }); - pass.set_bind_group(0, &bind, &[]); - pass.dispatch_workgroups(u32_of(wx)?, u32_of(wy)?, u32_of(wz)?); - if let Some(summing) = &summing { - pass.set_pipeline(&d.reduce); - pass.set_bind_group(0, summing, &[]); - pass.dispatch_workgroups(u32_of(rx)?, u32_of(ry)?, 1); - } + .collect::>(); + let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("faces"), timestamp_writes: None }); + for (x, bind) in list.iter().zip(&binds) { + pass.set_pipeline(x.pipeline); + pass.set_bind_group(0, bind, &[]); + let [gx, gy, gz] = x.groups; + pass.dispatch_workgroups(gx, gy, gz); } - enc.copy_buffer_to_buffer(&output, 0, &staging, 0, bytes); + Ok(()) + } + + /// Submit `enc`, which ends by copying into `staging` (`bytes` long), wait for it and read `staging` back. + fn finish(&self, enc: wgpu::CommandEncoder, staging: &wgpu::Buffer, scopes: [wgpu::ErrorScopeGuard; 2]) -> Result, String> { + let d = self.dev; let index = d.queue.submit([enc.finish()]); let slice = staging.slice(..); let (tx, rx) = std::sync::mpsc::channel(); @@ -389,7 +531,7 @@ impl FaceGpu { return Err(format!("the GPU reported an error: {e}")); } if let Err(e) = waited { - mark_broken(&format!("the GPU did not finish a convolution: {e}")); + mark_broken(&format!("the GPU did not finish a face model: {e}")); return Err(format!("the GPU did not finish: {e}")); } match rx.recv_timeout(Duration::from_secs(10)) { @@ -402,8 +544,244 @@ impl FaceGpu { data.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect::>() }; staging.unmap(); - Ok(Some(out)) + Ok(out) } + + fn scopes(&self) -> [wgpu::ErrorScopeGuard; 2] { + [self.dev.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), self.dev.device.push_error_scope(wgpu::ErrorFilter::Validation)] + } + + fn staging(&self, bytes: u64) -> wgpu::Buffer { + self.dev.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("faces read-back"), + size: bytes.max(4), + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }) + } + + fn scratch(&self, bytes: u64) -> wgpu::Buffer { + self.dev.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("faces values"), + size: bytes.max(4), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }) + } + + /// One convolution, uploaded and read back on its own. + fn run(&self, job: &ConvJob) -> Result>, String> { + let total = job.cout.checked_mul(job.oh).and_then(|v| v.checked_mul(job.ow)).ok_or("a layer too large")?; + let bytes = total as u64 * 4; + if bytes > self.largest_binding() || job.input.len() as u64 * 4 > self.largest_binding() { + return Ok(None); + } + let Some(_work) = crate::exit::enter() else { return Err("the app is closing".into()) }; + let _turn = Turn::take(); + let d = self.dev; + let scopes = self.scopes(); + let input = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("faces input"), + contents: bytemuck::cast_slice(job.input), + usage: wgpu::BufferUsages::STORAGE, + }); + let output = self.scratch(bytes); + let Some(list) = self.conv_dispatches(job, &input, &output)? else { + // leave the scopes balanced + let [oom, validation] = scopes; + let _ = (pollster::block_on(validation.pop()), pollster::block_on(oom.pop())); + return Ok(None); + }; + let staging = self.staging(bytes); + let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + self.record(&mut enc, &list)?; + enc.copy_buffer_to_buffer(&output, 0, &staging, 0, bytes); + self.finish(enc, &staging, scopes).map(Some) + } + + /// The dispatch of one non-convolution step. + fn op_dispatch(&self, op: &Op, out_len: usize, args: [wgpu::Buffer; 3], out: &wgpu::Buffer) -> Result, String> { + let d = self.dev; + let limit = d.limits.max_compute_workgroups_per_dimension as usize; + let u = u32_of; + let n = u(out_len)?; + let per_value = |pipeline: &'static wgpu::ComputePipeline, rest: &[u32]| -> Option { + let groups = flat(out_len.div_ceil(256), limit)?; + let mut all = vec![n, groups[2]]; + all.extend_from_slice(rest); + Some(Dispatch { + pipeline, + params: params24(&all), + buffers: [args[0].clone(), args[1].clone(), args[2].clone(), out.clone()], + groups: [groups[0], groups[1], 1], + }) + }; + let per_group = |pipeline: &'static wgpu::ComputePipeline, rest: &[u32]| -> Option { + let groups = flat(out_len, limit)?; + let mut all = vec![n, groups[2]]; + all.extend_from_slice(rest); + Some(Dispatch { + pipeline, + params: params24(&all), + buffers: [args[0].clone(), args[1].clone(), args[2].clone(), out.clone()], + groups: [groups[0], groups[1], 1], + }) + }; + let four = |v: &[usize; 4]| -> Result<[u32; 4], String> { Ok([u(v[0])?, u(v[1])?, u(v[2])?, u(v[3])?]) }; + let k = &d.ops; + Ok(match op { + Op::Conv(_) => None, + Op::Unary(f) => per_value(&k.unary, &[u32::from(*f == Unary::Sigmoid)]), + Op::ChannelAffine { channels, plane } => per_value(&k.affine, &[u(*plane)?, u(*channels)?]), + Op::ChannelPrelu { channels, plane } => per_value(&k.prelu, &[u(*plane)?, u(*channels)?]), + Op::Binary { op, dims, a, b } => { + let kind = match op { + Binary::Add => 0, + Binary::Sub => 1, + Binary::Mul => 2, + }; + let mut rest = vec![kind]; + rest.extend(four(dims)?); + rest.extend(four(a)?); + rest.extend(four(b)?); + per_value(&k.binary, &rest) + } + Op::Transpose { dims, strides } => { + let mut rest = vec![0]; + rest.extend(four(dims)?); + rest.extend(four(strides)?); + per_value(&k.transpose, &rest) + } + Op::MaxPool { h, w, kh, kw, sh, sw, oh, ow, .. } => { + per_value(&k.maxpool, &[u(*h)?, u(*w)?, u(*kh)?, u(*kw)?, u(*sh)?, u(*sw)?, u(*oh)?, u(*ow)?]) + } + Op::Resize { h, w, sh, sw, .. } => per_value(&k.resize, &[u(*h)?, u(*w)?, u(*sh)?, u(*sw)?]), + Op::GlobalAverage { plane, .. } => per_group(&k.gavg, &[u(*plane)?]), + Op::Gemv { k: kk, n: nn, trans_b, alpha } => per_group(&k.gemv, &[u(*kk)?, u(*nn)?, u32::from(*trans_b), alpha.to_bits()]), + }) + } + + /// A whole program for each input, every value staying on the device; one submission and one read-back. + fn program(&self, prog: &Program, inputs: &[&[f32]]) -> Result>>>, String> { + let d = self.dev; + let largest = self.largest_binding(); + if prog.lens.iter().any(|l| *l as u64 * 4 > largest) || inputs.iter().any(|i| Some(&i.len()) != prog.lens.get(INPUT)) { + return Ok(None); + } + let out_bytes: u64 = prog.outputs.iter().map(|(_, v, _)| prog.lens.get(*v).copied().unwrap_or(0) as u64 * 4).sum(); + let total_bytes = out_bytes.checked_mul(inputs.len() as u64).ok_or("too many faces at once")?; + if total_bytes > d.limits.max_buffer_size { + return Ok(None); + } + let Some(_work) = crate::exit::enter() else { return Err("the app is closing".into()) }; + let _turn = Turn::take(); + let scopes = self.scopes(); + let staging = self.staging(total_bytes); + let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("faces program") }); + // buffers nothing reads any more, reused for later values (best fit) + let mut pool: Vec<(u64, wgpu::Buffer)> = Vec::new(); + let take = |pool: &mut Vec<(u64, wgpu::Buffer)>, bytes: u64| -> wgpu::Buffer { + let best = pool + .iter() + .enumerate() + .filter(|(_, (size, _))| *size >= bytes && *size <= bytes.saturating_mul(2)) + .min_by_key(|(_, (size, _))| *size) + .map(|(i, _)| i); + match best { + Some(i) => pool.swap_remove(i).1, + None => self.scratch(bytes), + } + }; + let mut at = 0u64; + let mut fits = true; + for input in inputs { + let mut values: Vec> = vec![None; prog.lens.len()]; + // a fresh buffer: queued writes land before any of the encoder's work, so a pooled one could still be in use + let first = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("faces input"), + contents: bytemuck::cast_slice(input), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }); + if let Some(v) = values.get_mut(INPUT) { + *v = Some((input.len() as u64 * 4, first)); + } + let mut list = Vec::new(); + for step in &prog.steps { + let len = prog.lens.get(step.out).copied().ok_or("a step writes a value the program does not have")?; + let out = take(&mut pool, len as u64 * 4); + let mut args = Vec::with_capacity(3); + for a in &step.args { + args.push(match a { + Arg::Value(v) => { + values.get(*v).and_then(|v| v.as_ref()).map(|(_, b)| b.clone()).ok_or("a step reads a value before it is made")? + } + Arg::Const { id, data } => self.constant(*id, data), + }); + } + let dispatches = match &step.op { + Op::Conv(job) => { + let input = args.first().ok_or("a convolution without input")?; + self.conv_dispatches(job, input, &out)? + } + op => { + let arg = |i: usize| args.get(i).cloned().unwrap_or_else(|| d.dummy.clone()); + self.op_dispatch(op, len, [arg(0), arg(1), arg(2)], &out)?.map(|x| vec![x]) + } + }; + let Some(dispatches) = dispatches else { + fits = false; + break; + }; + list.extend(dispatches); + if let Some(v) = values.get_mut(step.out) { + *v = Some((len as u64 * 4, out)); + } + for f in &step.free { + if let Some(b) = values.get_mut(*f).and_then(Option::take) { + pool.push(b); + } + } + } + if !fits { + break; + } + self.record(&mut enc, &list)?; + for (_, v, _) in &prog.outputs { + let (bytes, buffer) = values.get(*v).and_then(|v| v.as_ref()).ok_or("an output was never made")?; + let len = prog.lens.get(*v).copied().unwrap_or(0) as u64 * 4; + enc.copy_buffer_to_buffer(buffer, 0, &staging, at, len.min(*bytes)); + at += len; + } + pool.extend(values.into_iter().flatten()); + } + if !fits { + let [oom, validation] = scopes; + let _ = (pollster::block_on(validation.pop()), pollster::block_on(oom.pop())); + return Ok(None); + } + let flat_out = self.finish(enc, &staging, scopes)?; + let mut rest = flat_out.as_slice(); + let mut per_input = Vec::with_capacity(inputs.len()); + for _ in inputs { + let mut outs = Vec::with_capacity(prog.outputs.len()); + for (_, v, _) in &prog.outputs { + let len = prog.lens.get(*v).copied().unwrap_or(0); + let (this, next) = rest.split_at_checked(len).ok_or("the read-back is short")?; + outs.push(this.to_vec()); + rest = next; + } + per_input.push(outs); + } + Ok(Some(per_input)) + } +} + +/// Run `f`, turning a panic into a broken device. +fn guarded(f: impl FnOnce() -> Result) -> Result { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)).unwrap_or_else(|_| { + mark_broken("a face model on the GPU panicked"); + Err("the GPU runner gave up".into()) + }) } impl Accelerator for FaceGpu { @@ -418,13 +796,23 @@ impl Accelerator for FaceGpu { if crate::switched_off().is_some() || job.macs() < MIN_MACS { return Ok(None); } - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.run(job))).unwrap_or_else(|_| { - mark_broken("a face convolution panicked"); - Err("the GPU runner gave up".into()) - }) + guarded(|| self.run(job)) + } + + fn run_program(&self, program: &Program, inputs: &[&[f32]]) -> Result>>>, String> { + if let Some(why) = broken_reason() { + return Err(format!("the GPU is not used any more: {why}")); + } + if crate::switched_off().is_some() || PER_CONV.load(Ordering::Relaxed) { + return Ok(None); + } + guarded(|| self.program(program, inputs)) } } +/// Tests and benchmarks only: run convolutions one by one (the first step's way) instead of whole programs. +static PER_CONV: AtomicBool = AtomicBool::new(false); + #[cfg(test)] mod tests { use super::*; @@ -513,6 +901,28 @@ mod tests { assert!(gpu.resident_bytes() > 0, "the weights stay on the device"); } + #[test] + fn every_operator_on_the_gpu_gives_what_the_cpu_gives() { + let Some(gpu) = gpu() else { return }; + let net = lightcraft_faces::net::Net::new(lightcraft_faces::synthetic::every_operator_graph()).unwrap(); + let inputs: Vec = (0..3) + .map(|s| { + lightcraft_faces::net::Tensor::new(vec![1, 4, 8, 8], (0..256).map(|i| (((i * 31 + s * 7) % 37) as f32 - 18.0) * 0.05).collect()) + .unwrap() + }) + .collect(); + let cpu: Vec<_> = inputs.iter().map(|x| net.run(x.clone()).unwrap()).collect(); + let program = net.program(&[1, 4, 8, 8]).unwrap().unwrap(); + let slices: Vec<&[f32]> = inputs.iter().map(|t| t.data.as_slice()).collect(); + let got = gpu.program(&program, &slices).unwrap().unwrap(); + for (want, got) in cpu.iter().zip(&got) { + for ((name, want), got) in want.iter().zip(got) { + let worst = want.data.iter().zip(got).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + assert!(worst < 1e-4, "{name}: off by {worst}"); + } + } + } + /// The real models on the GPU agree with the CPU interpreter: `LC_FACE_MODELS=` holding /// `face_detection_yunet_2023mar.onnx`, `face_recognition_sface_2021dec.onnx` and (optionally) `glintr100.onnx`. #[test] @@ -555,7 +965,8 @@ mod tests { } } - /// Timings, CPU against GPU: `LC_FACE_MODELS= cargo test --release -p lightcraft-gpu faces::tests::timings -- --ignored --nocapture`. + /// Timings, CPU against the GPU one convolution at a time and as whole programs (one face, and a batch of 16): + /// `LC_FACE_MODELS= cargo test --release -p lightcraft-gpu faces::tests::timings -- --ignored --nocapture`. #[test] #[ignore = "a benchmark: needs LC_FACE_MODELS and a GPU"] fn timings() { @@ -573,21 +984,35 @@ mod tests { let (w, h) = (1024, 768); let rgb: Vec = (0..w * h * 3).map(|i| (i * 7 % 251) as u8).collect(); let opts = lightcraft_faces::yunet::Options::default(); - let cpu = time(5, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap())); + let cpu = time(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap())); + PER_CONV.store(true, Ordering::Relaxed); det.accelerate(accelerator().unwrap(), false).unwrap(); - let fast = time(5, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap())); - eprintln!("YuNet 640: CPU {cpu:.1} ms, GPU {fast:.1} ms"); + let per_conv = time(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap())); + PER_CONV.store(false, Ordering::Relaxed); + det.accelerate(accelerator().unwrap(), false).unwrap(); + let whole = time(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap())); + assert!(det.net().accelerator().is_some()); + eprintln!("YuNet 640: CPU {cpu:.1} ms, GPU per convolution {per_conv:.1} ms, GPU whole network {whole:.1} ms"); for (file, sha) in [ ("face_recognition_sface_2021dec.onnx", lightcraft_faces::known::SFACE_SHA256), ("glintr100.onnx", lightcraft_faces::known::AURAFACE_SHA256), ] { let e = lightcraft_faces::runtime::Embedder::load(&dir.join(file), &lightcraft_faces::known::lookup(sha).unwrap()).unwrap(); let (ew, eh) = e.input_size(); - let face: Vec = (0..ew * eh * 3).map(|i| (i * 13 % 251) as u8).collect(); - let cpu = time(3, &mut || drop(e.embed(&face).unwrap())); + let faces: Vec> = (0..16).map(|s| (0..ew * eh * 3).map(|i| ((i + s * 101) * 13 % 251) as u8).collect()).collect(); + let refs: Vec<&[u8]> = faces.iter().map(Vec::as_slice).collect(); + let cpu = time(5, &mut || drop(e.embed(&faces[0]).unwrap())); + PER_CONV.store(true, Ordering::Relaxed); e.accelerate(accelerator().unwrap(), false).unwrap(); - let fast = time(3, &mut || drop(e.embed(&face).unwrap())); - eprintln!("{file}: CPU {cpu:.1} ms, GPU {fast:.1} ms"); + let per_conv = time(5, &mut || drop(e.embed(&faces[0]).unwrap())); + PER_CONV.store(false, Ordering::Relaxed); + e.accelerate(accelerator().unwrap(), false).unwrap(); + let whole = time(5, &mut || drop(e.embed(&faces[0]).unwrap())); + let batch = time(3, &mut || drop(e.embed_batch(&refs))) / 16.0; + assert!(e.net().accelerator().is_some()); + eprintln!( + "{file}: CPU {cpu:.1} ms, GPU per convolution {per_conv:.1} ms, GPU whole network {whole:.1} ms, batch of 16 {batch:.1} ms a face" + ); } } } diff --git a/crates/gpu/src/wgsl/faces_ops.wgsl b/crates/gpu/src/wgsl/faces_ops.wgsl new file mode 100644 index 000000000..aa7eba242 --- /dev/null +++ b/crates/gpu/src/wgsl/faces_ops.wgsl @@ -0,0 +1,188 @@ +// The face networks' other operators (lightcraft-faces `net::program`), so activations stay on the device between layers. +// Same bindings as faces_conv.wgsl: the parameters, three read-only operands and the output. Element kernels run one +// thread per output value; `gavg` and `gemv` one workgroup per output value. +// +// Parameters are 24 numbers: q(0) = values (or workgroups) to make, q(1) = workgroups along x, then per kernel. + +struct P { + v: array, 6>, +} + +@group(0) @binding(0) var p: P; +@group(0) @binding(1) var a: array; +@group(0) @binding(2) var b: array; +@group(0) @binding(3) var c: array; +@group(0) @binding(4) var dst: array; + +fn q(i: u32) -> u32 { + return p.v[i / 4u][i % 4u]; +} + +fn index(wid: vec3, lid: u32) -> u32 { + return (wid.y * q(1) + wid.x) * 256u + lid; +} + +// q(2): 0 relu, 1 sigmoid +@compute @workgroup_size(256) +fn unary(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + let x = a[i]; + if (q(2) == 0u) { + dst[i] = max(x, 0.0); + } else { + dst[i] = 1.0 / (1.0 + exp(-x)); + } +} + +// q(2) plane, q(3) channels: x * b[c] + c[c] +@compute @workgroup_size(256) +fn affine(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + let ch = (i / q(2)) % q(3); + dst[i] = a[i] * b[ch] + c[ch]; +} + +// q(2) plane, q(3) channels: x < 0 ? x * b[c] : x +@compute @workgroup_size(256) +fn prelu(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + let x = a[i]; + let ch = (i / q(2)) % q(3); + dst[i] = select(x, x * b[ch], x < 0.0); +} + +// q(2): 0 add, 1 sub, 2 mul; q(3..7) output dims; q(7..11) a's strides; q(11..15) b's strides +@compute @workgroup_size(256) +fn binary(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + var rest = i; + var ia = 0u; + var ib = 0u; + for (var d = 3; d >= 0; d = d - 1) { + let du = u32(d); + let size = q(3u + du); + let at = rest % size; + rest = rest / size; + ia = ia + at * q(7u + du); + ib = ib + at * q(11u + du); + } + let x = a[ia]; + let y = b[ib]; + if (q(2) == 0u) { + dst[i] = x + y; + } else if (q(2) == 1u) { + dst[i] = x - y; + } else { + dst[i] = x * y; + } +} + +// q(3..7) output dims; q(7..11) the input's strides for each output axis +@compute @workgroup_size(256) +fn transpose(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + var rest = i; + var at_in = 0u; + for (var d = 3; d >= 0; d = d - 1) { + let du = u32(d); + let size = q(3u + du); + at_in = at_in + (rest % size) * q(7u + du); + rest = rest / size; + } + dst[i] = a[at_in]; +} + +// q(2) h, q(3) w, q(4) kh, q(5) kw, q(6) sh, q(7) sw, q(8) oh, q(9) ow +@compute @workgroup_size(256) +fn maxpool(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + let h = q(2); + let w = q(3); + let oh = q(8); + let ow = q(9); + let ox = i % ow; + let oy = (i / ow) % oh; + let ch = i / (ow * oh); + var m = -3.402823e38; + for (var ky = 0u; ky < q(4); ky = ky + 1u) { + for (var kx = 0u; kx < q(5); kx = kx + 1u) { + let y = oy * q(6) + ky; + let x = ox * q(7) + kx; + if (y < h && x < w) { + m = max(m, a[(ch * h + y) * w + x]); + } + } + } + dst[i] = m; +} + +// q(2) h, q(3) w, q(4) sh, q(5) sw: nearest upscale by whole numbers +@compute @workgroup_size(256) +fn resize(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let i = index(wid, lid); + if (i >= q(0)) { return; } + let h = q(2); + let w = q(3); + let ow = w * q(5); + let oh = h * q(4); + let ox = i % ow; + let oy = (i / ow) % oh; + let ch = i / (ow * oh); + dst[i] = a[(ch * h + oy / q(4)) * w + ox / q(5)]; +} + +var partial: array; + +fn reduce(lid: u32) { + for (var s = 128u; s > 0u; s = s >> 1u) { + workgroupBarrier(); + if (lid < s) { + partial[lid] = partial[lid] + partial[lid + s]; + } + } + workgroupBarrier(); +} + +// one workgroup per channel; q(2) plane +@compute @workgroup_size(256) +fn gavg(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let ch = wid.y * q(1) + wid.x; + let plane = q(2); + var sum = 0.0; + if (ch < q(0)) { + for (var j = lid; j < plane; j = j + 256u) { + sum = sum + a[ch * plane + j]; + } + } + partial[lid] = sum; + reduce(lid); + if (lid == 0u && ch < q(0)) { + dst[ch] = partial[0] / f32(plane); + } +} + +// one workgroup per output; q(2) k, q(3) n, q(4) 1 when B is [n][k], q(5) alpha's bits +@compute @workgroup_size(256) +fn gemv(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_index) lid: u32) { + let o = wid.y * q(1) + wid.x; + let k = q(2); + let n = q(3); + var sum = 0.0; + if (o < q(0)) { + for (var j = lid; j < k; j = j + 256u) { + let wi = select(j * n + o, o * k + j, q(4) == 1u); + sum = sum + a[j] * b[wi]; + } + } + partial[lid] = sum; + reduce(lid); + if (lid == 0u && o < q(0)) { + dst[o] = bitcast(q(5)) * partial[0] + c[o]; + } +} diff --git a/crates/ui-egui/locales/de.json b/crates/ui-egui/locales/de.json index 25c411ce0..bfaf57fdf 100644 --- a/crates/ui-egui/locales/de.json +++ b/crates/ui-egui/locales/de.json @@ -2118,5 +2118,13 @@ "Adding to Apple Photos": "Wird zu Apple Fotos hinzugefügt", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Kein Grafikprozessor gefunden: LightCraft rendert auf der CPU, was viel langsamer ist. Details unter Hilfe → Systeminformationen.", "On NixOS, start it from `nix develop`.": "Unter NixOS starten Sie es aus `nix develop`.", - "Files with the same content as another file in this import: only the first is imported.": "Dateien mit demselben Inhalt wie eine andere Datei dieses Imports: Nur die erste wird importiert." + "Files with the same content as another file in this import: only the first is imported.": "Dateien mit demselben Inhalt wie eine andere Datei dieses Imports: Nur die erste wird importiert.", + "Run face models on the CPU only": "Gesichtsmodelle nur auf der CPU ausführen", + "Detection": "Erkennung", + "Recognition": "Wiedererkennung", + "Not loaded yet": "Noch nicht geladen", + "Checking the GPU…": "GPU wird geprüft …", + "CPU": "CPU", + "A face model uses the GPU only when a check shows it gives the same results faster.": "Ein Gesichtsmodell nutzt die GPU nur, wenn eine Prüfung zeigt, dass sie dieselben Ergebnisse schneller liefert.", + "Checking the GPU for faces": "GPU für Gesichter wird geprüft" } diff --git a/crates/ui-egui/locales/es.json b/crates/ui-egui/locales/es.json index ba3e1e331..d153c911e 100644 --- a/crates/ui-egui/locales/es.json +++ b/crates/ui-egui/locales/es.json @@ -1995,5 +1995,13 @@ "Adding to Apple Photos": "Añadiendo a Fotos de Apple", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "No se encontró ningún procesador gráfico: LightCraft renderiza en la CPU, lo que es mucho más lento. Detalles en Ayuda → Información del sistema.", "On NixOS, start it from `nix develop`.": "En NixOS, inícielo desde `nix develop`.", - "Files with the same content as another file in this import: only the first is imported.": "Archivos con el mismo contenido que otro archivo de esta importación: solo se importa el primero." + "Files with the same content as another file in this import: only the first is imported.": "Archivos con el mismo contenido que otro archivo de esta importación: solo se importa el primero.", + "Run face models on the CPU only": "Ejecutar los modelos de caras solo en la CPU", + "Detection": "Detección", + "Recognition": "Reconocimiento", + "Not loaded yet": "Aún no cargado", + "Checking the GPU…": "Comprobando la GPU…", + "CPU": "CPU", + "A face model uses the GPU only when a check shows it gives the same results faster.": "Un modelo de caras usa la GPU solo cuando una comprobación muestra que da los mismos resultados más rápido.", + "Checking the GPU for faces": "Comprobando la GPU para las caras" } diff --git a/crates/ui-egui/locales/fr.json b/crates/ui-egui/locales/fr.json index b0becb154..3827ce3ac 100644 --- a/crates/ui-egui/locales/fr.json +++ b/crates/ui-egui/locales/fr.json @@ -1957,5 +1957,13 @@ "Adding to Apple Photos": "Ajout à Apple Photos", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Aucun processeur graphique trouvé : LightCraft effectue le rendu sur le processeur, ce qui est beaucoup plus lent. Détails dans Aide → Informations système.", "On NixOS, start it from `nix develop`.": "Sous NixOS, lancez-le depuis `nix develop`.", - "Files with the same content as another file in this import: only the first is imported.": "Fichiers au contenu identique à un autre fichier de cette importation : seul le premier est importé." + "Files with the same content as another file in this import: only the first is imported.": "Fichiers au contenu identique à un autre fichier de cette importation : seul le premier est importé.", + "Run face models on the CPU only": "Exécuter les modèles de visages sur le processeur uniquement", + "Detection": "Détection", + "Recognition": "Reconnaissance", + "Not loaded yet": "Pas encore chargé", + "Checking the GPU…": "Vérification du GPU…", + "CPU": "Processeur", + "A face model uses the GPU only when a check shows it gives the same results faster.": "Un modèle de visages n’utilise le GPU que si une vérification montre qu’il donne les mêmes résultats plus vite.", + "Checking the GPU for faces": "Vérification du GPU pour les visages" } diff --git a/crates/ui-egui/locales/ja.json b/crates/ui-egui/locales/ja.json index 7a282e676..d9d83d109 100644 --- a/crates/ui-egui/locales/ja.json +++ b/crates/ui-egui/locales/ja.json @@ -1961,5 +1961,13 @@ "Adding to Apple Photos": "Appleの写真に追加中", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "グラフィックプロセッサが見つかりません。LightCraftはCPUで描画するため、かなり遅くなります。詳細は「ヘルプ → システム情報」を参照してください。", "On NixOS, start it from `nix develop`.": "NixOSでは `nix develop` から起動してください。", - "Files with the same content as another file in this import: only the first is imported.": "この読み込み内の別のファイルと内容が同じファイルです。最初のファイルだけが読み込まれます。" + "Files with the same content as another file in this import: only the first is imported.": "この読み込み内の別のファイルと内容が同じファイルです。最初のファイルだけが読み込まれます。", + "Run face models on the CPU only": "顔モデルを CPU だけで実行", + "Detection": "検出", + "Recognition": "認識", + "Not loaded yet": "未読み込み", + "Checking the GPU…": "GPU を確認中…", + "CPU": "CPU", + "A face model uses the GPU only when a check shows it gives the same results faster.": "顔モデルが GPU を使うのは、同じ結果をより速く出せると確認できた場合だけです。", + "Checking the GPU for faces": "顔用の GPU を確認中" } diff --git a/crates/ui-egui/locales/pt-br.json b/crates/ui-egui/locales/pt-br.json index 2e21c727c..f626e6433 100644 --- a/crates/ui-egui/locales/pt-br.json +++ b/crates/ui-egui/locales/pt-br.json @@ -1998,5 +1998,13 @@ "Recognition models": "Modelos de reconhecimento", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Nenhum processador gráfico encontrado: o LightCraft renderiza na CPU, o que é muito mais lento. Detalhes em Ajuda → Informações do sistema.", "On NixOS, start it from `nix develop`.": "No NixOS, inicie-o a partir de `nix develop`.", - "Files with the same content as another file in this import: only the first is imported.": "Arquivos com o mesmo conteúdo de outro arquivo desta importação: apenas o primeiro é importado." + "Files with the same content as another file in this import: only the first is imported.": "Arquivos com o mesmo conteúdo de outro arquivo desta importação: apenas o primeiro é importado.", + "Run face models on the CPU only": "Executar os modelos de rostos só na CPU", + "Detection": "Detecção", + "Recognition": "Reconhecimento", + "Not loaded yet": "Ainda não carregado", + "Checking the GPU…": "Verificando a GPU…", + "CPU": "CPU", + "A face model uses the GPU only when a check shows it gives the same results faster.": "Um modelo de rostos só usa a GPU quando uma verificação mostra que ela dá os mesmos resultados mais rápido.", + "Checking the GPU for faces": "Verificando a GPU para rostos" } diff --git a/crates/ui-egui/locales/ru.json b/crates/ui-egui/locales/ru.json index c840218a0..f08c3cb7d 100644 --- a/crates/ui-egui/locales/ru.json +++ b/crates/ui-egui/locales/ru.json @@ -2151,5 +2151,13 @@ "Adding to Apple Photos": "Добавление в Apple Фото", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Графический процессор не найден: LightCraft выполняет рендеринг на CPU, а это намного медленнее. Подробности: Справка → Сведения о системе.", "On NixOS, start it from `nix develop`.": "В NixOS запускайте его из `nix develop`.", - "Files with the same content as another file in this import: only the first is imported.": "Файлы с тем же содержимым, что и другой файл этого импорта: импортируется только первый." + "Files with the same content as another file in this import: only the first is imported.": "Файлы с тем же содержимым, что и другой файл этого импорта: импортируется только первый.", + "Run face models on the CPU only": "Запускать модели лиц только на ЦП", + "Detection": "Обнаружение", + "Recognition": "Распознавание", + "Not loaded yet": "Ещё не загружена", + "Checking the GPU…": "Проверка ГП…", + "CPU": "ЦП", + "A face model uses the GPU only when a check shows it gives the same results faster.": "Модель лиц использует ГП, только если проверка показала, что он даёт те же результаты быстрее.", + "Checking the GPU for faces": "Проверка ГП для лиц" } diff --git a/crates/ui-egui/locales/uk.json b/crates/ui-egui/locales/uk.json index 85a838098..a0783aa5f 100644 --- a/crates/ui-egui/locales/uk.json +++ b/crates/ui-egui/locales/uk.json @@ -2151,5 +2151,13 @@ "Adding to Apple Photos": "Додавання до Apple Фото", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Графічний процесор не знайдено: LightCraft виконує рендеринг на ЦП, а це набагато повільніше. Докладніше: Довідка → Інформація про систему.", "On NixOS, start it from `nix develop`.": "У NixOS запускайте його з `nix develop`.", - "Files with the same content as another file in this import: only the first is imported.": "Файли з тим самим вмістом, що й інший файл цього імпорту: імпортується лише перший." + "Files with the same content as another file in this import: only the first is imported.": "Файли з тим самим вмістом, що й інший файл цього імпорту: імпортується лише перший.", + "Run face models on the CPU only": "Запускати моделі облич лише на ЦП", + "Detection": "Виявлення", + "Recognition": "Розпізнавання", + "Not loaded yet": "Ще не завантажено", + "Checking the GPU…": "Перевірка ГП…", + "CPU": "ЦП", + "A face model uses the GPU only when a check shows it gives the same results faster.": "Модель облич використовує ГП, лише якщо перевірка показала, що він дає ті самі результати швидше.", + "Checking the GPU for faces": "Перевірка ГП для облич" } diff --git a/crates/ui-egui/locales/zh-hans.json b/crates/ui-egui/locales/zh-hans.json index ad379f1ef..c2ef31084 100644 --- a/crates/ui-egui/locales/zh-hans.json +++ b/crates/ui-egui/locales/zh-hans.json @@ -1997,5 +1997,13 @@ "Adding to Apple Photos": "正在添加到 Apple 照片", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "未找到图形处理器：LightCraft 改用 CPU 渲染，速度会慢很多。详情见“帮助 → 系统信息”。", "On NixOS, start it from `nix develop`.": "在 NixOS 上，请从 `nix develop` 启动。", - "Files with the same content as another file in this import: only the first is imported.": "与本次导入中另一个文件内容相同的文件：只导入第一个。" + "Files with the same content as another file in this import: only the first is imported.": "与本次导入中另一个文件内容相同的文件：只导入第一个。", + "Run face models on the CPU only": "仅在 CPU 上运行人脸模型", + "Detection": "检测", + "Recognition": "识别", + "Not loaded yet": "尚未加载", + "Checking the GPU…": "正在检查 GPU…", + "CPU": "CPU", + "A face model uses the GPU only when a check shows it gives the same results faster.": "只有检查表明 GPU 能更快给出相同结果时，人脸模型才会使用 GPU。", + "Checking the GPU for faces": "正在检查用于人脸的 GPU" } diff --git a/crates/ui-egui/locales/zh-hant.json b/crates/ui-egui/locales/zh-hant.json index 647d153d0..b830bf434 100644 --- a/crates/ui-egui/locales/zh-hant.json +++ b/crates/ui-egui/locales/zh-hant.json @@ -1959,5 +1959,13 @@ "Adding to Apple Photos": "正在加入 Apple 照片", "No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "找不到圖形處理器：LightCraft 改用 CPU 算繪，速度會慢很多。詳情請見「說明 → 系統資訊」。", "On NixOS, start it from `nix develop`.": "在 NixOS 上，請從 `nix develop` 啟動。", - "Files with the same content as another file in this import: only the first is imported.": "與本次匯入中另一個檔案內容相同的檔案：只匯入第一個。" + "Files with the same content as another file in this import: only the first is imported.": "與本次匯入中另一個檔案內容相同的檔案：只匯入第一個。", + "Run face models on the CPU only": "僅在 CPU 上執行臉孔模型", + "Detection": "偵測", + "Recognition": "辨識", + "Not loaded yet": "尚未載入", + "Checking the GPU…": "正在檢查 GPU…", + "CPU": "CPU", + "A face model uses the GPU only when a check shows it gives the same results faster.": "只有在檢查顯示 GPU 能更快得出相同結果時，臉孔模型才會使用 GPU。", + "Checking the GPU for faces": "正在檢查用於臉孔的 GPU" } diff --git a/crates/ui-egui/src/panels/faces.rs b/crates/ui-egui/src/panels/faces.rs index 2b3fe0329..27933f53d 100644 --- a/crates/ui-egui/src/panels/faces.rs +++ b/crates/ui-egui/src/panels/faces.rs @@ -157,6 +157,42 @@ pub fn setup_banner(app: &mut LightcraftApp, ui: &mut egui::Ui) { }); } +/// Where the face models run (the GPU only once a check shows it gives the CPU's results faster), and the switch that +/// keeps them on the CPU. +fn gpu_section(app: &mut LccApp, ui: &mut egui::Ui, t: &Tokens, gpu: &Value) { + heading(ui, t, "Processor"); + let mut cpu = gpu["cpuOnly"].as_bool().unwrap_or(false); + let r = ui.checkbox(&mut cpu, crate::i18n::tr("Run face models on the CPU only")); + register(ui.ctx(), "check:settings.facesCpuOnly", r.rect); + if r.changed() { + match app.run("faces.gpu", json!({"cpuOnly": cpu})) { + Ok(_) => app.caches.faces_epoch += 1, + Err(e) => app.toast(ui.ctx(), e), + } + } + for (what, label) in [("detection", "Detection"), ("recognition", "Recognition")] { + let st = &gpu[what]; + let (text, why) = if st.is_null() { + (crate::i18n::tr("Not loaded yet").to_string(), None) + } else if st["checking"] == true { + (crate::i18n::tr("Checking the GPU…").to_string(), None) + } else if let Some(device) = st["device"].as_str() { + (device.to_string(), None) + } else { + (crate::i18n::tr("CPU").to_string(), st["why"].as_str()) + }; + ui.horizontal(|ui| { + ui.label(RichText::new(crate::i18n::tr(label)).color(t.text_dim)); + let r = ui.label(text); + register(ui.ctx(), format!("faces:device:{what}"), r.rect); + if let Some(why) = why { + r.on_hover_text(sentence(why)); + } + }); + } + hint(ui, t, "A face model uses the GPU only when a check shows it gives the same results faster."); +} + /// A download-style size: decimal megabytes, as file managers and model pages show them. fn mb(bytes: Option) -> String { match bytes { @@ -231,6 +267,7 @@ pub fn settings_tab(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens) { let all: Vec = list["models"].as_array().cloned().unwrap_or_default(); let downloads = app.session.execute("faces.models.downloads", &json!({})).map(|v| v["downloads"].as_array().cloned().unwrap_or_default()).unwrap_or_default(); + gpu_section(app, ui, t, &list["gpu"]); heading(ui, t, "Face detector"); for m in all.iter().filter(|m| m["role"] == "detector") { model_row(app, ui, t, m, downloads.iter().find(|d| d["id"] == m["id"]), runtime); diff --git a/docs/faces.md b/docs/faces.md index f4f4a4757..c767ce913 100644 --- a/docs/faces.md +++ b/docs/faces.md @@ -145,31 +145,40 @@ names alone do not guarantee the same speed, and those additional model files we ### On the GPU -On a machine with a usable GPU, the convolutions of the detector and the recogniser (most of their time) run on it -through wgpu (`crates/gpu/src/faces.rs`, `wgsl/faces_conv.wgsl`: an implicit matrix product that gathers the input -straight into workgroup memory, cuts K into slices when a layer's map is small, and a direct kernel for depthwise -layers). Everything else in the graph stays on the CPU interpreter, which remains the reference. After a model loads, -a background check runs it once on the CPU and twice through the GPU on a fixed test input: the GPU is used only when +On a machine with a usable GPU, the whole detector and recogniser run on it through wgpu (`crates/gpu/src/faces.rs`). +The interpreter lowers the graph to a flat program for the input's shape (`crates/faces/src/net/program.rs`): each +convolution (`wgsl/faces_conv.wgsl`: an implicit matrix product that gathers the input straight into workgroup memory, +cuts K into slices when a map is small, and a direct kernel for depthwise layers) and every other operator (activations, +batch norm, PRelu, broadcasting add/sub/mul, max pooling, nearest upscaling, transposes, global pooling and the fully +connected layer, `wgsl/faces_ops.wgsl`) is a kernel, every value stays on the device between layers in reused +buffers, and a run is one upload, one submission and one read-back. Shape-only operators make no step. A model that +needs something the program cannot express runs its convolutions one by one on the GPU instead. Several faces of one +photo go together (`Embedder::embed_batch`, up to 16 per submission). + +The CPU interpreter remains the reference. After a model loads, a background check (a row in the activity stack, "Checking +the GPU for faces") runs it once on the CPU and twice through the GPU on a fixed test input: the GPU is used only when every output is within 1e-3 of the CPU's (relative to the output's largest value) **and** it was faster; until then, -and whenever the GPU is switched off (`LIGHTCRAFT_GPU=0`, `LIGHTCRAFT_GPU_BACKEND=off`, the GPU preference), missing, -wrong or slower, the model runs on the CPU. The first GPU error drops it for that model. The weights stay on a device of -the scan's own (separate from the editor's) while the model is loaded, at most 768 MB per model. The log says which -device each model runs on. +and whenever the GPU is switched off (`LIGHTCRAFT_GPU=0`, `LIGHTCRAFT_GPU_BACKEND=off`, the GPU preference, or **Settings ▸ +Faces ▸ Run face models on the CPU only**), missing, wrong or slower, the model runs on the CPU. The first GPU error, +or a result of the wrong size or with values that are not finite, drops it for that model. Settings ▸ Faces shows +which device each model runs on (`faces.gpu`). The weights stay on a device of the scan's own (separate from the +editor's) while the model is loaded, at most 768 MB per model. -Measured on Linux with a Radeon (Vulkan) and a 32-thread CPU, one call at a time, machine under heavy load from other -jobs (so the absolute numbers are noisy; the ratios held across three runs): +The background scan's pace also limits the GPU: paused or light, one model run at a time and the GPU busy at most a +quarter of the time (the editor renders on the same GPU); normal, two at once and three quarters; full, as fast as it +goes. -| Model | CPU | GPU | -| --- | ---: | ---: | -| YuNet (640², one photo) | 33–34 ms | 35–39 ms | -| SFace 2021dec | 25–28 ms | 22–27 ms | -| AuraFace v1 | 204–208 ms | 81–103 ms | +Measured on Linux with a Radeon (Vulkan) and a 32-thread CPU, one call at a time, machine under very heavy load from +other jobs (load average 100–250, so the absolute numbers are noisy): + +| Model | CPU | GPU, one layer at a time | GPU, whole network | Batch of 16, per face | +| --- | ---: | ---: | ---: | ---: | +| YuNet (640², one photo) | 58–60 ms | 85–89 ms | 21 ms | – | +| SFace 2021dec | 40–42 ms | 47–51 ms | 9.6 ms | 8.1 ms | +| AuraFace v1 | 295–308 ms | 165–172 ms | 70–73 ms | 71 ms | -Each layer still goes to the GPU and back: for YuNet and SFace that round trip costs about what the layer saves, so on -most machines they stay on the CPU by the speed check, while AuraFace runs 2–2.5× faster. Keeping activations on the GPU -between layers (the remaining operators as kernels) and batching faces per dispatch are the next steps. `LC_FACE_MODELS= cargo test --release -p lightcraft-gpu faces::tests -- --include-ignored --nocapture` checks -agreement on the real models and prints these timings. +agreement on the real models and every operator, and prints these timings. ### Agreement and regression coverage