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

crates/sam2/examples/op_bench.rs

@@ -47,6 +47,22 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     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(())
 }

crates/sam2/examples/sam2_bench.rs

@@ -46,6 +46,16 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
         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(())
 }

crates/sam2/src/hiera.rs

@@ -48,10 +48,16 @@ fn at<T: Copy>(v: &[T], i: usize, what: &str) -> Result<T> {
     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<Tensor> {
-    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::<f32>().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();

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<Tensor> {
+    static CACHE: OnceLock<Mutex<HashMap<(usize, usize, usize), Tensor>>> = 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<Tensor> {
     let eps = 1e-6f32;
     let scale = 2.0 * std::f32::consts::PI;
     let dim_t: Vec<f32> = (0..num_pos_features).map(|i| 10000f32.powf(2.0 * (i / 2) as f32 / num_pos_features as f32)).collect();