1use crate::{Backend, Executable, Plan, Tensor, kernels};
4use oxfuse_core::{Error, Graph};
5use std::sync::Arc;
6
7pub struct ReferenceBackend;
8
9impl Backend for ReferenceBackend {
10 fn compile(&self, graph: &Graph) -> Result<Box<dyn Executable>, Error> {
11 Ok(Box::new(ReferenceExecutable {
12 plan: Plan::build(graph),
13 }))
14 }
15}
16
17struct ReferenceExecutable {
18 plan: Plan,
19}
20
21enum Slot {
22 Empty,
23 Shared(Arc<[f32]>),
24 Owned(Vec<f32>),
25}
26
27impl Slot {
28 fn data(&self) -> &[f32] {
29 match self {
30 Slot::Empty => &[],
31 Slot::Shared(data) => data,
32 Slot::Owned(data) => data,
33 }
34 }
35}
36
37impl Executable for ReferenceExecutable {
38 fn run(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, Error> {
39 let plan = &self.plan;
40 if inputs.len() != plan.inputs.len() {
41 return Err(Error::InputCount {
42 expected: plan.inputs.len(),
43 found: inputs.len(),
44 });
45 }
46 let mut slots: Vec<Slot> = (0..plan.slot_shapes.len()).map(|_| Slot::Empty).collect();
47 for (binding, tensor) in plan.inputs.iter().zip(inputs) {
48 if tensor.shape() != binding.shape.as_slice() {
49 return Err(Error::InputShape {
50 name: binding.name.clone(),
51 expected: binding.shape.clone(),
52 found: tensor.shape().to_vec(),
53 });
54 }
55 slots[binding.slot] = Slot::Owned(tensor.data().to_vec());
56 }
57 for (slot, data) in &plan.constants {
58 slots[*slot] = Slot::Shared(Arc::clone(data));
59 }
60 for step in &plan.steps {
61 let out = {
62 let inputs: Vec<&[f32]> = step.inputs.iter().map(|&s| slots[s].data()).collect();
63 let shapes: Vec<&[usize]> = step
64 .inputs
65 .iter()
66 .map(|&s| plan.slot_shapes[s].as_slice())
67 .collect();
68 kernels::apply(&step.op, &inputs, &shapes)
69 };
70 slots[step.output] = Slot::Owned(out);
71 }
72 plan.outputs
73 .iter()
74 .map(|&slot| Tensor::new(plan.slot_shapes[slot].clone(), slots[slot].data().to_vec()))
75 .collect()
76 }
77}
78
79#[cfg(test)]
80mod tests {
81 use super::*;
82 use oxfuse_core::{GraphBuilder, Limits, Op};
83
84 fn relu_graph(shape: Vec<usize>) -> Graph {
85 let mut b = GraphBuilder::new(Limits::default());
86 let x = b.input("x", shape).unwrap();
87 let y = b.node("y", Op::Relu, &[x]).unwrap();
88 b.finish(&[y]).unwrap()
89 }
90
91 fn run(graph: &Graph, inputs: &[Tensor]) -> Result<Vec<Tensor>, Error> {
92 ReferenceBackend.compile(graph)?.run(inputs)
93 }
94
95 #[test]
96 fn relu_clamps_negatives() {
97 let input = Tensor::new(vec![2, 3], vec![-1.5, 0.0, 2.5, -3.0, 3.0, -0.5]).unwrap();
98 let out = run(&relu_graph(vec![2, 3]), &[input]).unwrap();
99 assert_eq!(out.len(), 1);
100 assert_eq!(out[0].shape(), &[2, 3]);
101 assert_eq!(out[0].data(), &[0.0, 0.0, 2.5, 0.0, 3.0, 0.0]);
102 }
103
104 #[test]
105 fn relu_of_negative_zero_is_positive_zero() {
106 let input = Tensor::new(vec![1], vec![-0.0]).unwrap();
107 let out = run(&relu_graph(vec![1]), &[input]).unwrap();
108 assert_eq!(out[0].data()[0].to_bits(), 0.0f32.to_bits());
109 }
110
111 #[test]
112 fn relu_propagates_nan() {
113 let input = Tensor::new(vec![2], vec![f32::NAN, f32::NEG_INFINITY]).unwrap();
114 let out = run(&relu_graph(vec![2]), &[input]).unwrap();
115 assert!(out[0].data()[0].is_nan());
116 assert_eq!(out[0].data()[1], 0.0);
117 }
118
119 #[test]
120 fn relu_passes_positive_infinity() {
121 let input = Tensor::new(vec![1], vec![f32::INFINITY]).unwrap();
122 let out = run(&relu_graph(vec![1]), &[input]).unwrap();
123 assert_eq!(out[0].data()[0], f32::INFINITY);
124 }
125
126 #[test]
127 fn scalar_input_works() {
128 let input = Tensor::new(vec![], vec![-4.0]).unwrap();
129 let out = run(&relu_graph(vec![]), &[input]).unwrap();
130 assert_eq!(out[0].shape(), &[] as &[usize]);
131 assert_eq!(out[0].data(), &[0.0]);
132 }
133
134 #[test]
135 fn constant_only_graph_runs_with_no_inputs() {
136 let mut b = GraphBuilder::new(Limits::default());
137 let c = b.constant("c", vec![2], vec![-1.0, 2.0]).unwrap();
138 let r = b.node("r", Op::Relu, &[c]).unwrap();
139 let g = b.finish(&[r]).unwrap();
140 let out = run(&g, &[]).unwrap();
141 assert_eq!(out[0].data(), &[0.0, 2.0]);
142 }
143
144 #[test]
145 fn output_can_be_an_input() {
146 let mut b = GraphBuilder::new(Limits::default());
147 let x = b.input("x", vec![2]).unwrap();
148 let g = b.finish(&[x]).unwrap();
149 let input = Tensor::new(vec![2], vec![-1.0, 1.0]).unwrap();
150 let out = run(&g, std::slice::from_ref(&input)).unwrap();
151 assert_eq!(out[0], input);
152 }
153
154 #[test]
155 fn same_value_can_be_output_twice() {
156 let mut b = GraphBuilder::new(Limits::default());
157 let x = b.input("x", vec![1]).unwrap();
158 let y = b.node("y", Op::Relu, &[x]).unwrap();
159 let g = b.finish(&[y, y]).unwrap();
160 let out = run(&g, &[Tensor::new(vec![1], vec![-2.0]).unwrap()]).unwrap();
161 assert_eq!(out.len(), 2);
162 assert_eq!(out[0], out[1]);
163 }
164
165 #[test]
166 fn wrong_input_count_is_rejected() {
167 let err = run(&relu_graph(vec![1]), &[]).unwrap_err();
168 assert_eq!(
169 err,
170 Error::InputCount {
171 expected: 1,
172 found: 0
173 }
174 );
175 }
176
177 #[test]
178 fn wrong_input_shape_is_rejected() {
179 let input = Tensor::new(vec![3], vec![1.0, 2.0, 3.0]).unwrap();
180 let err = run(&relu_graph(vec![2]), &[input]).unwrap_err();
181 assert_eq!(
182 err,
183 Error::InputShape {
184 name: "x".into(),
185 expected: vec![2],
186 found: vec![3]
187 }
188 );
189 }
190
191 #[test]
192 fn executable_is_reusable() {
193 let exe = ReferenceBackend.compile(&relu_graph(vec![1])).unwrap();
194 let a = exe
195 .run(&[Tensor::new(vec![1], vec![-1.0]).unwrap()])
196 .unwrap();
197 let b = exe
198 .run(&[Tensor::new(vec![1], vec![5.0]).unwrap()])
199 .unwrap();
200 assert_eq!(a[0].data(), &[0.0]);
201 assert_eq!(b[0].data(), &[5.0]);
202 }
203
204 #[test]
205 fn add_then_scale_by_a_scalar_constant() {
206 let mut b = GraphBuilder::new(Limits::default());
207 let a = b.input("a", vec![2, 2]).unwrap();
208 let c = b.input("b", vec![2, 2]).unwrap();
209 let two = b.constant("two", vec![], vec![2.0]).unwrap();
210 let sum = b.node("sum", Op::Add, &[a, c]).unwrap();
211 let scaled = b.node("scaled", Op::Mul, &[sum, two]).unwrap();
212 let g = b.finish(&[sum, scaled]).unwrap();
213 let a = Tensor::new(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
214 let c = Tensor::new(vec![2, 2], vec![10.0, 20.0, 30.0, 40.0]).unwrap();
215 let out = run(&g, &[a, c]).unwrap();
216 assert_eq!(out[0].data(), &[11.0, 22.0, 33.0, 44.0]);
217 assert_eq!(out[1].data(), &[22.0, 44.0, 66.0, 88.0]);
218 }
219
220 #[test]
221 fn matmul_graph_runs() {
222 let mut b = GraphBuilder::new(Limits::default());
223 let x = b.input("x", vec![2, 2]).unwrap();
224 let w = b
225 .constant("w", vec![2, 2], vec![5.0, 6.0, 7.0, 8.0])
226 .unwrap();
227 let y = b.node("y", Op::MatMul, &[x, w]).unwrap();
228 let g = b.finish(&[y]).unwrap();
229 let input = Tensor::new(vec![2, 2], vec![1.0, 2.0, 3.0, 4.0]).unwrap();
230 let out = run(&g, &[input]).unwrap();
231 assert_eq!(out[0].data(), &[19.0, 22.0, 43.0, 50.0]);
232 }
233}