taconite_sam3/
neck.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 Brishen Hawkins
2// SPDX-License-Identifier: Apache-2.0
3
4//! The FPN neck and the 3x3 convolutions it shares with the mask decoder,
5//! as `iron/applications/sam3/neck_npu.py`.
6//!
7//! A 3x3 (pad 1) conv is one GEMM over an overlapping view: the host writes
8//! `Y[p] = [X_pad[p] | X_pad[p + Wp] | X_pad[p + 2 Wp]]` (pixel `p` of the
9//! zero-bordered image at row pitch `Wp = W + 2`, next to the pixels one and
10//! two rows below), so output pixel `p`'s window is `Y[p..p + 3]`,
11//! contiguous; a GEMM row covers two pixels (`K = 4D` at row stride `2D`)
12//! against a banded B. The output is pixel-major `[H x Wp, OC]` with two
13//! junk columns per image row, dropped here.
14//!
15//! Levels 0-2 start from the same `[72 x 72, 1024]` map, so their ConvT /
16//! 1x1 first steps are one GEMM (columns `[level 0 | level 1 | level 2]`);
17//! level 0's GELU and pixel shuffle happen here, then its second ConvT
18//! (1x1 folded in) on the NPU; every level ends with its 3x3.
19
20use taconite::bf16_to_f32;
21
22use crate::cpu::{gelu_bf16, par_rows};
23use crate::npu::pull;
24use crate::{Error, Sam3, gemm, narrow};
25
26impl Sam3 {
27    /// 3x3 conv, pad 1: `x [side, side, C]` (bf16) -> `conv(x) + bias`
28    /// `[side, side, OC]` (f32), weights `self.w[wkey]`.
29    pub(crate) fn conv3x3(&mut self, x: &[u16], side: usize, wkey: &str, bias: &[f32]) -> Result<Vec<f32>, Error> {
30        let g = self.npu.spec("conv")?;
31        let (k, lda) = (g.k, g.lda);
32        let d = (k - lda) / 2;
33        let c = d / 3;
34        let oc = bias.len();
35        let wp = side + 2;
36        let n = side * wp;
37        let t0 = std::time::Instant::now();
38        let mut ybuf = vec![0u16; n * d];
39        par_rows(&mut ybuf, d, |p0, piece| {
40            for (pi, row) in piece.chunks_mut(d).enumerate() {
41                let p = p0 + pi;
42                for dy in 0..3 {
43                    let r = p + dy * wp;
44                    let (yy, xx) = (r / wp, r % wp);
45                    let dst = &mut row[dy * c..(dy + 1) * c];
46                    if (1..=side).contains(&yy) && (1..=side).contains(&xx) {
47                        let s = ((yy - 1) * side + xx - 1) * c;
48                        dst.copy_from_slice(&x[s..s + c]);
49                    } else {
50                        dst.fill(0);
51                    }
52                }
53            }
54        });
55        let io = self.io.conv.get_mut(&side).ok_or_else(|| Error::Input(format!("no conv buffers for {side} px")))?;
56        io.set_a(&ybuf)?;
57        self.timing.add("conv_in", t0.elapsed());
58        let io = &self.io.conv[&side];
59        gemm(&mut self.npu, io, &self.w[wkey], &mut self.timing)?;
60        let t0 = std::time::Instant::now();
61        let out_bf = pull(&io.c, n * oc)?;
62        let mut out = vec![0f32; side * side * oc];
63        par_rows(&mut out, side * oc, |y0, piece| {
64            for (yi, row) in piece.chunks_mut(side * oc).enumerate() {
65                let src = &out_bf[(y0 + yi) * wp * oc..][..side * oc];
66                for (j, o) in row.iter_mut().enumerate() {
67                    *o = bf16_to_f32(src[j]) + bias[j % oc];
68                }
69            }
70        });
71        self.timing.add("conv_out", t0.elapsed());
72        Ok(out)
73    }
74
75    /// The backbone's `[T, 1024]` (raster) -> FPN levels 0-2, channels-last
76    /// f32: `[288^2, 256]`, `[144^2, 256]`, `[72^2, 256]`.
77    pub fn neck(&mut self, vit: &[f32]) -> Result<[Vec<f32>; 3], Error> {
78        let cfg = self.cfg.clone();
79        let (g, t) = (cfg.grid, cfg.tokens());
80        let (s0, s1, s2) = (cfg.neck_splits[0], cfg.neck_splits[1], cfg.neck_splits[2]);
81        let width = s0 + s1 + s2;
82        let fpn = cfg.d_model;
83        let mut vb = vec![0u16; vit.len()];
84        narrow(vit, &mut vb);
85        self.io.n_in.set_a(&vb)?;
86        gemm(&mut self.npu, &self.io.n_in, &self.w["n.in"], &mut self.timing)?;
87        let t0 = std::time::Instant::now();
88        let y = self.io.n_in.get_c(t)?;
89        let y = &y[..];
90
91        // level 0: gelu, then shuffle the 2x2 taps onto the 144 grid
92        let mid = s0 / 4;
93        let g2 = 2 * g;
94        let mut a0 = vec![0u16; 4 * t * mid];
95        par_rows(&mut a0, mid, |r0, piece| {
96            for (ri, row) in piece.chunks_mut(mid).enumerate() {
97                let r = r0 + ri;
98                let (yy, xx) = (r / g2, r % g2);
99                let src = &y[((yy / 2) * g + xx / 2) * width + ((yy % 2) * 2 + xx % 2) * mid..][..mid];
100                gelu_bf16(src, row);
101            }
102        });
103        // levels 1 and 2 straight from the shared GEMM
104        let shuffle = |src: &[u16], h: usize, off: usize| -> Vec<u16> {
105            let mut out = vec![0u16; 4 * h * h * fpn];
106            par_rows(&mut out, fpn, |r0, piece| {
107                for (ri, row) in piece.chunks_mut(fpn).enumerate() {
108                    let r = r0 + ri;
109                    let (yy, xx) = (r / (2 * h), r % (2 * h));
110                    let s = ((yy / 2) * h + xx / 2) * width + off + ((yy % 2) * 2 + xx % 2) * fpn;
111                    row.copy_from_slice(&src[s..s + fpn]);
112                }
113            });
114            out
115        };
116        let u1 = shuffle(y, g, s0);
117        let mut u2 = vec![0u16; t * fpn];
118        for (p, row) in u2.chunks_mut(fpn).enumerate() {
119            row.copy_from_slice(&y[p * width + s0 + s1..][..fpn]);
120        }
121        self.io.n_up.set_a(&a0)?;
122        self.timing.add("neck_gelu_shuffle", t0.elapsed());
123        gemm(&mut self.npu, &self.io.n_up, &self.w["n.up"], &mut self.timing)?;
124        let t0 = std::time::Instant::now();
125        let up = self.io.n_up.get_c(4 * t)?;
126        let n_up = self.npu.spec("n_up")?.n;
127        let mut u0 = vec![0u16; 16 * t * fpn];
128        {
129            let g4 = 4 * g;
130            par_rows(&mut u0, fpn, |r0, piece| {
131                for (ri, row) in piece.chunks_mut(fpn).enumerate() {
132                    let r = r0 + ri;
133                    let (yy, xx) = (r / g4, r % g4);
134                    let s = ((yy / 2) * g2 + xx / 2) * n_up + ((yy % 2) * 2 + xx % 2) * fpn;
135                    row.copy_from_slice(&up[s..s + fpn]);
136                }
137            });
138        }
139        self.timing.add("neck_up_shuffle", t0.elapsed());
140        let b: Vec<Vec<f32>> =
141            (0..3).map(|i| self.store.f32(&format!("n.conv{i}.b")).map(<[f32]>::to_vec)).collect::<Result<_, _>>()?;
142        let f0 = self.conv3x3(&u0, 4 * g, "n.conv0", &b[0])?;
143        let f1 = self.conv3x3(&u1, 2 * g, "n.conv1", &b[1])?;
144        let f2 = self.conv3x3(&u2, g, "n.conv2", &b[2])?;
145        Ok([f0, f1, f2])
146    }
147}
taconite_sam3/neck.rs

taconite_sam3/
neck.rs
