S2.60: parallel exact GELU: encode 3.4 s -> 2.6 s per frame

Candle's gelu_erf runs on one thread and was a third of the Hiera block time.
A rayon CustomOp1 calls the same erf per element (bit-identical, tested). All
reference tests unchanged. op_bench example times the block's operations.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

Cargo.lock

@@ -1835,6 +1835,7 @@ dependencies = [
  "candle-core",
  "candle-nn",
  "effectcraft-segment",
+ "rayon",
 ]
 
 [[package]]

crates/sam2/Cargo.toml

@@ -11,6 +11,7 @@ workspace = true
 
 [dependencies]
 effectcraft-segment = { workspace = true }
+rayon = { workspace = true }
 
 # Native only (D1): candle is not built for wasm32, where this crate is empty.
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]

crates/sam2/examples/op_bench.rs

@@ -0,0 +1,52 @@
+//! Times the tensor operations a Hiera block is made of, at the sizes of its first (biggest) block.
+//! `cargo run --release -p effectcraft-sam2 --example op_bench`
+
+use std::time::Instant;
+
+use candle_core::{DType, Device, Tensor};
+use candle_nn::{Linear, Module};
+
+fn time<T>(name: &str, mut f: impl FnMut() -> T) {
+    let _ = f();
+    let n = 5;
+    let t = Instant::now();
+    for _ in 0..n {
+        let _ = f();
+    }
+    println!("{name:<44} {:>8.1} ms", t.elapsed().as_secs_f64() * 1e3 / n as f64);
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let d = Device::Cpu;
+    // block 0: 1 x 256 x 256 x 96 activations, windows of 8 x 8
+    let x = Tensor::randn(0f32, 1.0, (1, 256, 256, 96), &d)?;
+    let w = Tensor::randn(0f32, 0.1, (288, 96), &d)?;
+    let b = Tensor::zeros(288, DType::F32, &d)?;
+    let lin = Linear::new(w, Some(b));
+    let flat = x.reshape((65536, 96))?;
+    time("linear 65536x96 -> 288 (qkv)", || lin.forward(&flat).unwrap());
+    let big = Tensor::randn(0f32, 1.0, (65536, 384), &d)?;
+    time("gelu_erf on 65536 x 384", || big.gelu_erf().unwrap());
+    time("relu on 65536 x 384", || big.relu().unwrap());
+    time("add 65536 x 384", || (&big + &big).unwrap());
+    let xs = x.clone();
+    time("permute+contiguous (window partition)", || x.reshape((1, 32, 8, 32, 8, 96)).unwrap().permute((0, 1, 3, 2, 4, 5)).unwrap().contiguous().unwrap());
+    let lnw = Tensor::ones(96, DType::F32, &d)?;
+    let lnb = Tensor::zeros(96, DType::F32, &d)?;
+    time("fused layer_norm 256x256x96", || candle_nn::ops::layer_norm(&xs, &lnw, &lnb, 1e-6).unwrap());
+    // windowed attention: 1024 windows x 64 tokens x 96 dims, one head
+    let q = Tensor::randn(0f32, 1.0, (1024, 1, 64, 96), &d)?;
+    let kt = Tensor::randn(0f32, 1.0, (1024, 1, 96, 64), &d)?;
+    time("batched matmul 1024 x (64x96)(96x64)", || q.matmul(&kt).unwrap());
+    let sc = Tensor::randn(0f32, 1.0, (1024, 1, 64, 64), &d)?;
+    time("softmax_last_dim 1024x64x64", || candle_nn::ops::softmax_last_dim(&sc).unwrap());
+    let v = Tensor::randn(0f32, 1.0, (1024, 1, 64, 96), &d)?;
+    time("batched matmul 1024 x (64x64)(64x96)", || sc.matmul(&v).unwrap());
+    // the MLP matmuls
+    let w1 = Linear::new(Tensor::randn(0f32, 0.1, (384, 96), &d)?, Some(Tensor::zeros(384, DType::F32, &d)?));
+    time("linear 65536x96 -> 384 (mlp in)", || w1.forward(&flat).unwrap());
+    let w2 = Linear::new(Tensor::randn(0f32, 0.1, (96, 384), &d)?, Some(Tensor::zeros(96, DType::F32, &d)?));
+    time("linear 65536x384 -> 96 (mlp out)", || w2.forward(&big).unwrap());
+    println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get()));
+    Ok(())
+}

crates/sam2/src/decoder.rs

@@ -8,7 +8,7 @@
 use candle_core::{D, Tensor};
 use candle_nn::{Linear, ops::sigmoid, ops::softmax};
 
-use crate::nn::{LayerNorm, apply, layer_norm, linear};
+use crate::nn::{LayerNorm, apply, gelu, layer_norm, linear};
 use crate::weights::Weights;
 use crate::{Error, Result};
 
@@ -218,8 +218,10 @@ impl MaskDecoder {
         let (s0, s1) = high;
         let up = add_bias(emb.conv_transpose2d(&self.up1.0, 0, 0, 2, 1)?, &self.up1.1)?;
         let up = (up + s1)?;
-        let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?.gelu_erf()?;
-        let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?.gelu_erf()?;
+        let up = self.up_ln.forward(&up.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?;
+        let up = gelu(&up)?;
+        let up = (add_bias(up.conv_transpose2d(&self.up2.0, 0, 0, 2, 1)?, &self.up2.1)? + s0)?;
+        let up = gelu(&up)?;
 
         let hyper: Vec<Tensor> = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::<Result<_>>()?;
         let hyper = Tensor::stack(&hyper, 1)?;

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, layer_norm, linear};
+use crate::nn::{LayerNorm, apply, 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, &apply(&self.pw1, &h)?.gelu_erf()?)?;
+        let h = apply(&self.pw2, &gelu(&apply(&self.pw1, &h)?)?)?;
         let h = h.broadcast_mul(&self.scale)?.permute((0, 3, 1, 2))?.contiguous()?;
         Ok((x + h)?)
     }
