diff --git a/Cargo.lock b/Cargo.lock index db6efd50e..c22ec3f19 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3300,6 +3300,7 @@ dependencies = [ "lightcraft-color", "lightcraft-denoise", "lightcraft-develop", + "lightcraft-faces", "lightcraft-geom", "lightcraft-pipeline", "lightcraft-raster", diff --git a/crates/engine/src/cmd/face_detect.rs b/crates/engine/src/cmd/face_detect.rs index 4676d1c65..10e9bf7fe 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::{AiModel, EngineError, Result, Session}; const C: &str = "faces.detect"; @@ -54,10 +54,12 @@ impl NoDetector { } } +/// 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, NoDetector> { - static LOADED: Mutex)>> = Mutex::new(None); let home = s .face_models_dir .as_deref() @@ -82,10 +84,133 @@ pub(crate) fn detector(s: &Session) -> std::result::Result, NoDete let d = Arc::new( Detector::new(&bytes).map_err(|e| NoDetector::Unusable(format!("{e}: remove the detector in Settings > Faces and download it again")))?, ); + #[cfg(not(target_arch = "wasm32"))] + { + let net = d.clone(); + 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, 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) => { + 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}"); + 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}, @@ -173,15 +298,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 7e9e459cf..788f30cdb 100644 --- a/crates/engine/src/cmd/face_recognize.rs +++ b/crates/engine/src/cmd/face_recognize.rs @@ -321,6 +321,8 @@ mod imp { return Ok(e.clone()); } 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(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")), @@ -672,6 +674,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 96bd4f124..61dbc2544 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. @@ -201,7 +212,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)); @@ -218,20 +230,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 @@ -239,6 +249,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 a308f20fb..4a603e0c0 100644 --- a/crates/engine/src/tests_face_models.rs +++ b/crates/engine/src/tests_face_models.rs @@ -627,3 +627,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/lib.rs b/crates/faces/src/lib.rs index b4b7066d5..030cda088 100644 --- a/crates/faces/src/lib.rs +++ b/crates/faces/src/lib.rs @@ -5,7 +5,8 @@ //! //! Everything here treats model files and manifests as hostile input: sizes are capped, numbers must be //! finite and sane, and a malformed file is an error, never a panic. Detection and recognition run -//! on the same checked, single-thread pure-Rust CPU interpreter. +//! on the same checked, single-thread pure-Rust CPU interpreter; an [`net::Accelerator`] (the GPU, in `lightcraft-gpu`) can +//! take its convolutions once it agrees with it. #![forbid(unsafe_code)] #![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] diff --git a/crates/faces/src/net.rs b/crates/faces/src/net.rs index 6222314f2..430335840 100644 --- a/crates/faces/src/net.rs +++ b/crates/faces/src/net.rs @@ -9,8 +9,10 @@ use std::collections::HashMap; use std::rc::Rc; +use std::sync::{Arc, PoisonError, RwLock}; mod kernels; +pub mod program; mod shape; use crate::graph::{Data, Graph, Node}; @@ -64,9 +66,64 @@ impl Tensor { } } +/// One convolution, as handed to an [`Accelerator`]: batch 1, planar (NCHW) float32, already checked. +/// +/// `out[co][oy][ox] = bias[co] + Σ weights[co][ci][ky][kx] · input[g·cg + ci][oy·sh + ky − pt][ox·sw + kx − pl]` +/// with `g = co / (cout / groups)` and zero outside the input. +#[derive(Clone, Copy, Debug)] +pub struct ConvJob<'a> { + /// The node's index in the graph: the same for every run of a net, so its weights can stay on the device. + pub layer: usize, + pub cin: usize, + pub h: usize, + pub w: usize, + pub cout: usize, + pub groups: usize, + pub kh: usize, + pub kw: usize, + pub sh: usize, + pub sw: usize, + pub pt: usize, + pub pl: usize, + pub oh: usize, + pub ow: usize, + /// `[cout][cin / groups][kh][kw]`. + pub weights: &'a [f32], + /// `[cout]`, if the node has one. + pub bias: Option<&'a [f32]>, + /// `[cin][h][w]`. + pub input: &'a [f32], +} + +impl ConvJob<'_> { + /// Multiply-adds the convolution takes. + pub fn macs(&self) -> usize { + let per_group = (self.cin / self.groups.max(1)).saturating_mul(self.kh).saturating_mul(self.kw); + self.cout.saturating_mul(per_group).saturating_mul(self.oh).saturating_mul(self.ow) + } +} + +/// Something that can run a network's convolutions faster than the CPU kernels (the GPU, in `lightcraft-gpu`). +/// +/// The CPU interpreter stays the reference: a net only uses an accelerator that agreed with it on a test input +/// ([`Net::attach_checked`]), and the first error makes the net drop it and run on the CPU again. +pub trait Accelerator: Send + Sync { + /// What it runs on, for the user ("AMD Radeon RX 7900 (Vulkan)"). + 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. pub struct Net { graph: Graph, + accelerator: RwLock>>, } /// Validate semantics rather than silently ignoring unknown or mistyped attributes. @@ -80,7 +137,75 @@ impl Net { for n in &graph.nodes { check(n, &graph)?; } - Ok(Net { graph }) + Ok(Net { graph, accelerator: RwLock::new(None) }) + } + + /// The accelerator in use, if any. + pub fn accelerator(&self) -> Option> { + self.accelerator.read().unwrap_or_else(PoisonError::into_inner).clone() + } + + /// Use `accelerator` for this net's convolutions, but only after it gave the CPU's outputs on `probe` (every + /// output value within `tolerance` of the CPU's, relative to that output's largest magnitude) and, with + /// `require_faster`, ran it faster than the CPU did (a GPU's per-layer round trip can cost more than a small + /// model's layers). Returns why it was not used, so the caller can say so; the net then stays on the CPU. + pub fn attach_checked( + &self, + accelerator: Arc, + probe: &Tensor, + tolerance: f32, + require_faster: bool, + ) -> std::result::Result<(), String> { + self.detach(); + let started = web_time::Instant::now(); + let cpu = self.run(probe.clone()).map_err(|e| e.to_string())?; + let cpu_time = started.elapsed(); + *self.accelerator.write().unwrap_or_else(PoisonError::into_inner) = Some(accelerator); + // the first run also uploads the weights and builds the kernels: time a second one + let fast = self.run(probe.clone()).map_err(|e| e.to_string()); + let mut verdict = match fast { + Err(e) => Err(e), + Ok(_) if self.accelerator().is_none() => Err("the device failed on the test input".into()), + Ok(fast) => agree(&cpu, &fast, tolerance), + }; + if verdict.is_ok() && require_faster { + let started = web_time::Instant::now(); + let again = self.run(probe.clone()).map_err(|e| e.to_string()); + let fast_time = started.elapsed(); + verdict = match again { + Err(e) => Err(e), + Ok(_) if self.accelerator().is_none() => Err("the device failed on the test input".into()), + Ok(_) if fast_time >= cpu_time => Err(format!( + "the device is not faster for this model ({:.0} ms against {:.0} ms on the CPU)", + fast_time.as_secs_f64() * 1000.0, + cpu_time.as_secs_f64() * 1000.0 + )), + Ok(_) => Ok(()), + }; + } + if verdict.is_err() { + self.detach(); + } + verdict + } + + /// Run on the CPU only from now on. + pub fn detach(&self) { + *self.accelerator.write().unwrap_or_else(PoisonError::into_inner) = None; + } + + /// 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.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), + Ok(Some(_)) | Err(_) => { + self.detach(); + Ok(None) + } + } } /// The shape the model declares for its input (0 for a dynamic size). @@ -96,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(); @@ -127,7 +312,13 @@ impl Net { }; let started = web_time::Instant::now(); let out = match node.op.as_str() { - "Conv" => kernels::conv(node, &*get(0)?, g, &mut scratch)?, + "Conv" => { + let x = get(0)?; + match self.accelerated_conv(i, node, &x)? { + Some(out) => out, + None => kernels::conv(node, &x, g, &mut scratch)?, + } + } "Relu" => map(&*get(0)?, |v| v.max(0.0)), "Sigmoid" => map(&*get(0)?, |v| 1.0 / (1.0 + (-v).exp())), "Add" | "Sub" | "Mul" => kernels::binary(node.op.as_str(), &*get(0)?, &*get(1)?)?, @@ -168,6 +359,25 @@ impl Net { } } +/// Whether two runs' outputs agree: same names and shapes, values within `tolerance` of each other relative to the +/// output's largest magnitude. +fn agree(a: &[(String, Tensor)], b: &[(String, Tensor)], tolerance: f32) -> std::result::Result<(), String> { + if a.len() != b.len() { + return Err("the device gave a different number of outputs".into()); + } + for ((name, x), (other, y)) in a.iter().zip(b) { + if name != other || x.shape != y.shape { + return Err(format!("the device gave `{other}` a different shape")); + } + let scale = x.data.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6); + let worst = x.data.iter().zip(&y.data).fold(0f32, |m, (p, q)| m.max((p - q).abs())); + if worst.is_nan() || worst > tolerance * scale { + return Err(format!("the device's `{name}` differs from the CPU's by {:.2e} (of {scale:.2e})", worst)); + } + } + Ok(()) +} + fn map(x: &Tensor, f: impl Fn(f32) -> f32) -> Tensor { Tensor { shape: x.shape.clone(), data: x.data.iter().map(|v| f(*v)).collect() } } @@ -534,6 +744,244 @@ mod tests { assert_eq!(rs.shape, vec![1, 6, 2]); } + /// An accelerator built on the slow reference convolution, optionally wrong or failing. + struct Fake { + skew: f32, + fail_after: std::sync::atomic::AtomicUsize, + } + impl Accelerator for Fake { + fn device(&self) -> String { + "test".into() + } + fn conv(&self, j: &ConvJob) -> std::result::Result>, String> { + if self.fail_after.fetch_sub(1, std::sync::atomic::Ordering::Relaxed) == 0 { + return Err("device lost".into()); + } + assert!(j.pt == j.pl && j.sh == j.sw); + let x = Tensor::new(vec![1, j.cin, j.h, j.w], j.input.to_vec()).unwrap(); + let zero = vec![0.0; j.cout]; + let dims = [j.cout, j.cin / j.groups, j.kh, j.kw]; + let out = reference_conv(&x, j.weights, dims, j.bias.unwrap_or(&zero), j.sh, j.pt, j.groups); + Ok(Some(out.data.iter().map(|v| v + self.skew).collect())) + } + } + + fn two_convs() -> (Graph, Tensor) { + let attrs = [("strides", Attr::Ints(vec![2, 2])), ("pads", Attr::Ints(vec![1; 4]))]; + let mut a = node("Conv", &["x", "w1", "b1"], &attrs); + a.outputs = vec!["h".into()]; + let b = node("Conv", &["h", "w2"], &[("group", Attr::Int(2)), ("pads", Attr::Ints(vec![1; 4]))]); + let mut g = graph( + vec![a, b], + vec![ + ("w1", w(&[4, 3, 3, 3], (0..108).map(|i| ((i * 7 % 13) as f32 - 6.0) * 0.05).collect())), + ("b1", w(&[4], vec![0.1, -0.2, 0.3, 0.0])), + ("w2", w(&[4, 2, 3, 3], (0..72).map(|i| ((i * 5 % 11) as f32 - 5.0) * 0.05).collect())), + ], + ); + g.inputs = vec![("x".into(), vec![1, 3, 9, 8])]; + (g, input(3, 9, 8, |i| ((i * 13 % 17) as f32 - 8.0) * 0.1)) + } + + #[test] + fn an_accelerator_is_used_only_when_it_agrees_with_the_cpu() { + let (g, x) = two_convs(); + let net = Net::new(g).unwrap(); + let cpu = net.run(x.clone()).unwrap(); + let good = Arc::new(Fake { skew: 0.0, fail_after: usize::MAX.into() }); + net.attach_checked(good, &x, 1e-4, false).unwrap(); + assert!(net.accelerator().is_some()); + assert_eq!(net.run(x.clone()).unwrap()[0].1.shape, cpu[0].1.shape); + let wrong = Arc::new(Fake { skew: 0.5, fail_after: usize::MAX.into() }); + assert!(net.attach_checked(wrong, &x, 1e-4, false).is_err()); + assert!(net.accelerator().is_none(), "a device that disagrees is not used"); + assert_eq!(net.run(x).unwrap(), cpu); + } + + #[test] + fn a_failing_accelerator_falls_back_to_the_cpu_mid_run() { + let (g, x) = two_convs(); + let net = Net::new(g).unwrap(); + let cpu = net.run(x.clone()).unwrap(); + // agrees on the probe (two convs), then fails on the second conv of the next run + net.attach_checked(Arc::new(Fake { skew: 0.0, fail_after: 3.into() }), &x, 1e-4, false).unwrap(); + let out = net.run(x).unwrap(); + assert!(net.accelerator().is_none()); + for (a, b) in out[0].1.data.iter().zip(&cpu[0].1.data) { + assert!((a - b).abs() < 1e-4); + } + } + + /// 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 73cc72b93..71ae95874 100644 --- a/crates/faces/src/net/kernels.rs +++ b/crates/faces/src/net/kernels.rs @@ -59,6 +59,33 @@ 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, 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 { + return shape_err("invalid convolution weights"); + }; + let (cout, cg, kh, kw) = (*cout, *cg, *kh, *kw); + let groups = usize::try_from(n.int("group").unwrap_or(1)).map_err(|_| NetError::Shape("invalid group".into()))?; + let (sh, sw) = pair(n, "strides", 1)?; + let (pt, pl, pb, pr) = pads(n)?; + let ph = h.checked_add(pt).and_then(|v| v.checked_add(pb)).ok_or_else(|| NetError::Shape("padding overflow".into()))?; + let pw = w.checked_add(pl).and_then(|v| v.checked_add(pr)).ok_or_else(|| NetError::Shape("padding overflow".into()))?; + if ph < kh || pw < kw || groups == 0 || cg == 0 || kh == 0 || kw == 0 || cin != groups * cg || cout % groups != 0 { + return shape_err("invalid convolution shape"); + } + if cout.checked_mul(cg).and_then(|v| v.checked_mul(kh * kw)) != Some(weights.len()) { + return shape_err("short convolution weights"); + } + let bias = if n.inputs.get(2).is_some_and(|s| !s.is_empty()) { Some(shape::floats(n, g, 2)?) } else { None }; + if bias.is_some_and(|b| b.len() != cout) { + 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 }) +} pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Result { let (cin, h, w) = chw(x)?; let wt = shape::weight(n, g, 1)?; @@ -181,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 89ffaba7a..a1594ca6b 100644 --- a/crates/faces/src/net/shape.rs +++ b/crates/faces/src/net/shape.rs @@ -257,7 +257,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)?; @@ -480,6 +485,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 { @@ -488,5 +494,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 bbb1457eb..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 { @@ -99,6 +102,31 @@ impl Embedder { }) } + /// The network, e.g. to see which device it runs on. + pub fn net(&self) -> &Net { + &self.plan + } + + /// Run the model's convolutions on `accelerator` if it gives the CPU's embedding on a test face (and, with + /// `require_faster`, runs it faster than the CPU); otherwise (or on + /// its first error) recognition stays on the CPU. Returns why it was not used. + pub fn accelerate(&self, accelerator: Arc, require_faster: bool) -> Result<(), RuntimeError> { + let (mean, std) = (self.manifest.input.mean, self.manifest.input.std); + let plane = self.width * self.height; + let data: Vec = crate::yunet::probe_pixels(3 * plane) + .into_iter() + .enumerate() + .map(|(i, v)| { + let c = (i / plane.max(1)).min(2); + (v - mean.get(c).copied().unwrap_or(0.0)) / std.get(c).copied().filter(|s| *s != 0.0).unwrap_or(1.0) + }) + .collect(); + let probe = Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))?; + guarded("checking the device", || { + self.plan.attach_checked(accelerator, &probe, crate::yunet::ACCELERATOR_TOLERANCE, require_faster).map_err(RuntimeError::Run) + }) + } + /// The size of the aligned face picture the model wants (width, height). pub fn input_size(&self) -> (usize, usize) { (self.width, self.height) @@ -110,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())); } @@ -124,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/faces/src/yunet.rs b/crates/faces/src/yunet.rs index b77f318bf..d1c62f75c 100644 --- a/crates/faces/src/yunet.rs +++ b/crates/faces/src/yunet.rs @@ -60,6 +60,14 @@ pub struct Face { pub landmarks: [(f32, f32); 5], } +/// How far an accelerator's outputs may stray from the CPU's on the test input, relative to each output's largest value. +pub const ACCELERATOR_TOLERANCE: f32 = 1e-3; + +/// A deterministic, photo-like test input (smooth gradients plus texture, in 0..255). +pub fn probe_pixels(n: usize) -> Vec { + (0..n).map(|i| ((i % 640) as f32 * 0.3 + ((i * 7919) % 97) as f32 + (i / 640 % 480) as f32 * 0.2) % 255.0).collect() +} + pub struct Detector { net: Net, /// Side of the square network input. @@ -90,6 +98,20 @@ impl Detector { Ok(Detector { net, side }) } + /// The network, e.g. to see which device it runs on. + pub fn net(&self) -> &Net { + &self.net + } + + /// Run the network's convolutions on `accelerator` if it gives the CPU's outputs on a test picture (and, with `require_faster`, + /// runs it faster than the CPU); otherwise + /// (or on its first error) detection stays on the CPU. Returns why it was not used. + pub fn accelerate(&self, accelerator: std::sync::Arc, require_faster: bool) -> Result<(), String> { + let side = self.side; + let probe = Tensor::new(vec![1, 3, side, side], probe_pixels(3 * side * side)).map_err(|e| e.to_string())?; + self.net.attach_checked(accelerator, &probe, ACCELERATOR_TOLERANCE, require_faster) + } + /// The square input side, used to check an installed manifest. pub fn input_side(&self) -> usize { self.side diff --git a/crates/gpu/Cargo.toml b/crates/gpu/Cargo.toml index dc348ee88..060abb112 100644 --- a/crates/gpu/Cargo.toml +++ b/crates/gpu/Cargo.toml @@ -14,6 +14,7 @@ lightcraft-geom = { workspace = true } lightcraft-color = { workspace = true } lightcraft-raster = { workspace = true } lightcraft-develop = { workspace = true } +lightcraft-faces = { workspace = true } lightcraft-pipeline = { workspace = true } log = { workspace = true } diff --git a/crates/gpu/src/faces.rs b/crates/gpu/src/faces.rs new file mode 100644 index 000000000..bee9929ef --- /dev/null +++ b/crates/gpu/src/faces.rs @@ -0,0 +1,1018 @@ +//! Face networks on the GPU: an [`lightcraft_faces::net::Accelerator`] that runs the interpreter's convolutions (most of +//! YuNet's, SFace's and AuraFace's time) with a WGSL kernel (`wgsl/faces_conv.wgsl`). Everything else in the graph +//! stays on the CPU interpreter, which is the reference: a net attaches this only after it agreed with the CPU on a +//! test input (`Net::attach_checked`), and drops it on the first error. +//! +//! The device is a separate one (like denoise's), so a face scan and the interactive renders never share a queue, and +//! it respects the same switches as GPU rendering (`LIGHTCRAFT_GPU`, `LIGHTCRAFT_GPU_BACKEND`, the GPU preference) +//! and the crash sentinel around device creation. A model's weights are uploaded once, on first use, and stay on the +//! device while the accelerator (one per loaded model) lives. Every failure is an `Err`, never a panic. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +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; + +/// Longest wait for one convolution before the device is given up on. +const TIMEOUT: Duration = Duration::from_secs(30); +/// Below this many multiply-adds a convolution stays on the CPU: the upload, dispatch and read-back cost more. +const MIN_MACS: usize = 1 << 20; +/// Workgroups a convolution should make at least, cutting its K into slices when its tiles are too few. +const TARGET_WORKGROUPS: usize = 256; +/// Weights one model may keep on the device (AuraFace's are about 250 MB). +const MAX_RESIDENT: u64 = 768 << 20; + +struct Dev { + device: wgpu::Device, + queue: wgpu::Queue, + info: wgpu::AdapterInfo, + limits: wgpu::Limits, + layout: wgpu::BindGroupLayout, + 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(); +static BROKEN: AtomicBool = AtomicBool::new(false); +static BROKEN_WHY: Mutex> = Mutex::new(None); + +fn mark_broken(why: &str) { + let mut w = BROKEN_WHY.lock().unwrap_or_else(PoisonError::into_inner); + if w.is_none() { + *w = Some(why.to_string()); + } + BROKEN.store(true, Ordering::Relaxed); +} + +/// Why the GPU stopped being used for faces (a device error), if it did. +pub fn broken_reason() -> Option { + BROKEN.load(Ordering::Relaxed).then(|| BROKEN_WHY.lock().unwrap_or_else(PoisonError::into_inner).clone().unwrap_or_else(|| "device error".into())) +} + +fn dev() -> Result<&'static Dev, String> { + if let Some(why) = broken_reason() { + return Err(format!("stopped after a GPU failure: {why}")); + } + if let Some(why) = crate::switched_off() { + return Err(why); + } + let _work = match DEV.get() { + Some(_) => None, + None => Some(crate::exit::enter().ok_or("the app is closing")?), + }; + DEV.get_or_init(create_device).as_ref().map_err(|e| e.clone()) +} + +fn create_device() -> Result { + let Some(backends) = crate::backend::compute_backends() else { return Err("disabled by LIGHTCRAFT_GPU_BACKEND=off".into()) }; + crate::backend::with_init_marker(backends, || { + std::panic::catch_unwind(|| make_device(backends)).unwrap_or_else(|_| Err("device creation panicked".into())) + }) +} + +fn entry(binding: u32, ty: wgpu::BufferBindingType) -> wgpu::BindGroupLayoutEntry { + wgpu::BindGroupLayoutEntry { + binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { ty, has_dynamic_offset: false, min_binding_size: None }, + count: None, + } +} + +fn make_device(backends: wgpu::Backends) -> Result { + let mut desc = wgpu::InstanceDescriptor::new_without_display_handle(); + desc.backends = backends; + desc.backend_options = crate::backend::backend_options(); + desc.flags = crate::backend::instance_flags_from_env(desc.flags); + let instance = wgpu::Instance::new(desc); + let wanted = std::env::var("LIGHTCRAFT_GPU_ADAPTER").ok().map(|s| s.trim().to_lowercase()).filter(|s| !s.is_empty()); + let adapter = match &wanted { + Some(w) => pollster::block_on(instance.enumerate_adapters(backends)) + .into_iter() + .find(|a| { + let i = a.get_info(); + i.device_type != wgpu::DeviceType::Cpu && i.name.to_lowercase().contains(w.as_str()) + }) + .ok_or_else(|| format!("no GPU adapter named like `{w}`"))?, + None => pollster::block_on( + instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }), + ) + .map_err(|e| format!("no GPU adapter among {backends:?} ({e})"))?, + }; + let info = adapter.get_info(); + if info.device_type == wgpu::DeviceType::Cpu { + return Err(format!("software adapter ({}) skipped: the CPU is faster", info.name)); + } + let limits = adapter.limits(); + if limits.max_storage_buffers_per_shader_stage < 4 + || limits.max_compute_invocations_per_workgroup < 256 + || limits.max_compute_workgroup_size_x < 256 + || limits.max_compute_workgroup_storage_size < 8 * 1024 + { + return Err(format!("{} is too limited for the face kernels", info.name)); + } + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("lightcraft-faces"), + required_limits: limits.clone(), + ..Default::default() + })) + .map_err(|e| format!("{}: device creation failed ({e})", info.name))?; + device.on_uncaptured_error(Arc::new(|e| { + log::error!("gpu (faces): {e}"); + mark_broken(&format!("device error: {e}")); + })); + device.set_device_lost_callback(|reason, msg| { + log::error!("gpu (faces): device lost ({reason:?}): {msg}"); + mark_broken(&format!("device lost ({reason:?}): {msg}")); + }); + use wgpu::BufferBindingType::{Storage, Uniform}; + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("faces conv"), + entries: &[ + entry(0, Uniform), + entry(1, Storage { read_only: true }), + entry(2, Storage { read_only: true }), + entry(3, Storage { read_only: true }), + entry(4, Storage { read_only: false }), + ], + }); + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("faces conv"), + 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 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, + 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, ops, dummy }) +} + +/// The adapter face scans would use ("AMD Radeon RX 7900 XTX (Vulkan)"), or why they cannot use one. Creates the +/// device on first use. +pub fn adapter() -> Result { + dev().map(|d| format!("{} ({:?})", d.info.name, d.info.backend)) +} + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct Params { + m: u32, + n: u32, + k: u32, + cg: u32, + h: u32, + w: u32, + ow: u32, + kh: u32, + kw: u32, + sh: u32, + sw: u32, + pt: u32, + pl: u32, + bias: u32, + total: u32, + row: u32, + splits: u32, + kchunk: u32, + _a: u32, + _b: u32, +} + +/// A layer's weights and bias on the device. +struct Resident { + weights: wgpu::Buffer, + bias: wgpu::Buffer, +} + +/// The GPU accelerator for one loaded face model: its weights stay on the device while it lives. +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()), 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 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, + }) + } + + /// 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); + let n = job.oh.checked_mul(job.ow).ok_or("a layer too large")?; + let k = cg.checked_mul(job.kh).and_then(|v| v.checked_mul(job.kw)).ok_or("a layer too large")?; + let total = job.cout.checked_mul(n).ok_or("a layer too large")?; + let direct = m < 8; + // Too few 64 x 64 tiles to keep the GPU busy (a 7 x 7 map of 512 channels makes 8): cut K into slices that + // run side by side, and sum them in a second pass. + let tiles = n.div_ceil(64) * m.div_ceil(64) * groups; + let splits = if direct { 1 } else { (TARGET_WORKGROUPS / tiles.max(1)).clamp(1, (k / 128).max(1)) }; + 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 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; + 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 res = self.resident(job); + let params = Params { + m: u32_of(m)?, + n: u32_of(n)?, + k: u32_of(k)?, + cg: u32_of(cg)?, + h: u32_of(job.h)?, + w: u32_of(job.w)?, + ow: u32_of(job.ow)?, + kh: u32_of(job.kh)?, + kw: u32_of(job.kw)?, + sh: u32_of(job.sh)?, + sw: u32_of(job.sw)?, + pt: u32_of(job.pt)?, + pl: u32_of(job.pl)?, + bias: u32::from(job.bias.is_some()), + total: u32_of(total)?, + row: sum[2], + splits: u32_of(splits)?, + kchunk: u32_of(kchunk)?, + _a: 0, + _b: 0, + }; + 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: &blob, + usage: wgpu::BufferUsages::UNIFORM, + }); + 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() }, + ], + }) + }) + .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); + } + 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(); + slice.map_async(wgpu::MapMode::Read, move |r| { + let _ = tx.send(r); + }); + let waited = d.device.poll(wgpu::PollType::Wait { submission_index: Some(index), timeout: Some(TIMEOUT) }); + let [oom, validation] = scopes; + let (validation, oom) = (pollster::block_on(validation.pop()), pollster::block_on(oom.pop())); + if let Some(e) = oom { + return Err(format!("the GPU ran out of memory: {e}")); + } + if let Some(e) = validation { + mark_broken(&format!("device error: {e}")); + return Err(format!("the GPU reported an error: {e}")); + } + if let Err(e) = waited { + 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)) { + Ok(Ok(())) => {} + Ok(Err(e)) => return Err(format!("reading the result back failed: {e}")), + Err(e) => return Err(format!("reading the result back failed: {e}")), + } + let out = { + let data = slice.get_mapped_range().map_err(|e| format!("reading the result back failed: {e:?}"))?; + data.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect::>() + }; + staging.unmap(); + 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 { + fn device(&self) -> String { + format!("{} ({:?})", self.dev.info.name, self.dev.info.backend) + } + + fn conv(&self, job: &ConvJob) -> 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() || job.macs() < MIN_MACS { + return Ok(None); + } + 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::*; + + fn gpu() -> Option> { + match accelerator() { + Ok(a) => Some(a), + Err(e) => { + eprintln!("no GPU for the face kernels, skipping: {e}"); + None + } + } + } + + /// The plain definition of a convolution, to check the kernels against. + fn reference(j: &ConvJob) -> Vec { + let (m, cg) = (j.cout / j.groups, j.cin / j.groups); + let mut out = vec![0f32; j.cout * j.oh * j.ow]; + for co in 0..j.cout { + let g = co / m; + 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] + * j.input[((g * cg + ci) * j.h + iy as usize) * j.w + ix as usize]; + } + } + } + } + out[(co * j.oh + oy) * j.ow + ox] = v; + } + } + } + out + } + + #[test] + fn the_gpu_convolution_gives_what_the_definition_gives() { + let Some(gpu) = gpu() else { return }; + // (cin, cout, groups, k, stride, pad, h, w): dense, strided, depthwise, grouped, pointwise, odd sizes + for (layer, (cin, cout, groups, k, s, p, h, w)) in [ + (3, 16, 1, 3, 2, 1, 41, 37), + (64, 64, 1, 3, 1, 1, 28, 28), + (32, 32, 32, 3, 1, 1, 30, 31), + (16, 32, 4, 3, 2, 1, 33, 20), + (48, 72, 1, 1, 1, 0, 50, 45), + (8, 130, 1, 5, 1, 2, 23, 19), + ] + .into_iter() + .enumerate() + { + let cg = cin / groups; + let weights: Vec = (0..cout * cg * k * k).map(|i| ((i * 37 % 23) as f32 - 11.0) * 0.03).collect(); + let bias: Vec = (0..cout).map(|i| i as f32 * 0.01 - 0.3).collect(); + let input: Vec = (0..cin * h * w).map(|i| ((i * 13 % 29) as f32 - 14.0) * 0.07).collect(); + let (oh, ow) = ((h + 2 * p - k) / s + 1, (w + 2 * p - k) / s + 1); + let job = ConvJob { + layer, + cin, + h, + w, + cout, + groups, + kh: k, + kw: k, + sh: s, + sw: s, + pt: p, + pl: p, + oh, + ow, + weights: &weights, + bias: Some(&bias), + input: &input, + }; + let got = gpu.run(&job).unwrap().unwrap(); + let want = reference(&job); + let worst = got.iter().zip(&want).fold(0f32, |m, (a, b)| m.max((a - b).abs())); + assert!(worst < 1e-3, "case {layer}: off by {worst}"); + } + 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] + fn real_face_models_on_the_gpu_agree_with_the_cpu() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS").map(std::path::PathBuf::from) else { + eprintln!("LC_FACE_MODELS is not set: skipping"); + return; + }; + if gpu().is_none() { + return; + } + let det = lightcraft_faces::yunet::Detector::new(&std::fs::read(dir.join("face_detection_yunet_2023mar.onnx")).unwrap()).unwrap(); + let (w, h) = (900, 700); + let rgb: Vec = (0..w * h * 3).map(|i| ((i / 3 % w) as f32 * 0.25 + (i / 3 / w) as f32 * 0.3 + (i * 7919 % 41) as f32) as u8).collect(); + let opts = lightcraft_faces::yunet::Options { score: 0.0, ..Default::default() }; + let cpu = det.detect(&rgb, w, h, &opts).unwrap(); + det.accelerate(accelerator().unwrap(), false).unwrap(); + let fast = det.detect(&rgb, w, h, &opts).unwrap(); + assert!(det.net().accelerator().is_some()); + assert_eq!(cpu.len(), fast.len()); + for (a, b) in cpu.iter().zip(&fast) { + assert!((a.score - b.score).abs() < 1e-3, "{} vs {}", a.score, b.score); + } + for (file, sha) in [ + ("face_recognition_sface_2021dec.onnx", lightcraft_faces::known::SFACE_SHA256), + ("glintr100.onnx", lightcraft_faces::known::AURAFACE_SHA256), + ] { + let path = dir.join(file); + if !path.exists() { + continue; + } + let e = lightcraft_faces::runtime::Embedder::load(&path, &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 = e.embed(&face).unwrap(); + e.accelerate(accelerator().unwrap(), false).unwrap(); + let fast = e.embed(&face).unwrap(); + let cosine: f32 = cpu.iter().zip(&fast).map(|(a, b)| a * b).sum(); + assert!(cosine > 0.9999, "{file}: cosine {cosine}"); + } + } + + /// 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() { + let Some(dir) = std::env::var_os("LC_FACE_MODELS").map(std::path::PathBuf::from) else { return }; + let Some(_) = gpu() else { return }; + let time = |runs: usize, f: &mut dyn FnMut()| { + f(); + let t = std::time::Instant::now(); + for _ in 0..runs { + f(); + } + t.elapsed().as_secs_f64() * 1000.0 / runs as f64 + }; + let det = lightcraft_faces::yunet::Detector::new(&std::fs::read(dir.join("face_detection_yunet_2023mar.onnx")).unwrap()).unwrap(); + 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(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap())); + PER_CONV.store(true, Ordering::Relaxed); + det.accelerate(accelerator().unwrap(), false).unwrap(); + 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 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 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/lib.rs b/crates/gpu/src/lib.rs index 76fb0f2c1..b0409962d 100644 --- a/crates/gpu/src/lib.rs +++ b/crates/gpu/src/lib.rs @@ -35,6 +35,8 @@ pub mod backend; mod ctx; #[cfg(not(target_arch = "wasm32"))] mod exit; +#[cfg(not(target_arch = "wasm32"))] +pub mod faces; #[cfg(all(feature = "denoise", not(target_arch = "wasm32")))] pub mod nn; #[cfg(not(target_arch = "wasm32"))] @@ -239,9 +241,9 @@ pub fn wait_idle(timeout: std::time::Duration) -> bool { } } -/// Why GPU compute that has a device of its own (denoise) is switched off by the user's settings, or `None`: the +/// Why GPU compute that has a device of its own (denoise, faces) is switched off by the user's settings, or `None`: the /// environment switches and the rendering preference, but not a render-device failure (the devices are separate). -#[cfg(all(feature = "denoise", not(target_arch = "wasm32")))] +#[cfg(not(target_arch = "wasm32"))] pub(crate) fn switched_off() -> Option { if env_disabled() { return Some(if backend::env_off() { "disabled by LIGHTCRAFT_GPU_BACKEND=off" } else { "disabled by LIGHTCRAFT_GPU=0" }.into()); diff --git a/crates/gpu/src/wgsl/faces_conv.wgsl b/crates/gpu/src/wgsl/faces_conv.wgsl new file mode 100644 index 000000000..8384394e9 --- /dev/null +++ b/crates/gpu/src/wgsl/faces_conv.wgsl @@ -0,0 +1,171 @@ +// Convolution for the face networks (lightcraft-faces), planar (NCHW) float32, batch 1, any stride, padding and grouping. +// +// `tiled` is an implicit matrix product per group: rows are the group's output channels (M), columns its output +// pixels (N), and K runs over input channel x kernel tap. The input is gathered (im2col) straight into workgroup +// memory, so the lowered matrix is never made. 256 threads make a 64 x 64 block, 4 x 4 values each. +// `direct` gives one output value per thread, for depthwise convolutions (one output channel per group), where a +// 64-row tile would be almost all waste. +// +// out[co][p] = bias[co] + sum_k W[co][k] * X[k][p] X[k][p] = input[g*cg + k/taps][oy*sh + ky - pt][ox*sw + kx - pl] + +struct P { + m: u32, // output channels per group + n: u32, // output pixels (oh * ow) + k: u32, // cg * kh * kw + cg: u32, // input channels per group + h: u32, + w: u32, + ow: u32, + kh: u32, + kw: u32, + sh: u32, + sw: u32, + pt: u32, + pl: u32, + bias: u32, // 1 when there is a bias + total: u32, // cout * n + row: u32, // workgroups along x, for `direct` and `reduce` (so a large layer can use y too) + splits: u32, // K is cut into this many slices (z = group * splits + slice), summed by `reduce` + kchunk: u32, // K values per slice, a multiple of BK + _a: u32, + _b: u32, +} + +@group(0) @binding(0) var p: P; +@group(0) @binding(1) var src: array; +@group(0) @binding(2) var wgt: array; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var dst: array; + +const BM: u32 = 64u; +const BN: u32 = 64u; +const BK: u32 = 16u; + +var As: array; // [BK][BM] +var Bs: array; // [BK][BN] + +fn gather(g: u32, k: u32, pix: u32) -> f32 { + let taps = p.kh * p.kw; + let ci = g * p.cg + k / taps; + let tap = k % taps; + let oy = pix / p.ow; + let ox = pix % p.ow; + let iy = i32(oy * p.sh + tap / p.kw) - i32(p.pt); + let ix = i32(ox * p.sw + tap % p.kw) - i32(p.pl); + if (iy < 0 || ix < 0 || iy >= i32(p.h) || ix >= i32(p.w)) { + return 0.0; + } + return src[(ci * p.h + u32(iy)) * p.w + u32(ix)]; +} + +@compute @workgroup_size(256) +fn tiled(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_index) tid: u32) { + let n0 = wg.x * BN; + let m0 = wg.y * BM; + let g = wg.z / p.splits; + let slice = wg.z % p.splits; + let k_lo = slice * p.kchunk; + let k_hi = min(p.k, k_lo + p.kchunk); + let tx = tid % 16u; + let ty = tid / 16u; + var acc: array, 4>; + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { + acc[i][j] = 0.0; + } + } + for (var k0 = k_lo; k0 < k_hi; k0 += BK) { + // 1024 values of each tile, 4 per thread; every slot has exactly one writer + for (var e = tid; e < 1024u; e += 256u) { + // weights: neighbouring threads read neighbouring k of one row + let mm = e / BK; + let ka = k0 + e % BK; + let m = m0 + mm; + var a = 0.0; + if (ka < k_hi && m < p.m) { + a = wgt[(g * p.m + m) * p.k + ka]; + } + As[(e % BK) * BM + mm] = a; + // input: neighbouring threads read neighbouring pixels + let kk = e / BN; + let k = k0 + kk; + let pix = n0 + e % BN; + var b = 0.0; + if (k < k_hi && pix < p.n) { + b = gather(g, k, pix); + } + Bs[e] = b; + } + workgroupBarrier(); + for (var kk = 0u; kk < BK; kk++) { + var a: array; + var b: array; + for (var i = 0u; i < 4u; i++) { + a[i] = As[kk * BM + ty * 4u + i]; + b[i] = Bs[kk * BN + tx + 16u * i]; + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { + acc[i][j] += a[i] * b[j]; + } + } + } + workgroupBarrier(); + } + for (var i = 0u; i < 4u; i++) { + let m = m0 + ty * 4u + i; + if (m >= p.m) { + continue; + } + let co = g * p.m + m; + var b0 = 0.0; + if (p.bias != 0u && p.splits == 1u) { + b0 = bias[co]; + } + // with slices, each writes a partial sum of its own: [slice][cout][n] + let base = slice * p.total + co * p.n; + for (var j = 0u; j < 4u; j++) { + // columns tx, tx + 16, tx + 32, tx + 48: neighbouring threads write neighbouring pixels + let pix = n0 + tx + 16u * j; + if (pix < p.n) { + dst[base + pix] = acc[i][j] + b0; + } + } + } +} + +@compute @workgroup_size(256) +fn direct(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_index) tid: u32) { + let at = (wg.y * p.row + wg.x) * 256u + tid; + if (at >= p.total) { + return; + } + let co = at / p.n; + let pix = at % p.n; + let g = co / p.m; + var v = 0.0; + if (p.bias != 0u) { + v = bias[co]; + } + for (var k = 0u; k < p.k; k++) { + v += wgt[co * p.k + k] * gather(g, k, pix); + } + dst[at] = v; +} + +// The sum of the K slices' partial results, plus the bias. `src` holds [splits][cout][n]. +@compute @workgroup_size(256) +fn reduce(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_index) tid: u32) { + let at = (wg.y * p.row + wg.x) * 256u + tid; + if (at >= p.total) { + return; + } + var v = 0.0; + if (p.bias != 0u) { + v = bias[at / p.n]; + } + for (var s = 0u; s < p.splits; s++) { + v += src[s * p.total + at]; + } + dst[at] = v; +} 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 e1168f126..e3c86c567 100644 --- a/crates/ui-egui/locales/de.json +++ b/crates/ui-egui/locales/de.json @@ -2286,5 +2286,13 @@ "Tidying the thumbnail cache": "Miniaturcache wird aufgeräumt", "Reset all edits (Cmd+Shift+R)": "Alle Bearbeitungen zurücksetzen (⌘⇧R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "Intelligente Alben aktualisieren sich selbst: Ändern Sie ihre Regeln oder verwenden Sie Foto → Foto löschen", - "LightCraft stopped": "LightCraft wurde beendet" + "LightCraft stopped": "LightCraft wurde beendet", + "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 77a175721..3d3faabf2 100644 --- a/crates/ui-egui/locales/es.json +++ b/crates/ui-egui/locales/es.json @@ -2285,5 +2285,13 @@ "Tidying the thumbnail cache": "Ordenando la caché de miniaturas", "Reset all edits (Cmd+Shift+R)": "Restablecer todas las ediciones (Cmd+Mayús+R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "Los álbumes inteligentes se actualizan solos: cambie sus reglas o use Foto → Eliminar foto", - "LightCraft stopped": "LightCraft se detuvo" + "LightCraft stopped": "LightCraft se detuvo", + "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 d70878dfd..99f33a1f4 100644 --- a/crates/ui-egui/locales/fr.json +++ b/crates/ui-egui/locales/fr.json @@ -2283,5 +2283,13 @@ "Tidying the thumbnail cache": "Rangement du cache des miniatures", "Reset all edits (Cmd+Shift+R)": "Réinitialiser toutes les retouches (Cmd+Maj+R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "Les albums intelligents se mettent à jour d'eux-mêmes : modifiez leurs règles ou utilisez Photo → Supprimer la photo", - "LightCraft stopped": "LightCraft s’est arrêté" + "LightCraft stopped": "LightCraft s’est arrêté", + "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 db459dd70..651c18a62 100644 --- a/crates/ui-egui/locales/ja.json +++ b/crates/ui-egui/locales/ja.json @@ -2285,5 +2285,13 @@ "Tidying the thumbnail cache": "サムネールキャッシュを整理中", "Reset all edits (Cmd+Shift+R)": "すべての編集をリセット（⌘⇧R）", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "スマートアルバムは自動で更新されます。ルールを変更するか、写真 → 写真を削除を使ってください", - "LightCraft stopped": "LightCraftが終了しました" + "LightCraft stopped": "LightCraftが終了しました", + "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 14c30f937..1a6f33083 100644 --- a/crates/ui-egui/locales/pt-br.json +++ b/crates/ui-egui/locales/pt-br.json @@ -2125,5 +2125,13 @@ "Tidying the thumbnail cache": "Organizando o cache de miniaturas", "Reset all edits (Cmd+Shift+R)": "Redefinir todas as edições (Cmd+Shift+R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "Álbuns inteligentes se atualizam sozinhos: altere as regras ou use Foto → Excluir foto", - "LightCraft stopped": "O LightCraft parou" + "LightCraft stopped": "O LightCraft parou", + "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 29c723d03..649d13583 100644 --- a/crates/ui-egui/locales/ru.json +++ b/crates/ui-egui/locales/ru.json @@ -2278,5 +2278,13 @@ "Tidying the thumbnail cache": "Упорядочивание кэша миниатюр", "Reset all edits (Cmd+Shift+R)": "Сбросить все правки (Cmd+Shift+R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "Смарт-альбомы обновляются сами: измените их правила или выберите Фото → Удалить фото", - "LightCraft stopped": "LightCraft остановился" + "LightCraft stopped": "LightCraft остановился", + "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 8bb244ac2..60dc06e6a 100644 --- a/crates/ui-egui/locales/uk.json +++ b/crates/ui-egui/locales/uk.json @@ -2283,5 +2283,13 @@ "Tidying the thumbnail cache": "Упорядкування кешу мініатюр", "Reset all edits (Cmd+Shift+R)": "Скинути всі зміни (Cmd+Shift+R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "Розумні альбоми оновлюються самі: змініть їхні правила або виберіть Фото → Видалити фото", - "LightCraft stopped": "LightCraft зупинився" + "LightCraft stopped": "LightCraft зупинився", + "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 d60edc75b..d41cc6fa4 100644 --- a/crates/ui-egui/locales/zh-hans.json +++ b/crates/ui-egui/locales/zh-hans.json @@ -2124,5 +2124,13 @@ "Tidying the thumbnail cache": "正在整理缩略图缓存", "Reset all edits (Cmd+Shift+R)": "重置所有编辑 (Cmd+Shift+R)", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "智能相册会自动更新：请更改其规则，或使用“照片 → 删除照片”", - "LightCraft stopped": "LightCraft 已停止" + "LightCraft stopped": "LightCraft 已停止", + "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 ce3b17f32..9f1b5ff69 100644 --- a/crates/ui-egui/locales/zh-hant.json +++ b/crates/ui-egui/locales/zh-hant.json @@ -2086,5 +2086,13 @@ "Tidying the thumbnail cache": "正在整理縮圖快取", "Reset all edits (Cmd+Shift+R)": "重設所有編輯（Cmd+Shift+R）", "Smart albums update automatically: change their rules, or use Photo ▸ Delete Photo": "智慧型相簿會自動更新：請變更其規則，或使用「照片 → 刪除照片」", - "LightCraft stopped": "LightCraft 已停止" + "LightCraft stopped": "LightCraft 已停止", + "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 364d196d3..8705ebc43 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 LightcraftApp, 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 cd6640aa3..07909a18f 100644 --- a/docs/faces.md +++ b/docs/faces.md @@ -143,6 +143,43 @@ SFace's builtin speed metadata now uses the measured 6.9× ratio. User `catalog. Older published tables for MobileFaceNet/R50/R100 came from a different runtime and are not measurements of this runner; architecture names alone do not guarantee the same speed, and those additional model files were not verified in this task. +### On the GPU + +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, 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. + +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. + +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 | + +`LC_FACE_MODELS= cargo test --release -p lightcraft-gpu faces::tests -- --include-ignored --nocapture` checks +agreement on the real models and every operator, and prints these timings. + ### Agreement and regression coverage The public-crop embeddings match the reference runtime's outputs: