From aab00e43fecb5be09b68cb8c7cc007067a614d62 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:44:55 +0200 Subject: [PATCH] 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 --- Cargo.lock | 1 + crates/sam2/Cargo.toml | 1 + crates/sam2/examples/op_bench.rs | 52 +++++++++++++++++++++++++++++ crates/sam2/src/decoder.rs | 8 +++-- crates/sam2/src/memory.rs | 6 ++-- crates/sam2/src/nn.rs | 56 ++++++++++++++++++++++++++++++-- crates/sam2/src/prompt.rs | 6 ++-- 7 files changed, 119 insertions(+), 11 deletions(-) create mode 100644 crates/sam2/examples/op_bench.rs diff --git a/Cargo.lock b/Cargo.lock index 423abb85a..a9b24748d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1835,6 +1835,7 @@ dependencies = [ "candle-core", "candle-nn", "effectcraft-segment", + "rayon", ] [[package]] diff --git a/crates/sam2/Cargo.toml b/crates/sam2/Cargo.toml index 4b6dd7e9f..61afb4117 100644 --- a/crates/sam2/Cargo.toml +++ b/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] diff --git a/crates/sam2/examples/op_bench.rs b/crates/sam2/examples/op_bench.rs new file mode 100644 index 000000000..754248835 --- /dev/null +++ b/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(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> { + 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(()) +} diff --git a/crates/sam2/src/decoder.rs b/crates/sam2/src/decoder.rs index d7a7a6dcf..d6b25a7e4 100644 --- a/crates/sam2/src/decoder.rs +++ b/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 = self.hyper.iter().enumerate().map(|(i, m)| m.forward(&mask_tokens_out.narrow(1, i, 1)?.squeeze(1)?)).collect::>()?; let hyper = Tensor::stack(&hyper, 1)?; diff --git a/crates/sam2/src/memory.rs b/crates/sam2/src/memory.rs index 51e3d3c92..735c4a563 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, 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 { 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)?; diff --git a/crates/sam2/src/nn.rs b/crates/sam2/src/nn.rs index f90958eb1..76bafabac 100644 --- a/crates/sam2/src/nn.rs +++ b/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 { - 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 { 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 { + 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 = (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::().unwrap(), b.flatten_all().unwrap().to_vec1::().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); + } +} diff --git a/crates/sam2/src/prompt.rs b/crates/sam2/src/prompt.rs index 51aade1c7..91496bba3 100644 --- a/crates/sam2/src/prompt.rs +++ b/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 { 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) } }