taconite_clip/
lib.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 Brishen Hawkins
2// SPDX-License-Identifier: Apache-2.0
3
4//! CLIP ViT-H/14 (`laion/CLIP-ViT-H-14-laion2B-s32B-b79K`) image and text
5//! embeddings on an AMD XDNA NPU.
6//!
7//! The bundle `iron/applications/clip_vit_h14/export_clip.py` writes holds
8//! every compiled IRON kernel and the model's weights (the NPU ones
9//! pre-packed); this crate replays the forward the Python app runs:
10//!
11//! | | NPU | host (here) |
12//! |---|---|---|
13//! | image | all 32 layers (`tower.rs`) | resize / crop / normalise (`preprocess.rs`), the patch embedding (f32), CLS + position embedding, pre- and post-LayerNorm, the projection |
14//! | text | all 24 layers | BPE tokenizer, token + position embedding, final LayerNorm at the end token, the projection |
15//!
16//! [`Clip::encode_images`] / [`Clip::encode_texts`] give the projected
17//! embeddings, [`logits`] CLIP's scaled cosine similarities.
18
19use std::fmt;
20use std::path::Path;
21use std::time::Instant;
22
23use taconite_bundle::{Manifest, Store};
24use npu::Buffer;
25pub use taconite_sam3::Timing;
26use taconite_sam3::cpu::{dot, ln_row, par_rows};
27use taconite_sam3::tokenizer::Tokenizer;
28
29pub mod npu;
30pub mod preprocess;
31pub mod tower;
32
33use npu::Npu;
34use preprocess::Preprocess;
35use tower::Tower;
36
37pub const VERSION: u32 = 1;
38
39#[derive(Debug)]
40pub enum Error {
41    Bundle(String),
42    Npu(String),
43    Input(String),
44}
45
46impl fmt::Display for Error {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        match self {
49            Error::Bundle(m) => write!(f, "bundle: {m}"),
50            Error::Npu(m) => write!(f, "NPU: {m}"),
51            Error::Input(m) => write!(f, "input: {m}"),
52        }
53    }
54}
55
56impl std::error::Error for Error {}
57
58impl From<taconite_bundle::Error> for Error {
59    fn from(e: taconite_bundle::Error) -> Self {
60        Error::Bundle(e.to_string())
61    }
62}
63
64impl From<taconite::Error> for Error {
65    fn from(e: taconite::Error) -> Self {
66        Error::Npu(e.to_string())
67    }
68}
69
70/// Model constants (the manifest's `param` records).
71#[derive(Debug, Clone)]
72pub struct Config {
73    pub v_layers: usize,
74    pub v_dim: usize,
75    pub v_tokens: usize,
76    /// rows between images (tokens rounded up to 8; see the exporter)
77    pub v_stride: usize,
78    pub v_batch: usize,
79    pub v_rows: usize,
80    pub v_mha_rows: usize,
81    pub v_eps: f32,
82    pub patch: usize,
83    pub image_size: usize,
84    pub t_layers: usize,
85    pub t_dim: usize,
86    pub t_batch: usize,
87    pub t_seq_pad: usize,
88    pub t_rows: usize,
89    pub t_eps: f32,
90    pub context: usize,
91    pub embed_dim: usize,
92    pub logit_scale: f32,
93    pub eos_argmax: bool,
94    pub prompt_template: String,
95}
96
97impl Config {
98    fn load(m: &Manifest) -> Result<Self, Error> {
99        let u = |k: &str| -> Result<usize, Error> { Ok(m.param_as(k)?) };
100        let f = |k: &str| -> Result<f32, Error> { Ok(m.param_as(k)?) };
101        Ok(Config {
102            v_layers: u("v_layers")?,
103            v_dim: u("v_dim")?,
104            v_tokens: u("v_tokens")?,
105            v_stride: u("v_stride")?,
106            v_batch: u("v_batch")?,
107            v_rows: u("v_rows")?,
108            v_mha_rows: u("v_mha_rows")?,
109            v_eps: f("v_eps")?,
110            patch: u("patch")?,
111            image_size: u("image_size")?,
112            t_layers: u("t_layers")?,
113            t_dim: u("t_dim")?,
114            t_batch: u("t_batch")?,
115            t_seq_pad: u("t_seq_pad")?,
116            t_rows: u("t_rows")?,
117            t_eps: f("t_eps")?,
118            context: u("context")?,
119            embed_dim: u("embed_dim")?,
120            logit_scale: f("logit_scale")?,
121            eos_argmax: u("eos_argmax")? != 0,
122            prompt_template: m.param("prompt_template")?.to_string(),
123        })
124    }
125
126    /// Patches a side (16 for 224 / 14).
127    pub fn grid(&self) -> usize {
128        self.image_size / self.patch
129    }
130}
131
132/// Copies `src` (bf16 bits) into the start of `b` and syncs it.
133pub(crate) fn push(src: &[u16], b: &mut Buffer) -> Result<(), Error> {
134    const COPY: usize = 1 << 16;
135    let dst = &mut b.as_mut_slice::<u16>()[..src.len()];
136    par_rows(dst, COPY, |r0, piece| piece.copy_from_slice(&src[r0 * COPY..][..piece.len()]));
137    Ok(b.sync_to_device()?)
138}
139
140/// The first `n` bf16 of `b`, synced from the device, in cached memory.
141pub(crate) fn pull(b: &Buffer, n: usize) -> Result<Vec<u16>, Error> {
142    const COPY: usize = 1 << 16;
143    b.sync_from_device()?;
144    let src = &b.as_slice::<u16>()[..n];
145    let mut dst = vec![0u16; n];
146    par_rows(&mut dst, COPY, |r0, piece| piece.copy_from_slice(&src[r0 * COPY..][..piece.len()]));
147    Ok(dst)
148}
149
150pub struct Clip {
151    pub cfg: Config,
152    pub manifest: Manifest,
153    pub store: Store,
154    pub tokenizer: Tokenizer,
155    pub preprocess: Preprocess,
156    pub timing: Timing,
157    npu: Npu,
158    vision: Tower,
159    text: Tower,
160}
161
162impl Clip {
163    /// Loads a bundle: every kernel onto the NPU, every weight into device
164    /// buffers.
165    pub fn load(dir: &Path) -> Result<Self, Error> {
166        let manifest = Manifest::load(dir, VERSION)?;
167        let cfg = Config::load(&manifest)?;
168        let store = Store::load(dir)?;
169        let u = |k: &str| -> Result<u32, Error> { Ok(manifest.param_as(k)?) };
170        let tokenizer = Tokenizer::load(dir, u("bos")?, u("eos")?, u("pad")?, cfg.context)
171            .map_err(|e| Error::Bundle(e.to_string()))?;
172        let three = |k: &str| -> Result<[f32; 3], Error> {
173            let v: Vec<f32> = manifest.list(k)?;
174            v.try_into().map_err(|_| Error::Bundle(format!("param {k}: expected 3 values")))
175        };
176        let preprocess = Preprocess {
177            shortest: manifest.param_as("resize_shortest")?,
178            crop: manifest.param_as("crop")?,
179            mean: three("mean")?,
180            std: three("std")?,
181        };
182        if preprocess.crop != cfg.image_size {
183            return Err(Error::Bundle("crop size != image size".into()));
184        }
185
186        let npu = Npu::open(&manifest)?;
187        // one MHA dispatch per image: its tokens, padded (masked) to the
188        // kernel's rows, which run into the next image's -- rewritten by
189        // that image's dispatch, issued after
190        let v_mha: Vec<(usize, usize)> = (0..cfg.v_batch).map(|i| (i * cfg.v_stride, cfg.v_mha_rows)).collect();
191        let vision = Tower::new(&npu, &store, "v", cfg.v_layers, cfg.v_dim, cfg.v_rows, cfg.v_eps, &v_mha)?;
192        let text = Tower::new(
193            &npu,
194            &store,
195            "t",
196            cfg.t_layers,
197            cfg.t_dim,
198            cfg.t_rows,
199            cfg.t_eps,
200            &[(0, cfg.t_batch * cfg.t_seq_pad)],
201        )?;
202        Ok(Clip { cfg, manifest, store, tokenizer, preprocess, timing: Timing::default(), npu, vision, text })
203    }
204
205    /// Hardware contexts the bundle holds (of NPU2's 16).
206    pub fn contexts(&self) -> usize {
207        self.npu.contexts
208    }
209
210    /// RGB8 `[h, w, 3]` -> the model input `[3, 224, 224]`.
211    pub fn preprocess(&self, rgb: &[u8], w: usize, h: usize) -> Vec<f32> {
212        self.preprocess.run(rgb, w, h)
213    }
214
215    /// `n` images `[n, 3, 224, 224]` -> their embeddings `[n, embed_dim]`
216    /// (projected, not normalised: `get_image_features`).
217    pub fn encode_images(&mut self, pixels: &[f32]) -> Result<Vec<f32>, Error> {
218        let c = self.cfg.clone();
219        let (s, ps, g, d, t, b) = (c.image_size, c.patch, c.grid(), c.v_dim, c.v_tokens, c.v_batch);
220        let st = c.v_stride;
221        let img = 3 * s * s;
222        if pixels.is_empty() || pixels.len() % img != 0 {
223            return Err(Error::Input(format!("pixels must be [n, 3, {s}, {s}]")));
224        }
225        let n = pixels.len() / img;
226        let (cls, pos) = (self.store.f32("v.cls")?, self.store.f32("v.pos")?);
227        let (pre_w, pre_b) = (self.store.f32("v.ln_pre.w")?, self.store.f32("v.ln_pre.b")?);
228        let mut out = Vec::with_capacity(n * c.embed_dim);
229        for i0 in (0..n).step_by(b) {
230            let ng = (n - i0).min(b);
231            let t0 = Instant::now();
232            // patches, row (image, py, px), columns (channel, ky, kx)
233            let kp = 3 * ps * ps;
234            let mut patches = vec![0f32; ng * g * g * kp];
235            par_rows(&mut patches, kp, |r0, piece| {
236                for (ri, row) in piece.chunks_mut(kp).enumerate() {
237                    let r = r0 + ri;
238                    let (im, p) = (i0 + r / (g * g), r % (g * g));
239                    let (py, px) = (p / g, p % g);
240                    let src = &pixels[im * img..][..img];
241                    for ch in 0..3 {
242                        for ky in 0..ps {
243                            let base = ch * s * s + (py * ps + ky) * s + px * ps;
244                            row[(ch * ps + ky) * ps..][..ps].copy_from_slice(&src[base..][..ps]);
245                        }
246                    }
247                }
248            });
249            // f32, as the model's conv (an NPU GEMM here, its output bf16
250            // before the position embedding and pre-LayerNorm, cost the
251            // image embedding ~0.0025 cosine for ~25 ms)
252            let wp = self.store.f32("v.patch_w")?;
253            let mut emb = vec![0f32; ng * g * g * d];
254            par_rows(&mut emb, d, |r0, piece| {
255                for (ri, row) in piece.chunks_mut(d).enumerate() {
256                    let x = &patches[(r0 + ri) * kp..][..kp];
257                    for (j, o) in row.iter_mut().enumerate() {
258                        *o = dot(&wp[j * kp..][..kp], x);
259                    }
260                }
261            });
262            self.timing.add("v_patches", t0.elapsed());
263            let t1 = Instant::now();
264            // [CLS | patches] + position embedding, pre-LayerNorm
265            // image im's tokens at rows [im * stride, im * stride + t)
266            let mut x = vec![0f32; self.vision.rows * d];
267            par_rows(&mut x[..ng * st * d], d, |r0, piece| {
268                let mut tmp = vec![0f32; d];
269                for (ri, row) in piece.chunks_mut(d).enumerate() {
270                    let r = r0 + ri;
271                    let (im, tok) = (r / st, r % st);
272                    if tok >= t {
273                        continue;
274                    }
275                    for j in 0..d {
276                        let e = if tok == 0 { cls[j] } else { emb[(im * g * g + tok - 1) * d + j] };
277                        tmp[j] = e + pos[tok * d + j];
278                    }
279                    ln_row(&tmp, row, pre_w, pre_b, c.v_eps);
280                }
281            });
282            self.timing.add("v_embed_host", t1.elapsed());
283            let y = self.vision.run(&self.npu, &x, &mut self.timing)?;
284            let t2 = Instant::now();
285            let (post_w, post_b) = (self.store.f32("v.ln_post.w")?, self.store.f32("v.ln_post.b")?);
286            for im in 0..ng {
287                let mut pooled = vec![0f32; d];
288                ln_row(&y[im * st * d..][..d], &mut pooled, post_w, post_b, c.v_eps);
289                out.extend(project(&pooled, self.store.f32("v.proj")?, c.embed_dim));
290            }
291            self.timing.add("v_head", t2.elapsed());
292        }
293        Ok(out)
294    }
295
296    /// A prompt's token ids (`context` long, padded).
297    pub fn tokenize(&self, text: &str) -> Vec<u32> {
298        self.tokenizer.encode(text).0
299    }
300
301    /// Prompts -> their text embeddings `[n, embed_dim]`.
302    pub fn encode_texts(&mut self, texts: &[&str]) -> Result<Vec<f32>, Error> {
303        let t0 = Instant::now();
304        let ids: Vec<u32> = texts.iter().flat_map(|t| self.tokenize(t)).collect();
305        self.timing.add("t_tokenize", t0.elapsed());
306        self.encode_token_ids(&ids)
307    }
308
309    /// `n` token-id rows `[n, context]` -> the text embeddings
310    /// `[n, embed_dim]` (projected, not normalised: `get_text_features`).
311    pub fn encode_token_ids(&mut self, ids: &[u32]) -> Result<Vec<f32>, Error> {
312        let c = self.cfg.clone();
313        let (l, d, sp, b) = (c.context, c.t_dim, c.t_seq_pad, c.t_batch);
314        if ids.is_empty() || ids.len() % l != 0 {
315            return Err(Error::Input(format!("token ids must be [n, {l}]")));
316        }
317        let n = ids.len() / l;
318        let vocab = self.store.shape("t.tok_emb")?[0];
319        if let Some(&bad) = ids.iter().find(|&&i| i as usize >= vocab) {
320            return Err(Error::Input(format!("token id {bad} >= vocabulary size {vocab}")));
321        }
322        let eos = self.tokenizer.eos;
323        let mut out = Vec::with_capacity(n * c.embed_dim);
324        for p0 in (0..n).step_by(b) {
325            let ng = (n - p0).min(b);
326            let t0 = Instant::now();
327            let (tok, pos) = (self.store.f32("t.tok_emb")?, self.store.f32("t.pos")?);
328            // prompt p's tokens at rows [p * seq_pad, p * seq_pad + context)
329            let mut x = vec![0f32; self.text.rows * d];
330            par_rows(&mut x[..ng * sp * d], d, |r0, piece| {
331                for (ri, row) in piece.chunks_mut(d).enumerate() {
332                    let (p, s) = ((r0 + ri) / sp, (r0 + ri) % sp);
333                    if s < l {
334                        let id = ids[(p0 + p) * l + s] as usize;
335                        for j in 0..d {
336                            row[j] = tok[id * d + j] + pos[s * d + j];
337                        }
338                    }
339                }
340            });
341            self.timing.add("t_embed_host", t0.elapsed());
342            let y = self.text.run(&self.npu, &x, &mut self.timing)?;
343            let t1 = Instant::now();
344            let (fw, fb) = (self.store.f32("t.ln_final.w")?, self.store.f32("t.ln_final.b")?);
345            for p in 0..ng {
346                let row = &ids[(p0 + p) * l..][..l];
347                // the end token: the first eos (or, for legacy configs, the
348                // highest id -- the same token for CLIP's vocabulary)
349                let e = if c.eos_argmax {
350                    let m = *row.iter().max().unwrap();
351                    row.iter().position(|&i| i == m).unwrap()
352                } else {
353                    row.iter().position(|&i| i == eos).unwrap_or(0)
354                };
355                let mut pooled = vec![0f32; d];
356                ln_row(&y[(p * sp + e) * d..][..d], &mut pooled, fw, fb, c.t_eps);
357                out.extend(project(&pooled, self.store.f32("t.proj")?, c.embed_dim));
358            }
359            self.timing.add("t_head", t1.elapsed());
360        }
361        Ok(out)
362    }
363
364    /// A class label's prompt (`a photo of a {}`).
365    pub fn prompt(&self, label: &str) -> String {
366        self.cfg.prompt_template.replace("{}", label)
367    }
368}
369
370/// `w [out, in] x` (f32).
371fn project(x: &[f32], w: &[f32], out: usize) -> Vec<f32> {
372    let n = x.len();
373    (0..out).map(|o| dot(&w[o * n..][..n], x)).collect()
374}
375
376/// Rows of `x` scaled to unit length.
377pub fn normalize(x: &[f32], dim: usize) -> Vec<f32> {
378    x.chunks(dim)
379        .flat_map(|r| {
380            let s = 1.0 / dot(r, r).sqrt().max(1e-12);
381            r.iter().map(move |v| v * s)
382        })
383        .collect()
384}
385
386/// CLIP's logits `[n_images, n_texts]`: `logit_scale * cos(image, text)`.
387pub fn logits(images: &[f32], texts: &[f32], dim: usize, scale: f32) -> Vec<f32> {
388    let (a, b) = (normalize(images, dim), normalize(texts, dim));
389    a.chunks(dim).flat_map(|i| b.chunks(dim).map(move |t| scale * dot(i, t))).collect()
390}
391
392/// Row-wise softmax.
393pub fn softmax(x: &[f32], n: usize) -> Vec<f32> {
394    x.chunks(n)
395        .flat_map(|r| {
396            let m = r.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
397            let e: Vec<f32> = r.iter().map(|v| (v - m).exp()).collect();
398            let s: f32 = e.iter().sum();
399            e.into_iter().map(move |v| v / s)
400        })
401        .collect()
402}
taconite_clip/lib.rs

taconite_clip/
lib.rs
