From 921bb00ba6906b8209f043503bf9c52b1b8057f9 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 21:01:17 +0200 Subject: [PATCH] S2.61: parallel bias add fused with GELU; GPU matmul spike Candle adds a bias on one thread; the Linear layers now add it (and apply GELU for the MLP in) in one parallel pass. Encode 2.4 s -> 2.1 s, memory attention 2.4 -> 2.3 s; reference tests unchanged. gpu_matmul example: a tiled WGSL matmul on the RX 7900 GRE reaches 1.9-3.6 TFLOPS against 0.16-0.47 for the CPU. Co-Authored-By: Claude Sonnet 5.5 --- Cargo.lock | 3 + crates/sam2/Cargo.toml | 6 ++ crates/sam2/examples/gpu_matmul.rs | 164 +++++++++++++++++++++++++++++ crates/sam2/src/memory.rs | 4 +- crates/sam2/src/nn.rs | 76 +++++++++++-- docs/roto-brush-sam2.md | 15 +-- 6 files changed, 251 insertions(+), 17 deletions(-) create mode 100644 crates/sam2/examples/gpu_matmul.rs diff --git a/Cargo.lock b/Cargo.lock index a9b24748d..9b174e6cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1832,10 +1832,13 @@ dependencies = [ name = "effectcraft-sam2" version = "0.18.0" dependencies = [ + "bytemuck", "candle-core", "candle-nn", "effectcraft-segment", + "pollster 0.4.0", "rayon", + "wgpu", ] [[package]] diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index 61afb4117..d1d74ebc6 100644 --- a/crates/sam2/Cargo.toml +++ b/crates/sam2/Cargo.toml @@ -17,3 +17,9 @@ rayon = { workspace = true } [target.'cfg(not(target_arch = "wasm32"))'.dependencies] candle-core = "=0.9.2" candle-nn = "=0.9.2" + +# The GPU matmul spike (examples/gpu_matmul.rs) only. +[dev-dependencies] +wgpu = { workspace = true } +pollster = { workspace = true } +bytemuck = { workspace = true } diff --git a/crates/sam2/examples/gpu_matmul.rs b/crates/sam2/examples/gpu_matmul.rs new file mode 100644 index 000000000..1429ebf02 --- /dev/null +++ b/crates/sam2/examples/gpu_matmul.rs @@ -0,0 +1,164 @@ +//! A spike for the GPU path (plan phase 6): how fast is a plain tiled WGSL matrix multiply on this +//! machine's GPU, against the candle CPU matmul the encoder uses today? Candle has no AMD/Vulkan +//! backend; this measures what hand-written wgpu kernels could reach on the shapes SAM 2 needs. +//! +//! `cargo run --release -p effectcraft-sam2 --example gpu_matmul` + +use std::time::Instant; + +use candle_core::{Device, Tensor}; +use wgpu::util::DeviceExt; + +const SHADER: &str = r#" +struct P { m: u32, n: u32, k: u32, pad: u32 } +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; +@group(0) @binding(2) var c: array; +@group(0) @binding(3) var p: P; +var sa: array; +var sb: array; +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_id) li: vec3) { + let row0 = wg.y * 64u; + let col0 = wg.x * 64u; + let tid = li.y * 16u + li.x; + var acc: array; + for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; } + for (var kt = 0u; kt < p.k; kt += 16u) { + for (var i = 0u; i < 4u; i++) { + let idx = tid + i * 256u; + let r = idx / 16u; + let cc = idx % 16u; + let gr = row0 + r; + let gc = kt + cc; + var va = 0.0; + if (gr < p.m && gc < p.k) { va = a[gr * p.k + gc]; } + sa[idx] = va; + let r2 = idx / 64u; + let c2 = idx % 64u; + let gr2 = kt + r2; + let gc2 = col0 + c2; + var vb = 0.0; + if (gr2 < p.k && gc2 < p.n) { vb = b[gr2 * p.n + gc2]; } + sb[idx] = vb; + } + workgroupBarrier(); + for (var kk = 0u; kk < 16u; kk++) { + var ar: array; + var br: array; + for (var i = 0u; i < 4u; i++) { + ar[i] = sa[(li.y * 4u + i) * 16u + kk]; + br[i] = sb[kk * 64u + li.x * 4u + i]; + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { acc[i * 4u + j] += ar[i] * br[j]; } + } + } + workgroupBarrier(); + } + for (var i = 0u; i < 4u; i++) { + for (var j = 0u; j < 4u; j++) { + let gr = row0 + li.y * 4u + i; + let gc = col0 + li.x * 4u + j; + if (gr < p.m && gc < p.n) { c[gr * p.n + gc] = acc[i * 4u + j]; } + } + } +} +"#; + +fn main() -> Result<(), Box> { + 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() }), + )?; + println!("adapter: {:?}", adapter.get_info().name); + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { required_limits: adapter.limits(), ..Default::default() }))?; + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(SHADER.into()) }); + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + // (name, m, k, n): the matmul shapes of SAM 2's encoder and memory attention + let shapes = [ + ("hiera block 0 qkv", 65536usize, 96usize, 288usize), + ("hiera block 0 mlp in", 65536, 96, 384), + ("hiera global attention q.kT (one head)", 4096, 96, 4096), + ("memory attention q.kT (full bank)", 4096, 256, 28736), + ]; + for (name, m, k, n) in shapes { + let a = Tensor::randn(0f32, 1.0, (m, k), &Device::Cpu)?; + let b = Tensor::randn(0f32, 1.0, (k, n), &Device::Cpu)?; + let av = a.flatten_all()?.to_vec1::()?; + let bv = b.flatten_all()?.to_vec1::()?; + let ba = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&av), + usage: wgpu::BufferUsages::STORAGE, + }); + let bb = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&bv), + usage: wgpu::BufferUsages::STORAGE, + }); + let bc = device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: (m * n * 4) as u64, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let params = [m as u32, n as u32, k as u32, 0]; + let bp = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(¶ms), + usage: wgpu::BufferUsages::UNIFORM, + }); + let bg = device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { binding: 0, resource: ba.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 1, resource: bb.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 2, resource: bc.as_entire_binding() }, + wgpu::BindGroupEntry { binding: 3, resource: bp.as_entire_binding() }, + ], + }); + let run = |reps: usize| { + let mut enc = device.create_command_encoder(&Default::default()); + for _ in 0..reps { + let mut pass = enc.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bg, &[]); + pass.dispatch_workgroups(n.div_ceil(64) as u32, m.div_ceil(64) as u32, 1); + } + queue.submit([enc.finish()]); + let _ = device.poll(wgpu::PollType::wait_indefinitely()); + }; + run(1); + let reps = 5; + let t = Instant::now(); + run(reps); + let gpu = t.elapsed().as_secs_f64() / reps as f64; + let flops = 2.0 * m as f64 * k as f64 * n as f64; + // the CPU: candle's matmul, as the encoder uses it + let _ = a.matmul(&b)?; + let t = Instant::now(); + let cpu_reps = 3; + for _ in 0..cpu_reps { + let _ = a.matmul(&b)?; + } + let cpu = t.elapsed().as_secs_f64() / cpu_reps as f64; + println!( + "{name:<42} gpu {:>7.1} ms ({:>6.0} GFLOPS) cpu {:>7.1} ms ({:>5.0} GFLOPS) x{:.1}", + gpu * 1e3, + flops / gpu / 1e9, + cpu * 1e3, + flops / cpu / 1e9, + cpu / gpu + ); + } + Ok(()) +} diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs index 735c4a563..45b470a3f 100644 --- a/crates/sam2/src/memory.rs +++ b/crates/sam2/src/memory.rs @@ -8,7 +8,7 @@ use candle_core::Tensor; use crate::Result; use crate::neck::sine_position; -use crate::nn::{LayerNorm, apply, gelu, layer_norm, linear}; +use crate::nn::{LayerNorm, apply, apply_gelu, gelu, layer_norm, linear}; use crate::weights::Weights; use candle_nn::Linear; @@ -49,7 +49,7 @@ impl FuserBlock { fn forward(&self, x: &Tensor) -> Result { let h = norm_channels_first(&self.norm, &self.depthwise.forward(x)?)?; let h = h.permute((0, 2, 3, 1))?.contiguous()?; - let h = apply(&self.pw2, &gelu(&apply(&self.pw1, &h)?)?)?; + let h = apply(&self.pw2, &apply_gelu(&self.pw1, &h)?)?; let h = h.broadcast_mul(&self.scale)?.permute((0, 3, 1, 2))?.contiguous()?; Ok((x + h)?) } diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs index 76bafabac..256a60da9 100644 --- a/crates/sam2/src/nn.rs +++ b/crates/sam2/src/nn.rs @@ -1,8 +1,8 @@ // Modified work: small layers used by the ported SAM 2 modules (Apache-2.0, see NOTICE). //! Linear, layer norm and a few tensor helpers on candle. -use candle_core::{CpuStorage, CustomOp1, Layout, Shape, Tensor}; -use candle_nn::{Linear, Module}; +use candle_core::{CpuStorage, CustomOp1, CustomOp2, Layout, Shape, Tensor}; +use candle_nn::Linear; use crate::weights::Weights; use crate::{Error, Result}; @@ -41,23 +41,83 @@ impl Mlp { } pub fn forward(&self, x: &Tensor) -> Result { - apply(&self.proj_out, &gelu(&apply(&self.proj_in, x)?)?) + apply(&self.proj_out, &apply_gelu(&self.proj_in, x)?) } } -/// Apply a linear layer to `[.., in]` through a flat 2-D matmul. Candle's `Linear::forward` -/// broadcasts the weight over every leading dimension (copying it per window), which is slow -/// on the 4-D activations of the backbone. -pub fn apply(l: &Linear, x: &Tensor) -> Result { +/// `x + bias` (and, for [`bias_gelu`], then GELU) over rows of `x`, on all cores. Candle adds a bias +/// on one thread, which cost more than the matmul it followed on the wide activations of the early +/// Hiera blocks. +struct BiasAdd { + gelu: bool, +} + +impl CustomOp2 for BiasAdd { + fn name(&self) -> &'static str { + "bias_add_parallel" + } + + fn cpu_fwd(&self, xs: &CpuStorage, xl: &Layout, bs: &CpuStorage, bl: &Layout) -> candle_core::Result<(CpuStorage, Shape)> { + use rayon::prelude::*; + let (CpuStorage::F32(x), CpuStorage::F32(b)) = (xs, bs) else { candle_core::bail!("bias_add_parallel needs f32") }; + let (Some((x0, x1)), Some((b0, b1))) = (xl.contiguous_offsets(), bl.contiguous_offsets()) else { + candle_core::bail!("bias_add_parallel needs contiguous inputs") + }; + let (x, b) = match (x.get(x0..x1), b.get(b0..b1)) { + (Some(x), Some(b)) if !b.is_empty() && x.len() % b.len() == 0 => (x, b), + _ => candle_core::bail!("bias_add_parallel: sizes do not fit"), + }; + let cols = b.len(); + let mut out = vec![0f32; x.len()]; + let rows_per = (64 * 1024 / cols).max(1); + let gelu = self.gelu; + out.par_chunks_mut(rows_per * cols).zip(x.par_chunks(rows_per * cols)).for_each(|(o, i)| { + for (orow, irow) in o.chunks_mut(cols).zip(i.chunks(cols)) { + for ((a, v), bias) in orow.iter_mut().zip(irow).zip(b) { + let y = v + bias; + *a = if gelu { (candle_core::cpu::erf::erf_f32(y * std::f32::consts::FRAC_1_SQRT_2) + 1.0) * 0.5 * y } else { y }; + } + } + }); + Ok((CpuStorage::F32(out), xl.shape().clone())) + } +} + +fn bias_op(x: &Tensor, bias: &Tensor, gelu: bool) -> Result { + if x.dtype() != candle_core::DType::F32 || !x.device().is_cpu() { + let y = x.broadcast_add(bias)?; + return Ok(if gelu { y.gelu_erf()? } else { y }); + } + Ok(x.contiguous()?.apply_op2_no_bwd(&bias.contiguous()?, &BiasAdd { gelu })?) +} + +fn matmul_bias(l: &Linear, x: &Tensor, gelu: bool) -> Result { let dims = x.dims(); let (&last, lead) = dims.split_last().ok_or_else(|| Error::Compute("linear on a scalar".into()))?; let rows: usize = lead.iter().product(); - let y = l.forward(&x.reshape((rows, last))?)?; + let y = x.reshape((rows, last))?.matmul(&l.weight().t()?)?; + let y = match l.bias() { + Some(b) => bias_op(&y, b, gelu)?, + None if gelu => y.gelu_erf()?, + None => y, + }; let mut out: Vec = lead.to_vec(); out.push(y.dim(1)?); Ok(y.reshape(out)?) } +/// Apply a linear layer to `[.., in]` through a flat 2-D matmul and a parallel bias add. Candle's +/// `Linear::forward` broadcasts the weight over every leading dimension (copying it per window) and +/// adds the bias on one thread: both slow on the 4-D activations of the backbone. +pub fn apply(l: &Linear, x: &Tensor) -> Result { + matmul_bias(l, x, false) +} + +/// [`apply`] followed by the exact GELU, in the same pass over the output. +pub fn apply_gelu(l: &Linear, x: &Tensor) -> Result { + matmul_bias(l, x, true) +} + /// Exact (erf) GELU on all cores. Candle's own `gelu_erf` runs on one thread, which made it a third of /// the time of the Hiera blocks; this gives the same values (it calls the same `erf` per element). struct GeluErf; diff --git a/docs/roto-brush-sam2.md b/docs/roto-brush-sam2.md index 580553655..29a9e83a7 100644 --- a/docs/roto-brush-sam2.md +++ b/docs/roto-brush-sam2.md @@ -67,15 +67,16 @@ of clicks, and measured confidence collapses with many (0.99 with 1–7 points, ## Speed (this PC: 16 threads, CPU only, release build) | Step | Time | |---|---| -| Encode one 1024 × 1024 frame | about **2.4 s** (3.4 s before the speed work: a parallel GELU, cached position encodings, pooling without permutes) | -| A click on an encoded frame | about 105 ms | -| Memory attention on a tracked frame (full bank of 7 memory frames) | about **2.4 s** | -| A tracked frame in all | about **5 s** (the first six have fewer memories) | +| 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) | -A 300-frame clip therefore takes on the order of 25 minutes on this CPU. (An earlier version of this page said one encode +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 Burn -or hand-written wgpu route); `cargo run --release -p effectcraft-sam2 --example sam2_bench` and `--example op_bench` time 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. ## Benchmark