Skip to main content

laddu_runtime/cpu/
evaluation.rs

1use laddu_data::{data::EventBatch, schema::Schema};
2use laddu_expr::{BinaryOp, ExprId, ExprNode, P4Component, UnaryOp, parameters::ParamValues};
3use laddu_kernel::ir::{GradientKernelIr, KernelInstruction, KernelValue, KernelValueKind};
4use nalgebra::{DMatrix, DVector};
5use num::complex::{Complex32, Complex64};
6
7use super::layout::{
8    F32Value, Value, eval_binary, eval_unary, f32_matrix_at, f32_scalar_at, f32_vector_at,
9    matrix_at, matrix_at_optional, matrix_values_row_major, matrix_values_row_major_f32, scalar_at,
10    scalar_at_optional, vector_at, vector_at_optional,
11};
12use super::scalar::{
13    OperandRun, SCALAR_BLOCK_SIZE, ScalarEvaluationPlan, ScalarEventWorkspace, ScalarInstruction,
14    ScalarInvariantValues, ScalarSlot,
15};
16use super::{
17    CpuBatchCache, CpuPlan, EventLookup, Precision, RuntimeError, RuntimeResult, ValueGradient,
18};
19
20#[cfg(feature = "jit")]
21use crate::jit::{JitCacheView, JitScalarKernel};
22
23#[derive(Copy, Clone, Debug, PartialEq, Eq)]
24pub(super) enum EventColumn {
25    Scalar(usize),
26    P4Component { col: usize, component: P4Component },
27}
28
29#[derive(Clone, Copy)]
30pub(super) enum F32KernelInput<'a> {
31    Cache(Option<(&'a CpuBatchCache, usize)>),
32    Event(&'a dyn EventLookup),
33}
34
35impl<'a> F32KernelInput<'a> {
36    pub(super) fn cache(self) -> Option<(&'a CpuBatchCache, usize)> {
37        match self {
38            Self::Cache(cache) => cache,
39            Self::Event(_) => None,
40        }
41    }
42}
43
44impl CpuPlan {
45    pub(super) fn parameter_value(&self, params: &ParamValues, node: usize) -> RuntimeResult<f64> {
46        let id = self.parameter_slots[node].ok_or_else(|| RuntimeError::InvalidShape {
47            index: node,
48            message: "node is not a parameter".into(),
49        })?;
50        params
51            .get(id)
52            .map_err(|err| RuntimeError::Parameter(err.to_string()))
53    }
54
55    pub(super) fn evaluate_inner(
56        &self,
57        params: &ParamValues,
58        event: Option<&dyn EventLookup>,
59    ) -> RuntimeResult<Complex64> {
60        #[cfg(feature = "jit")]
61        if event.is_none()
62            && let Some(kernel) = self.scalar_jit_kernel()
63        {
64            if params.as_slice().len() != self.params.len() {
65                return Err(RuntimeError::Parameter(format!(
66                    "expected {} parameter values, got {}",
67                    self.params.len(),
68                    params.as_slice().len()
69                )));
70            }
71            return kernel.evaluate_invariant(params);
72        }
73        if self.precision == Precision::F32 {
74            let input = match event {
75                Some(event) => F32KernelInput::Event(event),
76                None => F32KernelInput::Cache(None),
77            };
78            return self.evaluate_f32_scalar(params, input);
79        }
80        let values = self.evaluate_values(params, event)?;
81        scalar_at(&values, self.graph.root().index())
82    }
83
84    /// Evaluates an event-independent model and its free-parameter gradient.
85    ///
86    /// # Errors
87    ///
88    /// Returns [`RuntimeError`] when parameters are incompatible, the model
89    /// requires event data, differentiation or evaluation fails, or a solve is
90    /// singular.
91    pub fn evaluate_with_gradient(&self, params: &ParamValues) -> RuntimeResult<ValueGradient> {
92        #[cfg(feature = "jit")]
93        if let (Some(value_kernel), Some(gradient_kernel)) =
94            (self.scalar_jit_kernel(), self.gradient_jit_kernel())
95        {
96            let value = value_kernel.evaluate_invariant(params)?;
97            let mut real = Vec::new();
98            let mut imag = Vec::new();
99            gradient_kernel.evaluate_invariant_component(params, 0, &mut real)?;
100            gradient_kernel.evaluate_invariant_component(params, 1, &mut imag)?;
101            let gradient = real
102                .into_iter()
103                .zip(imag)
104                .map(|(re, im)| Complex64::new(re, im))
105                .collect();
106            return Ok(ValueGradient { value, gradient });
107        }
108        if self.precision == Precision::F32 {
109            return self.evaluate_f32_gradient(params, F32KernelInput::Cache(None));
110        }
111        self.require_f64_gradient()?;
112        if let Some(interpreter) = self.gradient_interpreter() {
113            let (value, gradient) = interpreter.evaluate(params, None)?;
114            return Ok(ValueGradient { value, gradient });
115        }
116        let values = self.evaluate_values(params, None)?;
117        self.value_gradient(values, None)
118    }
119
120    pub(super) fn require_f64_gradient(&self) -> RuntimeResult<()> {
121        if self.precision == Precision::F32 {
122            return Err(crate::ExecutionError::UnsupportedCpuF32Gradient.into());
123        }
124        Ok(())
125    }
126
127    pub(super) fn evaluate_f32_scalar(
128        &self,
129        params: &ParamValues,
130        input: F32KernelInput<'_>,
131    ) -> RuntimeResult<Complex64> {
132        let kernel = self
133            .scalar_kernel
134            .as_ref()
135            .ok_or(crate::ExecutionError::UnsupportedCpuF32Model)?;
136        let values = self.evaluate_f32_kernel_values(kernel.values(), params, input)?;
137        let value = f32_scalar_at(&values, kernel.root())?;
138        Ok(Complex64::new(value.re as f64, value.im as f64))
139    }
140
141    pub(super) fn evaluate_f32_kernel_values(
142        &self,
143        kernel_values: &[KernelValue],
144        params: &ParamValues,
145        input: F32KernelInput<'_>,
146    ) -> RuntimeResult<Vec<F32Value>> {
147        let mut values = Vec::with_capacity(kernel_values.len());
148        for (index, value) in kernel_values.iter().enumerate() {
149            let result = match &value.instruction {
150                KernelInstruction::Cached(slot) => self.evaluate_f32_cached_value(
151                    *slot,
152                    params,
153                    input,
154                    crate::ExecutionError::UnsupportedCpuF32Model,
155                )?,
156                KernelInstruction::RealConstant(value) => {
157                    F32Value::Scalar(Complex32::from(*value as f32))
158                }
159                KernelInstruction::ComplexConstant(value) => {
160                    F32Value::Scalar(Complex32::new(value.re as f32, value.im as f32))
161                }
162                KernelInstruction::Parameter(id) => F32Value::Scalar(Complex32::from(
163                    params
164                        .get(*id)
165                        .map_err(|error| RuntimeError::Parameter(error.to_string()))?
166                        as f32,
167                )),
168                KernelInstruction::Unary { op, input } => {
169                    F32Value::Scalar(eval_unary(*op, f32_scalar_at(&values, *input)?))
170                }
171                KernelInstruction::Binary { op, lhs, rhs } => F32Value::Scalar(eval_binary(
172                    *op,
173                    f32_scalar_at(&values, *lhs)?,
174                    f32_scalar_at(&values, *rhs)?,
175                )),
176                KernelInstruction::Add(terms) => F32Value::Scalar(
177                    terms
178                        .iter()
179                        .map(|id| f32_scalar_at(&values, *id))
180                        .sum::<RuntimeResult<Complex32>>()?,
181                ),
182                KernelInstruction::Mul(factors) => {
183                    F32Value::Scalar(factors.iter().try_fold(Complex32::ONE, |product, id| {
184                        Ok::<_, RuntimeError>(product * f32_scalar_at(&values, *id)?)
185                    })?)
186                }
187                KernelInstruction::Complex { re, im } => F32Value::Scalar(Complex32::new(
188                    f32_scalar_at(&values, *re)?.re,
189                    f32_scalar_at(&values, *im)?.re,
190                )),
191                KernelInstruction::Vector(elements) => F32Value::Vector(
192                    elements
193                        .iter()
194                        .map(|id| f32_scalar_at(&values, *id))
195                        .collect::<RuntimeResult<_>>()?,
196                ),
197                KernelInstruction::Matrix {
198                    rows,
199                    cols,
200                    elements,
201                } => F32Value::Matrix {
202                    rows: *rows,
203                    cols: *cols,
204                    values: elements
205                        .iter()
206                        .map(|id| f32_scalar_at(&values, *id))
207                        .collect::<RuntimeResult<_>>()?,
208                },
209                KernelInstruction::Component {
210                    input,
211                    index: element,
212                } => {
213                    let vector = f32_vector_at(&values, *input)?;
214                    F32Value::Scalar(*vector.get(*element).ok_or_else(|| {
215                        RuntimeError::InvalidShape {
216                            index,
217                            message: format!(
218                                "component index {element} out of bounds for len {}",
219                                vector.len()
220                            ),
221                        }
222                    })?)
223                }
224                KernelInstruction::MatrixElement { input, row, col } => {
225                    let (rows, cols, matrix) = f32_matrix_at(&values, *input)?;
226                    let Some(offset) = (KernelValueKind::Matrix { rows, cols })
227                        .checked_row_major_index(*row, *col)
228                    else {
229                        return Err(RuntimeError::InvalidShape {
230                            index,
231                            message: format!(
232                                "matrix element ({row}, {col}) out of bounds for shape {rows}x{cols}"
233                            ),
234                        });
235                    };
236                    F32Value::Scalar(matrix[offset])
237                }
238                KernelInstruction::Dot { lhs, rhs } => {
239                    let lhs = f32_vector_at(&values, *lhs)?;
240                    let rhs = f32_vector_at(&values, *rhs)?;
241                    if lhs.len() != rhs.len() {
242                        return Err(RuntimeError::InvalidShape {
243                            index,
244                            message: format!(
245                                "cannot dot len {} vector with len {} vector",
246                                lhs.len(),
247                                rhs.len()
248                            ),
249                        });
250                    }
251                    F32Value::Scalar(lhs.iter().zip(rhs).map(|(lhs, rhs)| lhs * rhs).sum())
252                }
253                KernelInstruction::MatVec { matrix, vector } => {
254                    let (rows, cols, matrix) = f32_matrix_at(&values, *matrix)?;
255                    let vector = f32_vector_at(&values, *vector)?;
256                    if cols != vector.len() {
257                        return Err(RuntimeError::InvalidShape {
258                            index,
259                            message: format!(
260                                "cannot multiply {rows}x{cols} matrix by len {} vector",
261                                vector.len()
262                            ),
263                        });
264                    }
265                    let output = DMatrix::from_row_slice(rows, cols, matrix)
266                        * DVector::from_row_slice(vector);
267                    F32Value::Vector(output.iter().copied().collect())
268                }
269                KernelInstruction::MatMul { lhs, rhs } => {
270                    let (lhs_rows, lhs_cols, lhs) = f32_matrix_at(&values, *lhs)?;
271                    let (rhs_rows, rhs_cols, rhs) = f32_matrix_at(&values, *rhs)?;
272                    if lhs_cols != rhs_rows {
273                        return Err(RuntimeError::InvalidShape {
274                            index,
275                            message: format!(
276                                "cannot multiply {lhs_rows}x{lhs_cols} by {rhs_rows}x{rhs_cols}"
277                            ),
278                        });
279                    }
280                    let output = DMatrix::from_row_slice(lhs_rows, lhs_cols, lhs)
281                        * DMatrix::from_row_slice(rhs_rows, rhs_cols, rhs);
282                    F32Value::Matrix {
283                        rows: output.nrows(),
284                        cols: output.ncols(),
285                        values: matrix_values_row_major_f32(&output),
286                    }
287                }
288                KernelInstruction::Solve { matrix, rhs } => {
289                    let (rows, cols, matrix) = f32_matrix_at(&values, *matrix)?;
290                    let rhs = f32_vector_at(&values, *rhs)?;
291                    if rows != cols || rows != rhs.len() {
292                        return Err(RuntimeError::InvalidShape {
293                            index,
294                            message: format!(
295                                "cannot solve {rows}x{cols} matrix against len {} vector",
296                                rhs.len()
297                            ),
298                        });
299                    }
300                    let solution = DMatrix::from_row_slice(rows, cols, matrix)
301                        .lu()
302                        .solve(&DVector::from_row_slice(rhs))
303                        .ok_or(RuntimeError::SingularMatrix(index))?;
304                    F32Value::Vector(solution.iter().copied().collect())
305                }
306                KernelInstruction::SolveRow { row_slot, rhs } => {
307                    let (cache, row) = input
308                        .cache()
309                        .ok_or(crate::ExecutionError::UnsupportedCpuF32Model)?;
310                    let inverse = cache.solve_row(*row_slot, row)?;
311                    if inverse.len() != rhs.len() {
312                        return Err(RuntimeError::InvalidShape {
313                            index,
314                            message: format!(
315                                "specialized solve row has len {}, expected {}",
316                                inverse.len(),
317                                rhs.len()
318                            ),
319                        });
320                    }
321                    F32Value::Scalar(
322                        inverse
323                            .iter()
324                            .zip(rhs)
325                            .map(|(coefficient, rhs)| {
326                                Ok::<_, RuntimeError>(
327                                    Complex32::new(coefficient.re as f32, coefficient.im as f32)
328                                        * f32_scalar_at(&values, *rhs)?,
329                                )
330                            })
331                            .sum::<RuntimeResult<Complex32>>()?,
332                    )
333                }
334                KernelInstruction::SolveRowAdjointElement {
335                    row_slot,
336                    index: element,
337                    len,
338                    adjoint,
339                } => {
340                    let (cache, row) = input
341                        .cache()
342                        .ok_or(crate::ExecutionError::UnsupportedCpuF32Model)?;
343                    let inverse = cache.solve_row(*row_slot, row)?;
344                    if inverse.len() != *len {
345                        return Err(RuntimeError::InvalidShape {
346                            index,
347                            message: format!(
348                                "specialized solve row has len {}, expected {len}",
349                                inverse.len()
350                            ),
351                        });
352                    }
353                    let coefficient = inverse[*element];
354                    F32Value::Scalar(
355                        f32_scalar_at(&values, *adjoint)?
356                            * Complex32::new(coefficient.re as f32, coefficient.im as f32).conj(),
357                    )
358                }
359            };
360            values.push(result);
361        }
362        Ok(values)
363    }
364
365    pub(super) fn evaluate_f32_cached_value(
366        &self,
367        slot: usize,
368        params: &ParamValues,
369        input: F32KernelInput<'_>,
370        missing_input: crate::ExecutionError,
371    ) -> RuntimeResult<F32Value> {
372        match input {
373            F32KernelInput::Cache(Some((cache, row))) => {
374                Ok(F32Value::from_value(cache.value(slot, row)?))
375            }
376            F32KernelInput::Cache(None) => Err(missing_input.into()),
377            F32KernelInput::Event(event) => {
378                let entry =
379                    self.cache_plan
380                        .entries()
381                        .get(slot)
382                        .ok_or(RuntimeError::InvalidShape {
383                            index: self.graph.root().index(),
384                            message: format!("cache slot {slot} is out of bounds"),
385                        })?;
386                let values = self.evaluate_values(params, Some(event))?;
387                let value = values
388                    .get(entry.node().index())
389                    .ok_or(RuntimeError::InvalidShape {
390                        index: entry.node().index(),
391                        message: "cached node is out of bounds".into(),
392                    })?
393                    .clone();
394                Ok(F32Value::from_value(value))
395            }
396        }
397    }
398
399    pub(super) fn evaluate_f32_gradient(
400        &self,
401        params: &ParamValues,
402        input: F32KernelInput<'_>,
403    ) -> RuntimeResult<ValueGradient> {
404        let real_ir = self
405            .f32_gradient_fallback_real
406            .as_ref()
407            .ok_or(crate::ExecutionError::UnsupportedCpuF32Model)?;
408        let mut real = Vec::new();
409        let (value, _) =
410            self.evaluate_f32_gradient_component_prepared(real_ir, params, input, &mut real)?;
411        let imag = if let Some(imag_ir) = self.f32_gradient_fallback_imag.as_ref() {
412            let mut imag = Vec::new();
413            self.evaluate_f32_gradient_component_prepared(imag_ir, params, input, &mut imag)?;
414            imag
415        } else {
416            vec![0.0; real.len()]
417        };
418        Ok(ValueGradient {
419            value,
420            gradient: real
421                .into_iter()
422                .zip(imag)
423                .map(|(re, im)| Complex64::new(re as f64, im as f64))
424                .collect(),
425        })
426    }
427
428    pub(super) fn evaluate_f32_gradient_component_prepared<'a>(
429        &self,
430        ir: &GradientKernelIr,
431        params: &ParamValues,
432        input: F32KernelInput<'_>,
433        gradient: &'a mut Vec<f32>,
434    ) -> RuntimeResult<(Complex64, &'a [f32])> {
435        let values = self.evaluate_f32_kernel_values(ir.values(), params, input)?;
436        let value = f32_scalar_at(&values, ir.primal_root())?;
437        gradient.clear();
438        gradient.reserve(ir.outputs().len());
439        for output in ir.outputs() {
440            gradient.push(f32_scalar_at(&values, *output)?.re);
441        }
442        Ok((Complex64::new(value.re as f64, value.im as f64), gradient))
443    }
444
445    pub(super) fn solve_primal(
446        &self,
447        matrix_id: ExprId,
448        dimension: usize,
449        matrix: &[Complex64],
450        rhs: &DVector<Complex64>,
451        node_index: usize,
452        cached: Option<(&CpuBatchCache, usize)>,
453    ) -> RuntimeResult<DVector<Complex64>> {
454        let solution = if let (Some(slot), Some((cache, row))) =
455            (self.factor_matrix_slots[matrix_id.index()], cached)
456        {
457            cache.factor(slot, row)?.solve(rhs)
458        } else if let Some(slot) = self.constant_factor_slots[matrix_id.index()] {
459            self.constant_factors[slot]
460                .get_or_init(|| DMatrix::from_row_slice(dimension, dimension, matrix).lu())
461                .solve(rhs)
462        } else {
463            DMatrix::from_row_slice(dimension, dimension, matrix)
464                .lu()
465                .solve(rhs)
466        };
467        solution.ok_or(RuntimeError::SingularMatrix(node_index))
468    }
469
470    pub(super) fn event_columns(&self, schema: &Schema) -> RuntimeResult<Vec<Option<EventColumn>>> {
471        self.graph
472            .nodes()
473            .iter()
474            .map(|node| {
475                if let ExprNode::EventScalar(name) = node {
476                    Ok(Some(EventColumn::Scalar(
477                        schema
478                            .scalar_index(name)
479                            .ok_or_else(|| RuntimeError::MissingEventColumn(name.to_string()))?,
480                    )))
481                } else if let ExprNode::EventP4Component { name, component } = node {
482                    Ok(Some(EventColumn::P4Component {
483                        col: schema
484                            .p4_index(name)
485                            .ok_or_else(|| RuntimeError::MissingEventColumn(name.to_string()))?,
486                        component: *component,
487                    }))
488                } else {
489                    Ok(None)
490                }
491            })
492            .collect()
493    }
494
495    pub(super) fn evaluate_cache_values_for_row(
496        &self,
497        batch: &EventBatch,
498        row: usize,
499        event_columns: &[Option<EventColumn>],
500    ) -> RuntimeResult<Vec<Option<Value>>> {
501        let mut values = vec![None; self.graph.nodes().len()];
502
503        for id in &self.cache_materialization_nodes {
504            let index = id.index();
505            let node = &self.graph.nodes()[index];
506            let value = match node {
507                ExprNode::RealConst(value) => Value::Scalar(Complex64::from(*value)),
508                ExprNode::ComplexConst(value) => Value::Scalar(*value),
509                ExprNode::EventScalar(name) => {
510                    let col = event_columns[index]
511                        .ok_or_else(|| RuntimeError::MissingEventColumn(name.to_string()))?;
512                    let EventColumn::Scalar(col) = col else {
513                        return Err(RuntimeError::MissingEventColumn(name.to_string()));
514                    };
515                    Value::Scalar(Complex64::from(batch.scalar_at(col, row)))
516                }
517                ExprNode::EventP4Component { name, component } => {
518                    let col = event_columns[index]
519                        .ok_or_else(|| RuntimeError::MissingEventColumn(name.to_string()))?;
520                    let EventColumn::P4Component {
521                        col,
522                        component: actual,
523                    } = col
524                    else {
525                        return Err(RuntimeError::MissingEventColumn(name.to_string()));
526                    };
527                    debug_assert_eq!(actual, *component);
528                    let p4 = batch.p4_at(col, row);
529                    let value = match component {
530                        P4Component::Px => p4.px,
531                        P4Component::Py => p4.py,
532                        P4Component::Pz => p4.pz,
533                        P4Component::E => p4.e,
534                    };
535                    Value::Scalar(Complex64::from(value))
536                }
537                ExprNode::Unary { op, input } => {
538                    let input = scalar_at_optional(&values, input.index())?;
539                    Value::Scalar(eval_unary(*op, input))
540                }
541                ExprNode::Binary { op, lhs, rhs } => {
542                    let lhs = scalar_at_optional(&values, lhs.index())?;
543                    let rhs = scalar_at_optional(&values, rhs.index())?;
544                    Value::Scalar(eval_binary(*op, lhs, rhs))
545                }
546                ExprNode::NaryAdd { terms } => {
547                    let mut sum = Complex64::ZERO;
548                    for term in terms {
549                        sum += scalar_at_optional(&values, term.index())?;
550                    }
551                    Value::Scalar(sum)
552                }
553                ExprNode::NaryMul { factors } => {
554                    let mut product = Complex64::ONE;
555                    for factor in factors {
556                        product *= scalar_at_optional(&values, factor.index())?;
557                    }
558                    Value::Scalar(product)
559                }
560                ExprNode::Complex { re, im } => {
561                    let re = scalar_at_optional(&values, re.index())?;
562                    let im = scalar_at_optional(&values, im.index())?;
563                    Value::Scalar(Complex64::new(re.re, im.re))
564                }
565                ExprNode::Vector { elements } => Value::Vector(
566                    elements
567                        .iter()
568                        .map(|id| scalar_at_optional(&values, id.index()))
569                        .collect::<RuntimeResult<_>>()?,
570                ),
571                ExprNode::Matrix {
572                    rows,
573                    cols,
574                    elements,
575                } => {
576                    if elements.len() != rows * cols {
577                        return Err(RuntimeError::InvalidShape {
578                            index,
579                            message: format!(
580                                "matrix has {} elements for shape {rows}x{cols}",
581                                elements.len()
582                            ),
583                        });
584                    }
585                    Value::Matrix {
586                        rows: *rows,
587                        cols: *cols,
588                        values: elements
589                            .iter()
590                            .map(|id| scalar_at_optional(&values, id.index()))
591                            .collect::<RuntimeResult<_>>()?,
592                    }
593                }
594                ExprNode::Component { input, index: i } => {
595                    let vector = vector_at_optional(&values, input.index())?;
596                    Value::Scalar(*vector.get(*i).ok_or_else(|| RuntimeError::InvalidShape {
597                        index,
598                        message: format!(
599                            "component index {i} out of bounds for len {}",
600                            vector.len()
601                        ),
602                    })?)
603                }
604                ExprNode::MatrixElement { input, row, col } => {
605                    let (rows, cols, matrix) = matrix_at_optional(&values, input.index())?;
606                    if *row >= rows || *col >= cols {
607                        return Err(RuntimeError::InvalidShape {
608                            index,
609                            message: format!(
610                                "matrix element ({row}, {col}) out of bounds for shape {rows}x{cols}"
611                            ),
612                        });
613                    }
614                    Value::Scalar(matrix[row * cols + col])
615                }
616                ExprNode::MatMul { lhs, rhs } => {
617                    let (lhs_rows, lhs_cols, lhs) = matrix_at_optional(&values, lhs.index())?;
618                    let (rhs_rows, rhs_cols, rhs) = matrix_at_optional(&values, rhs.index())?;
619                    if lhs_cols != rhs_rows {
620                        return Err(RuntimeError::InvalidShape {
621                            index,
622                            message: format!(
623                                "cannot multiply {lhs_rows}x{lhs_cols} by {rhs_rows}x{rhs_cols}"
624                            ),
625                        });
626                    }
627                    let lhs = DMatrix::from_row_slice(lhs_rows, lhs_cols, lhs);
628                    let rhs = DMatrix::from_row_slice(rhs_rows, rhs_cols, rhs);
629                    let out = lhs * rhs;
630                    Value::Matrix {
631                        rows: out.nrows(),
632                        cols: out.ncols(),
633                        values: matrix_values_row_major(&out),
634                    }
635                }
636                ExprNode::MatVec { matrix, vector } => {
637                    let (rows, cols, matrix) = matrix_at_optional(&values, matrix.index())?;
638                    let vector = vector_at_optional(&values, vector.index())?;
639                    if cols != vector.len() {
640                        return Err(RuntimeError::InvalidShape {
641                            index,
642                            message: format!(
643                                "cannot multiply {rows}x{cols} matrix by len {} vector",
644                                vector.len()
645                            ),
646                        });
647                    }
648                    let matrix = DMatrix::from_row_slice(rows, cols, matrix);
649                    let vector = DVector::from_row_slice(vector);
650                    Value::Vector((matrix * vector).iter().copied().collect())
651                }
652                ExprNode::Dot { lhs, rhs } => {
653                    let lhs = vector_at_optional(&values, lhs.index())?;
654                    let rhs = vector_at_optional(&values, rhs.index())?;
655                    if lhs.len() != rhs.len() {
656                        return Err(RuntimeError::InvalidShape {
657                            index,
658                            message: format!(
659                                "cannot dot len {} vector with len {} vector",
660                                lhs.len(),
661                                rhs.len()
662                            ),
663                        });
664                    }
665                    Value::Scalar(lhs.iter().zip(rhs).map(|(lhs, rhs)| lhs * rhs).sum())
666                }
667                ExprNode::Solve { matrix, rhs } => {
668                    let matrix_id = *matrix;
669                    let (rows, cols, matrix) = matrix_at_optional(&values, matrix_id.index())?;
670                    let rhs = vector_at_optional(&values, rhs.index())?;
671                    if rows != cols || rows != rhs.len() {
672                        return Err(RuntimeError::InvalidShape {
673                            index,
674                            message: format!(
675                                "cannot solve {rows}x{cols} matrix against len {} vector",
676                                rhs.len()
677                            ),
678                        });
679                    }
680                    let rhs = DVector::from_row_slice(rhs);
681                    let solution = self.solve_primal(matrix_id, rows, matrix, &rhs, index, None)?;
682                    Value::Vector(solution.iter().copied().collect())
683                }
684                ExprNode::ScalarParam(_) => {
685                    return Err(RuntimeError::InvalidShape {
686                        index,
687                        message: "parameter-dependent node cannot be part of an event cache".into(),
688                    });
689                }
690            };
691            values[index] = Some(value);
692        }
693
694        Ok(values)
695    }
696
697    pub(super) fn evaluate_values(
698        &self,
699        params: &ParamValues,
700        event: Option<&dyn EventLookup>,
701    ) -> RuntimeResult<Vec<Value>> {
702        let mut values = Vec::with_capacity(self.graph.nodes().len());
703
704        for (index, node) in self.graph.nodes().iter().enumerate() {
705            let value = match node {
706                ExprNode::RealConst(value) => Value::Scalar(Complex64::from(*value)),
707                ExprNode::ComplexConst(value) => Value::Scalar(*value),
708                ExprNode::ScalarParam(_) => {
709                    Value::Scalar(Complex64::from(self.parameter_value(params, index)?))
710                }
711                ExprNode::EventScalar(name) => {
712                    let Some(event) = event else {
713                        return Err(RuntimeError::MissingEventScalar(name.to_string()));
714                    };
715                    Value::Scalar(Complex64::from(
716                        event
717                            .scalar(name)
718                            .ok_or_else(|| RuntimeError::MissingEventScalar(name.to_string()))?,
719                    ))
720                }
721                ExprNode::EventP4Component { name, component } => {
722                    let Some(event) = event else {
723                        return Err(RuntimeError::MissingEventScalar(format!(
724                            "{name}.{}",
725                            component.label()
726                        )));
727                    };
728                    Value::Scalar(Complex64::from(
729                        event.p4_component(name, *component).ok_or_else(|| {
730                            RuntimeError::MissingEventScalar(format!(
731                                "{name}.{}",
732                                component.label()
733                            ))
734                        })?,
735                    ))
736                }
737                ExprNode::Unary { op, input } => {
738                    let input = scalar_at(&values, input.index())?;
739                    Value::Scalar(eval_unary(*op, input))
740                }
741                ExprNode::Binary { op, lhs, rhs } => {
742                    let lhs = scalar_at(&values, lhs.index())?;
743                    let rhs = scalar_at(&values, rhs.index())?;
744                    Value::Scalar(eval_binary(*op, lhs, rhs))
745                }
746                ExprNode::NaryAdd { terms } => {
747                    let mut sum = Complex64::ZERO;
748                    for term in terms {
749                        sum += scalar_at(&values, term.index())?;
750                    }
751                    Value::Scalar(sum)
752                }
753                ExprNode::NaryMul { factors } => {
754                    let mut product = Complex64::ONE;
755                    for factor in factors {
756                        product *= scalar_at(&values, factor.index())?;
757                    }
758                    Value::Scalar(product)
759                }
760                ExprNode::Complex { re, im } => {
761                    let re = scalar_at(&values, re.index())?;
762                    let im = scalar_at(&values, im.index())?;
763                    Value::Scalar(Complex64::new(re.re, im.re))
764                }
765                ExprNode::Vector { elements } => Value::Vector(
766                    elements
767                        .iter()
768                        .map(|id| scalar_at(&values, id.index()))
769                        .collect::<RuntimeResult<_>>()?,
770                ),
771                ExprNode::Matrix {
772                    rows,
773                    cols,
774                    elements,
775                } => {
776                    if elements.len() != rows * cols {
777                        return Err(RuntimeError::InvalidShape {
778                            index,
779                            message: format!(
780                                "matrix has {} elements for shape {rows}x{cols}",
781                                elements.len()
782                            ),
783                        });
784                    }
785                    Value::Matrix {
786                        rows: *rows,
787                        cols: *cols,
788                        values: elements
789                            .iter()
790                            .map(|id| scalar_at(&values, id.index()))
791                            .collect::<RuntimeResult<_>>()?,
792                    }
793                }
794                ExprNode::Component { input, index: i } => {
795                    let vector = vector_at(&values, input.index())?;
796                    Value::Scalar(*vector.get(*i).ok_or_else(|| RuntimeError::InvalidShape {
797                        index,
798                        message: format!(
799                            "component index {i} out of bounds for len {}",
800                            vector.len()
801                        ),
802                    })?)
803                }
804                ExprNode::MatrixElement { input, row, col } => {
805                    let (rows, cols, matrix) = matrix_at(&values, input.index())?;
806                    if *row >= rows || *col >= cols {
807                        return Err(RuntimeError::InvalidShape {
808                            index,
809                            message: format!(
810                                "matrix element ({row}, {col}) out of bounds for shape {rows}x{cols}"
811                            ),
812                        });
813                    }
814                    Value::Scalar(matrix[row * cols + col])
815                }
816                ExprNode::MatMul { lhs, rhs } => {
817                    let (lhs_rows, lhs_cols, lhs) = matrix_at(&values, lhs.index())?;
818                    let (rhs_rows, rhs_cols, rhs) = matrix_at(&values, rhs.index())?;
819                    if lhs_cols != rhs_rows {
820                        return Err(RuntimeError::InvalidShape {
821                            index,
822                            message: format!(
823                                "cannot multiply {lhs_rows}x{lhs_cols} by {rhs_rows}x{rhs_cols}"
824                            ),
825                        });
826                    }
827                    let lhs = DMatrix::from_row_slice(lhs_rows, lhs_cols, lhs);
828                    let rhs = DMatrix::from_row_slice(rhs_rows, rhs_cols, rhs);
829                    let out = lhs * rhs;
830                    Value::Matrix {
831                        rows: out.nrows(),
832                        cols: out.ncols(),
833                        values: matrix_values_row_major(&out),
834                    }
835                }
836                ExprNode::MatVec { matrix, vector } => {
837                    let (rows, cols, matrix) = matrix_at(&values, matrix.index())?;
838                    let vector = vector_at(&values, vector.index())?;
839                    if cols != vector.len() {
840                        return Err(RuntimeError::InvalidShape {
841                            index,
842                            message: format!(
843                                "cannot multiply {rows}x{cols} matrix by len {} vector",
844                                vector.len()
845                            ),
846                        });
847                    }
848                    let matrix = DMatrix::from_row_slice(rows, cols, matrix);
849                    let vector = DVector::from_row_slice(vector);
850                    Value::Vector((matrix * vector).iter().copied().collect())
851                }
852                ExprNode::Dot { lhs, rhs } => {
853                    let lhs = vector_at(&values, lhs.index())?;
854                    let rhs = vector_at(&values, rhs.index())?;
855                    if lhs.len() != rhs.len() {
856                        return Err(RuntimeError::InvalidShape {
857                            index,
858                            message: format!(
859                                "cannot dot len {} vector with len {} vector",
860                                lhs.len(),
861                                rhs.len()
862                            ),
863                        });
864                    }
865                    Value::Scalar(lhs.iter().zip(rhs).map(|(lhs, rhs)| lhs * rhs).sum())
866                }
867                ExprNode::Solve { matrix, rhs } => {
868                    let matrix_id = *matrix;
869                    let (rows, cols, matrix) = matrix_at(&values, matrix_id.index())?;
870                    let rhs = vector_at(&values, rhs.index())?;
871                    if rows != cols || rows != rhs.len() {
872                        return Err(RuntimeError::InvalidShape {
873                            index,
874                            message: format!(
875                                "cannot solve {rows}x{cols} matrix against len {} vector",
876                                rhs.len()
877                            ),
878                        });
879                    }
880                    let rhs = DVector::from_row_slice(rhs);
881                    let solution = self.solve_primal(matrix_id, rows, matrix, &rhs, index, None)?;
882                    Value::Vector(solution.iter().copied().collect())
883                }
884            };
885            values.push(value);
886        }
887
888        Ok(values)
889    }
890
891    pub(super) fn evaluate_values_from_cache(
892        &self,
893        params: &ParamValues,
894        cache: &CpuBatchCache,
895        row: usize,
896    ) -> RuntimeResult<Vec<Value>> {
897        let mut values = Vec::with_capacity(self.cached_evaluation_nodes.len());
898
899        for id in &self.cached_evaluation_nodes {
900            let index = id.index();
901            let node = &self.graph.nodes()[index];
902            if let Some(slot) = self.cache_slots[index] {
903                values.push(cache.value(slot, row)?);
904                continue;
905            }
906            let value = match node {
907                ExprNode::RealConst(value) => Value::Scalar(Complex64::from(*value)),
908                ExprNode::ComplexConst(value) => Value::Scalar(*value),
909                ExprNode::ScalarParam(_) => {
910                    Value::Scalar(Complex64::from(self.parameter_value(params, index)?))
911                }
912                ExprNode::EventScalar(name) => {
913                    return Err(RuntimeError::MissingEventScalar(name.to_string()));
914                }
915                ExprNode::EventP4Component { name, component } => {
916                    return Err(RuntimeError::MissingEventScalar(format!(
917                        "{name}.{}",
918                        component.label()
919                    )));
920                }
921                ExprNode::Unary { op, input } => {
922                    let input = self.cached_scalar_at(&values, *input)?;
923                    Value::Scalar(eval_unary(*op, input))
924                }
925                ExprNode::Binary { op, lhs, rhs } => {
926                    let lhs = self.cached_scalar_at(&values, *lhs)?;
927                    let rhs = self.cached_scalar_at(&values, *rhs)?;
928                    Value::Scalar(eval_binary(*op, lhs, rhs))
929                }
930                ExprNode::NaryAdd { terms } => {
931                    let mut sum = Complex64::ZERO;
932                    for term in terms {
933                        sum += self.cached_scalar_at(&values, *term)?;
934                    }
935                    Value::Scalar(sum)
936                }
937                ExprNode::NaryMul { factors } => {
938                    let mut product = Complex64::ONE;
939                    for factor in factors {
940                        product *= self.cached_scalar_at(&values, *factor)?;
941                    }
942                    Value::Scalar(product)
943                }
944                ExprNode::Complex { re, im } => {
945                    let re = self.cached_scalar_at(&values, *re)?;
946                    let im = self.cached_scalar_at(&values, *im)?;
947                    Value::Scalar(Complex64::new(re.re, im.re))
948                }
949                ExprNode::Vector { elements } => Value::Vector(
950                    elements
951                        .iter()
952                        .map(|id| self.cached_scalar_at(&values, *id))
953                        .collect::<RuntimeResult<_>>()?,
954                ),
955                ExprNode::Matrix {
956                    rows,
957                    cols,
958                    elements,
959                } => {
960                    if elements.len() != rows * cols {
961                        return Err(RuntimeError::InvalidShape {
962                            index,
963                            message: format!(
964                                "matrix has {} elements for shape {rows}x{cols}",
965                                elements.len()
966                            ),
967                        });
968                    }
969                    Value::Matrix {
970                        rows: *rows,
971                        cols: *cols,
972                        values: elements
973                            .iter()
974                            .map(|id| self.cached_scalar_at(&values, *id))
975                            .collect::<RuntimeResult<_>>()?,
976                    }
977                }
978                ExprNode::Component { input, index: i } => {
979                    if let Some(plan) = self.solve_components[index] {
980                        let inverse_row = cache.solve_row(plan.row_slot(), row)?;
981                        if inverse_row.len() != plan.dimension() {
982                            return Err(RuntimeError::InvalidShape {
983                                index,
984                                message: format!(
985                                    "specialized solve expected row len {}, got {}",
986                                    plan.dimension(),
987                                    inverse_row.len()
988                                ),
989                            });
990                        }
991                        if let Some(elements) = &self.solve_rhs_elements[plan.rhs().index()] {
992                            if elements.len() != plan.dimension() {
993                                return Err(RuntimeError::InvalidShape {
994                                    index,
995                                    message: format!(
996                                        "specialized solve expected {} RHS elements, got {}",
997                                        plan.dimension(),
998                                        elements.len()
999                                    ),
1000                                });
1001                            }
1002                            Value::Scalar(
1003                                inverse_row
1004                                    .iter()
1005                                    .zip(elements)
1006                                    .map(|(lhs, rhs)| {
1007                                        Ok(lhs * self.cached_scalar_at(&values, *rhs)?)
1008                                    })
1009                                    .sum::<RuntimeResult<Complex64>>()?,
1010                            )
1011                        } else {
1012                            let rhs = self.cached_vector_at(&values, plan.rhs())?;
1013                            if rhs.len() != plan.dimension() {
1014                                return Err(RuntimeError::InvalidShape {
1015                                    index,
1016                                    message: format!(
1017                                        "specialized solve expected RHS len {}, got {}",
1018                                        plan.dimension(),
1019                                        rhs.len()
1020                                    ),
1021                                });
1022                            }
1023                            Value::Scalar(
1024                                inverse_row
1025                                    .iter()
1026                                    .zip(rhs)
1027                                    .map(|(lhs, rhs)| lhs * rhs)
1028                                    .sum(),
1029                            )
1030                        }
1031                    } else {
1032                        let vector = self.cached_vector_at(&values, *input)?;
1033                        Value::Scalar(*vector.get(*i).ok_or_else(|| {
1034                            RuntimeError::InvalidShape {
1035                                index,
1036                                message: format!(
1037                                    "component index {i} out of bounds for len {}",
1038                                    vector.len()
1039                                ),
1040                            }
1041                        })?)
1042                    }
1043                }
1044                ExprNode::MatrixElement { input, row, col } => {
1045                    let (rows, cols, matrix) = self.cached_matrix_at(&values, *input)?;
1046                    if *row >= rows || *col >= cols {
1047                        return Err(RuntimeError::InvalidShape {
1048                            index,
1049                            message: format!(
1050                                "matrix element ({row}, {col}) out of bounds for shape {rows}x{cols}"
1051                            ),
1052                        });
1053                    }
1054                    Value::Scalar(matrix[row * cols + col])
1055                }
1056                ExprNode::MatMul { lhs, rhs } => {
1057                    let (lhs_rows, lhs_cols, lhs) = self.cached_matrix_at(&values, *lhs)?;
1058                    let (rhs_rows, rhs_cols, rhs) = self.cached_matrix_at(&values, *rhs)?;
1059                    if lhs_cols != rhs_rows {
1060                        return Err(RuntimeError::InvalidShape {
1061                            index,
1062                            message: format!(
1063                                "cannot multiply {lhs_rows}x{lhs_cols} by {rhs_rows}x{rhs_cols}"
1064                            ),
1065                        });
1066                    }
1067                    let lhs = DMatrix::from_row_slice(lhs_rows, lhs_cols, lhs);
1068                    let rhs = DMatrix::from_row_slice(rhs_rows, rhs_cols, rhs);
1069                    let out = lhs * rhs;
1070                    Value::Matrix {
1071                        rows: out.nrows(),
1072                        cols: out.ncols(),
1073                        values: matrix_values_row_major(&out),
1074                    }
1075                }
1076                ExprNode::MatVec { matrix, vector } => {
1077                    let (rows, cols, matrix) = self.cached_matrix_at(&values, *matrix)?;
1078                    let vector = self.cached_vector_at(&values, *vector)?;
1079                    if cols != vector.len() {
1080                        return Err(RuntimeError::InvalidShape {
1081                            index,
1082                            message: format!(
1083                                "cannot multiply {rows}x{cols} matrix by len {} vector",
1084                                vector.len()
1085                            ),
1086                        });
1087                    }
1088                    let matrix = DMatrix::from_row_slice(rows, cols, matrix);
1089                    let vector = DVector::from_row_slice(vector);
1090                    Value::Vector((matrix * vector).iter().copied().collect())
1091                }
1092                ExprNode::Dot { lhs, rhs } => {
1093                    let lhs = self.cached_vector_at(&values, *lhs)?;
1094                    let rhs = self.cached_vector_at(&values, *rhs)?;
1095                    if lhs.len() != rhs.len() {
1096                        return Err(RuntimeError::InvalidShape {
1097                            index,
1098                            message: format!(
1099                                "cannot dot len {} vector with len {} vector",
1100                                lhs.len(),
1101                                rhs.len()
1102                            ),
1103                        });
1104                    }
1105                    Value::Scalar(lhs.iter().zip(rhs).map(|(lhs, rhs)| lhs * rhs).sum())
1106                }
1107                ExprNode::Solve { matrix, rhs } => {
1108                    let matrix_id = *matrix;
1109                    let (rows, cols, matrix) = self.cached_matrix_at(&values, matrix_id)?;
1110                    let rhs = self.cached_vector_at(&values, *rhs)?;
1111                    if rows != cols || rows != rhs.len() {
1112                        return Err(RuntimeError::InvalidShape {
1113                            index,
1114                            message: format!(
1115                                "cannot solve {rows}x{cols} matrix against len {} vector",
1116                                rhs.len()
1117                            ),
1118                        });
1119                    }
1120                    let rhs = DVector::from_row_slice(rhs);
1121                    let solution = self.solve_primal(
1122                        matrix_id,
1123                        rows,
1124                        matrix,
1125                        &rhs,
1126                        index,
1127                        Some((cache, row)),
1128                    )?;
1129                    Value::Vector(solution.iter().copied().collect())
1130                }
1131            };
1132            values.push(value);
1133        }
1134
1135        Ok(values)
1136    }
1137
1138    pub(super) fn cached_value_slot(&self, id: ExprId) -> RuntimeResult<usize> {
1139        self.cached_value_slots[id.index()].ok_or_else(|| RuntimeError::InvalidShape {
1140            index: id.index(),
1141            message: "node is not part of the cached evaluation schedule".into(),
1142        })
1143    }
1144
1145    pub(super) fn cached_scalar_at(
1146        &self,
1147        values: &[Value],
1148        id: ExprId,
1149    ) -> RuntimeResult<Complex64> {
1150        scalar_at(values, self.cached_value_slot(id)?)
1151    }
1152
1153    pub(super) fn cached_vector_at<'a>(
1154        &self,
1155        values: &'a [Value],
1156        id: ExprId,
1157    ) -> RuntimeResult<&'a [Complex64]> {
1158        vector_at(values, self.cached_value_slot(id)?)
1159    }
1160
1161    pub(super) fn cached_matrix_at<'a>(
1162        &self,
1163        values: &'a [Value],
1164        id: ExprId,
1165    ) -> RuntimeResult<(usize, usize, &'a [Complex64])> {
1166        matrix_at(values, self.cached_value_slot(id)?)
1167    }
1168
1169    pub(super) fn check_batch_cache(&self, cache: &CpuBatchCache) -> RuntimeResult<()> {
1170        if cache.nodes
1171            == self
1172                .cache_plan
1173                .entries()
1174                .iter()
1175                .map(|entry| entry.node())
1176                .collect::<Vec<_>>()
1177            && cache.factor_nodes
1178                == self
1179                    .factor_matrices
1180                    .iter()
1181                    .map(|(node, _)| *node)
1182                    .collect::<Vec<_>>()
1183            && cache.solve_row_keys == self.solve_row_keys
1184        {
1185            Ok(())
1186        } else {
1187            Err(RuntimeError::InvalidCacheLayout)
1188        }
1189    }
1190
1191    pub(super) fn scalar_invariant_values(
1192        &self,
1193        params: &ParamValues,
1194    ) -> RuntimeResult<Option<ScalarInvariantValues>> {
1195        if self.precision == Precision::F32 {
1196            return Ok(None);
1197        }
1198        let Some(plan) = self.scalar_interpreter_plan() else {
1199            return Ok(None);
1200        };
1201        let mut values = ScalarInvariantValues {
1202            real: vec![0.0; plan.invariant_real_slot_count],
1203            complex: vec![Complex64::ZERO; plan.invariant_complex_slot_count],
1204        };
1205        let event = ScalarEventWorkspace::default();
1206        for instruction in &plan.invariant_instructions {
1207            match instruction.output_slot {
1208                ScalarSlot::Real(slot) => {
1209                    values.real[slot] = instruction.instruction.evaluate_real(
1210                        Some(params),
1211                        None,
1212                        &values,
1213                        &event,
1214                    )?;
1215                }
1216                ScalarSlot::Complex(slot) => {
1217                    values.complex[slot] = instruction.instruction.evaluate_complex(
1218                        Some(params),
1219                        None,
1220                        &values,
1221                        &event,
1222                    )?;
1223                }
1224            }
1225        }
1226        Ok(Some(values))
1227    }
1228
1229    pub(super) fn evaluate_cache_row_unchecked(
1230        &self,
1231        params: &ParamValues,
1232        cache: &CpuBatchCache,
1233        row: usize,
1234    ) -> RuntimeResult<Complex64> {
1235        let invariant = self.scalar_invariant_values(params)?;
1236        self.evaluate_cache_row_prepared(
1237            params,
1238            cache,
1239            row,
1240            invariant.as_ref(),
1241            &mut ScalarEventWorkspace::default(),
1242        )
1243    }
1244
1245    /// Evaluates every row in a materialized batch cache.
1246    ///
1247    /// # Errors
1248    ///
1249    /// Returns [`RuntimeError`] when parameters or cache layout are
1250    /// incompatible, evaluation fails, or a matrix is singular.
1251    pub fn evaluate_cache(
1252        &self,
1253        params: &ParamValues,
1254        cache: &CpuBatchCache,
1255    ) -> RuntimeResult<Vec<Complex64>> {
1256        self.check_batch_cache(cache)?;
1257        #[cfg(feature = "jit")]
1258        if let Some(kernel) = self.scalar_jit_kernel() {
1259            let mut output = Vec::with_capacity(cache.len());
1260            kernel.evaluate(params, cache, 0, cache.len(), &mut output)?;
1261            return Ok(output);
1262        }
1263        let invariant = self.scalar_invariant_values(params)?;
1264        let mut out = Vec::with_capacity(cache.len());
1265        let mut workspace = ScalarEventWorkspace::default();
1266        for row in 0..cache.len() {
1267            out.push(self.evaluate_cache_row_prepared(
1268                params,
1269                cache,
1270                row,
1271                invariant.as_ref(),
1272                &mut workspace,
1273            )?);
1274        }
1275        Ok(out)
1276    }
1277
1278    pub(super) fn evaluate_cache_row_prepared(
1279        &self,
1280        params: &ParamValues,
1281        cache: &CpuBatchCache,
1282        row: usize,
1283        invariant: Option<&ScalarInvariantValues>,
1284        workspace: &mut ScalarEventWorkspace,
1285    ) -> RuntimeResult<Complex64> {
1286        #[cfg(feature = "jit")]
1287        if let Some(kernel) = self.scalar_jit_kernel() {
1288            let mut output = Vec::with_capacity(1);
1289            kernel.evaluate(params, cache, row, row + 1, &mut output)?;
1290            return Ok(output[0]);
1291        }
1292        if self.precision == Precision::F32 {
1293            return self.evaluate_f32_scalar(params, F32KernelInput::Cache(Some((cache, row))));
1294        }
1295        if let (Some(plan), Some(invariant)) = (self.scalar_interpreter_plan(), invariant) {
1296            return self.evaluate_scalar_cache_row(cache, row, plan, invariant, workspace);
1297        }
1298        let values = self.evaluate_values_from_cache(params, cache, row)?;
1299        self.cached_scalar_at(&values, self.graph.root())
1300    }
1301
1302    fn evaluate_scalar_cache_row(
1303        &self,
1304        cache: &CpuBatchCache,
1305        row: usize,
1306        plan: &ScalarEvaluationPlan,
1307        invariant: &ScalarInvariantValues,
1308        values: &mut ScalarEventWorkspace,
1309    ) -> RuntimeResult<Complex64> {
1310        values.real.clear();
1311        values
1312            .real
1313            .resize(plan.event_real_slot_count, [0.0; SCALAR_BLOCK_SIZE]);
1314        values.complex.clear();
1315        values.complex.resize(
1316            plan.event_complex_slot_count,
1317            [Complex64::ZERO; SCALAR_BLOCK_SIZE],
1318        );
1319        for event_instruction in &plan.event_instructions {
1320            match event_instruction.output_slot {
1321                ScalarSlot::Real(slot) => {
1322                    values.real[slot][0] = event_instruction.instruction.evaluate_real(
1323                        None,
1324                        Some((cache, row)),
1325                        invariant,
1326                        values,
1327                    )?;
1328                }
1329                ScalarSlot::Complex(slot) => {
1330                    values.complex[slot][0] = event_instruction.instruction.evaluate_complex(
1331                        None,
1332                        Some((cache, row)),
1333                        invariant,
1334                        values,
1335                    )?;
1336                }
1337            }
1338        }
1339        Ok(plan.root().complex_value(invariant, values))
1340    }
1341
1342    #[allow(clippy::too_many_arguments)]
1343    pub(super) fn evaluate_cache_block_prepared(
1344        &self,
1345        params: &ParamValues,
1346        cache: &CpuBatchCache,
1347        start: usize,
1348        end: usize,
1349        invariant: Option<&ScalarInvariantValues>,
1350        workspace: &mut ScalarEventWorkspace,
1351        output: &mut Vec<Complex64>,
1352        #[cfg(feature = "jit")] jit_cache: Option<&JitCacheView>,
1353    ) -> RuntimeResult<()> {
1354        #[cfg(feature = "jit")]
1355        if let Some(kernel) = self.scalar_jit_kernel() {
1356            let owned;
1357            let jit_cache = if let Some(jit_cache) = jit_cache {
1358                jit_cache
1359            } else {
1360                owned = JitScalarKernel::prepare_cache(cache);
1361                &owned
1362            };
1363            return kernel.evaluate_prepared(params, jit_cache, start, end, output);
1364        }
1365        if self.precision == Precision::F32 {
1366            output.clear();
1367            output.reserve(end - start);
1368            for row in start..end {
1369                output.push(
1370                    self.evaluate_f32_scalar(params, F32KernelInput::Cache(Some((cache, row))))?,
1371                );
1372            }
1373            return Ok(());
1374        }
1375        if let (Some(plan), Some(invariant)) = (self.scalar_interpreter_plan(), invariant) {
1376            return evaluate_scalar_cache_block(
1377                cache, start, end, plan, invariant, workspace, output,
1378            );
1379        }
1380        output.clear();
1381        for row in start..end {
1382            output
1383                .push(self.evaluate_cache_row_prepared(params, cache, row, invariant, workspace)?);
1384        }
1385        Ok(())
1386    }
1387
1388    pub(super) fn evaluate_cache_row_with_gradient_unchecked(
1389        &self,
1390        params: &ParamValues,
1391        cache: &CpuBatchCache,
1392        row: usize,
1393    ) -> RuntimeResult<ValueGradient> {
1394        #[cfg(feature = "jit")]
1395        if self.gradient_jit_kernel().is_some() {
1396            return self
1397                .evaluate_cache_gradient_jit(params, cache, row, row + 1)?
1398                .pop()
1399                .ok_or_else(|| RuntimeError::InvalidShape {
1400                    index: row,
1401                    message: "single-row JIT gradient produced no value".into(),
1402                });
1403        }
1404        if self.precision == Precision::F32 {
1405            return self.evaluate_f32_gradient(params, F32KernelInput::Cache(Some((cache, row))));
1406        }
1407        if self.autodiff.mode() == laddu_autodiff::AutodiffMode::Reverse {
1408            self.require_f64_gradient()?;
1409            let values = self.evaluate_values_from_cache(params, cache, row)?;
1410            return self.value_gradient(values, Some((cache, row)));
1411        }
1412        if let Some(interpreter) = self.gradient_interpreter() {
1413            let (value, gradient) = interpreter.evaluate(params, Some((cache, row)))?;
1414            return Ok(ValueGradient { value, gradient });
1415        }
1416        let values = self.evaluate_values_from_cache(params, cache, row)?;
1417        self.value_gradient(values, Some((cache, row)))
1418    }
1419
1420    /// Evaluates every cached row and its free-parameter gradient.
1421    ///
1422    /// # Errors
1423    ///
1424    /// Returns [`RuntimeError`] when parameters or cache layout are
1425    /// incompatible, or differentiation or evaluation fails.
1426    pub fn evaluate_cache_with_gradient(
1427        &self,
1428        params: &ParamValues,
1429        cache: &CpuBatchCache,
1430    ) -> RuntimeResult<Vec<ValueGradient>> {
1431        self.check_batch_cache(cache)?;
1432        #[cfg(feature = "jit")]
1433        if self.gradient_jit_kernel().is_some() {
1434            return self.evaluate_cache_gradient_jit(params, cache, 0, cache.len());
1435        }
1436        if self.precision == Precision::F32 {
1437            return (0..cache.len())
1438                .map(|row| {
1439                    self.evaluate_f32_gradient(params, F32KernelInput::Cache(Some((cache, row))))
1440                })
1441                .collect();
1442        }
1443        self.require_f64_gradient()?;
1444        (0..cache.len())
1445            .map(|row| self.evaluate_cache_row_with_gradient_unchecked(params, cache, row))
1446            .collect()
1447    }
1448
1449    #[cfg(feature = "jit")]
1450    fn evaluate_cache_gradient_jit(
1451        &self,
1452        params: &ParamValues,
1453        cache: &CpuBatchCache,
1454        start: usize,
1455        end: usize,
1456    ) -> RuntimeResult<Vec<ValueGradient>> {
1457        let (Some(value_kernel), Some(gradient_kernel)) =
1458            (self.scalar_jit_kernel(), self.gradient_jit_kernel())
1459        else {
1460            return Err(RuntimeError::InvalidShape {
1461                index: self.graph.root().index(),
1462                message: "JIT gradient evaluation requires both scalar and gradient kernels".into(),
1463            });
1464        };
1465        let view = JitScalarKernel::prepare_cache(cache);
1466        let mut values = Vec::new();
1467        let mut real = Vec::new();
1468        let mut imag = Vec::new();
1469        value_kernel.evaluate_prepared(params, &view, start, end, &mut values)?;
1470        gradient_kernel.evaluate_prepared(params, &view, start, end, 0, &mut real)?;
1471        gradient_kernel.evaluate_prepared(params, &view, start, end, 1, &mut imag)?;
1472        let parameter_count = self.free_parameter_count();
1473        Ok(values
1474            .into_iter()
1475            .enumerate()
1476            .map(|(row, value)| ValueGradient {
1477                value,
1478                gradient: (0..parameter_count)
1479                    .map(|parameter| {
1480                        let index = row * parameter_count + parameter;
1481                        Complex64::new(real[index], imag[index])
1482                    })
1483                    .collect(),
1484            })
1485            .collect())
1486    }
1487}
1488
1489pub(super) fn evaluate_scalar_cache_block(
1490    cache: &CpuBatchCache,
1491    start: usize,
1492    end: usize,
1493    plan: &ScalarEvaluationPlan,
1494    invariant: &ScalarInvariantValues,
1495    workspace: &mut ScalarEventWorkspace,
1496    output: &mut Vec<Complex64>,
1497) -> RuntimeResult<()> {
1498    let block_len = end - start;
1499    workspace
1500        .real
1501        .resize(plan.event_real_slot_count, [0.0; SCALAR_BLOCK_SIZE]);
1502    workspace.complex.resize(
1503        plan.event_complex_slot_count,
1504        [Complex64::ZERO; SCALAR_BLOCK_SIZE],
1505    );
1506
1507    for event_instruction in &plan.event_instructions {
1508        match event_instruction.output_slot {
1509            ScalarSlot::Real(slot) => {
1510                let output_slot = slot;
1511                match &event_instruction.instruction {
1512                    ScalarInstruction::Cached(slot) => {
1513                        workspace.real[output_slot][..block_len]
1514                            .copy_from_slice(cache.real_range(*slot, start, end)?);
1515                    }
1516                    ScalarInstruction::Unary { op, input } => {
1517                        for lane in 0..block_len {
1518                            workspace.real[output_slot][lane] = match op {
1519                                UnaryOp::Neg => -input.block_real_value(invariant, workspace, lane),
1520                                UnaryOp::Real | UnaryOp::Conj => {
1521                                    input.block_complex_value(invariant, workspace, lane).re
1522                                }
1523                                UnaryOp::Imag => {
1524                                    input.block_complex_value(invariant, workspace, lane).im
1525                                }
1526                                UnaryOp::NormSqr => input
1527                                    .block_complex_value(invariant, workspace, lane)
1528                                    .norm_sqr(),
1529                                UnaryOp::Sqrt => {
1530                                    input.block_real_value(invariant, workspace, lane).sqrt()
1531                                }
1532                                UnaryOp::Exp => {
1533                                    input.block_real_value(invariant, workspace, lane).exp()
1534                                }
1535                                UnaryOp::Sin => {
1536                                    input.block_real_value(invariant, workspace, lane).sin()
1537                                }
1538                                UnaryOp::Cos => {
1539                                    input.block_real_value(invariant, workspace, lane).cos()
1540                                }
1541                                UnaryOp::Log => {
1542                                    input.block_real_value(invariant, workspace, lane).ln()
1543                                }
1544                                UnaryOp::PowI(power) => input
1545                                    .block_real_value(invariant, workspace, lane)
1546                                    .powi(*power),
1547                            };
1548                        }
1549                    }
1550                    ScalarInstruction::Binary { op, lhs, rhs } => {
1551                        for lane in 0..block_len {
1552                            let lhs = lhs.block_real_value(invariant, workspace, lane);
1553                            let rhs = rhs.block_real_value(invariant, workspace, lane);
1554                            workspace.real[output_slot][lane] = match op {
1555                                BinaryOp::Add => lhs + rhs,
1556                                BinaryOp::Sub => lhs - rhs,
1557                                BinaryOp::Mul => lhs * rhs,
1558                                BinaryOp::Div => lhs / rhs,
1559                                BinaryOp::Atan2 => lhs.atan2(rhs),
1560                            };
1561                        }
1562                    }
1563                    ScalarInstruction::Add(runs) => {
1564                        workspace.real[output_slot][..block_len].fill(0.0);
1565                        for run in runs {
1566                            match run {
1567                                OperandRun::InvariantReal(slots) => {
1568                                    for slot in slots {
1569                                        let operand = invariant.real[*slot];
1570                                        for lane in 0..block_len {
1571                                            workspace.real[output_slot][lane] += operand;
1572                                        }
1573                                    }
1574                                }
1575                                OperandRun::EventReal(slots) => {
1576                                    for slot in slots {
1577                                        for lane in 0..block_len {
1578                                            workspace.real[output_slot][lane] +=
1579                                                workspace.real[*slot][lane];
1580                                        }
1581                                    }
1582                                }
1583                                OperandRun::InvariantComplex(_) | OperandRun::EventComplex(_) => {
1584                                    unreachable!("complex operand appeared in real add")
1585                                }
1586                            }
1587                        }
1588                    }
1589                    ScalarInstruction::Mul(runs) => {
1590                        workspace.real[output_slot][..block_len].fill(1.0);
1591                        for run in runs {
1592                            match run {
1593                                OperandRun::InvariantReal(slots) => {
1594                                    for slot in slots {
1595                                        let operand = invariant.real[*slot];
1596                                        for lane in 0..block_len {
1597                                            workspace.real[output_slot][lane] *= operand;
1598                                        }
1599                                    }
1600                                }
1601                                OperandRun::EventReal(slots) => {
1602                                    for slot in slots {
1603                                        for lane in 0..block_len {
1604                                            workspace.real[output_slot][lane] *=
1605                                                workspace.real[*slot][lane];
1606                                        }
1607                                    }
1608                                }
1609                                OperandRun::InvariantComplex(_) | OperandRun::EventComplex(_) => {
1610                                    unreachable!("complex operand appeared in real multiply")
1611                                }
1612                            }
1613                        }
1614                    }
1615                    ScalarInstruction::Constant(_)
1616                    | ScalarInstruction::Parameter(_)
1617                    | ScalarInstruction::Complex { .. }
1618                    | ScalarInstruction::SolveRow { .. }
1619                    | ScalarInstruction::SolveRowAdjointElement { .. } => {
1620                        unreachable!("non-real event instruction appeared in a real slot")
1621                    }
1622                }
1623            }
1624            ScalarSlot::Complex(slot) => {
1625                let output_slot = slot;
1626                match &event_instruction.instruction {
1627                    ScalarInstruction::Cached(slot) => {
1628                        workspace.complex[output_slot][..block_len]
1629                            .copy_from_slice(cache.complex_range(*slot, start, end)?);
1630                    }
1631                    ScalarInstruction::Unary { op, input } => {
1632                        for lane in 0..block_len {
1633                            let input = input.block_complex_value(invariant, workspace, lane);
1634                            workspace.complex[output_slot][lane] = eval_unary(*op, input);
1635                        }
1636                    }
1637                    ScalarInstruction::Binary { op, lhs, rhs } => {
1638                        for lane in 0..block_len {
1639                            let lhs = lhs.block_complex_value(invariant, workspace, lane);
1640                            let rhs = rhs.block_complex_value(invariant, workspace, lane);
1641                            workspace.complex[output_slot][lane] = eval_binary(*op, lhs, rhs);
1642                        }
1643                    }
1644                    ScalarInstruction::Add(runs) => {
1645                        workspace.complex[output_slot][..block_len].fill(Complex64::ZERO);
1646                        for run in runs {
1647                            match run {
1648                                OperandRun::InvariantReal(slots) => {
1649                                    for slot in slots {
1650                                        let operand = invariant.real[*slot];
1651                                        for lane in 0..block_len {
1652                                            workspace.complex[output_slot][lane] += operand;
1653                                        }
1654                                    }
1655                                }
1656                                OperandRun::InvariantComplex(slots) => {
1657                                    for slot in slots {
1658                                        let operand = invariant.complex[*slot];
1659                                        for lane in 0..block_len {
1660                                            workspace.complex[output_slot][lane] += operand;
1661                                        }
1662                                    }
1663                                }
1664                                OperandRun::EventReal(slots) => {
1665                                    for slot in slots {
1666                                        for lane in 0..block_len {
1667                                            workspace.complex[output_slot][lane] +=
1668                                                workspace.real[*slot][lane];
1669                                        }
1670                                    }
1671                                }
1672                                OperandRun::EventComplex(slots) => {
1673                                    for slot in slots {
1674                                        for lane in 0..block_len {
1675                                            let operand = workspace.complex[*slot][lane];
1676                                            workspace.complex[output_slot][lane] += operand;
1677                                        }
1678                                    }
1679                                }
1680                            }
1681                        }
1682                    }
1683                    ScalarInstruction::Mul(runs) => {
1684                        workspace.complex[output_slot][..block_len].fill(Complex64::ONE);
1685                        for run in runs {
1686                            match run {
1687                                OperandRun::InvariantReal(slots) => {
1688                                    for slot in slots {
1689                                        let operand = invariant.real[*slot];
1690                                        for lane in 0..block_len {
1691                                            workspace.complex[output_slot][lane] *= operand;
1692                                        }
1693                                    }
1694                                }
1695                                OperandRun::InvariantComplex(slots) => {
1696                                    for slot in slots {
1697                                        let operand = invariant.complex[*slot];
1698                                        for lane in 0..block_len {
1699                                            workspace.complex[output_slot][lane] *= operand;
1700                                        }
1701                                    }
1702                                }
1703                                OperandRun::EventReal(slots) => {
1704                                    for slot in slots {
1705                                        for lane in 0..block_len {
1706                                            workspace.complex[output_slot][lane] *=
1707                                                workspace.real[*slot][lane];
1708                                        }
1709                                    }
1710                                }
1711                                OperandRun::EventComplex(slots) => {
1712                                    for slot in slots {
1713                                        for lane in 0..block_len {
1714                                            let operand = workspace.complex[*slot][lane];
1715                                            workspace.complex[output_slot][lane] *= operand;
1716                                        }
1717                                    }
1718                                }
1719                            }
1720                        }
1721                    }
1722                    ScalarInstruction::Complex { re, im } => {
1723                        for lane in 0..block_len {
1724                            workspace.complex[output_slot][lane] = Complex64::new(
1725                                re.block_real_value(invariant, workspace, lane),
1726                                im.block_real_value(invariant, workspace, lane),
1727                            );
1728                        }
1729                    }
1730                    ScalarInstruction::SolveRow { row_slot, rhs } => {
1731                        for lane in 0..block_len {
1732                            let inverse_row = cache.solve_row(*row_slot, start + lane)?;
1733                            if inverse_row.len() != rhs.len() {
1734                                return Err(RuntimeError::InvalidShape {
1735                                    index: start + lane,
1736                                    message: format!(
1737                                        "specialized solve row has len {}, expected {}",
1738                                        inverse_row.len(),
1739                                        rhs.len()
1740                                    ),
1741                                });
1742                            }
1743                            workspace.complex[output_slot][lane] = inverse_row
1744                                .iter()
1745                                .zip(rhs)
1746                                .map(|(lhs, operand)| {
1747                                    lhs * operand.block_complex_value(invariant, workspace, lane)
1748                                })
1749                                .sum();
1750                        }
1751                    }
1752                    ScalarInstruction::SolveRowAdjointElement {
1753                        row_slot,
1754                        index,
1755                        len,
1756                        adjoint,
1757                    } => {
1758                        for lane in 0..block_len {
1759                            let inverse_row = cache.solve_row(*row_slot, start + lane)?;
1760                            if inverse_row.len() != *len {
1761                                return Err(RuntimeError::InvalidShape {
1762                                    index: start + lane,
1763                                    message: format!(
1764                                        "specialized solve row has len {}, expected {len}",
1765                                        inverse_row.len()
1766                                    ),
1767                                });
1768                            }
1769                            workspace.complex[output_slot][lane] = adjoint
1770                                .block_complex_value(invariant, workspace, lane)
1771                                * inverse_row[*index].conj();
1772                        }
1773                    }
1774                    ScalarInstruction::Constant(_) | ScalarInstruction::Parameter(_) => {
1775                        unreachable!("invariant instruction appeared in the event tape")
1776                    }
1777                }
1778            }
1779        }
1780    }
1781
1782    output.clear();
1783    output.reserve(block_len * plan.outputs.len());
1784    for lane in 0..block_len {
1785        for output_operand in &plan.outputs {
1786            output.push(output_operand.block_complex_value(invariant, workspace, lane));
1787        }
1788    }
1789    Ok(())
1790}