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 elementwise;
6mod kernels;
7mod matmul;
8mod vecmath;
9
10use crate::plan::Step;
11use crate::{Backend, Executable, Plan, Tensor};
12use oxfuse_core::{Error, Graph, Op};
13
14pub struct CpuBackend;
15
16impl Backend for CpuBackend {
17    fn compile(&self, graph: &Graph) -> Result<Box<dyn Executable>, Error> {
18        let plan = Plan::build(graph);
19        // Uses of each slot: once per step that reads it, once per listing as
20        // an output. Outputs are never released, because they are not consumed.
21        let mut uses = vec![0usize; plan.slot_shapes.len()];
22        for &slot in plan
23            .steps
24            .iter()
25            .flat_map(|s| &s.inputs)
26            .chain(&plan.outputs)
27        {
28            uses[slot] += 1;
29        }
30        Ok(Box::new(CpuExecutable { plan, uses }))
31    }
32}
33
34// Output buffers copied in a run on this thread. The bits are the same
35// whether an output is moved or copied, so a test counts the copies instead.
36#[cfg(test)]
37thread_local! {
38    static OUTPUT_COPIES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
39    // Steps that wrote their result into an operand's buffer.
40    static REUSED: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
41}
42
43struct CpuExecutable {
44    plan: Plan,
45    uses: Vec<usize>,
46}
47
48/// A buffer for one value during a run. Inputs and constants are borrowed,
49/// computed values are owned. Borrowed slots are never written.
50enum Slot<'a> {
51    Empty,
52    Borrowed(&'a [f32]),
53    Owned(Vec<f32>),
54}
55
56impl Slot<'_> {
57    fn data(&self) -> &[f32] {
58        match self {
59            Slot::Empty => &[],
60            Slot::Borrowed(data) => data,
61            Slot::Owned(data) => data,
62        }
63    }
64}
65
66/// An operand buffer this step may overwrite with its result: computed, read
67/// by no later step and not an output (its remaining use count is this step's
68/// one read), and the same size as the result. Returns the taken buffer and the
69/// slot of the other operand if any. A slot read twice by one step never
70/// qualifies.
71fn take_reusable(
72    step: &Step,
73    plan: &Plan,
74    remaining: &[usize],
75    slots: &mut [Slot],
76) -> Option<(Vec<f32>, Option<usize>)> {
77    let candidates: &[usize] = match step.op {
78        Op::Relu | Op::Sigmoid | Op::Reshape { .. } => &[0],
79        Op::Add | Op::Mul => &[0, 1],
80        Op::MatMul | Op::Transpose { .. } => return None,
81    };
82    let size = plan.slot_shapes[step.output].iter().product::<usize>();
83    for &at in candidates {
84        let slot = step.inputs[at];
85        let reusable = remaining[slot] == 1
86            && matches!(slots[slot], Slot::Owned(_))
87            && plan.slot_shapes[slot].iter().product::<usize>() == size;
88        if reusable {
89            let Slot::Owned(buf) = std::mem::replace(&mut slots[slot], Slot::Empty) else {
90                unreachable!("checked above")
91            };
92            #[cfg(test)]
93            REUSED.with(|c| c.set(c.get() + 1));
94            return Some((buf, step.inputs.get(1 - at).copied()));
95        }
96    }
97    None
98}
99
100impl Executable for CpuExecutable {
101    fn run(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, Error> {
102        let plan = &self.plan;
103        if inputs.len() != plan.inputs.len() {
104            return Err(Error::InputCount {
105                expected: plan.inputs.len(),
106                found: inputs.len(),
107            });
108        }
109        let mut slots: Vec<Slot> = (0..plan.slot_shapes.len()).map(|_| Slot::Empty).collect();
110        for (binding, tensor) in plan.inputs.iter().zip(inputs) {
111            if tensor.shape() != binding.shape.as_slice() {
112                return Err(Error::InputShape {
113                    name: binding.name.clone(),
114                    expected: binding.shape.clone(),
115                    found: tensor.shape().to_vec(),
116                });
117            }
118            slots[binding.slot] = Slot::Borrowed(tensor.data());
119        }
120        for (slot, data) in &plan.constants {
121            slots[*slot] = Slot::Borrowed(data);
122        }
123        let mut remaining = self.uses.clone();
124        for step in &plan.steps {
125            let out = match take_reusable(step, plan, &remaining, &mut slots) {
126                Some((mut buf, other)) => {
127                    if !matches!(step.op, Op::Reshape { .. }) {
128                        let other = other.map(|s| slots[s].data());
129                        elementwise::apply_in_place(&step.op, &mut buf, other);
130                    }
131                    buf
132                }
133                None => {
134                    let inputs: Vec<&[f32]> =
135                        step.inputs.iter().map(|&s| slots[s].data()).collect();
136                    let shapes: Vec<&[usize]> = step
137                        .inputs
138                        .iter()
139                        .map(|&s| plan.slot_shapes[s].as_slice())
140                        .collect();
141                    kernels::apply(&step.op, &inputs, &shapes)
142                }
143            };
144            // Release operands this step consumed last, so peak memory is the
145            // live set and not every intermediate.
146            for &s in &step.inputs {
147                remaining[s] -= 1;
148                if remaining[s] == 0 {
149                    slots[s] = Slot::Empty;
150                }
151            }
152            slots[step.output] = Slot::Owned(out);
153        }
154        // A computed buffer listed once as an output moves out; anything else
155        // (borrowed, or listed more than once) is copied.
156        let mut uses = vec![0usize; slots.len()];
157        for &slot in &plan.outputs {
158            uses[slot] += 1;
159        }
160        plan.outputs
161            .iter()
162            .map(|&slot| {
163                let data = match &mut slots[slot] {
164                    Slot::Owned(buffer) if uses[slot] == 1 => std::mem::take(buffer),
165                    other => {
166                        #[cfg(test)]
167                        OUTPUT_COPIES.with(|c| c.set(c.get() + 1));
168                        other.data().to_vec()
169                    }
170                };
171                Tensor::new(plan.slot_shapes[slot].clone(), data)
172            })
173            .collect()
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use oxfuse_core::{GraphBuilder, Limits, ValueId};
181
182    /// Runs the graph on `[1, -2, 3]` and returns the output buffers copied
183    /// and the steps that reused an operand's buffer.
184    fn counts_when_running(
185        build: impl FnOnce(&mut GraphBuilder) -> Vec<ValueId>,
186    ) -> (usize, usize) {
187        let mut b = GraphBuilder::new(Limits::default());
188        let outputs = build(&mut b);
189        let graph = b.finish(&outputs).unwrap();
190        let exe = CpuBackend.compile(&graph).unwrap();
191        let input = Tensor::new(vec![3], vec![1.0, -2.0, 3.0]).unwrap();
192        OUTPUT_COPIES.with(|c| c.set(0));
193        REUSED.with(|c| c.set(0));
194        exe.run(&[input]).unwrap();
195        (OUTPUT_COPIES.with(Cell::get), REUSED.with(Cell::get))
196    }
197
198    fn copies_when_running(build: impl FnOnce(&mut GraphBuilder) -> Vec<ValueId>) -> usize {
199        counts_when_running(build).0
200    }
201
202    fn reuses_when_running(build: impl FnOnce(&mut GraphBuilder) -> Vec<ValueId>) -> usize {
203        counts_when_running(build).1
204    }
205
206    use std::cell::Cell;
207
208    #[test]
209    fn a_computed_output_listed_once_moves_out_without_a_copy() {
210        let copies = copies_when_running(|b| {
211            let x = b.input("x", vec![3]).unwrap();
212            vec![b.node("y", Op::Relu, &[x]).unwrap()]
213        });
214        assert_eq!(copies, 0);
215    }
216
217    #[test]
218    fn a_computed_output_listed_twice_is_copied_each_time() {
219        let copies = copies_when_running(|b| {
220            let x = b.input("x", vec![3]).unwrap();
221            let y = b.node("y", Op::Relu, &[x]).unwrap();
222            vec![y, y]
223        });
224        assert_eq!(copies, 2);
225    }
226
227    #[test]
228    fn a_borrowed_input_listed_as_an_output_is_copied() {
229        let copies = copies_when_running(|b| vec![b.input("x", vec![3]).unwrap()]);
230        assert_eq!(copies, 1);
231    }
232
233    #[test]
234    fn an_output_that_is_also_read_later_is_still_moved_once_at_the_end() {
235        let copies = copies_when_running(|b| {
236            let x = b.input("x", vec![3]).unwrap();
237            let y = b.node("y", Op::Relu, &[x]).unwrap();
238            let z = b.node("z", Op::Add, &[y, y]).unwrap();
239            vec![y, z]
240        });
241        assert_eq!(copies, 0);
242    }
243
244    #[test]
245    fn a_chain_reuses_each_intermediate_buffer_but_not_the_borrowed_input() {
246        let reused = reuses_when_running(|b| {
247            let x = b.input("x", vec![3]).unwrap();
248            let a = b.node("a", Op::Relu, &[x]).unwrap();
249            let c = b.node("c", Op::Relu, &[a]).unwrap();
250            vec![b.node("d", Op::Relu, &[c]).unwrap()]
251        });
252        assert_eq!(reused, 2);
253    }
254
255    #[test]
256    fn a_value_that_is_also_an_output_is_never_overwritten() {
257        let reused = reuses_when_running(|b| {
258            let x = b.input("x", vec![3]).unwrap();
259            let y = b.node("y", Op::Relu, &[x]).unwrap();
260            let z = b.node("z", Op::Relu, &[y]).unwrap();
261            vec![y, z]
262        });
263        assert_eq!(reused, 0);
264    }
265
266    #[test]
267    fn a_value_read_twice_by_one_step_is_never_overwritten() {
268        let reused = reuses_when_running(|b| {
269            let x = b.input("x", vec![3]).unwrap();
270            let y = b.node("y", Op::Relu, &[x]).unwrap();
271            vec![b.node("z", Op::Add, &[y, y]).unwrap()]
272        });
273        assert_eq!(reused, 0);
274    }
275
276    #[test]
277    fn the_last_reader_of_a_shared_value_reuses_its_buffer() {
278        // y feeds a and then c: a must copy, c may overwrite y.
279        let reused = reuses_when_running(|b| {
280            let x = b.input("x", vec![3]).unwrap();
281            let y = b.node("y", Op::Relu, &[x]).unwrap();
282            let a = b.node("a", Op::Relu, &[y]).unwrap();
283            let c = b.node("c", Op::Relu, &[y]).unwrap();
284            vec![a, c]
285        });
286        assert_eq!(reused, 1);
287    }
288
289    #[test]
290    fn a_scalar_operand_is_not_taken_as_the_buffer_of_a_larger_result() {
291        // r is a computed rank-0 value of size 1; the sum has size 3.
292        let reused = reuses_when_running(|b| {
293            let x = b.input("x", vec![3]).unwrap();
294            let y = b.node("y", Op::Relu, &[x]).unwrap();
295            let s = b.constant("s", vec![], vec![0.5]).unwrap();
296            let r = b.node("r", Op::Relu, &[s]).unwrap();
297            vec![b.node("z", Op::Add, &[r, y]).unwrap()]
298        });
299        assert_eq!(reused, 1);
300    }
301}
oxfuse_runtime/cpu/mod.rs

oxfuse_runtime/cpu/
mod.rs
