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            // NaN propagates, matching sum; see ReduceOp::combine.
82            ReduceOp::Max if acc.is_nan() || x.is_nan() => f64::NAN,
83            ReduceOp::Min if acc.is_nan() || x.is_nan() => f64::NAN,
84            ReduceOp::Max => acc.max(x),
85            ReduceOp::Min => acc.min(x),
86        }
87    }
88}
89
90fn fill_par(out: &mut [f64], f: impl Fn(usize, &mut [f64]) + Sync) {
91    if out.len() >= PAR_THRESHOLD {
92        let chunk = PAR_THRESHOLD / 4;
93        out.par_chunks_mut(chunk)
94            .enumerate()
95            .for_each(|(i, c)| f(i * chunk, c));
96    } else {
97        f(0, out);
98    }
99}
100
101pub(crate) fn unary(op: UnaryOp, a: &Tensor) -> Result<Tensor> {
102    let src = f64_input(a, "unary")?;
103    let mut out = vec![0.0f64; a.numel()];
104    fill_par(&mut out, |start, chunk| {
105        let mut w = OffsetWalker::at(a.layout(), start);
106        for o in chunk.iter_mut() {
107            *o = op.eval_f64(src[w.next_offset()]);
108        }
109    });
110    Tensor::from_vec_f64(out, a.shape().clone())
111}
112
113pub(crate) fn binary(op: BinaryOp, a: &Tensor, b: &Tensor) -> Result<Tensor> {
114    let out_shape = broadcast_shapes(a.shape(), b.shape())?;
115    let av = a.broadcast_view(&out_shape)?;
116    let bv = b.broadcast_view(&out_shape)?;
117    let asrc = f64_input(&av, "binary")?;
118    let bsrc = f64_input(&bv, "binary")?;
119    let mut out = vec![0.0f64; out_shape.numel()];
120    fill_par(&mut out, |start, chunk| {
121        let mut wa = OffsetWalker::at(av.layout(), start);
122        let mut wb = OffsetWalker::at(bv.layout(), start);
123        for o in chunk.iter_mut() {
124            *o = op.eval_f64(asrc[wa.next_offset()], bsrc[wb.next_offset()]);
125        }
126    });
127    Tensor::from_vec_f64(out, out_shape)
128}
129
130pub(crate) fn matmul(a: &Tensor, b: &Tensor) -> Result<Tensor> {
131    let plan = crate::backend::plan_matmul(a.shape(), b.shape())?;
132    let av = a.to_vec_f64()?;
133    let bv = b.to_vec_f64()?;
134    let (m, k, n) = (plan.m, plan.k, plan.n);
135    let mut out = vec![0.0f64; plan.batch * m * n];
136    out.par_chunks_mut(m * n).enumerate().for_each(|(bi, c)| {
137        let ao = bi * plan.a_batch_stride;
138        let bo = bi * plan.b_batch_stride;
139        for i in 0..m {
140            let row = &mut c[i * n..(i + 1) * n];
141            for kk in 0..k {
142                let x = av[ao + i * k + kk];
143                if x == 0.0 {
144                    continue;
145                }
146                let brow = &bv[bo + kk * n..bo + kk * n + n];
147                for (o, &y) in row.iter_mut().zip(brow) {
148                    *o += x * y;
149                }
150            }
151        }
152    });
153    Tensor::from_vec_f64(out, plan.out_shape)
154}
155
156fn split_axes(dims: &[usize], axes: &[usize], keepdim: bool) -> (Vec<usize>, Vec<usize>) {
157    let mut out_dims = Vec::new();
158    let mut kept = Vec::new();
159    for (i, &d) in dims.iter().enumerate() {
160        if axes.contains(&i) {
161            if keepdim {
162                out_dims.push(1);
163            }
164        } else {
165            out_dims.push(d);
166            kept.push(i);
167        }
168    }
169    (out_dims, kept)
170}
171
172pub(crate) fn reduce(op: ReduceOp, a: &Tensor, axes: &[usize], keepdim: bool) -> Result<Tensor> {
173    let src = f64_input(a, "reduce")?;
174    let dims = a.dims().to_vec();
175    let strides = a.layout().strides.values().to_vec();
176    let base = a.layout().offset as isize;
177    let (out_dims, kept) = split_axes(&dims, axes, keepdim);
178    let reduced_dims: Vec<usize> = axes.iter().map(|&ax| dims[ax]).collect();
179    let reduced_strides: Vec<isize> = axes.iter().map(|&ax| strides[ax]).collect();
180    let kept_dims: Vec<usize> = kept.iter().map(|&ax| dims[ax]).collect();
181    let kept_strides: Vec<isize> = kept.iter().map(|&ax| strides[ax]).collect();
182    let out_numel: usize = out_dims.iter().product();
183    let empty = reduced_dims.contains(&0);
184    let reduce_one = |out_i: usize| -> f64 {
185        if empty {
186            return op.identity_f64();
187        }
188        let mut rem = out_i;
189        let mut offset = base;
190        for (i, &d) in kept_dims.iter().enumerate().rev() {
191            offset += (rem % d) as isize * kept_strides[i];
192            rem /= d;
193        }
194        // Sums are Neumaier-compensated even in f64: a research metric that
195        // reaches for f64 is usually one that cancels.
196        let mut acc = op.identity_f64();
197        let mut comp = 0.0f64;
198        let mut coords = vec![0usize; reduced_dims.len()];
199        let mut off = offset;
200        loop {
201            let x = src[off as usize];
202            if op == ReduceOp::Sum {
203                let t = acc + x;
204                comp += if acc.abs() >= x.abs() {
205                    (acc - t) + x
206                } else {
207                    (x - t) + acc
208                };
209                acc = t;
210            } else {
211                acc = op.combine_f64(acc, x);
212            }
213            let mut d = reduced_dims.len();
214            loop {
215                if d == 0 {
216                    return if op == ReduceOp::Sum { acc + comp } else { acc };
217                }
218                d -= 1;
219                coords[d] += 1;
220                off += reduced_strides[d];
221                if coords[d] < reduced_dims[d] {
222                    break;
223                }
224                off -= reduced_dims[d] as isize * reduced_strides[d];
225                coords[d] = 0;
226            }
227        }
228    };
229    let out: Vec<f64> = if out_numel >= 1024 {
230        (0..out_numel).into_par_iter().map(reduce_one).collect()
231    } else {
232        (0..out_numel).map(reduce_one).collect()
233    };
234    Tensor::from_vec_f64(out, Shape::new(out_dims))
235}
236
237pub(crate) fn argmax(a: &Tensor, dim: usize, keepdim: bool) -> Result<Tensor> {
238    let src = f64_input(a, "argmax")?;
239    let dims = a.dims().to_vec();
240    let strides = a.layout().strides.values().to_vec();
241    let base = a.layout().offset as isize;
242    let (out_dims, kept) = split_axes(&dims, &[dim], keepdim);
243    let kept_strides: Vec<isize> = kept.iter().map(|&ax| strides[ax]).collect();
244    let kept_dims: Vec<usize> = kept.iter().map(|&ax| dims[ax]).collect();
245    let (n, s) = (dims[dim], strides[dim]);
246    if n == 0 {
247        return Err(Error::InvalidArgument {
248            op: "argmax",
249            detail: format!("dimension {dim} has extent 0; argmax of nothing is undefined"),
250        });
251    }
252    let out_numel: usize = out_dims.iter().product();
253    let out: Vec<i64> = (0..out_numel)
254        .map(|out_i| {
255            let mut rem = out_i;
256            let mut offset = base;
257            for (i, &d) in kept_dims.iter().enumerate().rev() {
258                offset += (rem % d) as isize * kept_strides[i];
259                rem /= d;
260            }
261            let mut best = f64::NEG_INFINITY;
262            let mut best_i = 0i64;
263            for j in 0..n {
264                let v = src[(offset + j as isize * s) as usize];
265                if v.is_nan() {
266                    return Err(Error::InvalidArgument {
267                        op: "argmax",
268                        detail: "input contains NaN; the maximum is undefined".into(),
269                    });
270                }
271                if v > best {
272                    best = v;
273                    best_i = j as i64;
274                }
275            }
276            Ok(best_i)
277        })
278        .collect::<Result<Vec<i64>>>()?;
279    Tensor::from_vec_i64(out, Shape::new(out_dims))
280}
281
282pub(crate) fn index_select(a: &Tensor, dim: usize, indices: &Tensor) -> Result<Tensor> {
283    let idx = indices.to_vec_i64()?;
284    let dims = a.dims().to_vec();
285    if dim >= dims.len() {
286        return Err(Error::InvalidArgument {
287            op: "index_select",
288            detail: format!("dim {dim} out of range for rank {}", dims.len()),
289        });
290    }
291    for (pos, &i) in idx.iter().enumerate() {
292        // See the f32 path: a negative index is reported as negative.
293        if i < 0 {
294            return Err(Error::InvalidArgument {
295                op: "index_select",
296                detail: format!("index {i} at position {pos} is negative"),
297            });
298        }
299        if i as usize >= dims[dim] {
300            return Err(Error::IndexOutOfBounds {
301                index: vec![i as usize],
302                shape: a.shape().clone(),
303            });
304        }
305    }
306    // Logical order, then gather rows: correct for any layout.
307    let src = a.to_vec_f64()?;
308    let inner: usize = dims[dim + 1..].iter().product();
309    let outer: usize = dims[..dim].iter().product();
310    let mut out_dims = dims.clone();
311    out_dims[dim] = idx.len();
312    let mut out = Vec::with_capacity(outer * idx.len() * inner);
313    for o in 0..outer {
314        for &sel in &idx {
315            let start = (o * dims[dim] + sel as usize) * inner;
316            out.extend_from_slice(&src[start..start + inner]);
317        }
318    }
319    Tensor::from_vec_f64(out, Shape::new(out_dims))
320}
321
322pub(crate) fn index_add(a: &Tensor, dim: usize, indices: &Tensor, srct: &Tensor) -> Result<Tensor> {
323    let idx = indices.to_vec_i64()?;
324    let dims = a.dims().to_vec();
325    if dim >= dims.len() {
326        return Err(Error::InvalidArgument {
327            op: "index_add",
328            detail: format!("dim {dim} out of range for rank {}", dims.len()),
329        });
330    }
331    let mut expected = dims.clone();
332    expected[dim] = idx.len();
333    if srct.dims() != expected.as_slice() {
334        return Err(Error::ShapeMismatch {
335            expected: Shape::new(expected),
336            got: srct.shape().clone(),
337            op: "index_add",
338        });
339    }
340    let mut out = a.to_vec_f64()?;
341    let s = srct.to_vec_f64()?;
342    let inner: usize = dims[dim + 1..].iter().product();
343    let outer: usize = dims[..dim].iter().product();
344    for o in 0..outer {
345        for (k, &sel) in idx.iter().enumerate() {
346            if sel < 0 || sel as usize >= dims[dim] {
347                return Err(Error::IndexOutOfBounds {
348                    index: vec![sel.max(0) as usize],
349                    shape: a.shape().clone(),
350                });
351            }
352            let dst = (o * dims[dim] + sel as usize) * inner;
353            let src = (o * idx.len() + k) * inner;
354            for j in 0..inner {
355                out[dst + j] += s[src + j];
356            }
357        }
358    }
359    Tensor::from_vec_f64(out, a.shape().clone())
360}
361
362/// `cholesky` in f64 in and out (the reference routine is f64 internally
363/// already; this keeps the result in f64 instead of rounding it).
364pub(crate) fn cholesky(a: &Tensor) -> Result<Tensor> {
365    let d = a.dims();
366    let n = d[d.len() - 1];
367    let batch: usize = d[..d.len() - 2].iter().product();
368    let data = a.to_vec_f64()?;
369    let as_f32: Vec<f32> = data.iter().map(|&x| x as f32).collect();
370    // Reuse the reference for the pivot check, then redo the arithmetic in
371    // f64 for the result: small matrices, so the double pass is cheap.
372    crate::cpu_linalg::cholesky(&as_f32, batch, n)?;
373    let mut out = vec![0.0f64; batch * n * n];
374    for b in 0..batch {
375        let m = &data[b * n * n..(b + 1) * n * n];
376        let l = &mut out[b * n * n..(b + 1) * n * n];
377        for j in 0..n {
378            let mut dd = m[j * n + j];
379            for k in 0..j {
380                dd -= l[j * n + k] * l[j * n + k];
381            }
382            if dd.is_nan() || dd <= 0.0 || dd.is_infinite() {
383                return Err(Error::InvalidArgument {
384                    op: "cholesky",
385                    detail: format!(
386                        "matrix {b} is not positive definite (pivot {j} is {dd:e}); cholesky/logdet need an SPD input"
387                    ),
388                });
389            }
390            let ljj = dd.sqrt();
391            l[j * n + j] = ljj;
392            for i in j + 1..n {
393                let mut s = m[i * n + j];
394                for k in 0..j {
395                    s -= l[i * n + k] * l[j * n + k];
396                }
397                l[i * n + j] = s / ljj;
398            }
399        }
400    }
401    Tensor::from_vec_f64(out, a.shape().clone())
402}
oxmera_tensor/cpu_f64.rs

oxmera_tensor/
cpu_f64.rs
