From b7c78065dba4e0644281688a3163bb5ef8592f76 Mon Sep 17 00:00:00 2001 From: Kristian Myram Date: Thu, 15 Oct 2026 20:51:17 +0200 Subject: [PATCH] S2.61: cache the sine positions, pool without permutes Sine position encodings depend only on the sizes: built once (150 ms of a frame). do_pool takes the maximum over grouped cells instead of permute + max_pool2d + permute (tested equal, odd sizes included). Encode 2.6 s -> 2.4 s per frame in release on 16 threads (3.4 s at the start of the speed work); reference tests unchanged. Chunked global attention was tried and did not help. Co-Authored-By: Claude Sonnet 5.5 --- crates/sam2/examples/op_bench.rs | 16 ++++++++++++++++ crates/sam2/examples/sam2_bench.rs | 10 ++++++++++ crates/sam2/src/hiera.rs | 23 ++++++++++++++++++++--- crates/sam2/src/neck.rs | 19 +++++++++++++++++++ 4 files changed, 65 insertions(+), 3 deletions(-) diff --git a/crates/sam2/examples/op_bench.rs b/crates/sam2/examples/op_bench.rs index 754248835..4e2b5b4a3 100644 --- a/crates/sam2/examples/op_bench.rs +++ b/crates/sam2/examples/op_bench.rs @@ -47,6 +47,22 @@ fn main() -> Result<(), Box> { 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()); + // block 1 (the first stage change): pooling and the wide projections + let p1 = Tensor::randn(0f32, 1.0, (1, 192, 256, 256), &d)?; + time("max_pool2d 2x2 on 1 x 192 x 256 x 256", || p1.max_pool2d_with_stride(2, 2).unwrap()); + let p2 = Tensor::randn(0f32, 1.0, (1024, 192, 8, 8), &d)?; + time("max_pool2d 2x2 on 1024 x 192 x 8 x 8", || p2.max_pool2d_with_stride(2, 2).unwrap()); + let nhwc = Tensor::randn(0f32, 1.0, (1, 256, 256, 192), &d)?; + time("permute NHWC -> NCHW + contiguous (192ch)", || nhwc.permute((0, 3, 1, 2)).unwrap().contiguous().unwrap()); + let w3 = Linear::new(Tensor::randn(0f32, 0.1, (192, 96), &d)?, Some(Tensor::zeros(192, DType::F32, &d)?)); + time("linear 65536x96 -> 192 (block 1 proj)", || w3.forward(&flat).unwrap()); + time("NHWC 2x2 max via reshape + max(2), max(3)", || nhwc.reshape((1, 128, 2, 128, 2, 192)).unwrap().max(2).unwrap().max(3).unwrap()); + time("NHWC 2x2 max via two strided maxes of narrowed", || { + let a = nhwc.reshape((1, 128, 2, 256, 192)).unwrap(); + let m = a.narrow(2, 0, 1).unwrap().maximum(&a.narrow(2, 1, 1).unwrap()).unwrap().squeeze(2).unwrap(); + let b = m.reshape((1, 128, 128, 2, 192)).unwrap(); + b.narrow(3, 0, 1).unwrap().maximum(&b.narrow(3, 1, 1).unwrap()).unwrap().squeeze(3).unwrap() + }); println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/crates/sam2/examples/sam2_bench.rs b/crates/sam2/examples/sam2_bench.rs index 7bc678b7c..ea871a4f0 100644 --- a/crates/sam2/examples/sam2_bench.rs +++ b/crates/sam2/examples/sam2_bench.rs @@ -46,6 +46,16 @@ fn run() -> Result<(), Box> { x = b.forward(&x)?; println!("block {i}: {:.0} ms, out {:?}", t.elapsed().as_secs_f64() * 1e3, x.dims()); } + // The rest of the encode: the neck (1x1 convolutions, top-down sums) and its sine position encodings. + let stages = hiera.forward(&pixels)?; + let t = Instant::now(); + let _ = model.vision.neck.forward(&stages)?; + println!("neck (convolutions + positions): {:.0} ms", t.elapsed().as_secs_f64() * 1e3); + let t = Instant::now(); + for (h, w) in [(32, 32), (64, 64), (128, 128), (256, 256)] { + let _ = effectcraft_sam2::neck::sine_position(128, h, w)?; + } + println!("sine positions for the four levels: {:.0} ms", t.elapsed().as_secs_f64() * 1e3); println!("threads: {}", std::thread::available_parallelism().map_or(1, |n| n.get())); Ok(()) } diff --git a/crates/sam2/src/hiera.rs b/crates/sam2/src/hiera.rs index 90833f7da..61becd418 100644 --- a/crates/sam2/src/hiera.rs +++ b/crates/sam2/src/hiera.rs @@ -48,10 +48,16 @@ fn at(v: &[T], i: usize, what: &str) -> Result { v.get(i).copied().ok_or_else(|| Error::Weights(format!("Hiera config has no {what} for stage {i}"))) } -/// Max pool a `[B, H, W, C]` tensor by `stride` (kernel = stride, no padding, floor). +/// Max pool a `[B, H, W, C]` tensor by `stride` (kernel = stride, no padding, floor), without the +/// permutes a pooling layer would need: it groups the cells and takes the maximum over each group. fn do_pool(x: &Tensor, stride: usize) -> Result { - let x = x.permute((0, 3, 1, 2))?.contiguous()?.max_pool2d_with_stride(stride, stride)?; - Ok(x.permute((0, 2, 3, 1))?.contiguous()?) + let (b, h, w, c) = x.dims4()?; + if stride == 0 { + return Err(Error::Input("pool stride 0".into())); + } + let (oh, ow) = (h / stride, w / stride); + let x = if oh * stride != h || ow * stride != w { x.narrow(1, 0, oh * stride)?.narrow(2, 0, ow * stride)? } else { x.clone() }; + Ok(x.contiguous()?.reshape((b, oh, stride, ow, stride, c))?.max(2)?.max(3)?.contiguous()?) } /// Split `[B, H, W, C]` into `[B * nH * nW, ws, ws, C]` windows, zero padding the bottom/right. @@ -353,6 +359,17 @@ mod tests { assert_eq!(d, 0.0); } + #[test] + fn pooling_equals_a_max_pool_layer_including_odd_sizes() { + let x = Tensor::randn(0f32, 1.0, (2, 7, 9, 5), &Device::Cpu).unwrap(); + let want = x.permute((0, 3, 1, 2)).unwrap().contiguous().unwrap().max_pool2d_with_stride(2, 2).unwrap().permute((0, 2, 3, 1)).unwrap(); + let got = do_pool(&x, 2).unwrap(); + assert_eq!(got.dims(), want.dims()); + let d = (got - want).unwrap().abs().unwrap().sum_all().unwrap().to_scalar::().unwrap(); + assert_eq!(d, 0.0); + assert!(do_pool(&x, 0).is_err()); + } + #[test] fn zero_window_is_an_error_not_a_panic() { let x = Tensor::zeros((1, 4, 4, 2), candle_core::DType::F32, &Device::Cpu).unwrap(); diff --git a/crates/sam2/src/neck.rs b/crates/sam2/src/neck.rs index 0a82ec7a9..02a253db3 100644 --- a/crates/sam2/src/neck.rs +++ b/crates/sam2/src/neck.rs @@ -3,6 +3,9 @@ // HuggingFace Inc. team) to Rust on candle. See NOTICE. //! The feature pyramid neck on top of Hiera, and the vision encoder that joins them. +use std::collections::HashMap; +use std::sync::{Mutex, OnceLock, PoisonError}; + use candle_core::{Device, Tensor}; use crate::hiera::{Hiera, HieraConfig}; @@ -10,7 +13,23 @@ use crate::weights::Weights; use crate::{Error, Result}; /// Sine position encoding `[1, 2 * num_pos_features, h, w]`, normalised (`Sam2SinePositionEmbedding`). +/// It depends only on the sizes, so each is built once (it was 150 ms of a 2.7 s frame). pub fn sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { + static CACHE: OnceLock>> = OnceLock::new(); + let cache = CACHE.get_or_init(Default::default); + if let Some(t) = cache.lock().unwrap_or_else(PoisonError::into_inner).get(&(num_pos_features, h, w)) { + return Ok(t.clone()); + } + let t = build_sine_position(num_pos_features, h, w)?; + let mut map = cache.lock().unwrap_or_else(PoisonError::into_inner); + if map.len() >= 16 { + map.clear(); + } + map.insert((num_pos_features, h, w), t.clone()); + Ok(t) +} + +fn build_sine_position(num_pos_features: usize, h: usize, w: usize) -> Result { let eps = 1e-6f32; let scale = 2.0 * std::f32::consts::PI; let dim_t: Vec = (0..num_pos_features).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / num_pos_features as f32)).collect();