oxmera_tensor/
cpu_f64.rs

1//! The CPU backend's `f64` path (issue #27). Correct first, fast enough
2//! second: every op is the plain reference formula over strided views, in
3//! `f64` throughout, parallel over output chunks past the same threshold
4//! the `f32` path uses. Research metrics — set-likelihood normalisers,
5//! log-determinants, compensated sums — are what this serves; the `f32`
6//! path keeps the tuned kernels.
7
8use oxmera_core::shape::broadcast_shapes;
9use oxmera_core::{DType, Error, Result, Shape};
10use rayon::prelude::*;
11
12use crate::backend::{BinaryOp, ReduceOp, UnaryOp};
13use crate::cpu_iter::OffsetWalker;
14use crate::tensor::Tensor;
15
16const PAR_THRESHOLD: usize = 16 * 1024;
17
18pub(crate) fn f64_input<'t>(t: &'t Tensor, op: &'static str) -> Result<&'t [f64]> {
19    if t.dtype() != DType::F64 {
20        return Err(Error::UnsupportedDType {
21            dtype: t.dtype(),
22            op,
23        });
24    }
25    t.storage().cpu()?.f64s()
26}
27
28impl UnaryOp {
29    /// The `f64` reference semantics of the op.
30    pub fn eval_f64(self, x: f64) -> f64 {
31        match self {
32            UnaryOp::Neg => -x,
33            UnaryOp::Exp => x.exp(),
34            UnaryOp::Ln => x.ln(),
35            UnaryOp::Abs => x.abs(),
36            UnaryOp::Sqrt => x.sqrt(),
37            UnaryOp::Sin => x.sin(),
38            UnaryOp::Cos => x.cos(),
39            UnaryOp::Tanh => x.tanh(),
40            UnaryOp::Relu => x.max(0.0),
41            UnaryOp::Gelu => {
42                const SQRT_2_OVER_PI: f64 = 0.797_884_560_802_865_4;
43                0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044_715 * x * x * x)).tanh())
44            }
45            UnaryOp::Sigmoid => 1.0 / (1.0 + (-x).exp()),
46        }
47    }
48}
49
50impl BinaryOp {
51    /// The `f64` reference semantics of the op.
52    pub fn eval_f64(self, a: f64, b: f64) -> f64 {
53        match self {
54            BinaryOp::Add => a + b,
55            BinaryOp::Sub => a - b,
56            BinaryOp::Mul => a * b,
57            BinaryOp::Div => a / b,
58            BinaryOp::Pow => a.powf(b),
59            BinaryOp::Maximum => a.max(b),
60            BinaryOp::Minimum => a.min(b),
61            BinaryOp::Gt => f64::from(a > b),
62            BinaryOp::Eq => f64::from(a == b),
63        }
64    }
65}
66
67impl ReduceOp {
68    /// The identity element, in `f64`.
69    pub fn identity_f64(self) -> f64 {
70        match self {
71            ReduceOp::Sum => 0.0,
72            ReduceOp::Max => f64::NEG_INFINITY,
73            ReduceOp::Min => f64::INFINITY,
74        }
75    }
76
77    /// Combine an accumulator with one element, in `f64`.
78    pub fn combine_f64(self, acc: f64, x: f64) -> f64 {
79        match self {
80            ReduceOp::Sum => acc + x,
81            ReduceOp::Max => acc.max(x),
82            ReduceOp::Min => acc.min(x),
83        }
84    }
85}
86
87fn fill_par(out: &mut [f64], f: impl Fn(usize, &mut [f64]) + Sync) {
88    if out.len() >= PAR_THRESHOLD {
89        let chunk = PAR_THRESHOLD / 4;
90        out.par_chunks_mut(chunk)
91            .enumerate()
92            .for_each(|(i, c)| f(i * chunk, c));
93    } else {
94        f(0, out);
95    }
96}
97
98pub(crate) fn unary(op: UnaryOp, a: &Tensor) -> Result<Tensor> {
99    let src = f64_input(a, "unary")?;
100    let mut out = vec![0.0f64; a.numel()];
101    fill_par(&mut out, |start, chunk| {
102        let mut w = OffsetWalker::at(a.layout(), start);
103        for o in chunk.iter_mut() {
104            *o = op.eval_f64(src[w.next_offset()]);
105        }
106    });
107    Tensor::from_vec_f64(out, a.shape().clone())
108}
109
110pub(crate) fn binary(op: BinaryOp, a: &Tensor, b: &Tensor) -> Result<Tensor> {
111    let out_shape = broadcast_shapes(a.shape(), b.shape())?;
112    let av = a.broadcast_view(&out_shape)?;
113    let bv = b.broadcast_view(&out_shape)?;
114    let asrc = f64_input(&av, "binary")?;
115    let bsrc = f64_input(&bv, "binary")?;
116    let mut out = vec![0.0f64; out_shape.numel()];
117    fill_par(&mut out, |start, chunk| {
118        let mut wa = OffsetWalker::at(av.layout(), start);
119        let mut wb = OffsetWalker::at(bv.layout(), start);
120        for o in chunk.iter_mut() {
121            *o = op.eval_f64(asrc[wa.next_offset()], bsrc[wb.next_offset()]);
122        }
123    });
124    Tensor::from_vec_f64(out, out_shape)
125}
126
127pub(crate) fn matmul(a: &Tensor, b: &Tensor) -> Result<Tensor> {
128    let plan = crate::backend::plan_matmul(a.shape(), b.shape())?;
129    let av = a.to_vec_f64()?;
130    let bv = b.to_vec_f64()?;
131    let (m, k, n) = (plan.m, plan.k, plan.n);
132    let mut out = vec![0.0f64; plan.batch * m * n];
133    out.par_chunks_mut(m * n).enumerate().for_each(|(bi, c)| {
134        let ao = bi * plan.a_batch_stride;
135        let bo = bi * plan.b_batch_stride;
136        for i in 0..m {
137            let row = &mut c[i * n..(i + 1) * n];
138            for kk in 0..k {
139                let x = av[ao + i * k + kk];
140                if x == 0.0 {
141                    continue;
142                }
143                let brow = &bv[bo + kk * n..bo + kk * n + n];
144                for (o, &y) in row.iter_mut().zip(brow) {
145                    *o += x * y;
146                }
147            }
148        }
149    });
150    Tensor::from_vec_f64(out, plan.out_shape)
151}
152
153fn split_axes(dims: &[usize], axes: &[usize], keepdim: bool) -> (Vec<usize>, Vec<usize>) {
154    let mut out_dims = Vec::new();
155    let mut kept = Vec::new();
156    for (i, &d) in dims.iter().enumerate() {
157        if axes.contains(&i) {
158            if keepdim {
159                out_dims.push(1);
160            }
161        } else {
162            out_dims.push(d);
163            kept.push(i);
164        }
165    }
166    (out_dims, kept)
167}
168
169pub(crate) fn reduce(op: ReduceOp, a: &Tensor, axes: &[usize], keepdim: bool) -> Result<Tensor> {
170    let src = f64_input(a, "reduce")?;
171    let dims = a.dims().to_vec();
172    let strides = a.layout().strides.values().to_vec();
173    let base = a.layout().offset as isize;
174    let (out_dims, kept) = split_axes(&dims, axes, keepdim);
175    let reduced_dims: Vec<usize> = axes.iter().map(|&ax| dims[ax]).collect();
176    let reduced_strides: Vec<isize> = axes.iter().map(|&ax| strides[ax]).collect();
177    let kept_dims: Vec<usize> = kept.iter().map(|&ax| dims[ax]).collect();
178    let kept_strides: Vec<isize> = kept.iter().map(|&ax| strides[ax]).collect();
179    let out_numel: usize = out_dims.iter().product();
180    let empty = reduced_dims.contains(&0);
181    let reduce_one = |out_i: usize| -> f64 {
182        if empty {
183            return op.identity_f64();
184        }
185        let mut rem = out_i;
186        let mut offset = base;
187        for (i, &d) in kept_dims.iter().enumerate().rev() {
188            offset += (rem % d) as isize * kept_strides[i];
189            rem /= d;
190        }
191        // Sums are Neumaier-compensated even in f64: a research metric that
192        // reaches for f64 is usually one that cancels.
193        let mut acc = op.identity_f64();
194        let mut comp = 0.0f64;
195        let mut coords = vec![0usize; reduced_dims.len()];
196        let mut off = offset;
197        loop {
198            let x = src[off as usize];
199            if op == ReduceOp::Sum {
200                let t = acc + x;
201                comp += if acc.abs() >= x.abs() {
202                    (acc - t) + x
203                } else {
204                    (x - t) + acc
205                };
206                acc = t;
207            } else {
208                acc = op.combine_f64(acc, x);
209            }
210            let mut d = reduced_dims.len();
211            loop {
212                if d == 0 {
213                    return if op == ReduceOp::Sum { acc + comp } else { acc };
214                }
215                d -= 1;
216                coords[d] += 1;
217                off += reduced_strides[d];
218                if coords[d] < reduced_dims[d] {
219                    break;
220                }
221                off -= reduced_dims[d] as isize * reduced_strides[d];
222                coords[d] = 0;
223            }
224        }
225    };
226    let out: Vec<f64> = if out_numel >= 1024 {
227        (0..out_numel).into_par_iter().map(reduce_one).collect()
228    } else {
229        (0..out_numel).map(reduce_one).collect()
230    };
231    Tensor::from_vec_f64(out, Shape::new(out_dims))
232}
233
234pub(crate) fn argmax(a: &Tensor, dim: usize, keepdim: bool) -> Result<Tensor> {
235    let src = f64_input(a, "argmax")?;
236    let dims = a.dims().to_vec();
237    let strides = a.layout().strides.values().to_vec();
238    let base = a.layout().offset as isize;
239    let (out_dims, kept) = split_axes(&dims, &[dim], keepdim);
240    let kept_strides: Vec<isize> = kept.iter().map(|&ax| strides[ax]).collect();
241    let kept_dims: Vec<usize> = kept.iter().map(|&ax| dims[ax]).collect();
242    let (n, s) = (dims[dim], strides[dim]);
243    if n == 0 {
244        return Err(Error::InvalidArgument {
245            op: "argmax",
246            detail: format!("dimension {dim} has extent 0; argmax of nothing is undefined"),
247        });
248    }
249    let out_numel: usize = out_dims.iter().product();
250    let out: Vec<i64> = (0..out_numel)
251        .map(|out_i| {
252            let mut rem = out_i;
253            let mut offset = base;
254            for (i, &d) in kept_dims.iter().enumerate().rev() {
255                offset += (rem % d) as isize * kept_strides[i];
256                rem /= d;
257            }
258            let mut best = f64::NEG_INFINITY;
259            let mut best_i = 0i64;
260            for j in 0..n {
261                let v = src[(offset + j as isize * s) as usize];
262                if v > best {
263                    best = v;
264                    best_i = j as i64;
265                }
266            }
267            best_i
268        })
269        .collect();
270    Tensor::from_vec_i64(out, Shape::new(out_dims))
271}
272
273pub(crate) fn index_select(a: &Tensor, dim: usize, indices: &Tensor) -> Result<Tensor> {
274    let idx = indices.to_vec_i64()?;
275    let dims = a.dims().to_vec();
276    if dim >= dims.len() {
277        return Err(Error::InvalidArgument {
278            op: "index_select",
279            detail: format!("dim {dim} out of range for rank {}", dims.len()),
280        });
281    }
282    for &i in &idx {
283        if i < 0 || i as usize >= dims[dim] {
284            return Err(Error::IndexOutOfBounds {
285                index: vec![i.max(0) as usize],
286                shape: a.shape().clone(),
287            });
288        }
289    }
290    // Logical order, then gather rows: correct for any layout.
291    let src = a.to_vec_f64()?;
292    let inner: usize = dims[dim + 1..].iter().product();
293    let outer: usize = dims[..dim].iter().product();
294    let mut out_dims = dims.clone();
295    out_dims[dim] = idx.len();
296    let mut out = Vec::with_capacity(outer * idx.len() * inner);
297    for o in 0..outer {
298        for &sel in &idx {
299            let start = (o * dims[dim] + sel as usize) * inner;
300            out.extend_from_slice(&src[start..start + inner]);
301        }
302    }
303    Tensor::from_vec_f64(out, Shape::new(out_dims))
304}
305
306pub(crate) fn index_add(a: &Tensor, dim: usize, indices: &Tensor, srct: &Tensor) -> Result<Tensor> {
307    let idx = indices.to_vec_i64()?;
308    let dims = a.dims().to_vec();
309    if dim >= dims.len() {
310        return Err(Error::InvalidArgument {
311            op: "index_add",
312            detail: format!("dim {dim} out of range for rank {}", dims.len()),
313        });
314    }
315    let mut expected = dims.clone();
316    expected[dim] = idx.len();
317    if srct.dims() != expected.as_slice() {
318        return Err(Error::ShapeMismatch {
319            expected: Shape::new(expected),
320            got: srct.shape().clone(),
321            op: "index_add",
322        });
323    }
324    let mut out = a.to_vec_f64()?;
325    let s = srct.to_vec_f64()?;
326    let inner: usize = dims[dim + 1..].iter().product();
327    let outer: usize = dims[..dim].iter().product();
328    for o in 0..outer {
329        for (k, &sel) in idx.iter().enumerate() {
330            if sel < 0 || sel as usize >= dims[dim] {
331                return Err(Error::IndexOutOfBounds {
332                    index: vec![sel.max(0) as usize],
333                    shape: a.shape().clone(),
334                });
335            }
336            let dst = (o * dims[dim] + sel as usize) * inner;
337            let src = (o * idx.len() + k) * inner;
338            for j in 0..inner {
339                out[dst + j] += s[src + j];
340            }
341        }
342    }
343    Tensor::from_vec_f64(out, a.shape().clone())
344}
345
346/// `cholesky` in f64 in and out (the reference routine is f64 internally
347/// already; this keeps the result in f64 instead of rounding it).
348pub(crate) fn cholesky(a: &Tensor) -> Result<Tensor> {
349    let d = a.dims();
350    let n = d[d.len() - 1];
351    let batch: usize = d[..d.len() - 2].iter().product();
352    let data = a.to_vec_f64()?;
353    let as_f32: Vec<f32> = data.iter().map(|&x| x as f32).collect();
354    // Reuse the reference for the pivot check, then redo the arithmetic in
355    // f64 for the result: small matrices, so the double pass is cheap.
356    crate::cpu_linalg::cholesky(&as_f32, batch, n)?;
357    let mut out = vec![0.0f64; batch * n * n];
358    for b in 0..batch {
359        let m = &data[b * n * n..(b + 1) * n * n];
360        let l = &mut out[b * n * n..(b + 1) * n * n];
361        for j in 0..n {
362            let mut dd = m[j * n + j];
363            for k in 0..j {
364                dd -= l[j * n + k] * l[j * n + k];
365            }
366            if dd.is_nan() || dd <= 0.0 || dd.is_infinite() {
367                return Err(Error::InvalidArgument {
368                    op: "cholesky",
369                    detail: format!(
370                        "matrix {b} is not positive definite (pivot {j} is {dd:e}); cholesky/logdet need an SPD input"
371                    ),
372                });
373            }
374            let ljj = dd.sqrt();
375            l[j * n + j] = ljj;
376            for i in j + 1..n {
377                let mut s = m[i * n + j];
378                for k in 0..j {
379                    s -= l[i * n + k] * l[j * n + k];
380                }
381                l[i * n + j] = s / ljj;
382            }
383        }
384    }
385    Tensor::from_vec_f64(out, a.shape().clone())
386}
oxmera_tensor/cpu_f64.rs

oxmera_tensor/
cpu_f64.rs
