mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-24 10:09:01 -06:00
Brandon Thomas
Generated
+1
@@ -3300,6 +3300,7 @@ dependencies = [
|
||||
"lightcraft-color",
|
||||
"lightcraft-denoise",
|
||||
"lightcraft-develop",
|
||||
"lightcraft-faces",
|
||||
"lightcraft-geom",
|
||||
"lightcraft-pipeline",
|
||||
"lightcraft-raster",
|
||||
|
||||
@@ -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<Option<(PathBuf, Arc<Detector>)>> = 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<Arc<Detector>, NoDetector> {
|
||||
static LOADED: Mutex<Option<(PathBuf, Arc<Detector>)>> = Mutex::new(None);
|
||||
let home = s
|
||||
.face_models_dir
|
||||
.as_deref()
|
||||
@@ -82,10 +84,133 @@ pub(crate) fn detector(s: &Session) -> std::result::Result<Arc<Detector>, 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<String>,
|
||||
/// Why it runs on the CPU.
|
||||
pub why: Option<String>,
|
||||
}
|
||||
|
||||
/// Each model's [`GpuState`], by what it does ("detection", "recognition").
|
||||
static GPU_STATE: Mutex<Vec<(&'static str, GpuState)>> = Mutex::new(Vec::new());
|
||||
|
||||
pub(crate) fn gpu_state(what: &str) -> Option<GpuState> {
|
||||
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<dyn lightcraft_faces::net::Accelerator>) -> 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<Value> {
|
||||
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<Arc<lightcraft_faces::runtime::Embedder>> = 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<Value> {
|
||||
}
|
||||
|
||||
pub fn specs() -> Vec<CommandSpec> {
|
||||
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)]
|
||||
|
||||
@@ -27,6 +27,9 @@ pub(super) struct FaceSettings {
|
||||
pub(super) enabled: bool,
|
||||
/// The installed recogniser in use.
|
||||
pub(super) embedder: Option<String>,
|
||||
/// 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<Value> {
|
||||
"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,
|
||||
|
||||
@@ -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 };
|
||||
|
||||
@@ -127,6 +127,17 @@ fn pool_threads_for(cores: usize, pace: Pace, cap: Option<usize>) -> 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 {
|
||||
/// `<model id>@<sha256>` 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<Vec<f32>>, 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<Vec<u8>>, 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<Vec<f32>>, 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));
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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)]
|
||||
|
||||
|
||||
+451
-3
@@ -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<Option<Vec<f32>>, 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<Option<Vec<Vec<Vec<f32>>>>, String> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// A loaded network, ready to run.
|
||||
pub struct Net {
|
||||
graph: Graph,
|
||||
accelerator: RwLock<Option<Arc<dyn Accelerator>>>,
|
||||
}
|
||||
|
||||
/// 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<Arc<dyn Accelerator>> {
|
||||
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<dyn Accelerator>,
|
||||
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<Option<Tensor>> {
|
||||
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<Vec<(String, Tensor)>> {
|
||||
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<Tensor>) -> Result<Vec<Vec<(String, Tensor)>>> {
|
||||
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<Option<Vec<Vec<(String, Tensor)>>>> {
|
||||
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<Vec<Vec<(String, Tensor)>>> = (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::<Option<Vec<_>>>()
|
||||
})
|
||||
.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<Vec<(String, Tensor)>> {
|
||||
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<Option<Vec<f32>>, 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<Vec<f32>> {
|
||||
use program::{Arg, Binary, Op, Unary};
|
||||
let mut values: Vec<Vec<f32>> = vec![Vec::new(); p.lens.len()];
|
||||
values[program::INPUT] = input.to_vec();
|
||||
for step in &p.steps {
|
||||
let arg = |i: usize| -> Vec<f32> {
|
||||
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<f32> = 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::<f32>() / *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::<f32>() + 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<Option<Vec<f32>>, String> {
|
||||
Ok(None)
|
||||
}
|
||||
fn run_program(&self, p: &program::Program, inputs: &[&[f32]]) -> std::result::Result<Option<Vec<Vec<Vec<f32>>>>, String> {
|
||||
self.calls.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||||
let mut out: Vec<Vec<Vec<f32>>> = 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<Tensor> = (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);
|
||||
|
||||
@@ -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<ConvJob<'a>> {
|
||||
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<Tensor> {
|
||||
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() {
|
||||
|
||||
@@ -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<Arg<'a>>,
|
||||
pub out: usize,
|
||||
/// Values nothing reads after this step (never an output): their memory can be reused.
|
||||
pub free: Vec<usize>,
|
||||
}
|
||||
|
||||
/// 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<usize>,
|
||||
pub steps: Vec<Step<'a>>,
|
||||
/// Each graph output: its name, the value that holds it and its shape.
|
||||
pub outputs: Vec<(String, usize, Vec<usize>)>,
|
||||
}
|
||||
|
||||
/// 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<Option<Program<'_>>> {
|
||||
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<usize>)> = HashMap::new();
|
||||
names.insert(input_name.as_str(), (INPUT, input.to_vec()));
|
||||
let mut steps: Vec<Step> = 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<usize>)> { 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<i64> = (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::<Result<Vec<_>>>()?;
|
||||
// the last step that reads each value; a value nothing reads is free right after it is made
|
||||
let mut last: Vec<Option<usize>> = 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])")),
|
||||
}
|
||||
}
|
||||
@@ -257,7 +257,12 @@ fn image(s: &[usize]) -> Result<(usize, usize, usize)> {
|
||||
}
|
||||
}
|
||||
pub(super) fn infer(g: &Graph, input: &[usize]) -> Result<Vec<Vec<usize>>> {
|
||||
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<usize>>, Vec<Vec<usize>>)> {
|
||||
elements(input)?;
|
||||
let mut per_node = Vec::with_capacity(g.nodes.len());
|
||||
let mut shapes: HashMap<&str, Vec<usize>> = 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<Vec<Vec<usize>>> {
|
||||
.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<Vec<Vec<usize>>> {
|
||||
}
|
||||
}
|
||||
}
|
||||
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::<Result<_>>()?;
|
||||
Ok((per_node, outputs))
|
||||
}
|
||||
|
||||
@@ -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<dyn crate::net::Accelerator>, 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<f32> = 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<Vec<f32>, 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<Result<Vec<f32>, RuntimeError>> {
|
||||
if self.plan.accelerator().is_none() || faces.len() < 2 {
|
||||
return faces.iter().map(|f| self.embed(f)).collect();
|
||||
}
|
||||
let mut results: Vec<Option<Result<Vec<f32>, 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<usize>, Vec<Tensor>) = 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<Tensor, RuntimeError> {
|
||||
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<f32>) -> Result<Vec<f32>, 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)));
|
||||
}
|
||||
|
||||
@@ -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<f32> { (0..n).map(|i| (((i * k) % 23) as f32 - 11.0) * 0.04).collect() };
|
||||
let w = |dims: &[usize], v: Vec<f32>| 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()],
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<f32> {
|
||||
(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<dyn crate::net::Accelerator>, 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
|
||||
|
||||
@@ -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 }
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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<String> {
|
||||
if env_disabled() {
|
||||
return Some(if backend::env_off() { "disabled by LIGHTCRAFT_GPU_BACKEND=off" } else { "disabled by LIGHTCRAFT_GPU=0" }.into());
|
||||
|
||||
@@ -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<uniform> p: P;
|
||||
@group(0) @binding(1) var<storage, read> src: array<f32>;
|
||||
@group(0) @binding(2) var<storage, read> wgt: array<f32>;
|
||||
@group(0) @binding(3) var<storage, read> bias: array<f32>;
|
||||
@group(0) @binding(4) var<storage, read_write> dst: array<f32>;
|
||||
|
||||
const BM: u32 = 64u;
|
||||
const BN: u32 = 64u;
|
||||
const BK: u32 = 16u;
|
||||
|
||||
var<workgroup> As: array<f32, 1024>; // [BK][BM]
|
||||
var<workgroup> Bs: array<f32, 1024>; // [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<u32>, @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<array<f32, 4>, 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<f32, 4>;
|
||||
var b: array<f32, 4>;
|
||||
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<u32>, @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<u32>, @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;
|
||||
}
|
||||
@@ -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<vec4<u32>, 6>,
|
||||
}
|
||||
|
||||
@group(0) @binding(0) var<uniform> p: P;
|
||||
@group(0) @binding(1) var<storage, read> a: array<f32>;
|
||||
@group(0) @binding(2) var<storage, read> b: array<f32>;
|
||||
@group(0) @binding(3) var<storage, read> c: array<f32>;
|
||||
@group(0) @binding(4) var<storage, read_write> dst: array<f32>;
|
||||
|
||||
fn q(i: u32) -> u32 {
|
||||
return p.v[i / 4u][i % 4u];
|
||||
}
|
||||
|
||||
fn index(wid: vec3<u32>, 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<u32>, @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<u32>, @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<u32>, @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<u32>, @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<u32>, @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<u32>, @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<u32>, @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<workgroup> partial: array<f32, 256>;
|
||||
|
||||
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<u32>, @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<u32>, @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<f32>(q(5)) * partial[0] + c[o];
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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 を確認中"
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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": "Проверка ГП для лиц"
|
||||
}
|
||||
|
||||
@@ -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": "Перевірка ГП для облич"
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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<u64>) -> String {
|
||||
match bytes {
|
||||
@@ -231,6 +267,7 @@ pub fn settings_tab(app: &mut LightcraftApp, ui: &mut egui::Ui, t: &Tokens) {
|
||||
let all: Vec<Value> = 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);
|
||||
|
||||
@@ -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=<folder> 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:
|
||||
|
||||
Reference in New Issue
Block a user