taconite_sam3/
lib.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 Brishen Hawkins
2// SPDX-License-Identifier: Apache-2.0
3
4//! SAM3 (Segment Anything 3) text-prompted instance segmentation on an AMD
5//! XDNA NPU.
6//!
7//! The bundle `iron/applications/sam3/export_sam3.py` writes holds every
8//! compiled IRON kernel and the model's weights (NPU ones pre-packed); this
9//! crate replays the forward the Python app (`iron/applications/sam3`)
10//! runs, stage for stage:
11//!
12//! | stage | NPU | host (here) |
13//! |---|---|---|
14//! | CLIP text encoder | | all of it (`text.rs`) |
15//! | ViT backbone, 32 layers | all of it: patch embed, Linears, RoPE, attention, GELU, residual adds + LayerNorms | the first LayerNorm, once (`vit.rs`) |
16//! | FPN neck | ConvTs, 1x1s, 3x3s | GELU, pixel shuffles (`neck.rs`) |
17//! | DETR encoder, 6 layers | projections, self-attention, folded prompt cross-attention, MLP | LayerNorms, prompt softmax (`detr.rs`) |
18//! | DETR decoder, 6 layers | all six layers' vision keys/values (one GEMM) | the 201-query layers, box refinement, scoring (`detr.rs`) |
19//! | mask decoder | pixel-decoder 3x3s, folded mask head, prompt cross-attention | GroupNorms, upsampling (`mask.rs`) |
20//!
21//! [`Sam3::segment`] is the whole thing; the stage methods are public so
22//! `sam3 check` can test each on the bundle's reference inputs.
23
24use std::collections::HashMap;
25use std::fmt;
26use std::path::Path;
27use std::time::{Duration, Instant};
28
29use npu::Buffer;
30
31pub mod bundle;
32pub mod cpu;
33mod detr;
34mod mask;
35mod neck;
36pub mod npu;
37pub mod pack;
38pub mod post;
39mod text;
40pub mod tokenizer;
41mod vit;
42
43pub use detr::Decoded;
44pub use post::{Instance, instances, preprocess};
45pub use text::Text;
46
47use bundle::{Manifest, Store};
48use npu::{Io, MhaIo, Npu};
49use tokenizer::Tokenizer;
50
51#[derive(Debug)]
52pub enum Error {
53    Bundle(String),
54    Npu(String),
55    Input(String),
56}
57
58impl fmt::Display for Error {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        match self {
61            Error::Bundle(m) => write!(f, "bundle: {m}"),
62            Error::Npu(m) => write!(f, "NPU: {m}"),
63            Error::Input(m) => write!(f, "input: {m}"),
64        }
65    }
66}
67
68impl std::error::Error for Error {}
69
70impl From<taconite_bundle::Error> for Error {
71    fn from(e: taconite_bundle::Error) -> Self {
72        Error::Bundle(e.to_string())
73    }
74}
75
76impl From<taconite::Error> for Error {
77    fn from(e: taconite::Error) -> Self {
78        Error::Npu(e.to_string())
79    }
80}
81
82/// Wall time per stage, in first-seen order; NPU dispatch time is kept
83/// under `npu:<kernel>`.
84#[derive(Default, Debug, Clone)]
85pub struct Timing {
86    entries: Vec<(String, Duration)>,
87}
88
89impl Timing {
90    pub fn add(&mut self, key: &str, d: Duration) {
91        match self.entries.iter_mut().find(|(k, _)| k == key) {
92            Some((_, t)) => *t += d,
93            None => self.entries.push((key.to_string(), d)),
94        }
95    }
96
97    pub fn get(&self, key: &str) -> Duration {
98        self.entries.iter().find(|(k, _)| k == key).map_or(Duration::ZERO, |e| e.1)
99    }
100
101    pub fn clear(&mut self) {
102        self.entries.clear();
103    }
104
105    pub fn npu_total(&self) -> Duration {
106        self.entries.iter().filter(|(k, _)| k.starts_with("npu:")).map(|e| e.1).sum()
107    }
108}
109
110impl fmt::Display for Timing {
111    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
112        let ms = |d: &Duration| d.as_secs_f64() * 1e3;
113        let (npu, host): (Vec<_>, Vec<_>) = self.entries.iter().partition(|(k, _)| k.starts_with("npu:"));
114        write!(f, "stages:")?;
115        for (k, d) in &host {
116            write!(f, " {k} {:.0}", ms(d))?;
117        }
118        write!(f, " ms; npu {:.0} ms:", ms(&self.npu_total()))?;
119        for (k, d) in &npu {
120            write!(f, " {} {:.0}", &k[4..], ms(d))?;
121        }
122        Ok(())
123    }
124}
125
126/// Model constants (from the manifest's `param` lines).
127#[derive(Debug, Clone)]
128pub struct Config {
129    pub grid: usize,
130    pub window: usize,
131    pub vit_dim: usize,
132    pub vit_heads: usize,
133    pub vit_layers: usize,
134    pub vit_ffn_pad: usize,
135    pub vit_global: Vec<usize>,
136    pub patch: usize,
137    pub image_size: usize,
138    pub d_model: usize,
139    pub d_heads: usize,
140    pub d_layers: usize,
141    pub d_ffn: usize,
142    pub text_len: usize,
143    pub text_dim: usize,
144    pub text_heads: usize,
145    pub text_layers: usize,
146    pub text_ffn: usize,
147    pub text_eps: f32,
148    pub vit_eps: f32,
149    pub queries: usize,
150    pub dec_layers: usize,
151    pub neck_splits: Vec<usize>,
152    pub mask_size: usize,
153    /// every ViT layer on the device (RoPE, LayerNorms, residual adds);
154    /// bundles from before it host that glue
155    pub vit_device: bool,
156    /// the device ViT's residual stream in f32 (`AddLayerNorm(f32_residual)`);
157    /// bundles from before it keep it bf16
158    pub vit_res_f32: bool,
159}
160
161impl Config {
162    fn from(m: &Manifest) -> Result<Self, Error> {
163        Ok(Config {
164            grid: m.usize("grid")?,
165            window: m.usize("window")?,
166            vit_dim: m.usize("vit_dim")?,
167            vit_heads: m.usize("vit_heads")?,
168            vit_layers: m.usize("vit_layers")?,
169            vit_ffn_pad: m.usize("vit_ffn_pad")?,
170            vit_global: m.list("vit_global")?,
171            patch: m.usize("patch")?,
172            image_size: m.usize("image_size")?,
173            d_model: m.usize("d_model")?,
174            d_heads: m.usize("d_heads")?,
175            d_layers: m.usize("d_layers")?,
176            d_ffn: m.usize("d_ffn")?,
177            text_len: m.usize("text_len")?,
178            text_dim: m.usize("text_dim")?,
179            text_heads: m.usize("text_heads")?,
180            text_layers: m.usize("text_layers")?,
181            text_ffn: m.usize("text_ffn")?,
182            text_eps: m.f32("text_eps")?,
183            vit_eps: m.f32("vit_eps")?,
184            queries: m.usize("queries")?,
185            dec_layers: m.usize("dec_layers")?,
186            neck_splits: m.list("neck_splits")?,
187            mask_size: m.usize("mask_size")?,
188            vit_device: m.param("vit_device").is_ok_and(|v| v == "1"),
189            vit_res_f32: m.param("vit_res_f32").is_ok_and(|v| v == "1"),
190        })
191    }
192
193    /// ViT tokens (72 x 72).
194    pub fn tokens(&self) -> usize {
195        self.grid * self.grid
196    }
197}
198
199/// The model's raw outputs, as `Sam3Model` returns them (batch of one).
200pub struct Output {
201    /// `[Q]` classification logits.
202    pub logits: Vec<f32>,
203    /// `[Q, 4]` boxes, normalised xyxy.
204    pub boxes: Vec<f32>,
205    pub presence: f32,
206    /// `[Q, S, S]` mask logits, `S` = `mask_size` (288).
207    pub masks: Vec<f32>,
208    /// `[S, S]` semantic segmentation logits.
209    pub semantic: Vec<f32>,
210}
211
212/// NPU operands, allocated once per session.
213pub(crate) struct Ios {
214    v_embed: Io,
215    v_qkv: Io,
216    v_o: Io,
217    v_fc1: Io,
218    v_fc2: Io,
219    mha_win: MhaIo,
220    mha_glob: MhaIo,
221    n_in: Io,
222    n_up: Io,
223    conv: HashMap<usize, Io>, // by image side
224    d_qkv: Io,
225    d_o: Io,
226    d_s: Io,
227    d_c: Io,
228    d_fc1: Io,
229    d_fc2: Io,
230    mha_d: MhaIo,
231    dec_kv: Io,
232    m_head: Io,
233    /// device-resident ViT: RoPE output `[T, 2 D]`, the residual stream's
234    /// two ping-pong buffers `[T, D]`
235    pub(crate) rope_out: Option<Buffer>,
236    pub(crate) xres: Vec<Buffer>,
237}
238
239pub struct Sam3 {
240    pub manifest: Manifest,
241    pub store: Store,
242    pub cfg: Config,
243    pub tokenizer: Tokenizer,
244    npu: Npu,
245    /// packed NPU weights, by tensor name
246    w: HashMap<String, Buffer>,
247    /// per-prompt packed weights (folded cross-attentions, mask head)
248    slots: HashMap<String, Buffer>,
249    io: Ios,
250    pub timing: Timing,
251}
252
253impl Sam3 {
254    /// Loads the bundle, opens the NPU, loads every kernel and uploads the
255    /// packed weights.
256    pub fn load(dir: &Path) -> Result<Self, Error> {
257        let manifest = Manifest::load(dir)?;
258        let store = Store::load(dir)?;
259        let cfg = Config::from(&manifest)?;
260        let tokenizer = Tokenizer::load(
261            dir,
262            manifest.usize("bos")? as u32,
263            manifest.usize("eos")? as u32,
264            manifest.usize("pad")? as u32,
265            cfg.text_len,
266        )?;
267        let npu = Npu::open(&manifest)?;
268        let mut w = HashMap::new();
269        let mut names = vec!["v.embed".to_string(), "n.in".into(), "n.up".into(), "dec.kv".into()];
270        for i in 0..cfg.vit_layers {
271            for k in ["qkv", "o", "fc1", "fc2"] {
272                names.push(format!("v.{i}.{k}"));
273            }
274        }
275        for i in 0..cfg.d_layers {
276            for k in ["qkv", "o", "fc1", "fc2"] {
277                names.push(format!("d.{i}.{k}"));
278            }
279        }
280        for i in 0..3 {
281            names.push(format!("n.conv{i}"));
282        }
283        for i in 0..2 {
284            names.push(format!("m.conv{i}"));
285        }
286        if cfg.vit_device {
287            names.push("v.rope_tab.win".into());
288            names.push("v.rope_tab.glob".into());
289        }
290        for n in names {
291            let b = npu.upload(store.bytes(&n)?)?;
292            w.insert(n, b);
293        }
294        let mut slots = HashMap::new();
295        for i in 0..cfg.d_layers {
296            slots.insert(format!("d.{i}.s"), npu.weight_slot("d_s")?);
297            slots.insert(format!("d.{i}.c"), npu.weight_slot("d_c")?);
298        }
299        slots.insert("m.s".into(), npu.weight_slot("d_s")?);
300        slots.insert("m.c".into(), npu.weight_slot("d_c")?);
301        slots.insert("m.head".into(), npu.weight_slot("m_head")?);
302
303        let t = cfg.tokens();
304        let v_fc1 = npu.io("v_fc1", t)?;
305        let v_fc2 = npu.io_chained("v_fc2", &v_fc1)?;
306        let d_fc1 = npu.io("d_fc1", t)?;
307        let d_fc2 = npu.io_chained("d_fc2", &d_fc1)?;
308        let mut conv = HashMap::new();
309        for s in [cfg.grid, 2 * cfg.grid, 4 * cfg.grid] {
310            conv.insert(s, npu.conv_io("conv", s, s)?);
311        }
312        let s = cfg.mask_size;
313        let io = Ios {
314            v_embed: npu.io("v_embed", t)?,
315            v_qkv: npu.io("v_qkv", t)?,
316            v_o: npu.io("v_o", t)?,
317            v_fc1,
318            v_fc2,
319            mha_win: npu.mha_io("mha_win")?,
320            mha_glob: npu.mha_io("mha_glob")?,
321            n_in: npu.io("n_in", t)?,
322            n_up: npu.io("n_up", 4 * t)?,
323            conv,
324            d_qkv: npu.io("d_qkv", t)?,
325            d_o: npu.io("d_o", t)?,
326            d_s: npu.io("d_s", t)?,
327            d_c: npu.io("d_c", t)?,
328            d_fc1,
329            d_fc2,
330            mha_d: npu.mha_io("mha_d")?,
331            dec_kv: npu.io("dec_kv", t)?,
332            m_head: npu.io("m_head", s * s)?,
333            rope_out: if cfg.vit_device { Some(npu.session.alloc(t * 2 * cfg.vit_dim * 2)?) } else { None },
334            xres: if cfg.vit_device {
335                let bytes = t * cfg.vit_dim * if cfg.vit_res_f32 { 4 } else { 2 };
336                vec![npu.session.alloc(bytes)?, npu.session.alloc(bytes)?]
337            } else {
338                vec![]
339            },
340        };
341        Ok(Sam3 { manifest, store, cfg, tokenizer, npu, w, slots, io, timing: Timing::default() })
342    }
343
344    /// Hardware contexts created and evicted so far (evictions happen when
345    /// another process holds some of the NPU's contexts).
346    pub fn contexts(&self) -> (usize, usize) {
347        (self.npu.loads, self.npu.evictions)
348    }
349
350    fn time<T>(&mut self, key: &str, f: impl FnOnce(&mut Self) -> Result<T, Error>) -> Result<T, Error> {
351        let t0 = Instant::now();
352        let r = f(self)?;
353        self.timing.add(key, t0.elapsed());
354        Ok(r)
355    }
356
357    /// Everything for one image and prompt: `pixels` is the preprocessed
358    /// `[3, 1008, 1008]` image ([`preprocess`]). The text encoder (host)
359    /// runs on a side thread while the ViT has the NPU.
360    pub fn segment(&mut self, pixels: &[f32], prompt: &str) -> Result<Output, Error> {
361        let (ids, mask) = self.tokenizer.encode(prompt);
362        if !self.cfg.vit_device {
363            let text = self.time("text", |s| s.text(&ids, &mask))?;
364            return self.forward(pixels, &text);
365        }
366        let Sam3 { store, cfg, npu, io, w, timing, .. } = self;
367        let (text, vit) = std::thread::scope(|s| {
368            let encoder = s.spawn(|| {
369                let t0 = Instant::now();
370                let r = cpu::run_inline(|| text::encode(store, cfg, &ids, &mask));
371                (r, t0.elapsed())
372            });
373            let t0 = Instant::now();
374            let vit = vit::vit_device(npu, io, w, store, cfg, timing, pixels);
375            timing.add("vit", t0.elapsed());
376            let (text, dt) = encoder.join().expect("the text encoder thread");
377            timing.add("text", dt);
378            (text, vit)
379        });
380        let (text, vit) = (text?, vit?);
381        self.finish(&vit, &text)
382    }
383
384    /// The forward from preprocessed pixels and an encoded prompt.
385    pub fn forward(&mut self, pixels: &[f32], text: &Text) -> Result<Output, Error> {
386        let vit = self.time("vit", |s| s.vit(pixels))?;
387        self.finish(&vit, text)
388    }
389
390    /// The forward after the backbone.
391    fn finish(&mut self, vit: &[f32], text: &Text) -> Result<Output, Error> {
392        let fpn = self.time("neck", |s| s.neck(vit))?;
393        let enc = self.time("detr_enc", |s| s.detr_encoder(&fpn[2], text))?;
394        let dec = self.time("detr_dec", |s| s.detr_decoder(&enc, text))?;
395        let (masks, semantic) = self.time("mask_dec", |s| s.mask_decoder(&dec.hidden, &fpn, &enc, text))?;
396        Ok(Output { logits: dec.logits, boxes: dec.boxes, presence: dec.presence, masks, semantic })
397    }
398}
399
400/// Runs `io` (A synced first unless `io` is chained) with weights `w`,
401/// accounting the dispatch time under `npu:<kernel>`.
402fn gemm(npu: &mut Npu, io: &Io, w: &Buffer, timing: &mut Timing) -> Result<(), Error> {
403    let d = npu.run_synced(io, w)?;
404    timing.add(&format!("npu:{}", io.key), d);
405    Ok(())
406}
407
408/// Runs GEMM `io` whose A another kernel wrote (no host sync).
409fn gemm_dev(npu: &mut Npu, io: &Io, w: &Buffer, timing: &mut Timing) -> Result<(), Error> {
410    let d = npu.run(io, w)?;
411    timing.add(&format!("npu:{}", io.key), d);
412    Ok(())
413}
414
415/// Runs kernel `key` over device-resident `args`.
416fn op(npu: &mut Npu, key: &str, args: &[&Buffer], timing: &mut Timing) -> Result<(), Error> {
417    let d = npu.run_args(key, args)?;
418    timing.add(&format!("npu:{key}"), d);
419    Ok(())
420}
421
422fn mha(npu: &mut Npu, io: &MhaIo, timing: &mut Timing) -> Result<(), Error> {
423    let d = npu.run_mha(io)?;
424    timing.add(&format!("npu:{}", io.key), d);
425    Ok(())
426}
427
428/// f32 -> bf16 bits (round to nearest even), over the threads.
429pub(crate) fn narrow(src: &[f32], dst: &mut [u16]) {
430    let n = src.len();
431    cpu::par_rows(&mut dst[..n], 1 << 14, |i0, out| {
432        for (i, o) in out.iter_mut().enumerate() {
433            *o = taconite::f32_to_bf16(src[i0 * (1 << 14) + i]);
434        }
435    });
436}
taconite_sam3/lib.rs

taconite_sam3/
lib.rs
