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 <noreply@anthropic.com>

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]]

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 }

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<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> c: array<f32>;
+@group(0) @binding(3) var<uniform> p: P;
+var<workgroup> sa: array<f32, 1024>;
+var<workgroup> sb: array<f32, 1024>;
+@compute @workgroup_size(16, 16)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) li: vec3<u32>) {
+    let row0 = wg.y * 64u;
+    let col0 = wg.x * 64u;
+    let tid = li.y * 16u + li.x;
+    var acc: array<f32, 16>;
+    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<f32, 4>;
+            var br: array<f32, 4>;
+            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<dyn std::error::Error>> {
+    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::<f32>()?;
+        let bv = b.flatten_all()?.to_vec1::<f32>()?;
+        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(&params),
+            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(())
+}

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<Tensor> {
         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)?)
     }

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<Tensor> {
-        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<Tensor> {
+/// `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<Tensor> {
+    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<Tensor> {
     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<usize> = 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<Tensor> {
+    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<Tensor> {
+    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;

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