mirror of
https://github.com/storytold/lightcraft.git
synced 2026-10-29 20:48:33 -06:00
Faces: whole networks on the GPU, batched faces, CPU-only switch, paced GPU use
Activations stay on the device between layers (net/program.rs lowers the graph; every operator is a WGSL kernel, one submission and read-back per run); faces of a photo go in batches of up to 16; Settings > Faces shows each model's device and a CPU-only switch (faces.gpu); the scan's pace limits GPU use; the GPU check is an activity row. YuNet 58 -> 21 ms, SFace 40 -> 9.6 ms, AuraFace 295 -> 71 ms (heavily loaded machine). Sources: own design. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Andrew Frahn
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
0242a75b5a
commit
f625c74c7a
@@ -17,7 +17,7 @@ use lightcraft_faces::yunet::{Detector, Face, Options};
|
||||
use lightcraft_meta::{Region, RegionKind};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{CommandSpec, bad, cmd, has_selection};
|
||||
use super::{CommandSpec, always, bad, cmd, has_selection};
|
||||
use crate::{Result, Session};
|
||||
|
||||
const C: &str = "faces.detect";
|
||||
@@ -34,10 +34,12 @@ const EXISTING_IOU: f64 = 0.3;
|
||||
/// What to tell someone who has not downloaded the detector yet.
|
||||
const NOT_INSTALLED: &str = "finding faces needs the YuNet face detector (232 KB, MIT): download it in Settings > Faces";
|
||||
|
||||
/// The YuNet loaded last, and from where.
|
||||
static LOADED: Mutex<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>, String> {
|
||||
static LOADED: Mutex<Option<(PathBuf, Arc<Detector>)>> = Mutex::new(None);
|
||||
let home = s.face_models_dir.as_deref().ok_or("this build has nowhere to keep face models (the desktop app does)")?.join(YUNET_ID);
|
||||
let path = home.join("model.onnx");
|
||||
// installing writes the manifest last, so a folder without one is not (yet) a model
|
||||
@@ -59,35 +61,130 @@ pub(crate) fn detector(s: &Session) -> std::result::Result<Arc<Detector>, String
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
{
|
||||
let net = d.clone();
|
||||
offer_gpu("detection", move |gpu| net.accelerate(gpu, true));
|
||||
offer_gpu(s, "detection", move |gpu| net.accelerate(gpu, true));
|
||||
}
|
||||
*loaded = Some((path, d.clone()));
|
||||
Ok(d)
|
||||
}
|
||||
|
||||
/// Where a face model runs, for Settings ▸ Faces and `faces.gpu`.
|
||||
#[derive(Clone, Debug, Default, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct GpuState {
|
||||
/// The background check is still running.
|
||||
pub checking: bool,
|
||||
/// The GPU it runs on; `None` when it runs on the CPU.
|
||||
pub device: Option<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.
|
||||
/// 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) => return log::info!("faces: {what} runs on the CPU ({why})"),
|
||||
Err(why) => {
|
||||
log::info!("faces: {what} runs on the CPU ({why})");
|
||||
return set_gpu_state(what, GpuState { why: Some(why), ..GpuState::default() });
|
||||
}
|
||||
};
|
||||
let device = lightcraft_faces::net::Accelerator::device(&*gpu);
|
||||
match attach(gpu) {
|
||||
Ok(()) => log::info!("faces: {what} runs on {device}"),
|
||||
Err(why) => log::info!("faces: {what} stays on the CPU: {why}"),
|
||||
Ok(()) => {
|
||||
log::info!("faces: {what} runs on {device}");
|
||||
set_gpu_state(what, GpuState { device: Some(device), ..GpuState::default() });
|
||||
}
|
||||
Err(why) => {
|
||||
log::info!("faces: {what} stays on the CPU: {why}");
|
||||
set_gpu_state(what, GpuState { why: Some(why), ..GpuState::default() });
|
||||
}
|
||||
}
|
||||
});
|
||||
if let Err(e) = spawned {
|
||||
log::warn!("faces: could not check the GPU for {what}: {e}");
|
||||
set_gpu_state(what, GpuState { why: Some(format!("could not check the GPU: {e}")), ..GpuState::default() });
|
||||
}
|
||||
}
|
||||
|
||||
/// `faces.gpu`: where the face models run, and the switch that keeps them on the CPU.
|
||||
fn gpu(s: &mut Session, p: &Value) -> Result<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},
|
||||
@@ -175,15 +272,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,
|
||||
|
||||
@@ -235,7 +235,7 @@ mod imp {
|
||||
}
|
||||
let embedder = Arc::new(Embedder::load(&dir.join(&id).join("model.onnx"), &installed.manifest).map_err(|e| bad(cmd, e.to_string()))?);
|
||||
let net = embedder.clone();
|
||||
super::super::face_detect::offer_gpu("recognition", move |gpu| net.accelerate(gpu, true).map_err(|e| e.to_string()));
|
||||
super::super::face_detect::offer_gpu(s, "recognition", move |gpu| net.accelerate(gpu, true).map_err(|e| e.to_string()));
|
||||
let dim = match installed.manifest.output {
|
||||
lightcraft_faces::OutputSpec::Embedding { dim } => dim as usize,
|
||||
_ => return Err(bad(cmd, "the chosen model is not a recognition model")),
|
||||
@@ -584,6 +584,9 @@ mod imp {
|
||||
if let Some(w) = &s.faces.worker {
|
||||
w.allow(target, pace);
|
||||
}
|
||||
// the GPU is held back the same way: a light pace leaves it mostly free for the window's own renders
|
||||
let (runs, share) = crate::faces_worker::gpu_pace(pace);
|
||||
lightcraft_gpu::faces::set_pace(runs, share);
|
||||
while target > 0 && s.faces.in_flight.len() < target + PREFETCH {
|
||||
let Some(id) = s.faces.queue.pop() else { break };
|
||||
let Some(job) = prepare(s, id, true) else { continue };
|
||||
|
||||
@@ -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.
|
||||
@@ -197,7 +208,8 @@ pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done {
|
||||
let img = Rgb { data: &rgb, width: w, height: h };
|
||||
let (iw, ih) = embedder.input_size();
|
||||
let scaled = |f: &Face| f.landmarks.map(|(x, y)| (x * w as f32, y * h as f32));
|
||||
let mut results: Vec<(Rect, Option<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));
|
||||
@@ -214,20 +226,18 @@ pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done {
|
||||
ih,
|
||||
),
|
||||
};
|
||||
let v = crop.and_then(|c| embedder.embed(&c).ok());
|
||||
// shown by the detector's box where it found this face: every face then has the same tightness
|
||||
let view = best.and_then(|(_, f)| face_rect(f)).unwrap_or(rect);
|
||||
(rect, v, view)
|
||||
(rect, crop, view)
|
||||
})
|
||||
.collect();
|
||||
// a photo with no face regions: what the detector is sure of becomes its faces, embedded in the same pass
|
||||
let new_faces = if detect {
|
||||
detected.map(|faces| {
|
||||
detected.map(|found| {
|
||||
let mut rects = Vec::new();
|
||||
for f in faces.iter().filter(|f| f.score >= REGION_SCORE) {
|
||||
for f in found.iter().filter(|f| f.score >= REGION_SCORE) {
|
||||
let Some(rect) = face_rect(f) else { continue };
|
||||
let v = align_to_template(&img, &scaled(f), iw, ih).and_then(|c| embedder.embed(&c).ok());
|
||||
results.push((rect, v, rect));
|
||||
faces.push((rect, align_to_template(&img, &scaled(f), iw, ih), rect));
|
||||
rects.push(rect);
|
||||
}
|
||||
rects
|
||||
@@ -235,6 +245,10 @@ pub(crate) fn process(p: Prepared, embedder: &Embedder) -> Done {
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let crops: Vec<&[u8]> = faces.iter().filter_map(|(_, c, _)| c.as_deref()).collect();
|
||||
let mut embedded = embedder.embed_batch(&crops).into_iter();
|
||||
let results: Vec<(Rect, Option<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));
|
||||
|
||||
@@ -613,3 +613,17 @@ fn recognition_cannot_be_turned_on_without_a_chosen_model() {
|
||||
assert!(s.execute("faces.enable", &json!({"enabled": true})).is_err());
|
||||
let _ = std::fs::remove_dir_all(&d);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_cpu_only_switch_is_kept_and_bad_values_are_refused() {
|
||||
let d = temp("cpu-only");
|
||||
let mut s = session(&d);
|
||||
assert_eq!(s.execute("faces.gpu", &json!({})).unwrap()["cpuOnly"], false);
|
||||
assert_eq!(s.execute("faces.gpu", &json!({"cpuOnly": true})).unwrap()["cpuOnly"], true);
|
||||
// kept with the models, so a new session starts on the CPU too
|
||||
let mut again = session(&d);
|
||||
assert_eq!(again.execute("faces.models.list", &json!({})).unwrap()["gpu"]["cpuOnly"], true);
|
||||
assert!(again.execute("faces.gpu", &json!({"cpuOnly": "yes"})).is_err());
|
||||
assert_eq!(again.execute("faces.gpu", &json!({"cpuOnly": false})).unwrap()["cpuOnly"], false);
|
||||
let _ = std::fs::remove_dir_all(d);
|
||||
}
|
||||
|
||||
+239
-2
@@ -12,6 +12,7 @@ use std::rc::Rc;
|
||||
use std::sync::{Arc, PoisonError, RwLock};
|
||||
|
||||
mod kernels;
|
||||
pub mod program;
|
||||
mod shape;
|
||||
|
||||
use crate::graph::{Data, Graph, Node};
|
||||
@@ -111,6 +112,12 @@ pub trait Accelerator: Send + Sync {
|
||||
fn device(&self) -> String;
|
||||
/// `[cout][oh][ow]`, or `Ok(None)` to leave this one to the CPU (too small to be worth a dispatch, say).
|
||||
fn conv(&self, job: &ConvJob) -> std::result::Result<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.
|
||||
@@ -190,7 +197,7 @@ impl Net {
|
||||
/// A convolution on the accelerator, if one is attached and takes it. A failing accelerator is dropped.
|
||||
fn accelerated_conv(&self, layer: usize, node: &Node, x: &Tensor) -> Result<Option<Tensor>> {
|
||||
let Some(accelerator) = self.accelerator() else { return Ok(None) };
|
||||
let job = kernels::conv_job(layer, node, x, &self.graph)?;
|
||||
let job = kernels::conv_job(layer, node, &x.shape, &x.data, &self.graph)?;
|
||||
match accelerator.conv(&job) {
|
||||
Ok(Some(out)) if out.len() == job.cout * job.oh * job.ow => Ok(Some(Tensor::new(vec![1, job.cout, job.oh, job.ow], out)?)),
|
||||
Ok(None) => Ok(None),
|
||||
@@ -214,7 +221,67 @@ impl Net {
|
||||
|
||||
/// Run the network on `input` (its single input); returns each output by name.
|
||||
pub fn run(&self, input: Tensor) -> Result<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();
|
||||
@@ -745,6 +812,176 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// An accelerator that runs whole programs with plain loops (the definition of each step), counting its calls;
|
||||
/// `poison` makes it give a value that is not finite.
|
||||
struct ProgramFake {
|
||||
calls: std::sync::atomic::AtomicUsize,
|
||||
poison: bool,
|
||||
}
|
||||
|
||||
fn run_steps(p: &program::Program, input: &[f32]) -> Vec<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);
|
||||
|
||||
@@ -60,8 +60,8 @@ fn shifted_plane(plane: &[f32], width: usize, ky: usize, kx: usize, pt: usize, p
|
||||
Ok(())
|
||||
}
|
||||
/// The checked geometry of a convolution node, for an accelerator.
|
||||
pub(super) fn conv_job<'a>(layer: usize, n: &Node, x: &'a Tensor, g: &'a Graph) -> Result<ConvJob<'a>> {
|
||||
let (cin, h, w) = chw(x)?;
|
||||
pub(super) fn conv_job<'a>(layer: usize, n: &Node, shape: &[usize], input: &'a [f32], g: &'a Graph) -> Result<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 {
|
||||
@@ -84,7 +84,7 @@ pub(super) fn conv_job<'a>(layer: usize, n: &Node, x: &'a Tensor, g: &'a Graph)
|
||||
return shape_err("the convolution bias does not match");
|
||||
}
|
||||
let (oh, ow) = ((ph - kh) / sh + 1, (pw - kw) / sw + 1);
|
||||
Ok(ConvJob { layer, cin, h, w, cout, groups, kh, kw, sh, sw, pt, pl, oh, ow, weights, bias, input: &x.data })
|
||||
Ok(ConvJob { layer, cin, h, w, cout, groups, kh, kw, sh, sw, pt, pl, oh, ow, weights, bias, input })
|
||||
}
|
||||
pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Result<Tensor> {
|
||||
let (cin, h, w) = chw(x)?;
|
||||
@@ -208,7 +208,7 @@ pub(super) fn conv(n: &Node, x: &Tensor, g: &Graph, scratch: &mut Scratch) -> Re
|
||||
Ok(out)
|
||||
}
|
||||
/// Flat offset under ONNX's trailing-dimension broadcasting; shapes are checked before execution.
|
||||
fn broadcast_offset(mut at: usize, out: &[usize], input: &[usize]) -> usize {
|
||||
pub(super) fn broadcast_offset(mut at: usize, out: &[usize], input: &[usize]) -> usize {
|
||||
let mut offset = 0;
|
||||
let mut stride = 1;
|
||||
for (axis, &size) in out.iter().enumerate().rev() {
|
||||
|
||||
@@ -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])")),
|
||||
}
|
||||
}
|
||||
@@ -261,7 +261,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)?;
|
||||
@@ -484,6 +489,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 {
|
||||
@@ -492,5 +498,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 {
|
||||
@@ -135,6 +138,66 @@ impl Embedder {
|
||||
|
||||
/// The face as a unit-length vector. `rgb` is an aligned face at exactly [`Self::input_size`], 8-bit RGB.
|
||||
pub fn embed(&self, rgb: &[u8]) -> Result<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()));
|
||||
}
|
||||
@@ -149,14 +212,11 @@ impl Embedder {
|
||||
}
|
||||
}
|
||||
}
|
||||
let out = guarded("running it", || {
|
||||
let input = Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))?;
|
||||
let mut result = self.plan.run(input).map_err(|e| RuntimeError::Run(e.to_string()))?;
|
||||
if result.len() != 1 {
|
||||
return Err(RuntimeError::Run("the model must give one embedding".into()));
|
||||
}
|
||||
Ok(result.pop().ok_or_else(|| RuntimeError::Run("no output".into()))?.1.data)
|
||||
})?;
|
||||
Tensor::new(vec![1, 3, self.height, self.width], data).map_err(|e| RuntimeError::Input(e.to_string()))
|
||||
}
|
||||
|
||||
/// The model's output, checked and scaled to unit length.
|
||||
fn unit(&self, out: Vec<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()],
|
||||
}
|
||||
}
|
||||
|
||||
+548
-123
@@ -14,6 +14,7 @@ use std::sync::{Arc, Mutex, OnceLock, PoisonError};
|
||||
use std::time::Duration;
|
||||
|
||||
use bytemuck::{Pod, Zeroable};
|
||||
use lightcraft_faces::net::program::{Arg, Binary, INPUT, Op, Program, Unary};
|
||||
use lightcraft_faces::net::{Accelerator, ConvJob};
|
||||
use wgpu::util::DeviceExt;
|
||||
|
||||
@@ -35,6 +36,22 @@ struct Dev {
|
||||
tiled: wgpu::ComputePipeline,
|
||||
direct: wgpu::ComputePipeline,
|
||||
reduce: wgpu::ComputePipeline,
|
||||
ops: Ops,
|
||||
/// Bound where a kernel has fewer operands than bindings.
|
||||
dummy: wgpu::Buffer,
|
||||
}
|
||||
|
||||
/// The kernels of `wgsl/faces_ops.wgsl`.
|
||||
struct Ops {
|
||||
unary: wgpu::ComputePipeline,
|
||||
affine: wgpu::ComputePipeline,
|
||||
prelu: wgpu::ComputePipeline,
|
||||
binary: wgpu::ComputePipeline,
|
||||
transpose: wgpu::ComputePipeline,
|
||||
maxpool: wgpu::ComputePipeline,
|
||||
resize: wgpu::ComputePipeline,
|
||||
gavg: wgpu::ComputePipeline,
|
||||
gemv: wgpu::ComputePipeline,
|
||||
}
|
||||
|
||||
static DEV: OnceLock<Result<Dev, String>> = OnceLock::new();
|
||||
@@ -146,20 +163,43 @@ fn make_device(backends: wgpu::Backends) -> Result<Dev, String> {
|
||||
source: wgpu::ShaderSource::Wgsl(include_str!("wgsl/faces_conv.wgsl").into()),
|
||||
});
|
||||
let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: None, bind_group_layouts: &[Some(&layout)], immediate_size: 0 });
|
||||
let pipe = |entry: &str| {
|
||||
let ops_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
|
||||
label: Some("faces ops"),
|
||||
source: wgpu::ShaderSource::Wgsl(include_str!("wgsl/faces_ops.wgsl").into()),
|
||||
});
|
||||
let pipe_in = |module: &wgpu::ShaderModule, entry: &str| {
|
||||
device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
|
||||
label: Some(entry),
|
||||
layout: Some(&pl),
|
||||
module: &module,
|
||||
module,
|
||||
entry_point: Some(entry),
|
||||
// both kernels write all of their workgroup memory before reading it
|
||||
compilation_options: wgpu::PipelineCompilationOptions { zero_initialize_workgroup_memory: false, ..Default::default() },
|
||||
cache: None,
|
||||
})
|
||||
};
|
||||
let pipe = |entry: &str| pipe_in(&module, entry);
|
||||
let op = |entry: &str| pipe_in(&ops_module, entry);
|
||||
let (tiled, direct, reduce) = (pipe("tiled"), pipe("direct"), pipe("reduce"));
|
||||
let ops = Ops {
|
||||
unary: op("unary"),
|
||||
affine: op("affine"),
|
||||
prelu: op("prelu"),
|
||||
binary: op("binary"),
|
||||
transpose: op("transpose"),
|
||||
maxpool: op("maxpool"),
|
||||
resize: op("resize"),
|
||||
gavg: op("gavg"),
|
||||
gemv: op("gemv"),
|
||||
};
|
||||
let dummy = device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces unused"),
|
||||
size: 16,
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
log::info!("gpu (faces): {} ({:?})", info.name, info.backend);
|
||||
Ok(Dev { device, queue, info, limits, layout, tiled, direct, reduce })
|
||||
Ok(Dev { device, queue, info, limits, layout, tiled, direct, reduce, ops, dummy })
|
||||
}
|
||||
|
||||
/// The adapter face scans would use ("AMD Radeon RX 7900 XTX (Vulkan)"), or why they cannot use one. Creates the
|
||||
@@ -203,53 +243,151 @@ struct Resident {
|
||||
pub struct FaceGpu {
|
||||
dev: &'static Dev,
|
||||
layers: Mutex<HashMap<usize, Arc<Resident>>>,
|
||||
consts: Mutex<HashMap<usize, wgpu::Buffer>>,
|
||||
resident: AtomicU64,
|
||||
}
|
||||
|
||||
/// A GPU accelerator for one face model, or why there is none (GPU off, missing or broken).
|
||||
pub fn accelerator() -> Result<Arc<FaceGpu>, String> {
|
||||
Ok(Arc::new(FaceGpu { dev: dev()?, layers: Mutex::new(HashMap::new()), resident: AtomicU64::new(0) }))
|
||||
Ok(Arc::new(FaceGpu { dev: dev()?, layers: Mutex::new(HashMap::new()), consts: Mutex::new(HashMap::new()), resident: AtomicU64::new(0) }))
|
||||
}
|
||||
|
||||
fn u32_of(v: usize) -> Result<u32, String> {
|
||||
u32::try_from(v).ok().filter(|v| *v < (1 << 31)).ok_or_else(|| "a layer too large for the GPU kernel".to_string())
|
||||
}
|
||||
|
||||
/// How much of the GPU face work may take: at most `concurrent` runs at once, and busy for at most `share` of the time
|
||||
/// (after a run of `t`, the next one waits `t · (1 − share) / share`). Set by the background scan's pace.
|
||||
struct Pace {
|
||||
concurrent: usize,
|
||||
share: f32,
|
||||
running: usize,
|
||||
}
|
||||
|
||||
static PACE: Mutex<Pace> = Mutex::new(Pace { concurrent: usize::MAX, share: 1.0, running: 0 });
|
||||
static PACE_FREED: std::sync::Condvar = std::sync::Condvar::new();
|
||||
|
||||
/// Limit the GPU work face models do: `concurrent` runs at once (at least one) and `share` of the time (0.05 to 1).
|
||||
/// The background face scan sets it from its pace; anything else that runs a face model shares the same limit.
|
||||
pub fn set_pace(concurrent: usize, share: f32) {
|
||||
let mut p = PACE.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
p.concurrent = concurrent.max(1);
|
||||
p.share = if share.is_finite() { share.clamp(0.05, 1.0) } else { 1.0 };
|
||||
PACE_FREED.notify_all();
|
||||
}
|
||||
|
||||
/// A turn on the GPU under the pace; waits for one. Dropping it gives it back.
|
||||
struct Turn {
|
||||
share: f32,
|
||||
started: std::time::Instant,
|
||||
}
|
||||
|
||||
impl Turn {
|
||||
fn take() -> Turn {
|
||||
let mut p = PACE.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
while p.running >= p.concurrent {
|
||||
p = PACE_FREED.wait_timeout(p, Duration::from_millis(200)).map(|(g, _)| g).unwrap_or_else(|e| e.into_inner().0);
|
||||
}
|
||||
p.running += 1;
|
||||
Turn { share: p.share, started: std::time::Instant::now() }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Turn {
|
||||
fn drop(&mut self) {
|
||||
// rest while still holding the turn, so the GPU stays idle for the rest of the period
|
||||
if self.share < 1.0 {
|
||||
let rest = self.started.elapsed().mul_f32((1.0 - self.share) / self.share).min(Duration::from_secs(2));
|
||||
std::thread::sleep(rest);
|
||||
}
|
||||
let mut p = PACE.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
p.running = p.running.saturating_sub(1);
|
||||
PACE_FREED.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
/// One dispatch to record: a pipeline, its parameters, its four buffers (three read, the last written) and its
|
||||
/// workgroup counts.
|
||||
struct Dispatch {
|
||||
pipeline: &'static wgpu::ComputePipeline,
|
||||
params: Vec<u8>,
|
||||
buffers: [wgpu::Buffer; 4],
|
||||
groups: [u32; 3],
|
||||
}
|
||||
|
||||
/// Workgroups for `count` items of 256 (or `count` workgroups when `per_group`): x, y and the row length the kernel
|
||||
/// reads back as `q(1)`.
|
||||
fn flat(count: usize, limit: usize) -> Option<[u32; 3]> {
|
||||
let blocks = count.max(1);
|
||||
let row = blocks.clamp(1, limit.min(65_535));
|
||||
let rows = blocks.div_ceil(row);
|
||||
(rows <= limit).then_some([u32::try_from(row).ok()?, u32::try_from(rows).ok()?, u32::try_from(row).ok()?])
|
||||
}
|
||||
|
||||
fn params24(values: &[u32]) -> Vec<u8> {
|
||||
let mut all = [0u32; 24];
|
||||
for (to, v) in all.iter_mut().zip(values) {
|
||||
*to = *v;
|
||||
}
|
||||
bytemuck::cast_slice(&all).to_vec()
|
||||
}
|
||||
|
||||
impl FaceGpu {
|
||||
/// Bytes of weights this model keeps on the device.
|
||||
pub fn resident_bytes(&self) -> u64 {
|
||||
self.resident.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn resident(&self, job: &ConvJob) -> Result<Option<Arc<Resident>>, String> {
|
||||
let mut layers = self.layers.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(r) = layers.get(&job.layer) {
|
||||
return Ok(Some(r.clone()));
|
||||
}
|
||||
let bytes = (job.weights.len() as u64 + job.cout as u64) * 4;
|
||||
if self.resident_bytes().saturating_add(bytes) > MAX_RESIDENT {
|
||||
return Ok(None);
|
||||
}
|
||||
let d = self.dev;
|
||||
let usage = wgpu::BufferUsages::STORAGE;
|
||||
let weights = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("faces weights"),
|
||||
contents: bytemuck::cast_slice(job.weights),
|
||||
usage,
|
||||
});
|
||||
let zero = [0f32; 1];
|
||||
let bias = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("faces bias"),
|
||||
contents: bytemuck::cast_slice(job.bias.unwrap_or(&zero)),
|
||||
usage,
|
||||
});
|
||||
let r = Arc::new(Resident { weights, bias });
|
||||
self.resident.fetch_add(bytes, Ordering::Relaxed);
|
||||
layers.insert(job.layer, r.clone());
|
||||
Ok(Some(r))
|
||||
fn upload(&self, label: &str, data: &[f32]) -> wgpu::Buffer {
|
||||
let zero = [0f32; 4];
|
||||
self.dev.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some(label),
|
||||
contents: bytemuck::cast_slice(if data.is_empty() { &zero } else { data }),
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
})
|
||||
}
|
||||
|
||||
fn run(&self, job: &ConvJob) -> Result<Option<Vec<f32>>, String> {
|
||||
/// Whether `bytes` more may stay on the device; counts them if so.
|
||||
fn keep(&self, bytes: u64) -> bool {
|
||||
if self.resident_bytes().saturating_add(bytes) > MAX_RESIDENT {
|
||||
return false;
|
||||
}
|
||||
self.resident.fetch_add(bytes, Ordering::Relaxed);
|
||||
true
|
||||
}
|
||||
|
||||
/// A convolution's weights and bias: kept on the device while the budget lasts, uploaded for this run past it.
|
||||
fn resident(&self, job: &ConvJob) -> Arc<Resident> {
|
||||
let mut layers = self.layers.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(r) = layers.get(&job.layer) {
|
||||
return r.clone();
|
||||
}
|
||||
let r = Arc::new(Resident { weights: self.upload("faces weights", job.weights), bias: self.upload("faces bias", job.bias.unwrap_or(&[])) });
|
||||
if self.keep((job.weights.len() as u64 + job.cout as u64) * 4) {
|
||||
layers.insert(job.layer, r.clone());
|
||||
}
|
||||
r
|
||||
}
|
||||
|
||||
/// A program's constant, kept on the device like the weights.
|
||||
fn constant(&self, id: usize, data: &[f32]) -> wgpu::Buffer {
|
||||
let mut consts = self.consts.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
if let Some(b) = consts.get(&id) {
|
||||
return b.clone();
|
||||
}
|
||||
let b = self.upload("faces constant", data);
|
||||
if self.keep(data.len() as u64 * 4) {
|
||||
consts.insert(id, b.clone());
|
||||
}
|
||||
b
|
||||
}
|
||||
|
||||
fn largest_binding(&self) -> u64 {
|
||||
self.dev.limits.max_storage_buffer_binding_size.min(self.dev.limits.max_buffer_size)
|
||||
}
|
||||
|
||||
/// The dispatches of one convolution reading `input` and writing `output`, or `None` when it exceeds the device.
|
||||
fn conv_dispatches(&self, job: &ConvJob, input: &wgpu::Buffer, output: &wgpu::Buffer) -> Result<Option<Vec<Dispatch>>, String> {
|
||||
let d = self.dev;
|
||||
let groups = job.groups.max(1);
|
||||
let (m, cg) = (job.cout / groups, job.cin / groups);
|
||||
@@ -264,22 +402,15 @@ impl FaceGpu {
|
||||
let kchunk = k.div_ceil(splits).div_ceil(16) * 16;
|
||||
let splits = k.div_ceil(kchunk.max(1)).max(1);
|
||||
let partial = total.checked_mul(splits).ok_or("a layer too large")?;
|
||||
let largest = [job.input.len(), job.weights.len(), partial].into_iter().max().unwrap_or(0) as u64 * 4;
|
||||
if largest > d.limits.max_storage_buffer_binding_size.min(d.limits.max_buffer_size) {
|
||||
return Ok(None);
|
||||
}
|
||||
let flat = |count: usize| {
|
||||
let blocks = count.div_ceil(256);
|
||||
let row = blocks.clamp(1, 65_535);
|
||||
(row, blocks.div_ceil(row), 1, row)
|
||||
};
|
||||
let (wx, wy, wz, row) = if direct { flat(total) } else { (n.div_ceil(64), m.div_ceil(64), groups * splits, flat(total).3) };
|
||||
let (rx, ry, _, _) = flat(total);
|
||||
let cin_len = job.cin.checked_mul(job.h).and_then(|v| v.checked_mul(job.w)).ok_or("a layer too large")?;
|
||||
let largest = [cin_len, job.weights.len(), partial].into_iter().max().unwrap_or(0) as u64 * 4;
|
||||
let limit = d.limits.max_compute_workgroups_per_dimension as usize;
|
||||
if [wx, wy, wz, rx, ry].into_iter().any(|v| v > limit) {
|
||||
let Some(sum) = flat(total.div_ceil(256), limit) else { return Ok(None) };
|
||||
let (wx, wy, wz) = if direct { (sum[0] as usize, sum[1] as usize, 1) } else { (n.div_ceil(64), m.div_ceil(64), groups * splits) };
|
||||
if largest > self.largest_binding() || [wx, wy, wz].into_iter().any(|v| v > limit) {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(res) = self.resident(job)? else { return Ok(None) };
|
||||
let res = self.resident(job);
|
||||
let params = Params {
|
||||
m: u32_of(m)?,
|
||||
n: u32_of(n)?,
|
||||
@@ -296,82 +427,93 @@ impl FaceGpu {
|
||||
pl: u32_of(job.pl)?,
|
||||
bias: u32::from(job.bias.is_some()),
|
||||
total: u32_of(total)?,
|
||||
row: u32_of(row)?,
|
||||
row: sum[2],
|
||||
splits: u32_of(splits)?,
|
||||
kchunk: u32_of(kchunk)?,
|
||||
_a: 0,
|
||||
_b: 0,
|
||||
};
|
||||
let Some(_work) = crate::exit::enter() else { return Err("the app is closing".into()) };
|
||||
let scopes = [d.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), d.device.push_error_scope(wgpu::ErrorFilter::Validation)];
|
||||
let params = bytemuck::bytes_of(¶ms).to_vec();
|
||||
let main = |out: &wgpu::Buffer| Dispatch {
|
||||
pipeline: if direct { &d.direct } else { &d.tiled },
|
||||
params: params.clone(),
|
||||
buffers: [input.clone(), res.weights.clone(), res.bias.clone(), out.clone()],
|
||||
groups: [wx as u32, wy as u32, wz as u32],
|
||||
};
|
||||
if splits == 1 {
|
||||
return Ok(Some(vec![main(output)]));
|
||||
}
|
||||
let partials = d.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces partial sums"),
|
||||
size: partial as u64 * 4,
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
Ok(Some(vec![
|
||||
main(&partials),
|
||||
Dispatch {
|
||||
pipeline: &d.reduce,
|
||||
params,
|
||||
buffers: [partials, res.weights.clone(), res.bias.clone(), output.clone()],
|
||||
groups: [sum[0], sum[1], 1],
|
||||
},
|
||||
]))
|
||||
}
|
||||
|
||||
/// Record `list` into `enc` as one compute pass, with all their parameters in one uniform buffer.
|
||||
fn record(&self, enc: &mut wgpu::CommandEncoder, list: &[Dispatch]) -> Result<(), String> {
|
||||
let d = self.dev;
|
||||
if list.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let align = (d.limits.min_uniform_buffer_offset_alignment as usize).max(256);
|
||||
let mut blob = vec![0u8; list.len() * align];
|
||||
for (i, x) in list.iter().enumerate() {
|
||||
blob.get_mut(i * align..i * align + x.params.len()).ok_or("parameters too large")?.copy_from_slice(&x.params);
|
||||
}
|
||||
let uniform = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("faces params"),
|
||||
contents: bytemuck::bytes_of(¶ms),
|
||||
contents: &blob,
|
||||
usage: wgpu::BufferUsages::UNIFORM,
|
||||
});
|
||||
let input = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("faces input"),
|
||||
contents: bytemuck::cast_slice(job.input),
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
});
|
||||
let bytes = total as u64 * 4;
|
||||
let output = d.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces output"),
|
||||
size: bytes,
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
let partials = (splits > 1).then(|| {
|
||||
d.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces partial sums"),
|
||||
size: partial as u64 * 4,
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
mapped_at_creation: false,
|
||||
let binds = list
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, x)| {
|
||||
let [a, b, c, out] = &x.buffers;
|
||||
d.device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: None,
|
||||
layout: &d.layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry {
|
||||
binding: 0,
|
||||
resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
|
||||
buffer: &uniform,
|
||||
offset: (i * align) as u64,
|
||||
size: std::num::NonZeroU64::new(x.params.len() as u64),
|
||||
}),
|
||||
},
|
||||
wgpu::BindGroupEntry { binding: 1, resource: a.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 2, resource: b.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 3, resource: c.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 4, resource: out.as_entire_binding() },
|
||||
],
|
||||
})
|
||||
})
|
||||
});
|
||||
let staging = d.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces read-back"),
|
||||
size: bytes,
|
||||
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
|
||||
mapped_at_creation: false,
|
||||
});
|
||||
let bind = d.device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: None,
|
||||
layout: &d.layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 1, resource: input.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 2, resource: res.weights.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 3, resource: res.bias.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 4, resource: partials.as_ref().unwrap_or(&output).as_entire_binding() },
|
||||
],
|
||||
});
|
||||
let summing = partials.as_ref().map(|partials| {
|
||||
d.device.create_bind_group(&wgpu::BindGroupDescriptor {
|
||||
label: None,
|
||||
layout: &d.layout,
|
||||
entries: &[
|
||||
wgpu::BindGroupEntry { binding: 0, resource: uniform.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 1, resource: partials.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 2, resource: res.weights.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 3, resource: res.bias.as_entire_binding() },
|
||||
wgpu::BindGroupEntry { binding: 4, resource: output.as_entire_binding() },
|
||||
],
|
||||
})
|
||||
});
|
||||
let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
|
||||
{
|
||||
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("faces conv"), timestamp_writes: None });
|
||||
pass.set_pipeline(if direct { &d.direct } else { &d.tiled });
|
||||
pass.set_bind_group(0, &bind, &[]);
|
||||
pass.dispatch_workgroups(u32_of(wx)?, u32_of(wy)?, u32_of(wz)?);
|
||||
if let Some(summing) = &summing {
|
||||
pass.set_pipeline(&d.reduce);
|
||||
pass.set_bind_group(0, summing, &[]);
|
||||
pass.dispatch_workgroups(u32_of(rx)?, u32_of(ry)?, 1);
|
||||
}
|
||||
.collect::<Vec<_>>();
|
||||
let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: Some("faces"), timestamp_writes: None });
|
||||
for (x, bind) in list.iter().zip(&binds) {
|
||||
pass.set_pipeline(x.pipeline);
|
||||
pass.set_bind_group(0, bind, &[]);
|
||||
let [gx, gy, gz] = x.groups;
|
||||
pass.dispatch_workgroups(gx, gy, gz);
|
||||
}
|
||||
enc.copy_buffer_to_buffer(&output, 0, &staging, 0, bytes);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Submit `enc`, which ends by copying into `staging` (`bytes` long), wait for it and read `staging` back.
|
||||
fn finish(&self, enc: wgpu::CommandEncoder, staging: &wgpu::Buffer, scopes: [wgpu::ErrorScopeGuard; 2]) -> Result<Vec<f32>, String> {
|
||||
let d = self.dev;
|
||||
let index = d.queue.submit([enc.finish()]);
|
||||
let slice = staging.slice(..);
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
@@ -389,7 +531,7 @@ impl FaceGpu {
|
||||
return Err(format!("the GPU reported an error: {e}"));
|
||||
}
|
||||
if let Err(e) = waited {
|
||||
mark_broken(&format!("the GPU did not finish a convolution: {e}"));
|
||||
mark_broken(&format!("the GPU did not finish a face model: {e}"));
|
||||
return Err(format!("the GPU did not finish: {e}"));
|
||||
}
|
||||
match rx.recv_timeout(Duration::from_secs(10)) {
|
||||
@@ -402,8 +544,244 @@ impl FaceGpu {
|
||||
data.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect::<Vec<f32>>()
|
||||
};
|
||||
staging.unmap();
|
||||
Ok(Some(out))
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn scopes(&self) -> [wgpu::ErrorScopeGuard; 2] {
|
||||
[self.dev.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory), self.dev.device.push_error_scope(wgpu::ErrorFilter::Validation)]
|
||||
}
|
||||
|
||||
fn staging(&self, bytes: u64) -> wgpu::Buffer {
|
||||
self.dev.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces read-back"),
|
||||
size: bytes.max(4),
|
||||
usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
|
||||
mapped_at_creation: false,
|
||||
})
|
||||
}
|
||||
|
||||
fn scratch(&self, bytes: u64) -> wgpu::Buffer {
|
||||
self.dev.device.create_buffer(&wgpu::BufferDescriptor {
|
||||
label: Some("faces values"),
|
||||
size: bytes.max(4),
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||||
mapped_at_creation: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// One convolution, uploaded and read back on its own.
|
||||
fn run(&self, job: &ConvJob) -> Result<Option<Vec<f32>>, String> {
|
||||
let total = job.cout.checked_mul(job.oh).and_then(|v| v.checked_mul(job.ow)).ok_or("a layer too large")?;
|
||||
let bytes = total as u64 * 4;
|
||||
if bytes > self.largest_binding() || job.input.len() as u64 * 4 > self.largest_binding() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(_work) = crate::exit::enter() else { return Err("the app is closing".into()) };
|
||||
let _turn = Turn::take();
|
||||
let d = self.dev;
|
||||
let scopes = self.scopes();
|
||||
let input = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("faces input"),
|
||||
contents: bytemuck::cast_slice(job.input),
|
||||
usage: wgpu::BufferUsages::STORAGE,
|
||||
});
|
||||
let output = self.scratch(bytes);
|
||||
let Some(list) = self.conv_dispatches(job, &input, &output)? else {
|
||||
// leave the scopes balanced
|
||||
let [oom, validation] = scopes;
|
||||
let _ = (pollster::block_on(validation.pop()), pollster::block_on(oom.pop()));
|
||||
return Ok(None);
|
||||
};
|
||||
let staging = self.staging(bytes);
|
||||
let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
|
||||
self.record(&mut enc, &list)?;
|
||||
enc.copy_buffer_to_buffer(&output, 0, &staging, 0, bytes);
|
||||
self.finish(enc, &staging, scopes).map(Some)
|
||||
}
|
||||
|
||||
/// The dispatch of one non-convolution step.
|
||||
fn op_dispatch(&self, op: &Op, out_len: usize, args: [wgpu::Buffer; 3], out: &wgpu::Buffer) -> Result<Option<Dispatch>, String> {
|
||||
let d = self.dev;
|
||||
let limit = d.limits.max_compute_workgroups_per_dimension as usize;
|
||||
let u = u32_of;
|
||||
let n = u(out_len)?;
|
||||
let per_value = |pipeline: &'static wgpu::ComputePipeline, rest: &[u32]| -> Option<Dispatch> {
|
||||
let groups = flat(out_len.div_ceil(256), limit)?;
|
||||
let mut all = vec![n, groups[2]];
|
||||
all.extend_from_slice(rest);
|
||||
Some(Dispatch {
|
||||
pipeline,
|
||||
params: params24(&all),
|
||||
buffers: [args[0].clone(), args[1].clone(), args[2].clone(), out.clone()],
|
||||
groups: [groups[0], groups[1], 1],
|
||||
})
|
||||
};
|
||||
let per_group = |pipeline: &'static wgpu::ComputePipeline, rest: &[u32]| -> Option<Dispatch> {
|
||||
let groups = flat(out_len, limit)?;
|
||||
let mut all = vec![n, groups[2]];
|
||||
all.extend_from_slice(rest);
|
||||
Some(Dispatch {
|
||||
pipeline,
|
||||
params: params24(&all),
|
||||
buffers: [args[0].clone(), args[1].clone(), args[2].clone(), out.clone()],
|
||||
groups: [groups[0], groups[1], 1],
|
||||
})
|
||||
};
|
||||
let four = |v: &[usize; 4]| -> Result<[u32; 4], String> { Ok([u(v[0])?, u(v[1])?, u(v[2])?, u(v[3])?]) };
|
||||
let k = &d.ops;
|
||||
Ok(match op {
|
||||
Op::Conv(_) => None,
|
||||
Op::Unary(f) => per_value(&k.unary, &[u32::from(*f == Unary::Sigmoid)]),
|
||||
Op::ChannelAffine { channels, plane } => per_value(&k.affine, &[u(*plane)?, u(*channels)?]),
|
||||
Op::ChannelPrelu { channels, plane } => per_value(&k.prelu, &[u(*plane)?, u(*channels)?]),
|
||||
Op::Binary { op, dims, a, b } => {
|
||||
let kind = match op {
|
||||
Binary::Add => 0,
|
||||
Binary::Sub => 1,
|
||||
Binary::Mul => 2,
|
||||
};
|
||||
let mut rest = vec![kind];
|
||||
rest.extend(four(dims)?);
|
||||
rest.extend(four(a)?);
|
||||
rest.extend(four(b)?);
|
||||
per_value(&k.binary, &rest)
|
||||
}
|
||||
Op::Transpose { dims, strides } => {
|
||||
let mut rest = vec![0];
|
||||
rest.extend(four(dims)?);
|
||||
rest.extend(four(strides)?);
|
||||
per_value(&k.transpose, &rest)
|
||||
}
|
||||
Op::MaxPool { h, w, kh, kw, sh, sw, oh, ow, .. } => {
|
||||
per_value(&k.maxpool, &[u(*h)?, u(*w)?, u(*kh)?, u(*kw)?, u(*sh)?, u(*sw)?, u(*oh)?, u(*ow)?])
|
||||
}
|
||||
Op::Resize { h, w, sh, sw, .. } => per_value(&k.resize, &[u(*h)?, u(*w)?, u(*sh)?, u(*sw)?]),
|
||||
Op::GlobalAverage { plane, .. } => per_group(&k.gavg, &[u(*plane)?]),
|
||||
Op::Gemv { k: kk, n: nn, trans_b, alpha } => per_group(&k.gemv, &[u(*kk)?, u(*nn)?, u32::from(*trans_b), alpha.to_bits()]),
|
||||
})
|
||||
}
|
||||
|
||||
/// A whole program for each input, every value staying on the device; one submission and one read-back.
|
||||
fn program(&self, prog: &Program, inputs: &[&[f32]]) -> Result<Option<Vec<Vec<Vec<f32>>>>, String> {
|
||||
let d = self.dev;
|
||||
let largest = self.largest_binding();
|
||||
if prog.lens.iter().any(|l| *l as u64 * 4 > largest) || inputs.iter().any(|i| Some(&i.len()) != prog.lens.get(INPUT)) {
|
||||
return Ok(None);
|
||||
}
|
||||
let out_bytes: u64 = prog.outputs.iter().map(|(_, v, _)| prog.lens.get(*v).copied().unwrap_or(0) as u64 * 4).sum();
|
||||
let total_bytes = out_bytes.checked_mul(inputs.len() as u64).ok_or("too many faces at once")?;
|
||||
if total_bytes > d.limits.max_buffer_size {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(_work) = crate::exit::enter() else { return Err("the app is closing".into()) };
|
||||
let _turn = Turn::take();
|
||||
let scopes = self.scopes();
|
||||
let staging = self.staging(total_bytes);
|
||||
let mut enc = d.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("faces program") });
|
||||
// buffers nothing reads any more, reused for later values (best fit)
|
||||
let mut pool: Vec<(u64, wgpu::Buffer)> = Vec::new();
|
||||
let take = |pool: &mut Vec<(u64, wgpu::Buffer)>, bytes: u64| -> wgpu::Buffer {
|
||||
let best = pool
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, (size, _))| *size >= bytes && *size <= bytes.saturating_mul(2))
|
||||
.min_by_key(|(_, (size, _))| *size)
|
||||
.map(|(i, _)| i);
|
||||
match best {
|
||||
Some(i) => pool.swap_remove(i).1,
|
||||
None => self.scratch(bytes),
|
||||
}
|
||||
};
|
||||
let mut at = 0u64;
|
||||
let mut fits = true;
|
||||
for input in inputs {
|
||||
let mut values: Vec<Option<(u64, wgpu::Buffer)>> = vec![None; prog.lens.len()];
|
||||
// a fresh buffer: queued writes land before any of the encoder's work, so a pooled one could still be in use
|
||||
let first = d.device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
|
||||
label: Some("faces input"),
|
||||
contents: bytemuck::cast_slice(input),
|
||||
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
|
||||
});
|
||||
if let Some(v) = values.get_mut(INPUT) {
|
||||
*v = Some((input.len() as u64 * 4, first));
|
||||
}
|
||||
let mut list = Vec::new();
|
||||
for step in &prog.steps {
|
||||
let len = prog.lens.get(step.out).copied().ok_or("a step writes a value the program does not have")?;
|
||||
let out = take(&mut pool, len as u64 * 4);
|
||||
let mut args = Vec::with_capacity(3);
|
||||
for a in &step.args {
|
||||
args.push(match a {
|
||||
Arg::Value(v) => {
|
||||
values.get(*v).and_then(|v| v.as_ref()).map(|(_, b)| b.clone()).ok_or("a step reads a value before it is made")?
|
||||
}
|
||||
Arg::Const { id, data } => self.constant(*id, data),
|
||||
});
|
||||
}
|
||||
let dispatches = match &step.op {
|
||||
Op::Conv(job) => {
|
||||
let input = args.first().ok_or("a convolution without input")?;
|
||||
self.conv_dispatches(job, input, &out)?
|
||||
}
|
||||
op => {
|
||||
let arg = |i: usize| args.get(i).cloned().unwrap_or_else(|| d.dummy.clone());
|
||||
self.op_dispatch(op, len, [arg(0), arg(1), arg(2)], &out)?.map(|x| vec![x])
|
||||
}
|
||||
};
|
||||
let Some(dispatches) = dispatches else {
|
||||
fits = false;
|
||||
break;
|
||||
};
|
||||
list.extend(dispatches);
|
||||
if let Some(v) = values.get_mut(step.out) {
|
||||
*v = Some((len as u64 * 4, out));
|
||||
}
|
||||
for f in &step.free {
|
||||
if let Some(b) = values.get_mut(*f).and_then(Option::take) {
|
||||
pool.push(b);
|
||||
}
|
||||
}
|
||||
}
|
||||
if !fits {
|
||||
break;
|
||||
}
|
||||
self.record(&mut enc, &list)?;
|
||||
for (_, v, _) in &prog.outputs {
|
||||
let (bytes, buffer) = values.get(*v).and_then(|v| v.as_ref()).ok_or("an output was never made")?;
|
||||
let len = prog.lens.get(*v).copied().unwrap_or(0) as u64 * 4;
|
||||
enc.copy_buffer_to_buffer(buffer, 0, &staging, at, len.min(*bytes));
|
||||
at += len;
|
||||
}
|
||||
pool.extend(values.into_iter().flatten());
|
||||
}
|
||||
if !fits {
|
||||
let [oom, validation] = scopes;
|
||||
let _ = (pollster::block_on(validation.pop()), pollster::block_on(oom.pop()));
|
||||
return Ok(None);
|
||||
}
|
||||
let flat_out = self.finish(enc, &staging, scopes)?;
|
||||
let mut rest = flat_out.as_slice();
|
||||
let mut per_input = Vec::with_capacity(inputs.len());
|
||||
for _ in inputs {
|
||||
let mut outs = Vec::with_capacity(prog.outputs.len());
|
||||
for (_, v, _) in &prog.outputs {
|
||||
let len = prog.lens.get(*v).copied().unwrap_or(0);
|
||||
let (this, next) = rest.split_at_checked(len).ok_or("the read-back is short")?;
|
||||
outs.push(this.to_vec());
|
||||
rest = next;
|
||||
}
|
||||
per_input.push(outs);
|
||||
}
|
||||
Ok(Some(per_input))
|
||||
}
|
||||
}
|
||||
|
||||
/// Run `f`, turning a panic into a broken device.
|
||||
fn guarded<T>(f: impl FnOnce() -> Result<T, String>) -> Result<T, String> {
|
||||
std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)).unwrap_or_else(|_| {
|
||||
mark_broken("a face model on the GPU panicked");
|
||||
Err("the GPU runner gave up".into())
|
||||
})
|
||||
}
|
||||
|
||||
impl Accelerator for FaceGpu {
|
||||
@@ -418,13 +796,23 @@ impl Accelerator for FaceGpu {
|
||||
if crate::switched_off().is_some() || job.macs() < MIN_MACS {
|
||||
return Ok(None);
|
||||
}
|
||||
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| self.run(job))).unwrap_or_else(|_| {
|
||||
mark_broken("a face convolution panicked");
|
||||
Err("the GPU runner gave up".into())
|
||||
})
|
||||
guarded(|| self.run(job))
|
||||
}
|
||||
|
||||
fn run_program(&self, program: &Program, inputs: &[&[f32]]) -> Result<Option<Vec<Vec<Vec<f32>>>>, String> {
|
||||
if let Some(why) = broken_reason() {
|
||||
return Err(format!("the GPU is not used any more: {why}"));
|
||||
}
|
||||
if crate::switched_off().is_some() || PER_CONV.load(Ordering::Relaxed) {
|
||||
return Ok(None);
|
||||
}
|
||||
guarded(|| self.program(program, inputs))
|
||||
}
|
||||
}
|
||||
|
||||
/// Tests and benchmarks only: run convolutions one by one (the first step's way) instead of whole programs.
|
||||
static PER_CONV: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -513,6 +901,28 @@ mod tests {
|
||||
assert!(gpu.resident_bytes() > 0, "the weights stay on the device");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_operator_on_the_gpu_gives_what_the_cpu_gives() {
|
||||
let Some(gpu) = gpu() else { return };
|
||||
let net = lightcraft_faces::net::Net::new(lightcraft_faces::synthetic::every_operator_graph()).unwrap();
|
||||
let inputs: Vec<lightcraft_faces::net::Tensor> = (0..3)
|
||||
.map(|s| {
|
||||
lightcraft_faces::net::Tensor::new(vec![1, 4, 8, 8], (0..256).map(|i| (((i * 31 + s * 7) % 37) as f32 - 18.0) * 0.05).collect())
|
||||
.unwrap()
|
||||
})
|
||||
.collect();
|
||||
let cpu: Vec<_> = inputs.iter().map(|x| net.run(x.clone()).unwrap()).collect();
|
||||
let program = net.program(&[1, 4, 8, 8]).unwrap().unwrap();
|
||||
let slices: Vec<&[f32]> = inputs.iter().map(|t| t.data.as_slice()).collect();
|
||||
let got = gpu.program(&program, &slices).unwrap().unwrap();
|
||||
for (want, got) in cpu.iter().zip(&got) {
|
||||
for ((name, want), got) in want.iter().zip(got) {
|
||||
let worst = want.data.iter().zip(got).fold(0f32, |m, (a, b)| m.max((a - b).abs()));
|
||||
assert!(worst < 1e-4, "{name}: off by {worst}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The real models on the GPU agree with the CPU interpreter: `LC_FACE_MODELS=<folder>` holding
|
||||
/// `face_detection_yunet_2023mar.onnx`, `face_recognition_sface_2021dec.onnx` and (optionally) `glintr100.onnx`.
|
||||
#[test]
|
||||
@@ -555,7 +965,8 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
/// Timings, CPU against GPU: `LC_FACE_MODELS=<folder> cargo test --release -p lightcraft-gpu faces::tests::timings -- --ignored --nocapture`.
|
||||
/// Timings, CPU against the GPU one convolution at a time and as whole programs (one face, and a batch of 16):
|
||||
/// `LC_FACE_MODELS=<folder> cargo test --release -p lightcraft-gpu faces::tests::timings -- --ignored --nocapture`.
|
||||
#[test]
|
||||
#[ignore = "a benchmark: needs LC_FACE_MODELS and a GPU"]
|
||||
fn timings() {
|
||||
@@ -573,21 +984,35 @@ mod tests {
|
||||
let (w, h) = (1024, 768);
|
||||
let rgb: Vec<u8> = (0..w * h * 3).map(|i| (i * 7 % 251) as u8).collect();
|
||||
let opts = lightcraft_faces::yunet::Options::default();
|
||||
let cpu = time(5, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap()));
|
||||
let cpu = time(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap()));
|
||||
PER_CONV.store(true, Ordering::Relaxed);
|
||||
det.accelerate(accelerator().unwrap(), false).unwrap();
|
||||
let fast = time(5, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap()));
|
||||
eprintln!("YuNet 640: CPU {cpu:.1} ms, GPU {fast:.1} ms");
|
||||
let per_conv = time(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap()));
|
||||
PER_CONV.store(false, Ordering::Relaxed);
|
||||
det.accelerate(accelerator().unwrap(), false).unwrap();
|
||||
let whole = time(10, &mut || drop(det.detect(&rgb, w, h, &opts).unwrap()));
|
||||
assert!(det.net().accelerator().is_some());
|
||||
eprintln!("YuNet 640: CPU {cpu:.1} ms, GPU per convolution {per_conv:.1} ms, GPU whole network {whole:.1} ms");
|
||||
for (file, sha) in [
|
||||
("face_recognition_sface_2021dec.onnx", lightcraft_faces::known::SFACE_SHA256),
|
||||
("glintr100.onnx", lightcraft_faces::known::AURAFACE_SHA256),
|
||||
] {
|
||||
let e = lightcraft_faces::runtime::Embedder::load(&dir.join(file), &lightcraft_faces::known::lookup(sha).unwrap()).unwrap();
|
||||
let (ew, eh) = e.input_size();
|
||||
let face: Vec<u8> = (0..ew * eh * 3).map(|i| (i * 13 % 251) as u8).collect();
|
||||
let cpu = time(3, &mut || drop(e.embed(&face).unwrap()));
|
||||
let faces: Vec<Vec<u8>> = (0..16).map(|s| (0..ew * eh * 3).map(|i| ((i + s * 101) * 13 % 251) as u8).collect()).collect();
|
||||
let refs: Vec<&[u8]> = faces.iter().map(Vec::as_slice).collect();
|
||||
let cpu = time(5, &mut || drop(e.embed(&faces[0]).unwrap()));
|
||||
PER_CONV.store(true, Ordering::Relaxed);
|
||||
e.accelerate(accelerator().unwrap(), false).unwrap();
|
||||
let fast = time(3, &mut || drop(e.embed(&face).unwrap()));
|
||||
eprintln!("{file}: CPU {cpu:.1} ms, GPU {fast:.1} ms");
|
||||
let per_conv = time(5, &mut || drop(e.embed(&faces[0]).unwrap()));
|
||||
PER_CONV.store(false, Ordering::Relaxed);
|
||||
e.accelerate(accelerator().unwrap(), false).unwrap();
|
||||
let whole = time(5, &mut || drop(e.embed(&faces[0]).unwrap()));
|
||||
let batch = time(3, &mut || drop(e.embed_batch(&refs))) / 16.0;
|
||||
assert!(e.net().accelerator().is_some());
|
||||
eprintln!(
|
||||
"{file}: CPU {cpu:.1} ms, GPU per convolution {per_conv:.1} ms, GPU whole network {whole:.1} ms, batch of 16 {batch:.1} ms a face"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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];
|
||||
}
|
||||
}
|
||||
@@ -2118,5 +2118,13 @@
|
||||
"Adding to Apple Photos": "Wird zu Apple Fotos hinzugefügt",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Kein Grafikprozessor gefunden: LightCraft rendert auf der CPU, was viel langsamer ist. Details unter Hilfe → Systeminformationen.",
|
||||
"On NixOS, start it from `nix develop`.": "Unter NixOS starten Sie es aus `nix develop`.",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Dateien mit demselben Inhalt wie eine andere Datei dieses Imports: Nur die erste wird importiert."
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Dateien mit demselben Inhalt wie eine andere Datei dieses Imports: Nur die erste wird importiert.",
|
||||
"Run face models on the CPU only": "Gesichtsmodelle nur auf der CPU ausführen",
|
||||
"Detection": "Erkennung",
|
||||
"Recognition": "Wiedererkennung",
|
||||
"Not loaded yet": "Noch nicht geladen",
|
||||
"Checking the GPU…": "GPU wird geprüft …",
|
||||
"CPU": "CPU",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "Ein Gesichtsmodell nutzt die GPU nur, wenn eine Prüfung zeigt, dass sie dieselben Ergebnisse schneller liefert.",
|
||||
"Checking the GPU for faces": "GPU für Gesichter wird geprüft"
|
||||
}
|
||||
|
||||
@@ -1995,5 +1995,13 @@
|
||||
"Adding to Apple Photos": "Añadiendo a Fotos de Apple",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "No se encontró ningún procesador gráfico: LightCraft renderiza en la CPU, lo que es mucho más lento. Detalles en Ayuda → Información del sistema.",
|
||||
"On NixOS, start it from `nix develop`.": "En NixOS, inícielo desde `nix develop`.",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Archivos con el mismo contenido que otro archivo de esta importación: solo se importa el primero."
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Archivos con el mismo contenido que otro archivo de esta importación: solo se importa el primero.",
|
||||
"Run face models on the CPU only": "Ejecutar los modelos de caras solo en la CPU",
|
||||
"Detection": "Detección",
|
||||
"Recognition": "Reconocimiento",
|
||||
"Not loaded yet": "Aún no cargado",
|
||||
"Checking the GPU…": "Comprobando la GPU…",
|
||||
"CPU": "CPU",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "Un modelo de caras usa la GPU solo cuando una comprobación muestra que da los mismos resultados más rápido.",
|
||||
"Checking the GPU for faces": "Comprobando la GPU para las caras"
|
||||
}
|
||||
|
||||
@@ -1957,5 +1957,13 @@
|
||||
"Adding to Apple Photos": "Ajout à Apple Photos",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Aucun processeur graphique trouvé : LightCraft effectue le rendu sur le processeur, ce qui est beaucoup plus lent. Détails dans Aide → Informations système.",
|
||||
"On NixOS, start it from `nix develop`.": "Sous NixOS, lancez-le depuis `nix develop`.",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Fichiers au contenu identique à un autre fichier de cette importation : seul le premier est importé."
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Fichiers au contenu identique à un autre fichier de cette importation : seul le premier est importé.",
|
||||
"Run face models on the CPU only": "Exécuter les modèles de visages sur le processeur uniquement",
|
||||
"Detection": "Détection",
|
||||
"Recognition": "Reconnaissance",
|
||||
"Not loaded yet": "Pas encore chargé",
|
||||
"Checking the GPU…": "Vérification du GPU…",
|
||||
"CPU": "Processeur",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "Un modèle de visages n’utilise le GPU que si une vérification montre qu’il donne les mêmes résultats plus vite.",
|
||||
"Checking the GPU for faces": "Vérification du GPU pour les visages"
|
||||
}
|
||||
|
||||
@@ -1961,5 +1961,13 @@
|
||||
"Adding to Apple Photos": "Appleの写真に追加中",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "グラフィックプロセッサが見つかりません。LightCraftはCPUで描画するため、かなり遅くなります。詳細は「ヘルプ → システム情報」を参照してください。",
|
||||
"On NixOS, start it from `nix develop`.": "NixOSでは `nix develop` から起動してください。",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "この読み込み内の別のファイルと内容が同じファイルです。最初のファイルだけが読み込まれます。"
|
||||
"Files with the same content as another file in this import: only the first is imported.": "この読み込み内の別のファイルと内容が同じファイルです。最初のファイルだけが読み込まれます。",
|
||||
"Run face models on the CPU only": "顔モデルを CPU だけで実行",
|
||||
"Detection": "検出",
|
||||
"Recognition": "認識",
|
||||
"Not loaded yet": "未読み込み",
|
||||
"Checking the GPU…": "GPU を確認中…",
|
||||
"CPU": "CPU",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "顔モデルが GPU を使うのは、同じ結果をより速く出せると確認できた場合だけです。",
|
||||
"Checking the GPU for faces": "顔用の GPU を確認中"
|
||||
}
|
||||
|
||||
@@ -1998,5 +1998,13 @@
|
||||
"Recognition models": "Modelos de reconhecimento",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Nenhum processador gráfico encontrado: o LightCraft renderiza na CPU, o que é muito mais lento. Detalhes em Ajuda → Informações do sistema.",
|
||||
"On NixOS, start it from `nix develop`.": "No NixOS, inicie-o a partir de `nix develop`.",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Arquivos com o mesmo conteúdo de outro arquivo desta importação: apenas o primeiro é importado."
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Arquivos com o mesmo conteúdo de outro arquivo desta importação: apenas o primeiro é importado.",
|
||||
"Run face models on the CPU only": "Executar os modelos de rostos só na CPU",
|
||||
"Detection": "Detecção",
|
||||
"Recognition": "Reconhecimento",
|
||||
"Not loaded yet": "Ainda não carregado",
|
||||
"Checking the GPU…": "Verificando a GPU…",
|
||||
"CPU": "CPU",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "Um modelo de rostos só usa a GPU quando uma verificação mostra que ela dá os mesmos resultados mais rápido.",
|
||||
"Checking the GPU for faces": "Verificando a GPU para rostos"
|
||||
}
|
||||
|
||||
@@ -2151,5 +2151,13 @@
|
||||
"Adding to Apple Photos": "Добавление в Apple Фото",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Графический процессор не найден: LightCraft выполняет рендеринг на CPU, а это намного медленнее. Подробности: Справка → Сведения о системе.",
|
||||
"On NixOS, start it from `nix develop`.": "В NixOS запускайте его из `nix develop`.",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Файлы с тем же содержимым, что и другой файл этого импорта: импортируется только первый."
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Файлы с тем же содержимым, что и другой файл этого импорта: импортируется только первый.",
|
||||
"Run face models on the CPU only": "Запускать модели лиц только на ЦП",
|
||||
"Detection": "Обнаружение",
|
||||
"Recognition": "Распознавание",
|
||||
"Not loaded yet": "Ещё не загружена",
|
||||
"Checking the GPU…": "Проверка ГП…",
|
||||
"CPU": "ЦП",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "Модель лиц использует ГП, только если проверка показала, что он даёт те же результаты быстрее.",
|
||||
"Checking the GPU for faces": "Проверка ГП для лиц"
|
||||
}
|
||||
|
||||
@@ -2151,5 +2151,13 @@
|
||||
"Adding to Apple Photos": "Додавання до Apple Фото",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "Графічний процесор не знайдено: LightCraft виконує рендеринг на ЦП, а це набагато повільніше. Докладніше: Довідка → Інформація про систему.",
|
||||
"On NixOS, start it from `nix develop`.": "У NixOS запускайте його з `nix develop`.",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Файли з тим самим вмістом, що й інший файл цього імпорту: імпортується лише перший."
|
||||
"Files with the same content as another file in this import: only the first is imported.": "Файли з тим самим вмістом, що й інший файл цього імпорту: імпортується лише перший.",
|
||||
"Run face models on the CPU only": "Запускати моделі облич лише на ЦП",
|
||||
"Detection": "Виявлення",
|
||||
"Recognition": "Розпізнавання",
|
||||
"Not loaded yet": "Ще не завантажено",
|
||||
"Checking the GPU…": "Перевірка ГП…",
|
||||
"CPU": "ЦП",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "Модель облич використовує ГП, лише якщо перевірка показала, що він дає ті самі результати швидше.",
|
||||
"Checking the GPU for faces": "Перевірка ГП для облич"
|
||||
}
|
||||
|
||||
@@ -1997,5 +1997,13 @@
|
||||
"Adding to Apple Photos": "正在添加到 Apple 照片",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "未找到图形处理器：LightCraft 改用 CPU 渲染，速度会慢很多。详情见“帮助 → 系统信息”。",
|
||||
"On NixOS, start it from `nix develop`.": "在 NixOS 上，请从 `nix develop` 启动。",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "与本次导入中另一个文件内容相同的文件：只导入第一个。"
|
||||
"Files with the same content as another file in this import: only the first is imported.": "与本次导入中另一个文件内容相同的文件：只导入第一个。",
|
||||
"Run face models on the CPU only": "仅在 CPU 上运行人脸模型",
|
||||
"Detection": "检测",
|
||||
"Recognition": "识别",
|
||||
"Not loaded yet": "尚未加载",
|
||||
"Checking the GPU…": "正在检查 GPU…",
|
||||
"CPU": "CPU",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "只有检查表明 GPU 能更快给出相同结果时，人脸模型才会使用 GPU。",
|
||||
"Checking the GPU for faces": "正在检查用于人脸的 GPU"
|
||||
}
|
||||
|
||||
@@ -1959,5 +1959,13 @@
|
||||
"Adding to Apple Photos": "正在加入 Apple 照片",
|
||||
"No graphics processor found: LightCraft renders on the CPU, which is much slower. Details in Help → System Info.": "找不到圖形處理器：LightCraft 改用 CPU 算繪，速度會慢很多。詳情請見「說明 → 系統資訊」。",
|
||||
"On NixOS, start it from `nix develop`.": "在 NixOS 上，請從 `nix develop` 啟動。",
|
||||
"Files with the same content as another file in this import: only the first is imported.": "與本次匯入中另一個檔案內容相同的檔案：只匯入第一個。"
|
||||
"Files with the same content as another file in this import: only the first is imported.": "與本次匯入中另一個檔案內容相同的檔案：只匯入第一個。",
|
||||
"Run face models on the CPU only": "僅在 CPU 上執行臉孔模型",
|
||||
"Detection": "偵測",
|
||||
"Recognition": "辨識",
|
||||
"Not loaded yet": "尚未載入",
|
||||
"Checking the GPU…": "正在檢查 GPU…",
|
||||
"CPU": "CPU",
|
||||
"A face model uses the GPU only when a check shows it gives the same results faster.": "只有在檢查顯示 GPU 能更快得出相同結果時，臉孔模型才會使用 GPU。",
|
||||
"Checking the GPU for faces": "正在檢查用於臉孔的 GPU"
|
||||
}
|
||||
|
||||
@@ -157,6 +157,42 @@ pub fn setup_banner(app: &mut LightcraftApp, ui: &mut egui::Ui) {
|
||||
});
|
||||
}
|
||||
|
||||
/// Where the face models run (the GPU only once a check shows it gives the CPU's results faster), and the switch that
|
||||
/// keeps them on the CPU.
|
||||
fn gpu_section(app: &mut LccApp, ui: &mut egui::Ui, t: &Tokens, gpu: &Value) {
|
||||
heading(ui, t, "Processor");
|
||||
let mut cpu = gpu["cpuOnly"].as_bool().unwrap_or(false);
|
||||
let r = ui.checkbox(&mut cpu, crate::i18n::tr("Run face models on the CPU only"));
|
||||
register(ui.ctx(), "check:settings.facesCpuOnly", r.rect);
|
||||
if r.changed() {
|
||||
match app.run("faces.gpu", json!({"cpuOnly": cpu})) {
|
||||
Ok(_) => app.caches.faces_epoch += 1,
|
||||
Err(e) => app.toast(ui.ctx(), e),
|
||||
}
|
||||
}
|
||||
for (what, label) in [("detection", "Detection"), ("recognition", "Recognition")] {
|
||||
let st = &gpu[what];
|
||||
let (text, why) = if st.is_null() {
|
||||
(crate::i18n::tr("Not loaded yet").to_string(), None)
|
||||
} else if st["checking"] == true {
|
||||
(crate::i18n::tr("Checking the GPU…").to_string(), None)
|
||||
} else if let Some(device) = st["device"].as_str() {
|
||||
(device.to_string(), None)
|
||||
} else {
|
||||
(crate::i18n::tr("CPU").to_string(), st["why"].as_str())
|
||||
};
|
||||
ui.horizontal(|ui| {
|
||||
ui.label(RichText::new(crate::i18n::tr(label)).color(t.text_dim));
|
||||
let r = ui.label(text);
|
||||
register(ui.ctx(), format!("faces:device:{what}"), r.rect);
|
||||
if let Some(why) = why {
|
||||
r.on_hover_text(sentence(why));
|
||||
}
|
||||
});
|
||||
}
|
||||
hint(ui, t, "A face model uses the GPU only when a check shows it gives the same results faster.");
|
||||
}
|
||||
|
||||
/// A download-style size: decimal megabytes, as file managers and model pages show them.
|
||||
fn mb(bytes: Option<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);
|
||||
|
||||
+29
-20
@@ -145,31 +145,40 @@ names alone do not guarantee the same speed, and those additional model files we
|
||||
|
||||
### On the GPU
|
||||
|
||||
On a machine with a usable GPU, the convolutions of the detector and the recogniser (most of their time) run on it
|
||||
through wgpu (`crates/gpu/src/faces.rs`, `wgsl/faces_conv.wgsl`: an implicit matrix product that gathers the input
|
||||
straight into workgroup memory, cuts K into slices when a layer's map is small, and a direct kernel for depthwise
|
||||
layers). Everything else in the graph stays on the CPU interpreter, which remains the reference. After a model loads,
|
||||
a background check runs it once on the CPU and twice through the GPU on a fixed test input: the GPU is used only when
|
||||
On a machine with a usable GPU, the whole detector and recogniser run on it through wgpu (`crates/gpu/src/faces.rs`).
|
||||
The interpreter lowers the graph to a flat program for the input's shape (`crates/faces/src/net/program.rs`): each
|
||||
convolution (`wgsl/faces_conv.wgsl`: an implicit matrix product that gathers the input straight into workgroup memory,
|
||||
cuts K into slices when a map is small, and a direct kernel for depthwise layers) and every other operator (activations,
|
||||
batch norm, PRelu, broadcasting add/sub/mul, max pooling, nearest upscaling, transposes, global pooling and the fully
|
||||
connected layer, `wgsl/faces_ops.wgsl`) is a kernel, every value stays on the device between layers in reused
|
||||
buffers, and a run is one upload, one submission and one read-back. Shape-only operators make no step. A model that
|
||||
needs something the program cannot express runs its convolutions one by one on the GPU instead. Several faces of one
|
||||
photo go together (`Embedder::embed_batch`, up to 16 per submission).
|
||||
|
||||
The CPU interpreter remains the reference. After a model loads, a background check (a row in the activity stack, "Checking
|
||||
the GPU for faces") runs it once on the CPU and twice through the GPU on a fixed test input: the GPU is used only when
|
||||
every output is within 1e-3 of the CPU's (relative to the output's largest value) **and** it was faster; until then,
|
||||
and whenever the GPU is switched off (`LIGHTCRAFT_GPU=0`, `LIGHTCRAFT_GPU_BACKEND=off`, the GPU preference), missing,
|
||||
wrong or slower, the model runs on the CPU. The first GPU error drops it for that model. The weights stay on a device of
|
||||
the scan's own (separate from the editor's) while the model is loaded, at most 768 MB per model. The log says which
|
||||
device each model runs on.
|
||||
and whenever the GPU is switched off (`LIGHTCRAFT_GPU=0`, `LIGHTCRAFT_GPU_BACKEND=off`, the GPU preference, or **Settings ▸
|
||||
Faces ▸ Run face models on the CPU only**), missing, wrong or slower, the model runs on the CPU. The first GPU error,
|
||||
or a result of the wrong size or with values that are not finite, drops it for that model. Settings ▸ Faces shows
|
||||
which device each model runs on (`faces.gpu`). The weights stay on a device of the scan's own (separate from the
|
||||
editor's) while the model is loaded, at most 768 MB per model.
|
||||
|
||||
Measured on Linux with a Radeon (Vulkan) and a 32-thread CPU, one call at a time, machine under heavy load from other
|
||||
jobs (so the absolute numbers are noisy; the ratios held across three runs):
|
||||
The background scan's pace also limits the GPU: paused or light, one model run at a time and the GPU busy at most a
|
||||
quarter of the time (the editor renders on the same GPU); normal, two at once and three quarters; full, as fast as it
|
||||
goes.
|
||||
|
||||
| Model | CPU | GPU |
|
||||
| --- | ---: | ---: |
|
||||
| YuNet (640², one photo) | 33–34 ms | 35–39 ms |
|
||||
| SFace 2021dec | 25–28 ms | 22–27 ms |
|
||||
| AuraFace v1 | 204–208 ms | 81–103 ms |
|
||||
Measured on Linux with a Radeon (Vulkan) and a 32-thread CPU, one call at a time, machine under very heavy load from
|
||||
other jobs (load average 100–250, so the absolute numbers are noisy):
|
||||
|
||||
| Model | CPU | GPU, one layer at a time | GPU, whole network | Batch of 16, per face |
|
||||
| --- | ---: | ---: | ---: | ---: |
|
||||
| YuNet (640², one photo) | 58–60 ms | 85–89 ms | 21 ms | – |
|
||||
| SFace 2021dec | 40–42 ms | 47–51 ms | 9.6 ms | 8.1 ms |
|
||||
| AuraFace v1 | 295–308 ms | 165–172 ms | 70–73 ms | 71 ms |
|
||||
|
||||
Each layer still goes to the GPU and back: for YuNet and SFace that round trip costs about what the layer saves, so on
|
||||
most machines they stay on the CPU by the speed check, while AuraFace runs 2–2.5× faster. Keeping activations on the GPU
|
||||
between layers (the remaining operators as kernels) and batching faces per dispatch are the next steps.
|
||||
`LC_FACE_MODELS=<folder> cargo test --release -p lightcraft-gpu faces::tests -- --include-ignored --nocapture` checks
|
||||
agreement on the real models and prints these timings.
|
||||
agreement on the real models and every operator, and prints these timings.
|
||||
|
||||
### Agreement and regression coverage
|
||||
|
||||
|
||||
Reference in New Issue
Block a user