Skip to main content

runmat_runtime/
numeric_region.rs

1//! Executor-neutral fused evaluation for a deliberately small, exact numeric
2//! region subset. Executors build plans; Runtime owns MATLAB value semantics.
3
4use 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
111/// Evaluate a supported scalar/dense-double expression DAG in one element
112/// pass. Equal-shape dense arrays and scalar expansion are admitted; every
113/// other representation or broadcast shape fails closed before publication.
114pub 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}