ferric_tensor/
cpu.rs

1//! Plain-Rust reference on logical (row-major) data — the source of truth for the GPU tensor runtime.
2//! Operates on the same logical layout `Tensor::to_vec` returns, so comparisons are apples-to-apples.
3
4fn unravel(mut i: usize, shape: &[usize]) -> Vec<usize> {
5    let mut idx = vec![0usize; shape.len()];
6    for d in (0..shape.len()).rev() {
7        idx[d] = i % shape[d];
8        i /= shape[d];
9    }
10    idx
11}
12fn ravel(idx: &[usize], shape: &[usize]) -> usize {
13    let mut i = 0;
14    for d in 0..shape.len() {
15        i = i * shape[d] + idx[d];
16    }
17    i
18}
19pub fn broadcast_shapes(a: &[usize], b: &[usize]) -> Vec<usize> {
20    let r = a.len().max(b.len());
21    (0..r).map(|i| {
22        let da = if i + a.len() >= r { a[i + a.len() - r] } else { 1 };
23        let db = if i + b.len() >= r { b[i + b.len() - r] } else { 1 };
24        da.max(db)
25    }).collect()
26}
27// map an out multi-index to the linear index of a right-aligned (broadcast) operand
28fn proj(out_idx: &[usize], shape: &[usize]) -> usize {
29    let r = out_idx.len();
30    let idx: Vec<usize> = (0..shape.len()).map(|d| {
31        let od = out_idx[d + r - shape.len()];
32        if shape[d] == 1 { 0 } else { od }
33    }).collect();
34    ravel(&idx, shape)
35}
36
37pub fn binary(a: &[f32], ash: &[usize], b: &[f32], bsh: &[usize], op: &str) -> (Vec<f32>, Vec<usize>) {
38    let sh = broadcast_shapes(ash, bsh);
39    let n: usize = sh.iter().product();
40    let out = (0..n).map(|i| {
41        let idx = unravel(i, &sh);
42        let (x, y) = (a[proj(&idx, ash)], b[proj(&idx, bsh)]);
43        match op { "+" => x + y, "-" => x - y, "*" => x * y, "/" => x / y, "max" => x.max(y), _ => unreachable!() }
44    }).collect();
45    (out, sh)
46}
47
48pub fn unary(x: &[f32], op: &str) -> Vec<f32> {
49    x.iter().map(|&v| match op { "exp" => v.exp(), "neg" => -v, "relu" => v.max(0.0), "sqrt" => v.sqrt(), _ => v }).collect()
50}
51
52pub fn reduce(x: &[f32], shape: &[usize], axes: &[usize], op: &str, keepdim: bool) -> (Vec<f32>, Vec<usize>) {
53    let keep: Vec<usize> = (0..shape.len()).filter(|d| !axes.contains(d)).collect();
54    let oshape: Vec<usize> = if keepdim {
55        (0..shape.len()).map(|d| if axes.contains(&d) { 1 } else { shape[d] }).collect()
56    } else if keep.is_empty() { vec![1] } else { keep.iter().map(|&d| shape[d]).collect() };
57    let on: usize = oshape.iter().product();
58    let mut acc = vec![if op == "max" { f32::NEG_INFINITY } else { 0.0 }; on];
59    for i in 0..x.len() {
60        let idx = unravel(i, shape);
61        let oidx: Vec<usize> = keep.iter().map(|&d| idx[d]).collect();
62        let o = if oidx.is_empty() { 0 } else { ravel(&oidx, &keep.iter().map(|&d| shape[d]).collect::<Vec<_>>()) };
63        if op == "max" { acc[o] = acc[o].max(x[i]); } else { acc[o] += x[i]; }
64    }
65    if op == "mean" {
66        let red: usize = axes.iter().map(|&d| shape[d]).product();
67        for v in acc.iter_mut() { *v /= red as f32; }
68    }
69    (acc, oshape)
70}
71
72pub fn matmul(a: &[f32], ash: &[usize], b: &[f32], bsh: &[usize]) -> (Vec<f32>, Vec<usize>) {
73    let (ra, rb) = (ash.len(), bsh.len());
74    let (m, k, n) = (ash[ra - 2], ash[ra - 1], bsh[rb - 1]);
75    let batch = broadcast_shapes(&ash[..ra - 2], &bsh[..rb - 2]);
76    let bn: usize = batch.iter().product::<usize>().max(1);
77    let oshape: Vec<usize> = batch.iter().chain([m, n].iter()).copied().collect();
78    let mut out = vec![0.0f32; bn * m * n];
79    for bt in 0..bn {
80        let bidx = unravel(bt, &batch);
81        let ab = proj(&bidx, &ash[..ra - 2]) * m * k;
82        let bb = proj(&bidx, &bsh[..rb - 2]) * k * n;
83        for i in 0..m {
84            for j in 0..n {
85                let mut s = 0.0;
86                for l in 0..k { s += a[ab + i * k + l] * b[bb + l * n + j]; }
87                out[bt * m * n + i * n + j] = s;
88            }
89        }
90    }
91    (out, oshape)
92}
ferric_tensor/cpu.rs

ferric_tensor/
cpu.rs
