1use 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#[derive(Debug)]
27pub struct Tensor {
28 pub shape: Vec<usize>,
30 pub data: Vec<f32>,
32}
33
34#[derive(Debug, Clone)]
36pub struct Weights(Arc<[Tensor]>);
37
38#[derive(Debug)]
40pub struct CpuBackend {
41 pool: Pool,
42}
43
44impl CpuBackend {
45 #[must_use]
47 pub fn new(threads: usize) -> Self {
48 Self { pool: Pool::new(threads.max(1)) }
49 }
50
51 #[must_use]
53 pub fn threads(&self) -> usize {
54 self.pool.threads()
55 }
56}
57
58#[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
95struct PerWorker<T>(Vec<UnsafeCell<T>>);
97
98unsafe impl<T: Send> Sync for PerWorker<T> {}
101
102impl<T> PerWorker<T> {
103 #[allow(clippy::mut_from_ref)]
107 unsafe fn get(&self, worker: usize) -> &mut T {
108 unsafe { &mut *self.0[worker].get() }
110 }
111}
112
113pub 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 cu: Vec<usize>,
125 mcu: Vec<usize>,
126 row_seq: Vec<u32>,
127 blocks: Vec<(u32, u32)>,
128 scratch: PerWorker<Vec<f32>>,
129 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 #[must_use]
146 pub fn bucket(&self) -> Bucket {
147 self.bucket
148 }
149
150 #[must_use]
152 pub fn arena_bytes(&self) -> usize {
153 self.arena.len() * 4
154 }
155
156 pub fn profile(&mut self) {
158 self.profile = Some(vec![0; self.steps.len()]);
159 }
160
161 #[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#[derive(Clone, Copy)]
178struct Arena(*mut f32);
179
180unsafe impl Send for Arena {}
183unsafe impl Sync for Arena {}
185
186impl Arena {
187 unsafe fn slice<'a>(self, off: usize, len: usize) -> &'a [f32] {
191 unsafe { std::slice::from_raw_parts(self.0.add(off), len) }
193 }
194
195 #[allow(clippy::mut_from_ref)]
199 unsafe fn slice_mut<'a>(self, off: usize, len: usize) -> &'a mut [f32] {
200 unsafe { std::slice::from_raw_parts_mut(self.0.add(off), len) }
202 }
203}
204
205const 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 let logits = unsafe { ctx.arena.slice(plan.logits.off, m) };
411 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
421struct 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 unsafe fn get(&self, l: Loc) -> &[f32] {
451 unsafe { self.arena.slice(l.off, self.rows(l.rows) * l.width) }
453 }
454
455 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 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 match *step {
477 Step::Embed { table, out } => {
478 let (tab, ids, d) = (self.w(table), self.batch.ids, out.width);
479 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 let x = unsafe { self.get(x) };
492 let (nw, nb, d) = (self.w(w), b.map(|b| self.w(b)), out.width);
493 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 let x = unsafe { self.get(a) };
504 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 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 let x = unsafe { self.get(qkv) };
536 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 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 let (x, d) = (unsafe { self.get(x) }, out.width);
554 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 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 let (x, d) = (unsafe { self.get(h) }, h.width);
576 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 let (x, d) = (unsafe { self.get(h) }, h.width);
590 let l = unsafe { self.get(logits) };
592 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
608fn 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}