From ee3498b24df12cffa38c774c317b5c516a97bd9d Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 21:43:18 +0200 Subject: [PATCH] S2.63: GPU memory attention and memory encoder; one shared device; CPU fallback Memory attention (1.8e-5 vs PyTorch) and the memory encoder (2.6e-6) as wgpu kernels, convolution and layer-norm+GELU kernels, VideoModel::load_auto picks the GPU when there is one (EFFECTCRAFT_SAM2_DEVICE overrides) and falls back to the CPU on any GPU error; disable_gpu switches it off mid-clip. One process-wide device (several at once crashed the Vulkan loader). A tracked frame: 4.3 s CPU, 0.7 s on the RX 7900 GRE. Clips tracked on the GPU match Sam2VideoModel (worst logit error 0.0032). Co-Authored-By: Claude Sonnet 5.5 --- crates/engine/src/cmd/roto_cmds.rs | 2 +- crates/sam2/examples/sam2_bench.rs | 51 ++++++++ crates/sam2/src/gpu/encoder.rs | 3 +- crates/sam2/src/gpu/kernels_ops.wgsl | 37 +++++- crates/sam2/src/gpu/memattn.rs | 184 +++++++++++++++++++++++++++ crates/sam2/src/gpu/memenc.rs | 127 ++++++++++++++++++ crates/sam2/src/gpu/mod.rs | 39 +++++- crates/sam2/src/gpu/ops.rs | 56 +++++++- crates/sam2/src/gpu/tests.rs | 61 +++++++-- crates/sam2/src/mask.rs | 8 +- crates/sam2/src/memattn.rs | 2 +- crates/sam2/src/memory.rs | 16 ++- crates/sam2/src/session.rs | 155 +++++++++++++++++++++- crates/sam2/tests/reference.rs | 130 ++++++++++++++++++- crates/segment/src/lib.rs | 4 + docs/roto-brush-sam2.md | 34 +++-- 16 files changed, 861 insertions(+), 48 deletions(-) create mode 100644 crates/sam2/src/gpu/memattn.rs create mode 100644 crates/sam2/src/gpu/memenc.rs diff --git a/crates/engine/src/cmd/roto_cmds.rs b/crates/engine/src/cmd/roto_cmds.rs index 481a23b5d..5849fa2c8 100644 --- a/crates/engine/src/cmd/roto_cmds.rs +++ b/crates/engine/src/cmd/roto_cmds.rs @@ -478,7 +478,7 @@ fn status(s: &mut Session, p: &Value) -> Result { .keys() .filter_map(|f| chain.trust_at(*f).map(|t| (f.to_string(), json!({"confidence": t.confidence, "occluded": t.occluded})))) .collect(); - out["tracker"] = json!({"model": s.models_roto_active(), "saved": chain.saved_model(), "frames": tracked}); + out["tracker"] = json!({"model": s.models_roto_active(), "device": fx::installed_model().and_then(|m| m.device_name()), "saved": chain.saved_model(), "frames": tracked}); let notice = s.roto_model_notice(s.uses_roto_model(cid, lid, uid)).or_else(|| s.saved_tracking_notice(chain.saved_model())); out["notice"] = json!(notice); out["source"] = json!(if tracked.contains_key(&frame.to_string()) { diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs index a4e78feb0..0b51dd668 100644 --- a/crates/sam2/examples/sam2_bench.rs +++ b/crates/sam2/examples/sam2_bench.rs @@ -70,6 +70,57 @@ fn run() -> Result<(), Box> { let _ = attn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?; println!("memory attention (run {run}, {m} memory tokens): {:.0} ms", t.elapsed().as_secs_f64() * 1e3); } + // The same two costs on the GPU (wgpu kernels), when there is one. + match effectcraft_sam2::gpu::Gpu::shared() { + Ok(gpu) => { + println!("GPU: {}", gpu.adapter_name()); + let enc = effectcraft_sam2::gpu::encoder::GpuEncoder::load(gpu.clone(), &weights, model.vision.hiera.clone())?; + let gattn = effectcraft_sam2::gpu::memattn::GpuMemoryAttention::load(gpu, &weights)?; + for run in 0..3 { + let t = Instant::now(); + let _ = enc.encode(&pixels)?; + let encode = t.elapsed().as_secs_f64() * 1e3; + let t = Instant::now(); + let _ = gattn.forward(&cur, &cur_pos, &mem, &mem_pos, 64)?; + println!("GPU run {run}: encode {encode:.0} ms, memory attention (full bank) {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + } + } + Err(e) => println!("no GPU: {e}"), + } + // A whole tracked frame (encode, memory attention, decode, memory encoding) through the video session, on each backend. + struct Same(Tensor); + impl effectcraft_sam2::session::FrameSource for Same { + fn num_frames(&self) -> usize { + 12 + } + fn frame(&self, _i: usize) -> effectcraft_sam2::Result { + Ok(self.0.clone()) + } + } + for backend in [effectcraft_sam2::session::Backend::Cpu, effectcraft_sam2::session::Backend::Gpu] { + let Ok(vm) = effectcraft_sam2::session::VideoModel::load_with(&weights, backend) else { continue }; + let vm = std::sync::Arc::new(vm); + let mut tracker = effectcraft_sam2::session::VideoTracker::new(vm.clone(), Box::new(Same(pixels.clone())))?; + let obj = effectcraft_sam2::session::ObjectId(1); + let t = Instant::now(); + tracker.prompt(0, obj, &Prompt { points: vec![([512.0, 512.0], true)], bbox: None })?; + println!("{:?} ({}): prompt frame {:.0} ms", backend, vm.device_name(), t.elapsed().as_secs_f64() * 1e3); + let mut last = Instant::now(); + let mut times = vec![]; + tracker.propagate(0, effectcraft_sam2::session::Direction::Forward, None, &effectcraft_sam2::session::Cancel::new(), &mut |m| { + if m.frame > 0 { + times.push(last.elapsed().as_secs_f64() * 1e3); + } + last = Instant::now(); + })?; + println!( + "{:?}: tracked frames {:.0} ms each (first {:.0}, last {:.0})", + backend, + times.iter().sum::() / times.len().max(1) as f64, + times.first().copied().unwrap_or(0.0), + times.last().copied().unwrap_or(0.0) + ); + } println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/crates/sam2/src/gpu/encoder.rs b/crates/sam2/src/gpu/encoder.rs index 9a6caeeb4..630f4f09b 100644 --- a/crates/sam2/src/gpu/encoder.rs +++ b/crates/sam2/src/gpu/encoder.rs @@ -107,7 +107,7 @@ impl GpuEncoder { fn norm(&self, x: &Buf, rows: usize, dim: usize, n: &GNorm) -> Result { let out = self.gpu.zeros(rows * dim)?; - self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps)?; + self.gpu.layernorm(x, &n.w, &n.b, &out, rows, dim, self.eps, false)?; Ok(out) } @@ -227,6 +227,7 @@ impl GpuEncoder { /// Encode normalised `[1, 3, H, W]` pixels (`H` and `W` multiples of 32). pub fn encode(&self, pixels: &Tensor) -> Result { let g = &self.gpu; + let _turn = g.exclusive(); let (_, c, ph, pw) = pixels.dims4()?; let window0 = self.cfg.windows.first().copied().unwrap_or(8); if c != 3 || ph == 0 || pw == 0 || ph % 4 != 0 || pw % 4 != 0 || (ph / 4) % window0 != 0 || (pw / 4) % window0 != 0 { diff --git a/crates/sam2/src/gpu/kernels_ops.wgsl b/crates/sam2/src/gpu/kernels_ops.wgsl index aae0c8323..419476600 100644 --- a/crates/sam2/src/gpu/kernels_ops.wgsl +++ b/crates/sam2/src/gpu/kernels_ops.wgsl @@ -29,7 +29,7 @@ fn reduce_max(li: u32) -> f32 { } // ---------------------------------------------------------------- layer norm over the last dimension -// words: 0 rows, 1 cols, 2 eps (f32). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row. +// words: 0 rows, 1 cols, 2 eps (f32), 3 gelu (1: apply GELU to the result). in0 x, in1 weight, in2 bias, out0 y. One workgroup per row. @compute @workgroup_size(64) fn layernorm(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { let row = wg.y * nw.x + wg.x; @@ -44,7 +44,40 @@ fn layernorm(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: for (var c = li.x; c < cols; c += 64u) { let d = in0[base + c] - mean; q += d * d; } red[li.x] = q; let inv = inverseSqrt(reduce_sum(li.x) / f32(cols) + f(2u)); - for (var c = li.x; c < cols; c += 64u) { out0[base + c] = (in0[base + c] - mean) * inv * in1[c] + in2[c]; } + for (var c = li.x; c < cols; c += 64u) { + var y = (in0[base + c] - mean) * inv * in1[c] + in2[c]; + if (u(3u) == 1u) { y = gelu_(y); } + out0[base + c] = y; + } +} + +// ---------------------------------------------------------------- convolution over channels-last [H, W, Ci] +// words: 0 n (output elements), 1 H, 2 W, 3 Ci, 4 Co, 5 Ho, 6 Wo, 7 k, 8 stride, 9 pad, 10 groups. +// in0 input, in1 weight [Co, Ci / groups, k, k], in2 bias [Co]; output [Ho, Wo, Co]. +@compute @workgroup_size(256) +fn conv2d(@builtin(workgroup_id) wg: vec3, @builtin(num_workgroups) nw: vec3, @builtin(local_invocation_id) li: vec3) { + let i = gid1(wg, nw, li); + if (i >= u(0u)) { return; } + let H = u(1u); let W = u(2u); let Ci = u(3u); let Co = u(4u); let Wo = u(6u); + let k = u(7u); let stride = u(8u); let pad = i32(u(9u)); let groups = u(10u); + let co = i % Co; let t = i / Co; + let ox = t % Wo; let oy = t / Wo; + let cig = Ci / groups; + let g0 = co / (Co / groups); + var acc = in2[co]; + for (var c = 0u; c < cig; c++) { + let ci = g0 * cig + c; + for (var ky = 0u; ky < k; ky++) { + let iy = i32(oy * stride + ky) - pad; + if (iy < 0 || iy >= i32(H)) { continue; } + for (var kx = 0u; kx < k; kx++) { + let ix = i32(ox * stride + kx) - pad; + if (ix < 0 || ix >= i32(W)) { continue; } + acc += in0[(u32(iy) * W + u32(ix)) * Ci + ci] * in1[((co * cig + c) * k + ky) * k + kx]; + } + } + } + out0[i] = acc; } // ---------------------------------------------------------------- softmax over rows, in place on out0 diff --git a/crates/sam2/src/gpu/memattn.rs b/crates/sam2/src/gpu/memattn.rs new file mode 100644 index 000000000..38c5139a7 --- /dev/null +++ b/crates/sam2/src/gpu/memattn.rs @@ -0,0 +1,184 @@ +// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 video +// implementation (Apache-2.0, see NOTICE). +//! Memory attention on the GPU: four layers of self-attention over the current frame and cross-attention to +//! the memory bank, with axial rotary position encoding. The attention is one head of 256 channels; the +//! score matrix against a full memory bank is 4096 x 28736 floats, which stays on the GPU. + +use std::sync::Arc; + +use candle_core::{Device, Tensor}; + +use super::{Act, Buf, GLinear, GNorm, Gpu, Matmul, Strides}; +use crate::memattn::{GRID, rope_tables}; +use crate::weights::Weights; +use crate::{Error, Result}; + +const DIM: usize = 256; +const LAYERS: usize = 4; + +struct GAttn { + q: GLinear, + k: GLinear, + v: GLinear, + o: GLinear, +} + +struct GLayer { + self_attn: GAttn, + cross: GAttn, + ln1: GNorm, + ln2: GNorm, + ln3: GNorm, + lin1: GLinear, + lin2: GLinear, +} + +pub struct GpuMemoryAttention { + gpu: Arc, + layers: Vec, + final_ln: GNorm, + cos: Buf, + sin: Buf, +} + +fn attn(g: &Gpu, w: &Weights, prefix: &str, kv: usize) -> Result { + Ok(GAttn { + q: GLinear::load(g, w, &format!("{prefix}.q_proj"), DIM, DIM)?, + k: GLinear::load(g, w, &format!("{prefix}.k_proj"), kv, DIM)?, + v: GLinear::load(g, w, &format!("{prefix}.v_proj"), kv, DIM)?, + o: GLinear::load(g, w, &format!("{prefix}.o_proj"), DIM, DIM)?, + }) +} + +impl GpuMemoryAttention { + pub fn load(gpu: Arc, w: &Weights) -> Result { + let layers = (0..LAYERS) + .map(|i| { + let p = format!("memory_attention.layers.{i}"); + Ok(GLayer { + self_attn: attn(&gpu, w, &format!("{p}.self_attn"), DIM)?, + cross: attn(&gpu, w, &format!("{p}.cross_attn_image"), 64)?, + ln1: GNorm::load(&gpu, w, &format!("{p}.layer_norm1"), DIM)?, + ln2: GNorm::load(&gpu, w, &format!("{p}.layer_norm2"), DIM)?, + ln3: GNorm::load(&gpu, w, &format!("{p}.layer_norm3"), DIM)?, + lin1: GLinear::load(&gpu, w, &format!("{p}.linear1"), DIM, 2048)?, + lin2: GLinear::load(&gpu, w, &format!("{p}.linear2"), 2048, DIM)?, + }) + }) + .collect::>>()?; + let (cos, sin) = rope_tables(10000.0)?; + let flat = |t: Tensor| -> Result> { Ok(t.flatten_all()?.to_vec1::()?) }; + Ok(GpuMemoryAttention { + final_ln: GNorm::load(&gpu, w, "memory_attention.layer_norm", DIM)?, + cos: gpu.upload(&flat(cos)?), + sin: gpu.upload(&flat(sin)?), + layers, + gpu, + }) + } + + fn norm(&self, x: &Buf, rows: usize, n: &GNorm) -> Result { + let out = self.gpu.zeros(rows * DIM)?; + self.gpu.layernorm(x, &n.w, &n.b, &out, rows, DIM, 1e-5, false)?; + Ok(out) + } + + fn rope(&self, x: &Buf, rows: usize, rotated: usize) -> Result { + let out = self.gpu.zeros(rows * DIM)?; + self.gpu.rope(x, &self.cos, &self.sin, &out, rows, DIM, GRID * GRID, rotated)?; + Ok(out) + } + + /// `softmax(q . k^T / 16) . v` for `[nq, 256]` queries and `[nk, 256]` keys and values, through `scores` + /// (at least `nq * nk` elements); returns `[nq, 256]`. + fn attention(&self, q: &Buf, k: &Buf, v: &Buf, nq: usize, nk: usize, scores: &Buf) -> Result { + let g = &self.gpu; + g.matmul(&Matmul { + m: nq, + n: nk, + k: DIM, + batch: 1, + a: (q, Strides::row_major(DIM)), + b: (k, Strides { offset: 0, batch: 0, row: 1, col: DIM }), + bias: None, + c: (scores, Strides::row_major(nk)), + residual: None, + alpha: (DIM as f32).powf(-0.5), + act: Act::None, + })?; + g.softmax(scores, nq, nk)?; + let out = g.zeros(nq * DIM)?; + g.matmul(&Matmul { + m: nq, + n: DIM, + k: nk, + batch: 1, + a: (scores, Strides::row_major(nk)), + b: (v, Strides::row_major(DIM)), + bias: None, + c: (&out, Strides::row_major(DIM)), + residual: None, + alpha: 1.0, + act: Act::None, + })?; + Ok(out) + } + + /// As [`crate::memattn::MemoryAttention::forward`]: `current` and `current_pos` `[4096, 256]`, `memory` and + /// `memory_pos` `[M, 64]`, the last `num_pointers` memory rows being object pointers (not rotated). + pub fn forward(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, num_pointers: usize) -> Result { + let g = &self.gpu; + let _turn = g.exclusive(); + let n = GRID * GRID; + if current.dims() != [n, DIM] || current_pos.dims() != [n, DIM] { + return Err(Error::Input(format!("current features must be [{n}, {DIM}], got {:?}", current.dims()))); + } + let (m, md) = memory.dims2()?; + if md != 64 || memory_pos.dims() != [m, 64] { + return Err(Error::Input(format!("memory must be [M, 64] with a matching position encoding, got {:?}", memory.dims()))); + } + let rotated = m.checked_sub(num_pointers).ok_or_else(|| Error::Input("more object pointers than memory tokens".into()))?; + if rotated % n != 0 { + return Err(Error::Input(format!("memory length {rotated} is not a multiple of {n}"))); + } + let flat = |t: &Tensor| -> Result> { Ok(t.contiguous()?.flatten_all()?.to_vec1::()?) }; + let mut x = g.upload(&flat(&(current + (current_pos * 0.1)?)?)?); + let mem = g.upload(&flat(memory)?); + let keys_in = g.upload(&flat(&(memory + memory_pos)?)?); + let scores = g.zeros(n * n.max(m))?; + for l in &self.layers { + // self-attention over the frame + let q_in = self.norm(&x, n, &l.ln1)?; + let (q, k, v) = (g.zeros(n * DIM)?, g.zeros(n * DIM)?, g.zeros(n * DIM)?); + g.linear(&q_in, n, &l.self_attn.q, &q, Act::None, None)?; + g.linear(&q_in, n, &l.self_attn.k, &k, Act::None, None)?; + g.linear(&q_in, n, &l.self_attn.v, &v, Act::None, None)?; + let (q, k) = (self.rope(&q, n, n)?, self.rope(&k, n, n)?); + let a = self.attention(&q, &k, &v, n, n, &scores)?; + let x1 = g.zeros(n * DIM)?; + g.linear(&a, n, &l.self_attn.o, &x1, Act::None, Some(&x))?; + // cross-attention to the memory bank + let q_in = self.norm(&x1, n, &l.ln2)?; + let q = g.zeros(n * DIM)?; + g.linear(&q_in, n, &l.cross.q, &q, Act::None, None)?; + let q = self.rope(&q, n, n)?; + let (k, v) = (g.zeros(m * DIM)?, g.zeros(m * DIM)?); + g.linear(&keys_in, m, &l.cross.k, &k, Act::None, None)?; + g.linear(&mem, m, &l.cross.v, &v, Act::None, None)?; + let k = self.rope(&k, m, rotated)?; + let a = self.attention(&q, &k, &v, n, m, &scores)?; + let x2 = g.zeros(n * DIM)?; + g.linear(&a, n, &l.cross.o, &x2, Act::None, Some(&x1))?; + // feed-forward + let f_in = self.norm(&x2, n, &l.ln3)?; + let hidden = g.zeros(n * 2048)?; + g.linear(&f_in, n, &l.lin1, &hidden, Act::Relu, None)?; + let x3 = g.zeros(n * DIM)?; + g.linear(&hidden, n, &l.lin2, &x3, Act::None, Some(&x2))?; + x = x3; + } + let out = self.norm(&x, n, &self.final_ln)?; + g.flush()?; + Ok(Tensor::from_vec(g.download(&out)?, (n, DIM), &Device::Cpu)?) + } +} diff --git a/crates/sam2/src/gpu/memenc.rs b/crates/sam2/src/gpu/memenc.rs new file mode 100644 index 000000000..2cfa693d0 --- /dev/null +++ b/crates/sam2/src/gpu/memenc.rs @@ -0,0 +1,127 @@ +// Original work for EffectCraft; the maths follows the Hugging Face `transformers` SAM 2 video +// implementation (Apache-2.0, see NOTICE). +//! The memory encoder on the GPU: the mask is down-sampled by four strided convolutions, fused with the +//! frame's features by two ConvNeXt-style blocks, and projected to the 64 channels of a memory. + +use std::sync::Arc; + +use candle_core::{Device, Tensor}; + +use super::{Act, Buf, GLinear, GNorm, Gpu}; +use crate::weights::Weights; +use crate::{Error, Result}; + +struct GConv { + w: Buf, + b: Buf, + ci: usize, + co: usize, + k: usize, +} + +struct GFuser { + dw: GConv, + norm: GNorm, + pw1: GLinear, + /// The layer scale is folded into this layer's weights. + pw2: GLinear, +} + +pub struct GpuMemoryEncoder { + gpu: Arc, + down: Vec<(GConv, GNorm)>, + final_conv: GLinear, + feature_projection: GLinear, + fuser: Vec, + projection: GLinear, +} + +fn conv(g: &Gpu, w: &Weights, name: &str, ci: usize, co: usize, k: usize, groups: usize) -> Result { + let flat = |t: Tensor| -> Result> { Ok(t.flatten_all()?.to_vec1::()?) }; + Ok(GConv { + w: g.upload(&flat(w.get(&format!("{name}.weight"), &[co, ci / groups, k, k])?)?), + b: g.upload(&flat(w.get(&format!("{name}.bias"), &[co])?)?), + ci, + co, + k, + }) +} + +impl GpuMemoryEncoder { + pub fn load(gpu: Arc, w: &Weights) -> Result { + let mut down = vec![]; + let mut ch = 1; + for i in 0..4 { + let out = ch * 4; + let p = format!("memory_encoder.mask_downsampler.layers.{i}"); + down.push((conv(&gpu, w, &format!("{p}.conv"), ch, out, 3, 1)?, GNorm::load(&gpu, w, &format!("{p}.layer_norm"), out)?)); + ch = out; + } + let fuser = (0..2) + .map(|i| { + let p = format!("memory_encoder.memory_fuser.layers.{i}"); + let scale = w.get(&format!("{p}.scale"), &[256])?.flatten_all()?.to_vec1::()?; + Ok(GFuser { + dw: conv(&gpu, w, &format!("{p}.depthwise_conv"), 256, 256, 7, 256)?, + norm: GNorm::load(&gpu, w, &format!("{p}.layer_norm"), 256)?, + pw1: GLinear::load(&gpu, w, &format!("{p}.pointwise_conv1"), 256, 1024)?, + pw2: GLinear::load_scaled(&gpu, w, &format!("{p}.pointwise_conv2"), 1024, 256, &scale)?, + }) + }) + .collect::>>()?; + Ok(GpuMemoryEncoder { + final_conv: GLinear::load_conv1x1(&gpu, w, "memory_encoder.mask_downsampler.final_conv", 256, 256)?, + feature_projection: GLinear::load_conv1x1(&gpu, w, "memory_encoder.feature_projection", 256, 256)?, + projection: GLinear::load_conv1x1(&gpu, w, "memory_encoder.projection", 256, 64)?, + down, + fuser, + gpu, + }) + } + + /// `features`: `[1, 256, 64, 64]`; `mask`: `[1, 1, 1024, 1024]` already scaled (`20 m - 10`). Returns the raw + /// memory tokens `[4096, 64]` (the occlusion embedding and bfloat16 rounding are the caller's). + pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result { + let g = &self.gpu; + let _turn = g.exclusive(); + let flat = |t: &Tensor| -> Result> { Ok(t.contiguous()?.flatten_all()?.to_vec1::()?) }; + if features.dims() != [1, 256, 64, 64] || mask.dims() != [1, 1, 1024, 1024] { + return Err(Error::Input(format!("memory encoder input shapes {:?} / {:?}", features.dims(), mask.dims()))); + } + let tokens = features.permute((0, 2, 3, 1))?.contiguous()?; + let ftok = g.upload(&flat(&tokens)?); + let mut cur = g.upload(&flat(mask)?); + let (mut h, mut w) = (1024usize, 1024usize); + for (c, ln) in &self.down { + let (ho, wo) = (h / 2, w / 2); + let raw = g.zeros(ho * wo * c.co)?; + g.conv2d(&cur, &c.w, &c.b, &raw, (h, w, c.ci, c.co), c.k, 2, 1, 1)?; + let normed = g.zeros(ho * wo * c.co)?; + g.layernorm(&raw, &ln.w, &ln.b, &normed, ho * wo, c.co, 1e-6, true)?; + (cur, h, w) = (normed, ho, wo); + } + let n = h * w; + if n != 4096 { + return Err(Error::Compute(format!("the mask down-sampler ended at {w}x{h}, not 64x64"))); + } + let m = g.zeros(n * 256)?; + g.linear(&cur, n, &self.final_conv, &m, Act::None, None)?; + let mut x = g.zeros(n * 256)?; + g.linear(&ftok, n, &self.feature_projection, &x, Act::None, Some(&m))?; + for b in &self.fuser { + let d = g.zeros(n * 256)?; + g.conv2d(&x, &b.dw.w, &b.dw.b, &d, (h, w, 256, 256), 7, 1, 3, 256)?; + let normed = g.zeros(n * 256)?; + g.layernorm(&d, &b.norm.w, &b.norm.b, &normed, n, 256, 1e-6, false)?; + let hidden = g.zeros(n * 1024)?; + g.linear(&normed, n, &b.pw1, &hidden, Act::Gelu, None)?; + let next = g.zeros(n * 256)?; + g.linear(&hidden, n, &b.pw2, &next, Act::None, Some(&x))?; + x = next; + } + let out = g.zeros(n * 64)?; + g.linear(&x, n, &self.projection, &out, Act::None, None)?; + g.flush()?; + Ok(Tensor::from_vec(g.download(&out)?, (n, 64), &Device::Cpu)?) + } +} diff --git a/crates/sam2/src/gpu/mod.rs b/crates/sam2/src/gpu/mod.rs index a14df4b16..9ca180932 100644 --- a/crates/sam2/src/gpu/mod.rs +++ b/crates/sam2/src/gpu/mod.rs @@ -72,6 +72,7 @@ pub struct Gpu { module: wgpu::ShaderModule, pipelines: Mutex>>, encoder: Mutex>, + turn: Mutex<()>, error: Arc>>, adapter: String, } @@ -89,8 +90,16 @@ fn u32_of(v: usize, what: &str) -> Result { } impl Gpu { + /// The process's one GPU device, created on first use. Several devices at once upset some Vulkan loaders + /// (a crash when two threads create or drop them together), and one is all SAM 2 needs: models come and go, + /// the device stays. An error (no GPU) is remembered too. + pub fn shared() -> Result> { + static DEVICE: std::sync::OnceLock, String>> = std::sync::OnceLock::new(); + DEVICE.get_or_init(|| Gpu::new().map(Arc::new).map_err(|e| e.to_string())).clone().map_err(Error::Compute) + } + /// A device on the best real GPU (a software adapter is refused: the CPU path is faster). - pub fn new() -> Result { + fn new() -> Result { let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle()); let adapter = pollster::block_on( instance.request_adapter(&wgpu::RequestAdapterOptions { power_preference: wgpu::PowerPreference::HighPerformance, ..Default::default() }), @@ -117,7 +126,13 @@ impl Gpu { sink.lock().unwrap_or_else(PoisonError::into_inner).get_or_insert_with(|| format!("GPU device lost ({reason:?}): {message}")); }); let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some("sam2 kernels"), source: wgpu::ShaderSource::Wgsl(SHADER.into()) }); - Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), error, adapter: info.name }) + Ok(Gpu { device, queue, module, pipelines: Mutex::default(), encoder: Mutex::new(None), turn: Mutex::new(()), error, adapter: info.name }) + } + + /// One GPU pass at a time: a caller holds this while it queues, flushes and reads back a whole job, so + /// two threads never interleave their dispatches in one command stream. + pub fn exclusive(&self) -> std::sync::MutexGuard<'_, ()> { + self.turn.lock().unwrap_or_else(PoisonError::into_inner) } pub fn adapter_name(&self) -> &str { @@ -304,6 +319,8 @@ impl Gpu { } pub mod encoder; +pub mod memattn; +pub mod memenc; mod ops; /// A linear layer's weight `[out, in]` and bias `[out]` on the GPU. @@ -335,6 +352,24 @@ impl GLinear { }) } + /// As [`GLinear::load`], with output channel `j` scaled by `scale[j]` (the memory fuser's layer scale). + pub fn load_scaled(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize, scale: &[f32]) -> Result { + if scale.len() != out { + return Err(Error::Weights(format!("`{name}`: {} scales for {out} outputs", scale.len()))); + } + let mut wt = flat(w.get(&format!("{name}.weight"), &[out, inp])?)?; + let mut bias = flat(w.get(&format!("{name}.bias"), &[out])?)?; + for (j, s) in scale.iter().enumerate() { + if let Some(row) = wt.get_mut(j * inp..(j + 1) * inp) { + row.iter_mut().for_each(|v| *v *= s); + } + if let Some(b) = bias.get_mut(j) { + *b *= s; + } + } + Ok(GLinear { w: g.upload(&wt), b: g.upload(&bias), inp, out }) + } + /// A 1 x 1 convolution (`[out, in, 1, 1]`) as a linear layer over channels-last tokens. pub fn load_conv1x1(g: &Gpu, w: &crate::weights::Weights, name: &str, inp: usize, out: usize) -> Result { Ok(GLinear { diff --git a/crates/sam2/src/gpu/ops.rs b/crates/sam2/src/gpu/ops.rs index 7bb99c495..3d9b8f184 100644 --- a/crates/sam2/src/gpu/ops.rs +++ b/crates/sam2/src/gpu/ops.rs @@ -27,12 +27,18 @@ fn need(buf: &Buf, len: usize, what: &str) -> Result<()> { impl Gpu { /// `y = layer_norm(x) * weight + bias` over the last dimension of a `[rows, cols]` buffer. - pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32) -> Result<()> { + #[allow(clippy::too_many_arguments)] + pub fn layernorm(&self, x: &Buf, weight: &Buf, bias: &Buf, y: &Buf, rows: usize, cols: usize, eps: f32, gelu: bool) -> Result<()> { need(x, rows * cols, "layernorm x")?; need(y, rows * cols, "layernorm y")?; need(weight, cols, "layernorm weight")?; need(bias, cols, "layernorm bias")?; - self.run("layernorm", &[(0, x), (1, weight), (2, bias), (3, y)], &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits()], grid(rows)?); + self.run( + "layernorm", + &[(0, x), (1, weight), (2, bias), (3, y)], + &[n32(rows, "rows")?, n32(cols, "cols")?, eps.to_bits(), u32::from(gelu)], + grid(rows)?, + ); Ok(()) } @@ -43,6 +49,52 @@ impl Gpu { Ok(()) } + /// A convolution over channels-last `[h, w, ci]` with weight `[co, ci / groups, k, k]` and bias `[co]`, + /// output `[ho, wo, co]` (`groups == ci == co` is depthwise). + #[allow(clippy::too_many_arguments)] + pub fn conv2d( + &self, + x: &Buf, + weight: &Buf, + bias: &Buf, + out: &Buf, + dims: (usize, usize, usize, usize), + k: usize, + stride: usize, + pad: usize, + groups: usize, + ) -> Result<()> { + let (h, w, ci, co) = dims; + if k == 0 || stride == 0 || groups == 0 || ci % groups != 0 || co % groups != 0 || h + 2 * pad < k || w + 2 * pad < k { + return Err(Error::Input(format!("a {k}x{k} convolution (stride {stride}, groups {groups}) does not fit {w}x{h}x{ci} -> {co}"))); + } + let (ho, wo) = ((h + 2 * pad - k) / stride + 1, (w + 2 * pad - k) / stride + 1); + need(x, h * w * ci, "conv2d input")?; + need(weight, co * (ci / groups) * k * k, "conv2d weight")?; + need(bias, co, "conv2d bias")?; + need(out, ho * wo * co, "conv2d output")?; + let n = ho * wo * co; + self.run( + "conv2d", + &[(0, x), (1, weight), (2, bias), (3, out)], + &[ + n32(n, "n")?, + n32(h, "h")?, + n32(w, "w")?, + n32(ci, "ci")?, + n32(co, "co")?, + n32(ho, "ho")?, + n32(wo, "wo")?, + n32(k, "k")?, + n32(stride, "stride")?, + n32(pad, "pad")?, + n32(groups, "groups")?, + ], + items(n)?, + ); + Ok(()) + } + /// `out = a + b` over `n` elements. pub fn add(&self, a: &Buf, b: &Buf, out: &Buf, n: usize) -> Result<()> { need(a, n, "add a")?; diff --git a/crates/sam2/src/gpu/tests.rs b/crates/sam2/src/gpu/tests.rs index 92ad20f41..ed5b3528a 100644 --- a/crates/sam2/src/gpu/tests.rs +++ b/crates/sam2/src/gpu/tests.rs @@ -6,27 +6,25 @@ use super::*; /// The one device the tests share (several devices at once upset the Vulkan loader), held for the whole /// test so the tests take turns. -pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, &'static Gpu); +pub(crate) struct Shared(#[allow(dead_code)] std::sync::MutexGuard<'static, ()>, Arc); impl std::ops::Deref for Shared { type Target = Gpu; fn deref(&self) -> &Gpu { - self.1 + &self.1 } } pub(crate) fn gpu() -> Option { static TURN: Mutex<()> = Mutex::new(()); - static DEVICE: std::sync::OnceLock> = std::sync::OnceLock::new(); let guard = TURN.lock().unwrap_or_else(PoisonError::into_inner); - let device = DEVICE.get_or_init(|| match Gpu::new() { - Ok(g) => Some(g), + match Gpu::shared() { + Ok(g) => Some(Shared(guard, g)), Err(e) => { eprintln!("skipped: {e}"); None } - }); - device.as_ref().map(|g| Shared(guard, g)) + } } /// Deterministic pseudo-random values in -1..1. @@ -176,7 +174,7 @@ fn layernorm_and_softmax_match_the_definition() { let x = noise(rows * cols, 7); let (w, b) = (noise(cols, 8), noise(cols, 9)); let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&w), g.upload(&b), g.zeros(rows * cols).unwrap()); - g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6).unwrap(); + g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6, false).unwrap(); let want: Vec = x .chunks(cols) .flat_map(|r| { @@ -299,3 +297,50 @@ fn patch_embed_rope_and_upsample_match_the_cpu_definitions() { let want: Vec = (0..h * w * c).map(|i| lat[i] + prev[((i / c / w / 2) * (w / 2) + (i / c % w) / 2) * c + i % c]).collect(); assert_close(&run_and_read(&g, &go), &want, 0.0, "upsample_add"); } + +#[test] +fn convolutions_and_the_layernorm_gelu_epilogue_match_candle() { + use candle_core::{Device, Tensor}; + let Some(g) = gpu() else { return }; + // (h, w, ci, co, k, stride, pad, groups): a strided 3 x 3, a depthwise 7 x 7, and a 1-channel input + for (h, w, ci, co, k, stride, pad, groups) in + [(9usize, 13usize, 3usize, 5usize, 3usize, 2usize, 1usize, 1usize), (10, 10, 4, 4, 7, 1, 3, 4), (16, 16, 1, 4, 3, 2, 1, 1)] + { + let (x, wt, b) = (noise(h * w * ci, 30), noise(co * (ci / groups) * k * k, 31), noise(co, 32)); + let (gx, gw, gb) = (g.upload(&x), g.upload(&wt), g.upload(&b)); + let (ho, wo) = ((h + 2 * pad - k) / stride + 1, (w + 2 * pad - k) / stride + 1); + let out = g.zeros(ho * wo * co).unwrap(); + g.conv2d(&gx, &gw, &gb, &out, (h, w, ci, co), k, stride, pad, groups).unwrap(); + // candle works on NCHW: bring the channels-last input over, convolve, bring the result back + let t = Tensor::from_vec(x, (1, h, w, ci), &Device::Cpu).unwrap().permute((0, 3, 1, 2)).unwrap().contiguous().unwrap(); + let kern = Tensor::from_vec(wt, (co, ci / groups, k, k), &Device::Cpu).unwrap(); + let want = t + .conv2d(&kern, pad, stride, 1, groups) + .unwrap() + .broadcast_add(&Tensor::from_vec(b, (1, co, 1, 1), &Device::Cpu).unwrap()) + .unwrap() + .permute((0, 2, 3, 1)) + .unwrap(); + assert_close( + &run_and_read(&g, &out), + &want.flatten_all().unwrap().to_vec1::().unwrap(), + 2e-4, + &format!("conv {k}x{k} stride {stride} groups {groups}"), + ); + } + assert!(g.conv2d(&g.upload(&[0.0; 4]), &g.upload(&[0.0; 4]), &g.upload(&[0.0; 1]), &g.zeros(4).unwrap(), (2, 2, 1, 1), 5, 1, 0, 1).is_err()); + // layer norm followed by GELU in the same pass + let (rows, cols) = (50usize, 96usize); + let (x, wt, b) = (noise(rows * cols, 33), noise(cols, 34), noise(cols, 35)); + let (gx, gw, gb, gy) = (g.upload(&x), g.upload(&wt), g.upload(&b), g.zeros(rows * cols).unwrap()); + g.layernorm(&gx, &gw, &gb, &gy, rows, cols, 1e-6, true).unwrap(); + let want: Vec = x + .chunks(cols) + .flat_map(|r| { + let mean = r.iter().sum::() / cols as f32; + let var = r.iter().map(|v| (v - mean).powi(2)).sum::() / cols as f32; + r.iter().enumerate().map(|(i, v)| gelu((v - mean) / (var + 1e-6).sqrt() * wt[i] + b[i])).collect::>() + }) + .collect(); + assert_close(&run_and_read(&g, &gy), &want, 1e-3, "layernorm + gelu"); +} diff --git a/crates/sam2/src/mask.rs b/crates/sam2/src/mask.rs index f917e29ef..94cf8b40a 100644 --- a/crates/sam2/src/mask.rs +++ b/crates/sam2/src/mask.rs @@ -79,7 +79,7 @@ impl Sam2Mask { { return Ok(e.clone()); } - let enc = Arc::new(self.model.image.encode(&preprocess(rgb, w, h)?)?); + let enc = Arc::new(self.model.encode(&preprocess(rgb, w, h)?)?); *self.last.lock().unwrap_or_else(PoisonError::into_inner) = Some((key, enc.clone())); Ok(enc) } @@ -117,6 +117,10 @@ impl MaskModel for Sam2Mask { Some(self) } + fn device_name(&self) -> Option { + Some(self.model.device_name()) + } + fn idle_secs(&self) -> Option { Some(self.used.lock().unwrap_or_else(PoisonError::into_inner).elapsed().as_secs_f64()) } @@ -129,7 +133,7 @@ pub fn load(bytes: &[u8]) -> std::result::Result, String> { return Err(format!("{}: the file does not match (SHA-256 {got}, expected {})", SAM2_TINY.name, SAM2_TINY.sha256)); } let w = crate::weights::Weights::from_bytes(bytes).map_err(|e| e.to_string())?; - let model = VideoModel::load(&w).map_err(|e| e.to_string())?; + let model = VideoModel::load_auto(&w).map_err(|e| e.to_string())?; Ok(Arc::new(Sam2Mask::new(Arc::new(model)))) } diff --git a/crates/sam2/src/memattn.rs b/crates/sam2/src/memattn.rs index c3bfc0893..d455f1ed1 100644 --- a/crates/sam2/src/memattn.rs +++ b/crates/sam2/src/memattn.rs @@ -18,7 +18,7 @@ const LAYERS: usize = 4; pub const GRID: usize = 64; /// `cos` and `sin` of the axial rotary encoding for the `GRID x GRID` grid, each `[GRID², 256]`. -fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> { +pub(crate) fn rope_tables(theta: f32) -> Result<(Tensor, Tensor)> { let freqs = DIM / 2 / 2; // per axis: head_dim / 4 let inv: Vec = (0..freqs).map(|i| 1.0 / theta.powf(i as f32 / freqs as f32)).collect(); let n = GRID * GRID; diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs index 45b470a3f..7a73fd5a3 100644 --- a/crates/sam2/src/memory.rs +++ b/crates/sam2/src/memory.rs @@ -123,10 +123,16 @@ impl MemoryEncoder { /// The memory kept for a frame: the encoder output, plus the occlusion embedding when the /// object is absent, rounded through bfloat16 like the reference. pub fn memory(&self, features: &Tensor, mask: &Tensor, object_present: bool) -> Result { - let (x, pos) = self.encode(features, mask)?; - let x = if object_present { x } else { x.broadcast_add(&self.occlusion.reshape((1, 64, 1, 1))?)? }; - let flat = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; - let features = flat(&x)?.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?; - Ok(Memory { features, pos: flat(&pos)? }) + let (x, _) = self.encode(features, mask)?; + self.finish(&x.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?, object_present) + } + + /// A memory from the encoder's raw tokens `[HW, 64]`: the occlusion embedding when the object is absent, + /// rounded through bfloat16, with the position encoding (shared by the CPU and GPU encoders). + pub fn finish(&self, tokens: &Tensor, object_present: bool) -> Result { + let tokens = if object_present { tokens.clone() } else { tokens.broadcast_add(&self.occlusion)? }; + let features = tokens.to_dtype(candle_core::DType::BF16)?.to_dtype(candle_core::DType::F32)?; + let pos = sine_position(32, 64, 64)?; + Ok(Memory { features, pos: pos.flatten_from(2)?.squeeze(0)?.t()?.contiguous()? }) } } diff --git a/crates/sam2/src/session.rs b/crates/sam2/src/session.rs index 2e45cc407..9de8fedc2 100644 --- a/crates/sam2/src/session.rs +++ b/crates/sam2/src/session.rs @@ -14,9 +14,14 @@ use candle_core::{DType, Tensor}; use candle_nn::{Linear, ops::sigmoid}; use crate::decoder::FeedForward; +use crate::gpu::Gpu; +use crate::gpu::encoder::GpuEncoder; +use crate::gpu::memattn::GpuMemoryAttention; +use crate::gpu::memenc::GpuMemoryEncoder; use crate::memattn::MemoryAttention; use crate::memory::{Memory, MemoryEncoder}; -use crate::model::ImageModel; +use crate::model::{Encoded, ImageModel}; +use crate::neck::ImageFeatures; use crate::nn::{apply, linear}; use crate::prompt::{INPUT_SIZE, Prompt}; use crate::resize::antialias_bilinear; @@ -79,8 +84,31 @@ pub struct FrameMask { pub occluded: bool, } +/// Which device runs the heavy parts (the backbone, the neck and the memory attention). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Backend { + /// Candle on the CPU: the reference path. + Cpu, + /// wgpu compute on the GPU (any vendor); an error if there is none. + Gpu, + /// The GPU when there is a real one and it loads, else the CPU. The environment variable + /// `EFFECTCRAFT_SAM2_DEVICE` (`cpu`, `gpu` or `auto`) overrides the choice. + Auto, +} + +/// The GPU path: kernels for the backbone and the memory attention. Any failure switches it off for good +/// and the CPU path takes over. +struct GpuParts { + encoder: GpuEncoder, + attention: GpuMemoryAttention, + memory_encoder: GpuMemoryEncoder, + ok: AtomicBool, + adapter: String, +} + /// The weights of everything the video tracker needs. pub struct VideoModel { + gpu: Option, pub image: ImageModel, memory_encoder: MemoryEncoder, memory_attention: MemoryAttention, @@ -93,9 +121,37 @@ pub struct VideoModel { } impl VideoModel { + /// On the CPU: the reference path every test compares against. pub fn load(w: &Weights) -> Result { + VideoModel::load_with(w, Backend::Cpu) + } + + /// On the GPU when there is one (see [`Backend::Auto`]). + pub fn load_auto(w: &Weights) -> Result { + VideoModel::load_with(w, Backend::Auto) + } + + pub fn load_with(w: &Weights, backend: Backend) -> Result { + let backend = match (backend, std::env::var("EFFECTCRAFT_SAM2_DEVICE").ok().as_deref()) { + (Backend::Auto, Some("cpu")) => Backend::Cpu, + (Backend::Auto, Some("gpu")) => Backend::Gpu, + (b, _) => b, + }; + let image = ImageModel::load(w)?; + let gpu = match backend { + Backend::Cpu => None, + Backend::Gpu => Some(Self::load_gpu(w, &image)?), + Backend::Auto => match Self::load_gpu(w, &image) { + Ok(g) => Some(g), + Err(e) => { + log::info!("SAM 2 runs on the CPU: {e}"); + None + } + }, + }; Ok(VideoModel { - image: ImageModel::load(w)?, + gpu, + image, memory_encoder: MemoryEncoder::load(w)?, memory_attention: MemoryAttention::load(w)?, temporal_pos: w.get("memory_temporal_positional_encoding", &[MEMORY_FRAMES, 1, 1, 64])?.flatten_all()?.to_vec1::()?, @@ -106,6 +162,79 @@ impl VideoModel { }) } + fn load_gpu(w: &Weights, image: &ImageModel) -> Result { + let gpu = Gpu::shared()?; + Ok(GpuParts { + encoder: GpuEncoder::load(gpu.clone(), w, image.vision.hiera.clone())?, + attention: GpuMemoryAttention::load(gpu.clone(), w)?, + memory_encoder: GpuMemoryEncoder::load(gpu.clone(), w)?, + ok: AtomicBool::new(true), + adapter: gpu.adapter_name().to_string(), + }) + } + + /// The device the heavy parts run on now: the GPU's name, or "CPU". + pub fn device_name(&self) -> String { + match &self.gpu { + Some(g) if g.ok.load(Ordering::Relaxed) => g.adapter.clone(), + _ => "CPU".to_string(), + } + } + + /// Run on the CPU from now on (a GPU that misbehaves, or the user's choice). Weights and results are + /// unchanged; only the device differs. + pub fn disable_gpu(&self) { + if let Some(g) = &self.gpu { + g.ok.store(false, Ordering::Relaxed); + } + } + + fn gpu_failed(&self, what: &str, e: &Error) { + if let Some(g) = &self.gpu + && g.ok.swap(false, Ordering::Relaxed) + { + log::warn!("SAM 2 {what} failed on the GPU ({e}); the CPU takes over"); + } + } + + /// Encode a normalised `[1, 3, 1024, 1024]` frame on the GPU if it is working, else on the CPU. + pub fn encode(&self, pixels: &Tensor) -> Result { + if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { + match g.encoder.encode(pixels) { + Ok(f) => { + let none = || Tensor::zeros((1, 0, 0, 0), DType::F32, &candle_core::Device::Cpu); + let pos = crate::neck::sine_position(128, GRID, GRID)?; + let features = ImageFeatures { features: vec![none()?, none()?, f.low], positions: vec![none()?, none()?, pos] }; + return Ok(Encoded { features, high: f.high }); + } + Err(e) => self.gpu_failed("encoding", &e), + } + } + self.image.encode(pixels) + } + + /// The memory of a frame (its features and the scaled mask), on the GPU if it is working, else on the CPU. + fn memorize(&self, features: &Tensor, mask: &Tensor, present: bool) -> Result { + if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { + match g.memory_encoder.encode(features, mask).and_then(|t| self.memory_encoder.finish(&t, present)) { + Ok(m) => return Ok(m), + Err(e) => self.gpu_failed("memory encoding", &e), + } + } + self.memory_encoder.memory(features, mask, present) + } + + /// Memory attention on the GPU if it is working, else on the CPU. + fn attend(&self, current: &Tensor, current_pos: &Tensor, memory: &Tensor, memory_pos: &Tensor, pointers: usize) -> Result { + if let Some(g) = self.gpu.as_ref().filter(|g| g.ok.load(Ordering::Relaxed)) { + match g.attention.forward(current, current_pos, memory, memory_pos, pointers) { + Ok(t) => return Ok(t), + Err(e) => self.gpu_failed("memory attention", &e), + } + } + self.memory_attention.forward(current, current_pos, memory, memory_pos, pointers) + } + pub fn open(path: &Path) -> Result { VideoModel::load(&Weights::open(path)?) } @@ -221,7 +350,7 @@ impl<'a> VideoTracker<'a> { return Ok(f.clone()); } let pixels = self.frames.frame(frame)?; - let enc = self.model.image.encode(&pixels)?; + let enc = self.model.encode(&pixels)?; let (low, pos) = match (enc.features.features.get(2), enc.features.positions.get(2)) { (Some(l), Some(p)) => (l.clone(), p.clone()), _ => return Err(Error::Compute("vision encoder returned too few levels".into())), @@ -295,7 +424,17 @@ impl<'a> VideoTracker<'a> { /// Run the model on one frame for one object and store the result. fn step(&mut self, id: ObjectId, frame: usize, prompt: Option<&Prompt>, reverse: bool, conditioning: bool) -> Result { + // EFFECTCRAFT_SAM2_PROFILE=1 prints where a frame's time goes (stderr). + let profile = std::env::var_os("EFFECTCRAFT_SAM2_PROFILE").is_some(); + let mut lap = std::time::Instant::now(); + let mut mark = |what: &str| { + if profile { + eprintln!("sam2 frame {frame}: {what} {:.0} ms", lap.elapsed().as_secs_f64() * 1e3); + } + lap = std::time::Instant::now(); + }; let feats = self.features(frame)?; + mark("encode"); let model = self.model.clone(); let tokens = |t: &Tensor| -> Result { Ok(t.flatten_from(2)?.squeeze(0)?.t()?.contiguous()?) }; let conditioned = if conditioning { @@ -303,9 +442,10 @@ impl<'a> VideoTracker<'a> { } else { let obj = self.objects.get(&id).ok_or_else(|| Error::Input("unknown object".into()))?; let (mem, mem_pos, n_ptr) = self.memory_bank(obj, frame, reverse)?; - let out = model.memory_attention.forward(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?; + let out = model.attend(&tokens(&feats.low)?, &tokens(&feats.pos)?, &mem, &mem_pos, n_ptr)?; out.t()?.contiguous()?.reshape((1, 256, GRID, GRID))? }; + mark("memory attention"); let num_points = prompt.map_or(0, |p| p.points.len() + 2 * usize::from(p.bbox.is_some())); let multimask = num_points <= 1; let sparse = match prompt { @@ -315,6 +455,7 @@ impl<'a> VideoTracker<'a> { let dense = model.image.prompt.dense_empty()?; let out = model.image.decode_embedded(&conditioned, &feats.high, &sparse, &dense, multimask)?; + mark("decoder"); let score = scalar(&out.object_score)?; let present = score > 0.0; let ious = out.iou.flatten_all()?.to_vec1::()?; @@ -327,7 +468,9 @@ impl<'a> VideoTracker<'a> { let high = low.upsample_bilinear2d(INPUT_SIZE, INPUT_SIZE, false)?; let for_memory = if prompt.is_some() { high.gt(0f32)?.to_dtype(DType::F32)? } else { sigmoid(&high)? }; let for_memory = ((for_memory * 20.0)? - 10.0)?; - let memory = model.memory_encoder.memory(&feats.low, &for_memory, present)?; + mark("upsample and mask"); + let memory = model.memorize(&feats.low, &for_memory, present)?; + mark("memory encoder"); let iou = ious.get(best).copied().unwrap_or(0.0); let result = FrameMask { frame, object: id, logits: low.flatten_all()?.to_vec1::()?, iou, object_score: score, occluded: !present }; @@ -414,7 +557,7 @@ impl<'a> VideoTracker<'a> { let decoder_present = scalar(&out.object_score)? > 0.0; let pointer = if decoder_present { model.pointer_proj.forward(&token)? } else { model.no_object_pointer.clone() }; let pointer = if present { pointer } else { model.no_object_pointer.clone() }; - let memory = model.memory_encoder.memory(&feats.low, &high, present)?; + let memory = model.memorize(&feats.low, &high, present)?; let score = if present { 10.0 } else { -10.0 }; let result = FrameMask { frame, object: id, logits: low, iou: 1.0, object_score: score, occluded: !present }; let state = self.objects.entry(id).or_default(); diff --git a/crates/sam2/tests/reference.rs b/crates/sam2/tests/reference.rs index 59b09b7f1..16f5cfd24 100644 --- a/crates/sam2/tests/reference.rs +++ b/crates/sam2/tests/reference.rs @@ -223,12 +223,17 @@ fn iou(a: &[f32], b: &[f32]) -> f64 { const VIDEO_LOGIT_TOL: f32 = 0.02; const VIDEO_MIN_IOU: f64 = 0.995; -#[test] -fn video_tracking_matches_sam2_video_model() { +fn check_video_tracking(backend: effectcraft_sam2::session::Backend) { use effectcraft_sam2::prompt::Prompt; use effectcraft_sam2::session::{Cancel, Direction, ObjectId, VideoModel, VideoTracker}; let Some((w, r)) = load_video() else { return }; - let model = std::sync::Arc::new(VideoModel::load(&w).unwrap()); + let model = std::sync::Arc::new(match VideoModel::load_with(&w, backend) { + Ok(m) => m, + Err(e) => { + eprintln!("skipped: {e}"); + return; + } + }); let clip = r.get("clip.pixel_values").unwrap().clone(); let mut tracker = VideoTracker::new(model, Box::new(Clip(clip))).unwrap(); let obj = ObjectId(1); @@ -261,6 +266,19 @@ fn video_tracking_matches_sam2_video_model() { } } +#[test] +fn video_tracking_matches_sam2_video_model() { + check_video_tracking(effectcraft_sam2::session::Backend::Cpu); +} + +/// The same clip on the GPU: every frame's mask against `Sam2VideoModel`, forwards and backwards. The GPU path +/// is within 1e-5 of PyTorch stage by stage; through the clip the bfloat16 memories can round the other way +/// on a few values, so the bound on logits is the video one, and the masks must agree. +#[test] +fn video_tracking_on_the_gpu_matches_sam2_video_model() { + check_video_tracking(effectcraft_sam2::session::Backend::Gpu); +} + #[test] fn tracker_rejects_hostile_input_and_honours_cancel() { use effectcraft_sam2::Error; @@ -603,8 +621,8 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() { use effectcraft_sam2::gpu::{Gpu, encoder::GpuEncoder}; use effectcraft_sam2::model::ImageModel; let Some((w, r)) = load() else { return }; - let gpu = match Gpu::new() { - Ok(g) => std::sync::Arc::new(g), + let gpu = match Gpu::shared() { + Ok(g) => g, Err(e) => { eprintln!("skipped: {e}"); return; @@ -632,3 +650,105 @@ fn gpu_encoder_matches_the_cpu_path_and_pytorch() { eprintln!("skip features vs CPU: {e0:e} / {e1:e}"); assert!(e0 < GPU_TOL && e1 < GPU_TOL, "skip feature errors {e0} / {e1}"); } + +/// The GPU memory attention against PyTorch (the fixture's frame: one memory frame and four pointer tokens). +/// Measured on the RX 7900 GRE, 2026-10-15; the CPU path is at 2.6e-5. +const GPU_ATTN_TOL: f32 = 1e-3; + +#[test] +fn gpu_memory_attention_matches_pytorch() { + use effectcraft_sam2::gpu::{Gpu, memattn::GpuMemoryAttention}; + let Some((w, r)) = load_video() else { return }; + let gpu = match Gpu::shared() { + Ok(g) => g, + Err(e) => { + eprintln!("skipped: {e}"); + return; + } + }; + let att = GpuMemoryAttention::load(gpu, &w).unwrap(); + let get = |k: &str| r.get(k).unwrap().squeeze(1).unwrap(); + let pointers = r.get("mem.attn.num_pointers").unwrap().to_vec1::().unwrap()[0] as usize; + let args = ( + get("mem.attn.current_vision_features"), + get("mem.attn.current_vision_position_embeddings"), + get("mem.attn.memory"), + get("mem.attn.memory_posision_embeddings"), + ); + let out = att.forward(&args.0, &args.1, &args.2, &args.3, pointers).unwrap(); + let want = r.get("mem.attn.out").unwrap().squeeze(0).unwrap().squeeze(0).unwrap(); + let e = max_abs(&out, &want); + eprintln!("GPU memory attention vs PyTorch: max abs error {e:e}"); + assert!(e < GPU_ATTN_TOL, "GPU memory attention error {e}"); + // a full bank (7 memory frames and 16 pointers of 4 tokens): the same maths on a bigger problem, against the CPU path + let cpu = effectcraft_sam2::memattn::MemoryAttention::load(&w).unwrap(); + let m = 7 * 4096 + 64; + let mem = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu).unwrap(); + let mem_pos = Tensor::randn(0f32, 1.0, (m, 64), &Device::Cpu).unwrap(); + let t = std::time::Instant::now(); + let big = att.forward(&args.0, &args.1, &mem, &mem_pos, 64).unwrap(); + eprintln!("GPU memory attention, full bank: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let want = cpu.forward(&args.0, &args.1, &mem, &mem_pos, 64).unwrap(); + let e = max_abs(&big, &want); + eprintln!("full bank vs CPU: max abs error {e:e}"); + assert!(e < GPU_ATTN_TOL, "full-bank error {e}"); + // bad shapes are errors + assert!(att.forward(&args.0, &args.1, &mem, &mem_pos, m + 1).is_err()); + assert!(att.forward(&args.0.narrow(0, 0, 10).unwrap(), &args.1, &mem, &mem_pos, 64).is_err()); +} + +/// The GPU can be switched off in the middle of a clip: the CPU takes over, the device name changes, and +/// the masks still match `Sam2VideoModel` (the memory banks built on the GPU are read by the CPU path). +#[test] +fn switching_the_gpu_off_mid_clip_keeps_the_tracking_right() { + use effectcraft_sam2::prompt::Prompt; + use effectcraft_sam2::session::{Backend, Cancel, Direction, ObjectId, VideoModel, VideoTracker}; + let Some((w, r)) = load_video() else { return }; + let Ok(model) = VideoModel::load_with(&w, Backend::Gpu) else { + eprintln!("skipped: no GPU"); + return; + }; + let model = std::sync::Arc::new(model); + assert_ne!(model.device_name(), "CPU"); + let mut t = VideoTracker::new(model.clone(), Box::new(Clip(r.get("clip.pixel_values").unwrap().clone()))).unwrap(); + t.prompt(2, ObjectId(1), &Prompt { points: vec![([416.0, 512.0], true)], bbox: None }).unwrap(); + let mut seen = vec![]; + t.propagate(2, Direction::Forward, Some(1), &Cancel::new(), &mut |m| seen.push(m)).unwrap(); + model.disable_gpu(); + assert_eq!(model.device_name(), "CPU"); + t.propagate(2, Direction::Forward, None, &Cancel::new(), &mut |m| seen.push(m)).unwrap(); + let last = seen.last().unwrap(); + assert_eq!(last.frame, 6); + let want = r.get("video.fwd.6.pred_masks").unwrap().flatten_all().unwrap().to_vec1::().unwrap(); + let j = iou(&last.logits, &want); + assert!(j >= VIDEO_MIN_IOU, "frame 6 after the switch: IoU {j}"); +} + +/// The GPU memory encoder against the CPU one on the fixture's frame (a mask, the frame's features). The CPU +/// path matches PyTorch to 3.5e-6; the GPU is held to the same output. +#[test] +fn gpu_memory_encoder_matches_the_cpu_one() { + use effectcraft_sam2::gpu::{Gpu, memenc::GpuMemoryEncoder}; + let Some((w, r)) = load_video() else { return }; + let Ok(gpu) = Gpu::shared() else { + eprintln!("skipped: no GPU"); + return; + }; + let cpu = effectcraft_sam2::memory::MemoryEncoder::load(&w).unwrap(); + let enc = GpuMemoryEncoder::load(gpu, &w).unwrap(); + let (features, mask) = (r.get("mem.enc.vision_features").unwrap(), r.get("mem.enc.masks").unwrap()); + let t = std::time::Instant::now(); + let got = enc.encode(features, mask).unwrap(); + eprintln!("GPU memory encoder: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let want = r.get("mem.enc.out_features").unwrap().flatten_from(2).unwrap().squeeze(0).unwrap().t().unwrap().contiguous().unwrap(); + let e = max_abs(&got, &want); + eprintln!("GPU memory encoder vs PyTorch: max abs error {e:e}"); + assert!(e < 1e-3, "memory encoder error {e}"); + // the whole memory (with the occlusion embedding and bfloat16 rounding) equals the CPU one + let (a, b) = (cpu.finish(&got, false).unwrap(), cpu.memory(features, mask, false).unwrap()); + let e = max_abs(&a.features, &b.features); + assert!(e < 0.05, "memory features differ by {e} (bfloat16 steps are 2^-8 of the value)"); + assert_eq!(max_abs(&a.pos, &b.pos), 0.0); + // wrong shapes are errors + assert!(enc.encode(&features.narrow(1, 0, 8).unwrap(), mask).is_err()); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index bd0b1e672..3fff1eec4 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -49,6 +49,10 @@ pub trait MaskModel: Send + Sync { fn tracker(&self) -> Option<&dyn ObjectTracker> { None } + /// What runs the model: the GPU's name or "CPU" (`None`: it does not say). + fn device_name(&self) -> Option { + None + } /// Seconds since the model was last used (`None`: it does not track use, so it is never unloaded /// for being idle). Large models free their memory after a quiet spell and load again on use. fn idle_secs(&self) -> Option { diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 29a9e83a7..b3ceba2ad 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -64,20 +64,28 @@ the CPU. SAM 2 is prompted with at most three foreground and three background points per frame: it was trained on a handful of clicks, and measured confidence collapses with many (0.99 with 1–7 points, 0.25 with 16, 0.00 with 32). -## Speed (this PC: 16 threads, CPU only, release build) -| Step | Time | -|---|---| -| Encode one 1024 × 1024 frame | about **2.1 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes, a parallel bias add) | -| A click on an encoded frame | about 100 ms | -| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.3 s** | -| A tracked frame in all | about **4.6 s** (the first six have fewer memories) | +## Speed (this PC: AMD RX 7900 GRE and a 16-thread CPU, release build) +SAM 2 runs on the GPU when there is one (wgpu compute kernels, so AMD, NVIDIA and Intel alike; candle itself has no AMD +backend) and on the CPU otherwise. The model loads on the GPU automatically; `EFFECTCRAFT_SAM2_DEVICE=cpu` (or `gpu`) +overrides it. A GPU that fails (no adapter, out of memory, a lost device) makes the CPU take over for the rest of the +session, and `roto.status` reports `tracker.device` (the GPU's name, or `CPU`). -A 300-frame clip therefore takes on the order of 23 minutes on this CPU. (An earlier version of this page said one encode -per frame; the memory attention is as large as the encode.) What was tried and did not help: chunking the queries of the -global and memory attention to keep scores in cache. Candle has no AMD/Vulkan backend, so a GPU path is future work. A spike (`--example gpu_matmul`) measured a plain tiled -WGSL matmul on the RX 7900 GRE at 1,900–3,600 GFLOPS against 160–470 for candle's CPU matmul on the same shapes (6–13×), -so hand-written wgpu kernels could bring a tracked frame to a second or less; `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time the -pieces. +| Step | CPU | GPU | +|---|---|---| +| Encode one 1024 × 1024 frame | 2.05 s | **0.24 s** | +| Memory attention (full bank of 7 memory frames) | 2.3 s | **0.29 s** | +| Memory encoder | 0.24 s | 0.05 s | +| Decoder (CPU on both) | 0.11 s | 0.11 s | +| **A tracked frame in all** | **4.3 s** | **0.7 s** | +| A click on an encoded frame | 0.10 s | 0.10 s | + +A 300-frame clip takes about 21 minutes on the CPU and about 3.5 minutes on this GPU; the decoder, still on the CPU, is the +next 16 % of a GPU frame. The CPU numbers came down from 3.4 s per encode at the start of the speed +work (parallel GELU and bias add, cached position encodings, pooling without permutes); chunking the attention queries was +tried and did not help. The GPU kernels (`crates/sam2/src/gpu`) are tested against CPU definitions and PyTorch: the +encoder's 64 × 64 level matches PyTorch to 7e-6, the memory attention to 1.8e-5, and a clip tracked on the GPU matches +`Sam2VideoModel` as the CPU does (worst logit error 0.0032). `cargo run --release -p effectcraft-sam2 --example sam2_bench` +times the pieces on both devices (`EFFECTCRAFT_SAM2_PROFILE=1` prints a frame's breakdown). ## Benchmark `effectcraft-cli bench --roto [--frames N] [--sam2 model.safetensors] [--mobilesam mobile_sam.pt]` runs synthetic clips