@@ -108,7 +108,7 @@ impl MemoryEncoder {
     pub fn encode(&self, features: &Tensor, mask: &Tensor) -> Result<(Tensor, Tensor)> {
         let mut m = mask.clone();
         for (conv, ln) in &self.down {
-            m = norm_channels_first(ln, &conv.forward(&m)?)?.gelu_erf()?;
+            m = gelu(&norm_channels_first(ln, &conv.forward(&m)?)?)?;
         }
         let m = self.final_conv.forward(&m)?;
         let mut x = (self.feature_projection.forward(features)? + m)?;

crates/sam2/src/nn.rs

@@ -1,7 +1,7 @@
 // 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::Tensor;
+use candle_core::{CpuStorage, CustomOp1, Layout, Shape, Tensor};
 use candle_nn::{Linear, Module};
 
 use crate::weights::Weights;
@@ -41,7 +41,7 @@ impl Mlp {
     }
 
     pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
-        apply(&self.proj_out, &apply(&self.proj_in, x)?.gelu_erf()?)
+        apply(&self.proj_out, &gelu(&apply(&self.proj_in, x)?)?)
     }
 }
 
@@ -57,3 +57,55 @@ pub fn apply(l: &Linear, x: &Tensor) -> Result<Tensor> {
     out.push(y.dim(1)?);
     Ok(y.reshape(out)?)
 }
+
+/// 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;
+
+impl CustomOp1 for GeluErf {
+    fn name(&self) -> &'static str {
+        "gelu_erf_parallel"
+    }
+
+    fn cpu_fwd(&self, storage: &CpuStorage, layout: &Layout) -> candle_core::Result<(CpuStorage, Shape)> {
+        use rayon::prelude::*;
+        let CpuStorage::F32(data) = storage else { candle_core::bail!("gelu_erf_parallel needs f32") };
+        let Some((start, end)) = layout.contiguous_offsets() else { candle_core::bail!("gelu_erf_parallel needs a contiguous input") };
+        let src = data.get(start..end).ok_or_else(|| candle_core::Error::Msg("gelu_erf_parallel: layout outside the storage".into()))?;
+        let mut out = vec![0f32; src.len()];
+        const CHUNK: usize = 16 * 1024;
+        out.par_chunks_mut(CHUNK).zip(src.par_chunks(CHUNK)).for_each(|(o, i)| {
+            for (a, v) in o.iter_mut().zip(i) {
+                *a = (candle_core::cpu::erf::erf_f32(v * std::f32::consts::FRAC_1_SQRT_2) + 1.0) * 0.5 * v;
+            }
+        });
+        Ok((CpuStorage::F32(out), layout.shape().clone()))
+    }
+}
+
+/// Exact GELU of `x`, in parallel for f32 tensors.
+pub fn gelu(x: &Tensor) -> Result<Tensor> {
+    if x.dtype() != candle_core::DType::F32 || !x.device().is_cpu() {
+        return Ok(x.gelu_erf()?);
+    }
+    Ok(x.contiguous()?.apply_op1_no_bwd(&GeluErf)?)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use candle_core::Device;
+
+    #[test]
+    fn the_parallel_gelu_equals_candles_bit_for_bit() {
+        let v: Vec<f32> = (0..101_000).map(|i| ((i as f32) * 0.00037).sin() * 6.0 - 0.5).collect();
+        let x = Tensor::from_vec(v, (101, 1000), &Device::Cpu).unwrap();
+        let (a, b) = (gelu(&x).unwrap(), x.gelu_erf().unwrap());
+        assert_eq!(a.flatten_all().unwrap().to_vec1::<f32>().unwrap(), b.flatten_all().unwrap().to_vec1::<f32>().unwrap());
+        // a transposed (non-contiguous) input works, and an empty one does too
+        let t = x.t().unwrap();
+        assert_eq!(gelu(&t).unwrap().dims(), t.dims());
+        let e = Tensor::zeros((0, 4), candle_core::DType::F32, &Device::Cpu).unwrap();
+        assert_eq!(gelu(&e).unwrap().elem_count(), 0);
+    }
+}

crates/sam2/src/prompt.rs

@@ -5,7 +5,7 @@
 
 use candle_core::{Device, Tensor};
 
-use crate::nn::{LayerNorm, layer_norm};
+use crate::nn::{LayerNorm, gelu, layer_norm};
 use crate::weights::Weights;
 use crate::{Error, Result};
 
@@ -166,8 +166,8 @@ impl PromptEncoder {
             Ok(x.conv2d(&c.0, 0, c.3, 1, 1)?.broadcast_add(&c.1.reshape((1, c.1.dim(0)?, 1, 1))?)?)
         };
         let cl = |x: &Tensor, ln: &LayerNorm| -> Result<Tensor> { Ok(ln.forward(&x.permute((0, 2, 3, 1))?)?.permute((0, 3, 1, 2))?.contiguous()?) };
-        let x = cl(&step(mask, c1)?, &self.mask.ln1)?.gelu_erf()?;
-        let x = cl(&step(&x, c2)?, &self.mask.ln2)?.gelu_erf()?;
+        let x = gelu(&cl(&step(mask, c1)?, &self.mask.ln1)?)?;
+        let x = gelu(&cl(&step(&x, c2)?, &self.mask.ln2)?)?;
         step(&x, c3)
     }
 }