Skip to main content

kime_cpu/
plan.rs

1//! The CPU backend: graphs lowered to a list of steps over one arena, run on a persistent pool.
2//!
3//! Lowering resolves every value to an arena offset, every weight to its converted tensor and every
4//! RoPE base to its table, and checks every shape, so running a step is a match and a call. The
5//! arena, the per worker attention scratch and the per batch index tables are sized for the bucket
6//! when the plan is built, and nothing on the run path allocates.
7//!
8//! Each step runs on the rows the batch has rather than the bucket's padded count. Padding only
9//! matters to a backend that captures a fixed shape, and on the CPU it would be wasted work. The
10//! kernels are the reference ones, so a plan gives the same bits as [`Compat`](crate::Compat).
11
12use std::cell::UnsafeCell;
13use std::sync::Arc;
14use std::time::Instant;
15
16use kime_tensor::plan::{Epilogue, Graph, Op, Rows, Val, layout};
17use kime_tensor::{Backend, Batch, Bucket, Caps, Error, HostTensor, Outputs, Result};
18
19use crate::attention::{self, HEAD, QB};
20use crate::gemm::Gemm;
21use crate::ops::{Rope, geglu, layer_norm};
22use crate::par::{self, Shared};
23use crate::pool::Pool;
24
25/// A weight converted to f32.
26#[derive(Debug)]
27pub struct Tensor {
28    /// Shape.
29    pub shape: Vec<usize>,
30    /// Row major values.
31    pub data: Vec<f32>,
32}
33
34/// Every weight of a checkpoint, shared by all the plans built from it.
35#[derive(Debug, Clone)]
36pub struct Weights(Arc<[Tensor]>);
37
38/// The CPU backend, which owns its threads.
39#[derive(Debug)]
40pub struct CpuBackend {
41    pool: Pool,
42}
43
44impl CpuBackend {
45    /// A backend on `threads` threads, the calling thread included.
46    #[must_use]
47    pub fn new(threads: usize) -> Self {
48        Self { pool: Pool::new(threads.max(1)) }
49    }
50
51    /// Threads.
52    #[must_use]
53    pub fn threads(&self) -> usize {
54        self.pool.threads()
55    }
56}
57
58/// A value resolved to its place in the arena.
59#[derive(Debug, Clone, Copy)]
60struct Loc {
61    off: usize,
62    rows: Rows,
63    width: usize,
64}
65
66#[derive(Debug, Clone, Copy)]
67enum Step {
68    Embed { table: usize, out: Loc },
69    LayerNorm { x: Loc, w: usize, b: Option<usize>, eps: f64, out: Loc },
70    Gemm { a: Loc, w: usize, b: Option<usize>, ep: Epilogue, out: Loc },
71    Rope { qkv: Loc, rope: usize },
72    Attention { qkv: Loc, window: Option<usize>, out: Loc },
73    GeGlu { x: Loc, out: Loc },
74    AddType { h: Loc, table: usize },
75    Gather { h: Loc, out: Loc },
76    ActFeatures { h: Loc, logits: Loc, out: Loc },
77}
78
79impl Step {
80    fn name(&self) -> &'static str {
81        match self {
82            Step::Embed { .. } => "embed",
83            Step::LayerNorm { .. } => "layer norm",
84            Step::Gemm { .. } => "gemm",
85            Step::Rope { .. } => "rope",
86            Step::Attention { .. } => "attention",
87            Step::GeGlu { .. } => "geglu",
88            Step::AddType { .. } => "type embedding",
89            Step::Gather { .. } => "gather markers",
90            Step::ActFeatures { .. } => "act features",
91        }
92    }
93}
94
95/// One slot per worker, each touched only by its own worker.
96struct PerWorker<T>(Vec<UnsafeCell<T>>);
97
98// SAFETY: slot w is only reached through `get(w)` from the task running on worker w, and the pool
99// never runs two tasks on one worker at once.
100unsafe impl<T: Send> Sync for PerWorker<T> {}
101
102impl<T> PerWorker<T> {
103    /// # Safety
104    ///
105    /// Only the task running on `worker` may call this, and only for its own worker.
106    #[allow(clippy::mut_from_ref)]
107    unsafe fn get(&self, worker: usize) -> &mut T {
108        // SAFETY: the caller is the only user of slot `worker` right now.
109        unsafe { &mut *self.0[worker].get() }
110    }
111}
112
113/// A graph lowered for one bucket.
114pub struct CpuPlan {
115    w: Weights,
116    bucket: Bucket,
117    steps: Vec<Step>,
118    arena: Vec<f32>,
119    ropes: Vec<Rope>,
120    logits: Loc,
121    act: Loc,
122    /// Token starts of each sequence, then marker starts, then the sequence of each token, then
123    /// the query blocks attention runs. Rebuilt per batch in place.
124    cu: Vec<usize>,
125    mcu: Vec<usize>,
126    row_seq: Vec<u32>,
127    blocks: Vec<(u32, u32)>,
128    scratch: PerWorker<Vec<f32>>,
129    /// Nanoseconds per step, summed over runs, when profiling.
130    profile: Option<Vec<u64>>,
131}
132
133impl std::fmt::Debug for CpuPlan {
134    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135        f.debug_struct("CpuPlan")
136            .field("bucket", &self.bucket)
137            .field("steps", &self.steps.len())
138            .field("arena", &self.arena.len())
139            .finish_non_exhaustive()
140    }
141}
142
143impl CpuPlan {
144    /// The bucket it was built for.
145    #[must_use]
146    pub fn bucket(&self) -> Bucket {
147        self.bucket
148    }
149
150    /// Arena size in bytes.
151    #[must_use]
152    pub fn arena_bytes(&self) -> usize {
153        self.arena.len() * 4
154    }
155
156    /// Starts timing every step, which costs two clock reads per step.
157    pub fn profile(&mut self) {
158        self.profile = Some(vec![0; self.steps.len()]);
159    }
160
161    /// Time per kind of step since [`CpuPlan::profile`], in nanoseconds, largest first.
162    #[must_use]
163    pub fn timings(&self) -> Vec<(&'static str, u64)> {
164        let mut by: Vec<(&'static str, u64)> = Vec::new();
165        for (s, &ns) in self.steps.iter().zip(self.profile.iter().flatten()) {
166            match by.iter_mut().find(|b| b.0 == s.name()) {
167                Some(b) => b.1 += ns,
168                None => by.push((s.name(), ns)),
169            }
170        }
171        by.sort_by_key(|b| std::cmp::Reverse(b.1));
172        by
173    }
174}
175
176/// A pointer to an arena that steps carve disjoint slices from.
177#[derive(Clone, Copy)]
178struct Arena(*mut f32);
179
180// SAFETY: the plan hands out slices of the arena only as the layout allows: values live at the
181// same time never overlap, and within a step each task writes rows no other task touches.
182unsafe impl Send for Arena {}
183// SAFETY: as above.
184unsafe impl Sync for Arena {}
185
186impl Arena {
187    /// # Safety
188    ///
189    /// The range must be in the arena and not written by anyone else while the slice lives.
190    unsafe fn slice<'a>(self, off: usize, len: usize) -> &'a [f32] {
191        // SAFETY: by the caller.
192        unsafe { std::slice::from_raw_parts(self.0.add(off), len) }
193    }
194
195    /// # Safety
196    ///
197    /// The range must be in the arena and not touched by anyone else while the slice lives.
198    #[allow(clippy::mut_from_ref)]
199    unsafe fn slice_mut<'a>(self, off: usize, len: usize) -> &'a mut [f32] {
200        // SAFETY: by the caller.
201        unsafe { std::slice::from_raw_parts_mut(self.0.add(off), len) }
202    }
203}
204
205/// Rows per task for the row wise steps.
206const ROWS: usize = 16;
207
208impl Backend for CpuBackend {
209    type Weights = Weights;
210    type Plan = CpuPlan;
211
212    fn caps(&self) -> Caps {
213        Caps { name: "cpu", threads: self.threads(), graphs: false, unified_memory: true }
214    }
215
216    fn upload(&self, tensors: &[HostTensor<'_>]) -> Result<Weights> {
217        let t = par::map(tensors.len(), self.threads(), |i| {
218            let h = &tensors[i];
219            let n = h.bytes.len() / h.dtype.size();
220            if n != h.shape.iter().product::<usize>() {
221                return Err(Error::Unsupported(format!(
222                    "tensor {i} has {n} values for {:?}",
223                    h.shape
224                )));
225            }
226            let data = (0..n).map(|j| h.dtype.read_f32(h.bytes, j)).collect();
227            Ok(Tensor { shape: h.shape.to_vec(), data })
228        });
229        Ok(Weights(t.into_iter().collect::<Result<Vec<_>>>()?.into()))
230    }
231
232    fn lower(&self, w: &Weights, graph: &Graph, bucket: Bucket) -> Result<CpuPlan> {
233        let lay = layout(graph, |r| bucket.rows(r));
234        let loc = |v: Val| {
235            let s = graph.shape(v);
236            Loc { off: lay.offsets[v.0 as usize], rows: s.rows, width: s.width }
237        };
238        let bad = |m: String| Err(Error::Unsupported(m));
239        let shape = |i: usize| -> Result<&[usize]> {
240            match w.0.get(i) {
241                Some(t) => Ok(&t.shape),
242                None => Err(Error::Unsupported(format!("weight {i} is not in the checkpoint"))),
243            }
244        };
245        let mut ropes: Vec<(u64, Rope)> = Vec::new();
246        let mut steps = Vec::with_capacity(graph.ops.len());
247        for (i, op) in graph.ops.iter().enumerate() {
248            let step = match *op {
249                Op::Embed { table, out } => {
250                    let out = loc(out);
251                    if shape(table)?.get(1) != Some(&out.width) || out.rows != Rows::Tokens {
252                        return bad(format!("op {i}: embedding table does not match its output"));
253                    }
254                    Step::Embed { table, out }
255                }
256                Op::LayerNorm { x, w: nw, b, eps, out } => {
257                    let (x, out) = (loc(x), loc(out));
258                    let ok = shape(nw)? == [x.width]
259                        && b.map_or(Ok(true), |b| shape(b).map(|s| s == [x.width]))?
260                        && x.width == out.width
261                        && x.rows == out.rows;
262                    if !ok {
263                        return bad(format!("op {i}: layer norm shapes do not match"));
264                    }
265                    Step::LayerNorm { x, w: nw, b, eps, out }
266                }
267                Op::Gemm { a, w: gw, b, epilogue, out } => {
268                    let (a, out) = (loc(a), loc(out));
269                    let ok = shape(gw)? == [out.width, a.width]
270                        && b.map_or(Ok(true), |b| shape(b).map(|s| s == [out.width]))?
271                        && a.rows == out.rows;
272                    if !ok {
273                        return bad(format!("op {i}: gemm shapes do not match"));
274                    }
275                    Step::Gemm { a, w: gw, b, ep: epilogue, out }
276                }
277                Op::Rope { qkv, theta } => {
278                    let qkv = loc(qkv);
279                    if !qkv.width.is_multiple_of(3 * HEAD) || qkv.rows != Rows::Tokens {
280                        return bad(format!("op {i}: rope needs token rows of 3 heads 64"));
281                    }
282                    let at = match ropes.iter().position(|r| r.0 == theta.to_bits()) {
283                        Some(at) => at,
284                        None => {
285                            ropes.push((theta.to_bits(), Rope::new(theta, HEAD, bucket.tokens)));
286                            ropes.len() - 1
287                        }
288                    };
289                    Step::Rope { qkv, rope: at }
290                }
291                Op::Attention { qkv, window, out } => {
292                    let (qkv, out) = (loc(qkv), loc(out));
293                    let ok = qkv.width.is_multiple_of(3 * HEAD)
294                        && out.width * 3 == qkv.width
295                        && qkv.rows == Rows::Tokens
296                        && out.rows == Rows::Tokens;
297                    if !ok {
298                        return bad(format!("op {i}: attention shapes do not match"));
299                    }
300                    Step::Attention { qkv, window, out }
301                }
302                Op::GeGlu { x, out } => {
303                    let (x, out) = (loc(x), loc(out));
304                    if x.width != 2 * out.width || x.rows != out.rows {
305                        return bad(format!("op {i}: geglu input is not twice its output"));
306                    }
307                    Step::GeGlu { x, out }
308                }
309                Op::AddType { h, table } => {
310                    let h = loc(h);
311                    if shape(table)?.get(1) != Some(&h.width) || h.rows != Rows::Tokens {
312                        return bad(format!("op {i}: type table does not match"));
313                    }
314                    Step::AddType { h, table }
315                }
316                Op::GatherMarkers { h, out } => {
317                    let (h, out) = (loc(h), loc(out));
318                    if h.width != out.width || h.rows != Rows::Tokens || out.rows != Rows::Markers {
319                        return bad(format!("op {i}: gather shapes do not match"));
320                    }
321                    Step::Gather { h, out }
322                }
323                Op::ActFeatures { h, logits, out } => {
324                    let (h, logits, out) = (loc(h), loc(logits), loc(out));
325                    let ok = out.width == h.width + 4
326                        && logits.width == 1
327                        && logits.rows == Rows::Markers
328                        && out.rows == Rows::Seqs;
329                    if !ok {
330                        return bad(format!("op {i}: act feature shapes do not match"));
331                    }
332                    Step::ActFeatures { h, logits, out }
333                }
334            };
335            steps.push(step);
336        }
337        let (Some(logits), Some(act)) = (graph.logits, graph.act) else {
338            return bad("the graph has no logits or act output".into());
339        };
340        let (logits, act) = (loc(logits), loc(act));
341        if logits.width != 1
342            || logits.rows != Rows::Markers
343            || act.width != 2
344            || act.rows != Rows::Seqs
345        {
346            return bad(
347                "outputs must be one logit per marker and two act logits per sequence".into()
348            );
349        }
350        let threads = self.threads();
351        Ok(CpuPlan {
352            w: w.clone(),
353            bucket,
354            steps,
355            arena: vec![0.0; lay.len],
356            ropes: ropes.into_iter().map(|r| r.1).collect(),
357            logits,
358            act,
359            cu: Vec::with_capacity(bucket.seqs + 1),
360            mcu: Vec::with_capacity(bucket.seqs + 1),
361            row_seq: Vec::with_capacity(bucket.tokens),
362            blocks: Vec::with_capacity(bucket.tokens.div_ceil(QB) + bucket.seqs),
363            scratch: PerWorker(
364                (0..threads).map(|_| UnsafeCell::new(Vec::with_capacity(bucket.tokens))).collect(),
365            ),
366            profile: None,
367        })
368    }
369
370    fn run(&self, plan: &mut CpuPlan, batch: &Batch<'_>, out: &mut Outputs) -> Result<()> {
371        let (t, s, m) = (batch.ids.len(), batch.seqs(), batch.markers.len());
372        if !plan.bucket.holds(t, s, m) {
373            return Err(Error::Batch(format!("batch does not fit bucket {}", plan.bucket)));
374        }
375        plan.cu.clear();
376        plan.cu.extend(batch.cu.iter().map(|&c| c as usize));
377        plan.mcu.clear();
378        plan.mcu.extend(batch.mcu.iter().map(|&c| c as usize));
379        plan.row_seq.clear();
380        plan.blocks.clear();
381        for q in 0..s {
382            let (lo, hi) = (plan.cu[q], plan.cu[q + 1]);
383            plan.row_seq.extend(std::iter::repeat_n(q as u32, hi - lo));
384            plan.blocks.extend((lo..hi).step_by(QB).map(|q0| (q as u32, q0 as u32)));
385        }
386        let ctx = Ctx {
387            pool: &self.pool,
388            w: &plan.w.0,
389            arena: Arena(plan.arena.as_mut_ptr()),
390            ropes: &plan.ropes,
391            batch,
392            cu: &plan.cu,
393            mcu: &plan.mcu,
394            row_seq: &plan.row_seq,
395            blocks: &plan.blocks,
396            scratch: &plan.scratch,
397            counts: [t, s, m],
398        };
399        for (i, step) in plan.steps.iter().enumerate() {
400            match plan.profile.as_mut() {
401                None => ctx.step(step),
402                Some(p) => {
403                    let at = Instant::now();
404                    ctx.step(step);
405                    p[i] += u64::try_from(at.elapsed().as_nanos()).unwrap_or(u64::MAX);
406                }
407            }
408        }
409        // SAFETY: the steps are done, so nothing else touches the arena.
410        let logits = unsafe { ctx.arena.slice(plan.logits.off, m) };
411        // SAFETY: as above.
412        let act = unsafe { ctx.arena.slice(plan.act.off, 2 * s) };
413        out.logits.clear();
414        out.logits.extend_from_slice(logits);
415        out.act.clear();
416        out.act.extend_from_slice(act.as_chunks::<2>().0);
417        Ok(())
418    }
419}
420
421/// Everything a step needs for one batch.
422struct Ctx<'a> {
423    pool: &'a Pool,
424    w: &'a [Tensor],
425    arena: Arena,
426    ropes: &'a [Rope],
427    batch: &'a Batch<'a>,
428    cu: &'a [usize],
429    mcu: &'a [usize],
430    row_seq: &'a [u32],
431    blocks: &'a [(u32, u32)],
432    scratch: &'a PerWorker<Vec<f32>>,
433    counts: [usize; 3],
434}
435
436impl Ctx<'_> {
437    fn rows(&self, r: Rows) -> usize {
438        self.counts[r as usize]
439    }
440
441    fn w(&self, i: usize) -> &[f32] {
442        &self.w[i].data
443    }
444
445    /// The live rows of `l`.
446    ///
447    /// # Safety
448    ///
449    /// Nobody may write `l` while the slice lives.
450    unsafe fn get(&self, l: Loc) -> &[f32] {
451        // SAFETY: in the arena by the layout, and by the caller.
452        unsafe { self.arena.slice(l.off, self.rows(l.rows) * l.width) }
453    }
454
455    /// Runs `f(first row, rows of out)` over the live rows of `out` in blocks of [`ROWS`].
456    ///
457    /// # Safety
458    ///
459    /// Nobody else may touch `out` meanwhile, and `f` must not reach `out` any other way.
460    unsafe fn rows_of(&self, out: Loc, f: &(dyn Fn(usize, &mut [f32]) + Sync)) {
461        let n = self.rows(out.rows);
462        let arena = self.arena;
463        self.pool.run(n.div_ceil(ROWS), &|task, _| {
464            let r0 = task * ROWS;
465            let r1 = (r0 + ROWS).min(n);
466            // SAFETY: rows r0..r1 of out belong to this task alone.
467            let rows = unsafe { arena.slice_mut(out.off + r0 * out.width, (r1 - r0) * out.width) };
468            f(r0, rows);
469        });
470    }
471
472    fn step(&self, step: &Step) {
473        // The layout gives the inputs and the output of a step disjoint arena ranges unless the
474        // step is in place, and an in place step takes one slice only. That is the argument behind
475        // every SAFETY comment below.
476        match *step {
477            Step::Embed { table, out } => {
478                let (tab, ids, d) = (self.w(table), self.batch.ids, out.width);
479                // SAFETY: the layout keeps the inputs and the output of a step apart.
480                unsafe {
481                    self.rows_of(out, &|r0, rows| {
482                        for (i, row) in rows.chunks_exact_mut(d).enumerate() {
483                            let id = ids[r0 + i] as usize;
484                            row.copy_from_slice(&tab[id * d..(id + 1) * d]);
485                        }
486                    });
487                }
488            }
489            Step::LayerNorm { x, w, b, eps, out } => {
490                // SAFETY: the layout keeps the inputs and the output of a step apart.
491                let x = unsafe { self.get(x) };
492                let (nw, nb, d) = (self.w(w), b.map(|b| self.w(b)), out.width);
493                // SAFETY: the layout keeps the inputs and the output of a step apart.
494                unsafe {
495                    self.rows_of(out, &|r0, rows| {
496                        layer_norm(&x[r0 * d..r0 * d + rows.len()], d, nw, nb, eps, rows);
497                    });
498                }
499            }
500            Step::Gemm { a, w, b, ep, out } => {
501                let rows = self.rows(a.rows);
502                // SAFETY: the layout keeps the inputs and the output of a step apart.
503                let x = unsafe { self.get(a) };
504                // SAFETY: the layout keeps the inputs and the output of a step apart.
505                let y = unsafe { self.arena.slice_mut(out.off, rows * out.width) };
506                let g = Gemm {
507                    x,
508                    m: rows,
509                    k: a.width,
510                    w: self.w(w),
511                    n: out.width,
512                    b: b.map(|b| self.w(b)),
513                    ep,
514                };
515                g.run(y, self.pool.threads(), |n, f| self.pool.run(n, &|i, _| f(i)));
516            }
517            Step::Rope { qkv, rope } => {
518                let (rope, d, cu, seq) = (&self.ropes[rope], qkv.width / 3, self.cu, self.row_seq);
519                // SAFETY: the layout keeps the inputs and the output of a step apart.
520                unsafe {
521                    self.rows_of(qkv, &|r0, rows| {
522                        for (i, row) in rows.chunks_exact_mut(qkv.width).enumerate() {
523                            let r = r0 + i;
524                            let pos = r - cu[seq[r] as usize];
525                            for head in row[..2 * d].as_chunks_mut::<HEAD>().0 {
526                                rope.apply(head, pos);
527                            }
528                        }
529                    });
530                }
531            }
532            Step::Attention { qkv, window, out } => {
533                let heads = out.width / HEAD;
534                // SAFETY: the layout keeps the inputs and the output of a step apart.
535                let x = unsafe { self.get(qkv) };
536                // SAFETY: the layout keeps the inputs and the output of a step apart.
537                let y = unsafe { self.arena.slice_mut(out.off, self.rows(out.rows) * out.width) };
538                let shared = Shared::new(y);
539                let (blocks, cu, scratch) = (self.blocks, self.cu, self.scratch);
540                self.pool.run(blocks.len() * heads, &|task, worker| {
541                    let (s, q0) = blocks[task / heads];
542                    let (s, q0, h) = (s as usize, q0 as usize, task % heads);
543                    // SAFETY: the scratch is this worker's, and this task owns rows q0 to q0 + QB
544                    // of head h.
545                    unsafe {
546                        let p = scratch.get(worker);
547                        attention::block(x, heads, (cu[s], cu[s + 1]), q0, h, window, p, &shared);
548                    }
549                });
550            }
551            Step::GeGlu { x, out } => {
552                // SAFETY: the layout keeps the inputs and the output of a step apart.
553                let (x, d) = (unsafe { self.get(x) }, out.width);
554                // SAFETY: the layout keeps the inputs and the output of a step apart.
555                unsafe {
556                    self.rows_of(out, &|r0, rows| {
557                        geglu(&x[2 * r0 * d..2 * (r0 * d + rows.len())], d, rows);
558                    });
559                }
560            }
561            Step::AddType { h, table } => {
562                let (tab, d, seq, qt) = (self.w(table), h.width, self.row_seq, self.batch.qtype);
563                // SAFETY: the layout keeps the inputs and the output of a step apart.
564                unsafe {
565                    self.rows_of(h, &|r0, rows| {
566                        for (i, row) in rows.chunks_exact_mut(d).enumerate() {
567                            let q = usize::from(qt[seq[r0 + i] as usize]);
568                            row.iter_mut().zip(&tab[q * d..(q + 1) * d]).for_each(|(a, b)| *a += b);
569                        }
570                    });
571                }
572            }
573            Step::Gather { h, out } => {
574                // SAFETY: the layout keeps the inputs and the output of a step apart.
575                let (x, d) = (unsafe { self.get(h) }, h.width);
576                // SAFETY: the layout keeps the inputs and the output of a step apart.
577                let y = unsafe { self.arena.slice_mut(out.off, self.rows(out.rows) * d) };
578                let mut at = 0;
579                for s in 0..self.cu.len() - 1 {
580                    for &p in &self.batch.markers[self.mcu[s]..self.mcu[s + 1]] {
581                        let r = self.cu[s] + p as usize;
582                        y[at * d..(at + 1) * d].copy_from_slice(&x[r * d..(r + 1) * d]);
583                        at += 1;
584                    }
585                }
586            }
587            Step::ActFeatures { h, logits, out } => {
588                // SAFETY: the layout keeps the inputs and the output of a step apart.
589                let (x, d) = (unsafe { self.get(h) }, h.width);
590                // SAFETY: the layout keeps the inputs and the output of a step apart.
591                let l = unsafe { self.get(logits) };
592                // SAFETY: the layout keeps the inputs and the output of a step apart.
593                let y = unsafe { self.arena.slice_mut(out.off, self.rows(out.rows) * out.width) };
594                for (s, row) in y.chunks_exact_mut(out.width).enumerate() {
595                    let (lo, hi) = (self.cu[s], self.cu[s + 1]);
596                    if hi > lo {
597                        row[..d].copy_from_slice(&x[lo * d..(lo + 1) * d]);
598                    } else {
599                        row[..d].fill(0.0);
600                    }
601                    row[d..].copy_from_slice(&act_features(&l[self.mcu[s]..self.mcu[s + 1]]));
602                }
603            }
604        }
605    }
606}
607
608/// `[top1, top1 - top2, entropy / ln k, k / 255]` over the softmax of the logits, with `k` at
609/// least 2, the same arithmetic as [`crate::compat::act_features`] without its buffers.
610fn act_features(logits: &[f32]) -> [f32; 4] {
611    let kf = logits.len().max(2) as f32;
612    if logits.is_empty() {
613        return [0.0, 0.0, 0.0, kf / 255.0];
614    }
615    let mx = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
616    let sum: f32 = logits.iter().map(|&l| (l - mx).exp()).sum();
617    let (mut top1, mut top2, mut ent) = (f32::NEG_INFINITY, 0f32, 0f32);
618    let mut first = true;
619    for &l in logits {
620        let q = (l - mx).exp() / sum;
621        ent += q * q.max(1e-9).ln();
622        if q > top1 || first {
623            if !first {
624                top2 = top1;
625            }
626            top1 = q;
627            first = false;
628        } else if q > top2 {
629            top2 = q;
630        }
631    }
632    let ent = -ent / kf.ln();
633    [top1, top1 - top2, ent, kf / 255.0]
634}
635
636#[cfg(test)]
637mod tests {
638    use super::*;
639
640    #[test]
641    fn act_features_match_the_reference() {
642        let cases: [&[f32]; 6] =
643            [&[], &[0.3], &[1.0, 1.0], &[2.0, -1.0, 2.0], &[5.0, 0.1, -3.0, 4.9], &[-1e3, 0.0]];
644        for l in cases {
645            assert_eq!(
646                act_features(l).map(f32::to_bits),
647                crate::compat::act_features(l).map(f32::to_bits),
648                "{l:?}"
649            );
650        }
651    }
652}