1use std::sync::{
5 atomic::{AtomicBool, Ordering},
6 Arc,
7};
8
9use runmat_value::{Tensor, Value};
10
11const MAX_NODES: usize = 256;
12const MAX_OUTPUTS: usize = 64;
13
14#[derive(Clone, Copy, Debug, Eq, PartialEq)]
15pub enum NumericUnaryOperation {
16 Plus,
17 Minus,
18}
19
20#[derive(Clone, Copy, Debug, Eq, PartialEq)]
21pub enum NumericBinaryOperation {
22 Add,
23 Subtract,
24 Multiply,
25 Divide,
26 LeftDivide,
27}
28
29#[derive(Clone, Debug, PartialEq)]
30pub enum NumericRegionNode {
31 Input(usize),
32 Constant(f64),
33 Unary {
34 operation: NumericUnaryOperation,
35 input: usize,
36 },
37 Binary {
38 operation: NumericBinaryOperation,
39 left: usize,
40 right: usize,
41 },
42}
43
44#[derive(Clone, Debug, PartialEq)]
45pub struct NumericRegionProgram {
46 pub nodes: Vec<NumericRegionNode>,
47 pub outputs: Vec<usize>,
48}
49
50#[derive(Clone, Debug, PartialEq)]
51pub enum NumericRegionExecution {
52 Completed(Vec<Value>),
53 Ineligible,
54 Cancelled,
55}
56
57#[derive(Clone, Copy, Debug, Eq, PartialEq)]
58pub struct NumericRegionWorkload {
59 pub elements: usize,
60 pub output_bytes_per_value: u64,
61}
62
63pub fn workload(inputs: &[&Value]) -> Option<NumericRegionWorkload> {
64 let (_, shape) = analyze_inputs(inputs)?;
65 let elements = shape
66 .as_ref()
67 .map(|shape| shape.iter().try_fold(1_usize, |n, dim| n.checked_mul(*dim)))
68 .unwrap_or(Some(1))?;
69 Some(NumericRegionWorkload {
70 elements,
71 output_bytes_per_value: u64::try_from(elements).ok()?.checked_mul(8)?,
72 })
73}
74
75impl NumericRegionProgram {
76 pub fn validate(&self, input_count: usize) -> Result<(), &'static str> {
77 if self.nodes.is_empty()
78 || self.nodes.len() > MAX_NODES
79 || self.outputs.is_empty()
80 || self.outputs.len() > MAX_OUTPUTS
81 {
82 return Err("numeric region plan exceeds its structural bounds");
83 }
84 for (index, node) in self.nodes.iter().enumerate() {
85 match node {
86 NumericRegionNode::Input(input) if *input >= input_count => {
87 return Err("numeric region input is out of bounds")
88 }
89 NumericRegionNode::Unary { input, .. } if *input >= index => {
90 return Err("numeric region unary dependency is not topological")
91 }
92 NumericRegionNode::Binary { left, right, .. }
93 if *left >= index || *right >= index =>
94 {
95 return Err("numeric region binary dependency is not topological")
96 }
97 _ => {}
98 }
99 }
100 if self
101 .outputs
102 .iter()
103 .any(|output| *output >= self.nodes.len())
104 {
105 return Err("numeric region output is out of bounds");
106 }
107 Ok(())
108 }
109}
110
111pub fn execute(
115 program: &NumericRegionProgram,
116 inputs: &[&Value],
117 cancellation: &Arc<AtomicBool>,
118) -> Result<NumericRegionExecution, &'static str> {
119 program.validate(inputs.len())?;
120 let Some((input_views, output_shape)) = analyze_inputs(inputs) else {
121 return Ok(NumericRegionExecution::Ineligible);
122 };
123 let element_count = output_shape
124 .as_ref()
125 .map(|shape| shape.iter().try_fold(1_usize, |n, dim| n.checked_mul(*dim)))
126 .unwrap_or(Some(1))
127 .ok_or("numeric region output shape overflows this host")?;
128 let total_output_elements = element_count
129 .checked_mul(program.outputs.len())
130 .ok_or("numeric region output allocation overflows this host")?;
131 let mut output_values = vec![Vec::with_capacity(element_count); program.outputs.len()];
132 let mut values = vec![0.0; program.nodes.len()];
133 for element in 0..element_count {
134 if element % 1_024 == 0 && cancellation.load(Ordering::Relaxed) {
135 return Ok(NumericRegionExecution::Cancelled);
136 }
137 for (index, node) in program.nodes.iter().enumerate() {
138 values[index] = match *node {
139 NumericRegionNode::Input(input) => input_views[input].at(element),
140 NumericRegionNode::Constant(value) => value,
141 NumericRegionNode::Unary { operation, input } => match operation {
142 NumericUnaryOperation::Plus => values[input],
143 NumericUnaryOperation::Minus => -values[input],
144 },
145 NumericRegionNode::Binary {
146 operation,
147 left,
148 right,
149 } => match operation {
150 NumericBinaryOperation::Add => values[left] + values[right],
151 NumericBinaryOperation::Subtract => values[left] - values[right],
152 NumericBinaryOperation::Multiply => values[left] * values[right],
153 NumericBinaryOperation::Divide => values[left] / values[right],
154 NumericBinaryOperation::LeftDivide => values[right] / values[left],
155 },
156 };
157 }
158 for (output, node) in output_values.iter_mut().zip(&program.outputs) {
159 output.push(values[*node]);
160 }
161 }
162 debug_assert_eq!(
163 output_values.iter().map(Vec::len).sum::<usize>(),
164 total_output_elements
165 );
166 let outputs = if let Some(shape) = output_shape {
167 output_values
168 .into_iter()
169 .map(|values| Tensor::new(values, shape.clone()).map(Value::Tensor))
170 .collect::<Result<Vec<_>, _>>()
171 .map_err(|_| "numeric region produced an invalid dense tensor")?
172 } else {
173 output_values
174 .into_iter()
175 .map(|values| Value::Num(values[0]))
176 .collect()
177 };
178 Ok(NumericRegionExecution::Completed(outputs))
179}
180
181fn analyze_inputs<'a>(inputs: &[&'a Value]) -> Option<(Vec<NumericInput<'a>>, Option<Vec<usize>>)> {
182 let mut output_shape: Option<Vec<usize>> = None;
183 let mut tensor_shape: Option<Vec<usize>> = None;
184 let mut input_views = Vec::with_capacity(inputs.len());
185 for &input in inputs {
186 match input {
187 Value::Num(value) => input_views.push(NumericInput::Scalar(*value)),
188 Value::Tensor(tensor) => {
189 let values = tensor.as_f64_slice()?;
190 if tensor.len() == 1 {
191 output_shape.get_or_insert_with(|| tensor.shape.clone());
192 input_views.push(NumericInput::Scalar(values[0]));
193 } else {
194 if tensor_shape
195 .as_ref()
196 .is_some_and(|shape| shape != &tensor.shape)
197 {
198 return None;
199 }
200 tensor_shape = Some(tensor.shape.clone());
201 output_shape = Some(tensor.shape.clone());
202 input_views.push(NumericInput::Dense(values));
203 }
204 }
205 _ => return None,
206 }
207 }
208 Some((input_views, output_shape))
209}
210
211#[derive(Clone, Copy)]
212enum NumericInput<'a> {
213 Scalar(f64),
214 Dense(&'a [f64]),
215}
216
217impl NumericInput<'_> {
218 fn at(self, index: usize) -> f64 {
219 match self {
220 Self::Scalar(value) => value,
221 Self::Dense(values) => values[index],
222 }
223 }
224}
225
226#[cfg(test)]
227mod tests {
228 use super::*;
229
230 fn execute_values(
231 program: &NumericRegionProgram,
232 inputs: &[Value],
233 cancellation: &Arc<AtomicBool>,
234 ) -> Result<NumericRegionExecution, &'static str> {
235 let inputs = inputs.iter().collect::<Vec<_>>();
236 execute(program, &inputs, cancellation)
237 }
238
239 #[test]
240 fn fuses_dense_double_chain_with_scalar_expansion() {
241 let program = NumericRegionProgram {
242 nodes: vec![
243 NumericRegionNode::Input(0),
244 NumericRegionNode::Input(1),
245 NumericRegionNode::Binary {
246 operation: NumericBinaryOperation::Add,
247 left: 0,
248 right: 1,
249 },
250 NumericRegionNode::Constant(2.0),
251 NumericRegionNode::Binary {
252 operation: NumericBinaryOperation::Multiply,
253 left: 2,
254 right: 3,
255 },
256 ],
257 outputs: vec![4],
258 };
259 let tensor = Tensor::new(vec![1.0, 2.0, 3.0], vec![1, 3]).unwrap();
260 let result = execute_values(
261 &program,
262 &[Value::Tensor(tensor), Value::Num(1.0)],
263 &Arc::new(AtomicBool::new(false)),
264 )
265 .unwrap();
266 let NumericRegionExecution::Completed(outputs) = result else {
267 panic!("eligible numeric region did not execute")
268 };
269 let Value::Tensor(output) = &outputs[0] else {
270 panic!("expected dense output")
271 };
272 assert_eq!(output.materialize_f64(), vec![4.0, 6.0, 8.0]);
273 }
274
275 #[test]
276 fn rejects_incompatible_dense_shapes_before_publication() {
277 let program = NumericRegionProgram {
278 nodes: vec![
279 NumericRegionNode::Input(0),
280 NumericRegionNode::Input(1),
281 NumericRegionNode::Binary {
282 operation: NumericBinaryOperation::Add,
283 left: 0,
284 right: 1,
285 },
286 ],
287 outputs: vec![2],
288 };
289 let left = Tensor::new(vec![1.0, 2.0], vec![1, 2]).unwrap();
290 let right = Tensor::new(vec![1.0, 2.0], vec![2, 1]).unwrap();
291 assert_eq!(
292 execute_values(
293 &program,
294 &[Value::Tensor(left), Value::Tensor(right)],
295 &Arc::new(AtomicBool::new(false)),
296 )
297 .unwrap(),
298 NumericRegionExecution::Ineligible
299 );
300 }
301
302 #[test]
303 fn cancellation_discards_the_transactional_result() {
304 let program = NumericRegionProgram {
305 nodes: vec![
306 NumericRegionNode::Input(0),
307 NumericRegionNode::Constant(1.0),
308 NumericRegionNode::Binary {
309 operation: NumericBinaryOperation::Add,
310 left: 0,
311 right: 1,
312 },
313 ],
314 outputs: vec![2],
315 };
316 let input = Tensor::new(vec![0.0; 2_048], vec![1, 2_048]).unwrap();
317 let cancellation = Arc::new(AtomicBool::new(true));
318 assert_eq!(
319 execute_values(&program, &[Value::Tensor(input)], &cancellation).unwrap(),
320 NumericRegionExecution::Cancelled
321 );
322 }
323
324 #[test]
325 fn preserves_scalar_array_shape() {
326 let program = NumericRegionProgram {
327 nodes: vec![
328 NumericRegionNode::Input(0),
329 NumericRegionNode::Constant(1.0),
330 NumericRegionNode::Binary {
331 operation: NumericBinaryOperation::Add,
332 left: 0,
333 right: 1,
334 },
335 ],
336 outputs: vec![2],
337 };
338 let input = Tensor::new(vec![2.0], vec![1, 1]).unwrap();
339 let result = execute_values(
340 &program,
341 &[Value::Tensor(input)],
342 &Arc::new(AtomicBool::new(false)),
343 )
344 .unwrap();
345 let NumericRegionExecution::Completed(outputs) = result else {
346 panic!("eligible numeric region did not execute")
347 };
348 let Value::Tensor(output) = &outputs[0] else {
349 panic!("scalar arrays must remain arrays")
350 };
351 assert_eq!(output.shape, vec![1, 1]);
352 assert_eq!(output.materialize_f64(), vec![3.0]);
353 }
354}