oxfuse_runtime/cpu/
mod.rs

1//! The fast CPU backend. It runs the same `Plan` as the reference backend and
2//! returns the same bits; it exists to be faster. It has its own executor so
3//! the reference executor stays a plain, obviously correct oracle.
4
5mod kernels;
6
7use crate::{Backend, Executable, Plan, Tensor};
8use oxfuse_core::{Error, Graph};
9
10pub struct CpuBackend;
11
12impl Backend for CpuBackend {
13    fn compile(&self, graph: &Graph) -> Result<Box<dyn Executable>, Error> {
14        Ok(Box::new(CpuExecutable {
15            plan: Plan::build(graph),
16        }))
17    }
18}
19
20struct CpuExecutable {
21    plan: Plan,
22}
23
24/// A buffer for one value during a run. Inputs and constants are borrowed,
25/// computed values are owned. Borrowed slots are never written.
26enum Slot<'a> {
27    Empty,
28    Borrowed(&'a [f32]),
29    Owned(Vec<f32>),
30}
31
32impl Slot<'_> {
33    fn data(&self) -> &[f32] {
34        match self {
35            Slot::Empty => &[],
36            Slot::Borrowed(data) => data,
37            Slot::Owned(data) => data,
38        }
39    }
40}
41
42impl Executable for CpuExecutable {
43    fn run(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, Error> {
44        let plan = &self.plan;
45        if inputs.len() != plan.inputs.len() {
46            return Err(Error::InputCount {
47                expected: plan.inputs.len(),
48                found: inputs.len(),
49            });
50        }
51        let mut slots: Vec<Slot> = (0..plan.slot_shapes.len()).map(|_| Slot::Empty).collect();
52        for (binding, tensor) in plan.inputs.iter().zip(inputs) {
53            if tensor.shape() != binding.shape.as_slice() {
54                return Err(Error::InputShape {
55                    name: binding.name.clone(),
56                    expected: binding.shape.clone(),
57                    found: tensor.shape().to_vec(),
58                });
59            }
60            slots[binding.slot] = Slot::Borrowed(tensor.data());
61        }
62        for (slot, data) in &plan.constants {
63            slots[*slot] = Slot::Borrowed(data);
64        }
65        for step in &plan.steps {
66            let out = {
67                let inputs: Vec<&[f32]> = step.inputs.iter().map(|&s| slots[s].data()).collect();
68                let shapes: Vec<&[usize]> = step
69                    .inputs
70                    .iter()
71                    .map(|&s| plan.slot_shapes[s].as_slice())
72                    .collect();
73                kernels::apply(&step.op, &inputs, &shapes)
74            };
75            slots[step.output] = Slot::Owned(out);
76        }
77        // A computed buffer listed once as an output moves out; anything else
78        // (borrowed, or listed more than once) is copied.
79        let mut uses = vec![0usize; slots.len()];
80        for &slot in &plan.outputs {
81            uses[slot] += 1;
82        }
83        plan.outputs
84            .iter()
85            .map(|&slot| {
86                let data = match &mut slots[slot] {
87                    Slot::Owned(buffer) if uses[slot] == 1 => std::mem::take(buffer),
88                    other => other.data().to_vec(),
89                };
90                Tensor::new(plan.slot_shapes[slot].clone(), data)
91            })
92            .collect()
93    }
94}
oxfuse_runtime/cpu/mod.rs

oxfuse_runtime/cpu/
mod.rs
