laddu_kernel/
ir.rs

1pub use crate::KernelIrError;
2use laddu_expr::{BinaryOp, UnaryOp, parameters::ParamId};
3use num::complex::Complex64;
4
5/// Stable identifier for a value in kernel IR.
6#[derive(Copy, Clone, Debug, PartialEq, Eq)]
7pub struct KernelValueId(usize);
8
9impl KernelValueId {
10    /// Creates an identifier from a zero-based value index.
11    pub fn from_index(index: usize) -> Self {
12        Self(index)
13    }
14
15    /// Returns the zero-based value index.
16    pub fn index(self) -> usize {
17        self.0
18    }
19}
20
21/// Runtime shape and scalar representation of a kernel value.
22#[derive(Copy, Clone, Debug, PartialEq, Eq)]
23pub enum KernelValueKind {
24    /// A real scalar.
25    Real,
26    /// A complex scalar.
27    Complex,
28    /// A complex vector.
29    Vector {
30        /// Number of elements.
31        len: usize,
32    },
33    /// A complex matrix.
34    Matrix {
35        /// Number of rows.
36        rows: usize,
37        /// Number of columns.
38        cols: usize,
39    },
40}
41
42impl KernelValueKind {
43    /// Returns the number of logical scalar elements.
44    pub fn width(self) -> usize {
45        match self {
46            Self::Real | Self::Complex => 1,
47            Self::Vector { len } => len,
48            Self::Matrix { rows, cols } => rows * cols,
49        }
50    }
51
52    fn scalar_combine(self, rhs: Self) -> Option<Self> {
53        match (self, rhs) {
54            (Self::Real, Self::Real) => Some(Self::Real),
55            (Self::Real | Self::Complex, Self::Real | Self::Complex) => Some(Self::Complex),
56            _ => None,
57        }
58    }
59
60    fn is_scalar(self) -> bool {
61        matches!(self, Self::Real | Self::Complex)
62    }
63}
64
65/// Whether a kernel value is constant across events or event-dependent.
66#[derive(Copy, Clone, Debug, PartialEq, Eq)]
67pub enum KernelValueClass {
68    /// The value depends only on constants and parameters.
69    Invariant,
70    /// The value depends on event data or cache inputs.
71    Event,
72}
73
74/// Operation that produces one value in kernel IR.
75#[derive(Clone, Debug)]
76pub enum KernelInstruction {
77    /// Reads a precomputed cache slot.
78    Cached(usize),
79    /// Emits a real constant.
80    RealConstant(f64),
81    /// Emits a complex constant.
82    ComplexConstant(Complex64),
83    /// Reads a scalar parameter.
84    Parameter(ParamId),
85    /// Applies a unary scalar operation.
86    Unary {
87        /// Operation to apply.
88        op: UnaryOp,
89        /// Input value.
90        input: KernelValueId,
91    },
92    /// Applies a binary scalar operation.
93    Binary {
94        /// Operation to apply.
95        op: BinaryOp,
96        /// Left operand.
97        lhs: KernelValueId,
98        /// Right operand.
99        rhs: KernelValueId,
100    },
101    /// Adds scalar operands.
102    Add(Vec<KernelValueId>),
103    /// Multiplies scalar operands.
104    Mul(Vec<KernelValueId>),
105    /// Constructs a complex scalar from real components.
106    Complex {
107        /// Real component.
108        re: KernelValueId,
109        /// Imaginary component.
110        im: KernelValueId,
111    },
112    /// Constructs a vector from scalar elements.
113    Vector(Vec<KernelValueId>),
114    /// Constructs a row-major matrix.
115    Matrix {
116        /// Number of rows.
117        rows: usize,
118        /// Number of columns.
119        cols: usize,
120        /// Row-major scalar elements.
121        elements: Vec<KernelValueId>,
122    },
123    /// Selects a vector component.
124    Component {
125        /// Vector input.
126        input: KernelValueId,
127        /// Zero-based component index.
128        index: usize,
129    },
130    /// Selects a matrix element.
131    MatrixElement {
132        /// Matrix input.
133        input: KernelValueId,
134        /// Zero-based row index.
135        row: usize,
136        /// Zero-based column index.
137        col: usize,
138    },
139    /// Multiplies two matrices.
140    MatMul {
141        /// Left matrix.
142        lhs: KernelValueId,
143        /// Right matrix.
144        rhs: KernelValueId,
145    },
146    /// Multiplies a matrix by a vector.
147    MatVec {
148        /// Matrix operand.
149        matrix: KernelValueId,
150        /// Vector operand.
151        vector: KernelValueId,
152    },
153    /// Computes a vector dot product.
154    Dot {
155        /// Left vector.
156        lhs: KernelValueId,
157        /// Right vector.
158        rhs: KernelValueId,
159    },
160    /// Solves a linear system.
161    Solve {
162        /// Coefficient matrix.
163        matrix: KernelValueId,
164        /// Right-hand-side vector.
165        rhs: KernelValueId,
166    },
167    /// Evaluates one row of a specialized cached solve.
168    SolveRow {
169        /// Cache slot containing the decomposed matrix row data.
170        row_slot: usize,
171        /// Right-hand-side scalar values.
172        rhs: Vec<KernelValueId>,
173    },
174    /// Evaluates one adjoint element for a specialized solve row.
175    SolveRowAdjointElement {
176        /// Cache slot containing the decomposed matrix row data.
177        row_slot: usize,
178        /// Element index within the row.
179        index: usize,
180        /// Row length.
181        len: usize,
182        /// Incoming scalar adjoint.
183        adjoint: KernelValueId,
184    },
185}
186
187impl KernelInstruction {
188    /// Returns the direct input value identifiers.
189    pub fn operands(&self) -> Vec<KernelValueId> {
190        match self {
191            Self::Cached(_)
192            | Self::RealConstant(_)
193            | Self::ComplexConstant(_)
194            | Self::Parameter(_) => Vec::new(),
195            Self::Unary { input, .. }
196            | Self::Component { input, .. }
197            | Self::MatrixElement { input, .. }
198            | Self::SolveRowAdjointElement { adjoint: input, .. } => vec![*input],
199            Self::Binary { lhs, rhs, .. } | Self::MatMul { lhs, rhs } | Self::Dot { lhs, rhs } => {
200                vec![*lhs, *rhs]
201            }
202            Self::MatVec { matrix, vector }
203            | Self::Solve {
204                matrix,
205                rhs: vector,
206            } => vec![*matrix, *vector],
207            Self::Add(values)
208            | Self::Mul(values)
209            | Self::Vector(values)
210            | Self::SolveRow { rhs: values, .. } => values.clone(),
211            Self::Complex { re, im } => vec![*re, *im],
212            Self::Matrix { elements, .. } => elements.clone(),
213        }
214    }
215
216    fn shape_error(
217        value: usize,
218        operation: &'static str,
219        message: impl Into<String>,
220    ) -> KernelIrError {
221        KernelIrError::InvalidShape {
222            value,
223            operation,
224            message: message.into(),
225        }
226    }
227
228    fn expected_kind(
229        &self,
230        values: &[KernelValue],
231        value: usize,
232    ) -> Result<Option<KernelValueKind>, KernelIrError> {
233        let kind = |id: KernelValueId| values[id.index()].kind;
234        Ok(Some(match self {
235            Self::Cached(_) => return Ok(None),
236            Self::RealConstant(_) | Self::Parameter(_) => KernelValueKind::Real,
237            Self::ComplexConstant(value) if value.im == 0.0 => KernelValueKind::Real,
238            Self::ComplexConstant(_) => KernelValueKind::Complex,
239            Self::Unary { op, input } => {
240                if !kind(*input).is_scalar() {
241                    return Err(Self::shape_error(
242                        value,
243                        "unary operation",
244                        "input is not scalar",
245                    ));
246                }
247                match op {
248                    UnaryOp::Real | UnaryOp::Imag | UnaryOp::NormSqr => KernelValueKind::Real,
249                    _ => kind(*input),
250                }
251            }
252            Self::Binary { op, lhs, rhs } => {
253                if *op == BinaryOp::Atan2 {
254                    if kind(*lhs) != KernelValueKind::Real || kind(*rhs) != KernelValueKind::Real {
255                        return Err(Self::shape_error(
256                            value,
257                            "atan2",
258                            "both inputs must be real",
259                        ));
260                    }
261                    KernelValueKind::Real
262                } else {
263                    kind(*lhs).scalar_combine(kind(*rhs)).ok_or_else(|| {
264                        Self::shape_error(value, "binary operation", "both inputs must be scalar")
265                    })?
266                }
267            }
268            Self::Add(terms) | Self::Mul(terms) => {
269                let operation = if matches!(self, Self::Add(_)) {
270                    "addition"
271                } else {
272                    "multiplication"
273                };
274                let mut terms = terms.iter();
275                let first = terms
276                    .next()
277                    .ok_or(KernelIrError::EmptyOperands { value, operation })?;
278                terms.try_fold(kind(*first), |acc, term| {
279                    acc.scalar_combine(kind(*term)).ok_or_else(|| {
280                        Self::shape_error(value, operation, "all inputs must be scalar")
281                    })
282                })?
283            }
284            Self::Complex { re, im } => {
285                if kind(*re) != KernelValueKind::Real || kind(*im) != KernelValueKind::Real {
286                    return Err(Self::shape_error(
287                        value,
288                        "complex construction",
289                        "components must be real",
290                    ));
291                }
292                KernelValueKind::Complex
293            }
294            Self::Vector(elements) => {
295                if elements.iter().any(|element| !kind(*element).is_scalar()) {
296                    return Err(Self::shape_error(
297                        value,
298                        "vector construction",
299                        "elements must be scalar",
300                    ));
301                }
302                KernelValueKind::Vector {
303                    len: elements.len(),
304                }
305            }
306            Self::Matrix {
307                rows,
308                cols,
309                elements,
310            } => {
311                if elements.len() != rows * cols
312                    || elements.iter().any(|element| !kind(*element).is_scalar())
313                {
314                    return Err(Self::shape_error(
315                        value,
316                        "matrix construction",
317                        format!("expected {} scalar elements", rows * cols),
318                    ));
319                }
320                KernelValueKind::Matrix {
321                    rows: *rows,
322                    cols: *cols,
323                }
324            }
325            Self::Component { input, index } => match kind(*input) {
326                KernelValueKind::Vector { len } if *index < len => KernelValueKind::Complex,
327                actual => {
328                    return Err(Self::shape_error(
329                        value,
330                        "component",
331                        format!("index {index} is invalid for {actual:?}"),
332                    ));
333                }
334            },
335            Self::MatrixElement { input, row, col } => match kind(*input) {
336                KernelValueKind::Matrix { rows, cols } if *row < rows && *col < cols => {
337                    KernelValueKind::Complex
338                }
339                actual => {
340                    return Err(Self::shape_error(
341                        value,
342                        "matrix element",
343                        format!("index ({row}, {col}) is invalid for {actual:?}"),
344                    ));
345                }
346            },
347            Self::MatMul { lhs, rhs } => match (kind(*lhs), kind(*rhs)) {
348                (
349                    KernelValueKind::Matrix { rows, cols: inner },
350                    KernelValueKind::Matrix {
351                        rows: rhs_rows,
352                        cols,
353                    },
354                ) if inner == rhs_rows => KernelValueKind::Matrix { rows, cols },
355                shapes => {
356                    return Err(Self::shape_error(
357                        value,
358                        "matrix multiplication",
359                        format!("incompatible operands {shapes:?}"),
360                    ));
361                }
362            },
363            Self::MatVec { matrix, vector } => match (kind(*matrix), kind(*vector)) {
364                (KernelValueKind::Matrix { rows, cols }, KernelValueKind::Vector { len })
365                    if cols == len =>
366                {
367                    KernelValueKind::Vector { len: rows }
368                }
369                shapes => {
370                    return Err(Self::shape_error(
371                        value,
372                        "matrix-vector multiplication",
373                        format!("incompatible operands {shapes:?}"),
374                    ));
375                }
376            },
377            Self::Dot { lhs, rhs } => match (kind(*lhs), kind(*rhs)) {
378                (KernelValueKind::Vector { len }, KernelValueKind::Vector { len: rhs_len })
379                    if len == rhs_len =>
380                {
381                    KernelValueKind::Complex
382                }
383                shapes => {
384                    return Err(Self::shape_error(
385                        value,
386                        "dot product",
387                        format!("incompatible operands {shapes:?}"),
388                    ));
389                }
390            },
391            Self::Solve { matrix, rhs } => match (kind(*matrix), kind(*rhs)) {
392                (KernelValueKind::Matrix { rows, cols }, KernelValueKind::Vector { len })
393                    if rows == cols && rows == len =>
394                {
395                    KernelValueKind::Vector { len }
396                }
397                shapes => {
398                    return Err(Self::shape_error(
399                        value,
400                        "linear solve",
401                        format!("incompatible operands {shapes:?}"),
402                    ));
403                }
404            },
405            Self::SolveRow { rhs, .. } => {
406                if rhs.is_empty() || rhs.iter().any(|entry| !kind(*entry).is_scalar()) {
407                    return Err(Self::shape_error(
408                        value,
409                        "specialized solve row",
410                        "right-hand side must contain scalars",
411                    ));
412                }
413                KernelValueKind::Complex
414            }
415            Self::SolveRowAdjointElement {
416                index,
417                len,
418                adjoint,
419                ..
420            } => {
421                if *len == 0 || *index >= *len || !kind(*adjoint).is_scalar() {
422                    return Err(Self::shape_error(
423                        value,
424                        "specialized solve-row adjoint",
425                        "adjoint must be scalar and index must be within a non-empty row",
426                    ));
427                }
428                KernelValueKind::Complex
429            }
430        }))
431    }
432
433    fn expected_class(&self, values: &[KernelValue]) -> KernelValueClass {
434        match self {
435            Self::Cached(_) | Self::SolveRow { .. } | Self::SolveRowAdjointElement { .. } => {
436                KernelValueClass::Event
437            }
438            Self::RealConstant(_) | Self::ComplexConstant(_) | Self::Parameter(_) => {
439                KernelValueClass::Invariant
440            }
441            _ if self
442                .operands()
443                .iter()
444                .any(|operand| values[operand.index()].class == KernelValueClass::Event) =>
445            {
446                KernelValueClass::Event
447            }
448            _ => KernelValueClass::Invariant,
449        }
450    }
451}
452
453/// Typed instruction and evaluation class for one kernel IR value.
454#[derive(Clone, Debug)]
455pub struct KernelValue {
456    /// Value shape and scalar representation.
457    pub kind: KernelValueKind,
458    /// Event-dependency class.
459    pub class: KernelValueClass,
460    /// Instruction that produces the value.
461    pub instruction: KernelInstruction,
462}
463
464/// Validated IR for a kernel with one scalar output.
465#[derive(Clone, Debug)]
466pub struct ScalarKernelIr {
467    values: Vec<KernelValue>,
468    root: KernelValueId,
469}
470
471/// Validated IR for a kernel that populates multiple cache outputs.
472#[derive(Clone, Debug)]
473pub struct CacheKernelIr {
474    values: Vec<KernelValue>,
475    outputs: Vec<KernelValueId>,
476}
477
478/// Scalar component of a complex primal output to differentiate.
479#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
480pub enum OutputComponent {
481    /// Differentiate the real component.
482    Real,
483    /// Differentiate the imaginary component.
484    Imag,
485}
486
487/// Validated IR containing a primal computation and real gradient outputs.
488#[derive(Clone, Debug)]
489pub struct GradientKernelIr {
490    values: Vec<KernelValue>,
491    primal_root: KernelValueId,
492    outputs: Vec<KernelValueId>,
493    component: OutputComponent,
494}
495
496/// Builder for appending type-checked instructions to existing scalar IR.
497#[derive(Clone, Debug)]
498pub struct KernelIrBuilder {
499    values: Vec<KernelValue>,
500}
501
502impl ScalarKernelIr {
503    /// Validates values and constructs a scalar kernel rooted at `root`.
504    ///
505    /// # Errors
506    ///
507    /// Returns [`KernelIrError`] when the value list is empty, `root` or an
508    /// operand is out of bounds, values are not topologically ordered, a
509    /// value's kind or class is inconsistent with its instruction, or the
510    /// root is not scalar.
511    pub fn new(values: Vec<KernelValue>, root: KernelValueId) -> Result<Self, KernelIrError> {
512        let ir = Self { values, root };
513        ir.validate()?;
514        if !ir.values[ir.root.index()].kind.is_scalar() {
515            return Err(KernelIrError::InvalidShape {
516                value: ir.root.index(),
517                operation: "kernel root",
518                message: "root must be scalar".into(),
519            });
520        }
521        Ok(ir)
522    }
523
524    /// Revalidates ordering, types, classes, and the scalar root.
525    ///
526    /// # Errors
527    ///
528    /// Returns [`KernelIrError`] when the IR is empty, its root or an operand
529    /// is out of bounds, its values are not topologically ordered, or a
530    /// value's kind, class, or shape is inconsistent with its instruction.
531    pub fn validate(&self) -> Result<(), KernelIrError> {
532        Self::validate_values(&self.values, self.root)
533    }
534
535    fn validate_values(values: &[KernelValue], root: KernelValueId) -> Result<(), KernelIrError> {
536        if values.is_empty() {
537            return Err(KernelIrError::Empty);
538        }
539        if root.index() >= values.len() {
540            return Err(KernelIrError::RootOutOfBounds {
541                root: root.index(),
542                len: values.len(),
543            });
544        }
545        for (index, value) in values.iter().enumerate() {
546            for operand in value.instruction.operands() {
547                if operand.index() >= index {
548                    return Err(KernelIrError::InvalidOperand {
549                        value: index,
550                        operand: operand.index(),
551                    });
552                }
553            }
554            if let Some(expected) = value.instruction.expected_kind(values, index)?
555                && value.kind != expected
556            {
557                return Err(KernelIrError::KindMismatch {
558                    value: index,
559                    expected,
560                    actual: value.kind,
561                });
562            }
563            let expected = value.instruction.expected_class(values);
564            if value.class != expected {
565                return Err(KernelIrError::ClassMismatch {
566                    value: index,
567                    expected,
568                    actual: value.class,
569                });
570            }
571        }
572        Ok(())
573    }
574
575    /// Returns all IR values in topological order.
576    pub fn values(&self) -> &[KernelValue] {
577        &self.values
578    }
579    /// Returns the scalar output identifier.
580    pub fn root(&self) -> KernelValueId {
581        self.root
582    }
583}
584
585impl CacheKernelIr {
586    /// Validates values and constructs a cache kernel with the given outputs.
587    ///
588    /// # Errors
589    ///
590    /// Returns [`KernelIrError`] when `outputs` is empty, an output or operand
591    /// is out of bounds, values are not topologically ordered, or a value's
592    /// kind, class, or shape is inconsistent with its instruction.
593    pub fn new(
594        values: Vec<KernelValue>,
595        outputs: Vec<KernelValueId>,
596    ) -> Result<Self, KernelIrError> {
597        let Some(first) = outputs.first().copied() else {
598            return Err(KernelIrError::EmptyCacheOutputs);
599        };
600        ScalarKernelIr::validate_values(&values, first)?;
601        for output in &outputs {
602            if output.index() >= values.len() {
603                return Err(KernelIrError::CacheOutputOutOfBounds {
604                    output: output.index(),
605                    len: values.len(),
606                });
607            }
608        }
609        Ok(Self { values, outputs })
610    }
611
612    /// Returns all IR values in topological order.
613    pub fn values(&self) -> &[KernelValue] {
614        &self.values
615    }
616
617    /// Returns cache output identifiers in storage order.
618    pub fn outputs(&self) -> &[KernelValueId] {
619        &self.outputs
620    }
621}
622
623impl GradientKernelIr {
624    /// Validates and constructs a gradient kernel.
625    ///
626    /// # Errors
627    ///
628    /// Returns [`KernelIrError`] when the primal IR is invalid, the primal
629    /// root is not scalar, a gradient output is out of bounds, or a gradient
630    /// output is not real-valued.
631    pub fn new(
632        values: Vec<KernelValue>,
633        primal_root: KernelValueId,
634        outputs: Vec<KernelValueId>,
635        component: OutputComponent,
636    ) -> Result<Self, KernelIrError> {
637        let ir = Self {
638            values,
639            primal_root,
640            outputs,
641            component,
642        };
643        ir.validate()?;
644        Ok(ir)
645    }
646
647    /// Revalidates the primal root and real gradient outputs.
648    ///
649    /// # Errors
650    ///
651    /// Returns [`KernelIrError`] when the primal IR is invalid, the primal
652    /// root is not scalar, a gradient output is out of bounds, or a gradient
653    /// output is not real-valued.
654    pub fn validate(&self) -> Result<(), KernelIrError> {
655        ScalarKernelIr::validate_values(&self.values, self.primal_root)?;
656        if !self.values[self.primal_root.index()].kind.is_scalar() {
657            return Err(KernelIrError::InvalidShape {
658                value: self.primal_root.index(),
659                operation: "gradient primal root",
660                message: "primal root must be scalar".into(),
661            });
662        }
663        for output in &self.outputs {
664            let Some(value) = self.values.get(output.index()) else {
665                return Err(KernelIrError::GradientOutOfBounds {
666                    output: output.index(),
667                    len: self.values.len(),
668                });
669            };
670            if value.kind != KernelValueKind::Real {
671                return Err(KernelIrError::GradientKindMismatch {
672                    output: output.index(),
673                    actual: value.kind,
674                });
675            }
676        }
677        Ok(())
678    }
679
680    /// Returns all primal and derivative IR values in topological order.
681    pub fn values(&self) -> &[KernelValue] {
682        &self.values
683    }
684
685    /// Returns the primal scalar output identifier.
686    pub fn primal_root(&self) -> KernelValueId {
687        self.primal_root
688    }
689
690    /// Returns derivative output identifiers.
691    pub fn outputs(&self) -> &[KernelValueId] {
692        &self.outputs
693    }
694
695    /// Returns the differentiated component of the complex primal.
696    pub fn component(&self) -> OutputComponent {
697        self.component
698    }
699}
700
701impl KernelIrBuilder {
702    /// Creates a builder initialized with a scalar kernel's values.
703    pub fn from_scalar(ir: &ScalarKernelIr) -> Self {
704        Self {
705            values: ir.values.clone(),
706        }
707    }
708
709    /// Appends an instruction after validating its operands and inferred type.
710    ///
711    /// # Errors
712    ///
713    /// Returns [`KernelIrError`] when an operand does not precede the new
714    /// instruction, the operand shapes are incompatible, or the instruction's
715    /// value kind cannot be inferred.
716    pub fn push(&mut self, instruction: KernelInstruction) -> Result<KernelValueId, KernelIrError> {
717        let index = self.values.len();
718        for operand in instruction.operands() {
719            if operand.index() >= index {
720                return Err(KernelIrError::InvalidOperand {
721                    value: index,
722                    operand: operand.index(),
723                });
724            }
725        }
726        let kind = instruction
727            .expected_kind(&self.values, index)?
728            .ok_or_else(|| KernelIrError::InvalidShape {
729                value: index,
730                operation: "derived instruction",
731                message: "instruction requires an explicitly supplied value kind".into(),
732            })?;
733        let class = instruction.expected_class(&self.values);
734        let id = KernelValueId::from_index(index);
735        self.values.push(KernelValue {
736            kind,
737            class,
738            instruction,
739        });
740        Ok(id)
741    }
742
743    /// Finishes the builder as a validated gradient kernel.
744    ///
745    /// # Errors
746    ///
747    /// Returns [`KernelIrError`] when the accumulated primal IR is invalid,
748    /// `primal_root` is not scalar, a gradient output is out of bounds, or a
749    /// gradient output is not real-valued.
750    pub fn finish_gradient(
751        self,
752        primal_root: KernelValueId,
753        outputs: Vec<KernelValueId>,
754        component: OutputComponent,
755    ) -> Result<GradientKernelIr, KernelIrError> {
756        GradientKernelIr::new(self.values, primal_root, outputs, component)
757    }
758
759    /// Returns the values accumulated so far.
760    pub fn values(&self) -> &[KernelValue] {
761        &self.values
762    }
763}
764
765#[cfg(test)]
766mod tests {
767    use super::*;
768
769    #[test]
770    fn validates_aggregate_operations() {
771        let values = vec![
772            KernelValue {
773                kind: KernelValueKind::Complex,
774                class: KernelValueClass::Invariant,
775                instruction: KernelInstruction::ComplexConstant(Complex64::new(1.0, 2.0)),
776            },
777            KernelValue {
778                kind: KernelValueKind::Vector { len: 1 },
779                class: KernelValueClass::Invariant,
780                instruction: KernelInstruction::Vector(vec![KernelValueId::from_index(0)]),
781            },
782            KernelValue {
783                kind: KernelValueKind::Complex,
784                class: KernelValueClass::Invariant,
785                instruction: KernelInstruction::Dot {
786                    lhs: KernelValueId::from_index(1),
787                    rhs: KernelValueId::from_index(1),
788                },
789            },
790        ];
791        ScalarKernelIr::new(values, KernelValueId::from_index(2)).unwrap();
792    }
793
794    #[test]
795    fn rejects_forward_references() {
796        let error = ScalarKernelIr::new(
797            vec![
798                KernelValue {
799                    kind: KernelValueKind::Real,
800                    class: KernelValueClass::Invariant,
801                    instruction: KernelInstruction::Unary {
802                        op: UnaryOp::Neg,
803                        input: KernelValueId::from_index(1),
804                    },
805                },
806                KernelValue {
807                    kind: KernelValueKind::Real,
808                    class: KernelValueClass::Invariant,
809                    instruction: KernelInstruction::RealConstant(1.0),
810                },
811            ],
812            KernelValueId::from_index(0),
813        )
814        .unwrap_err();
815        assert_eq!(
816            error,
817            KernelIrError::InvalidOperand {
818                value: 0,
819                operand: 1
820            }
821        );
822    }
823
824    #[test]
825    fn rejects_invalid_matrix_shapes() {
826        let error = ScalarKernelIr::new(
827            vec![
828                KernelValue {
829                    kind: KernelValueKind::Real,
830                    class: KernelValueClass::Invariant,
831                    instruction: KernelInstruction::RealConstant(1.0),
832                },
833                KernelValue {
834                    kind: KernelValueKind::Matrix { rows: 2, cols: 2 },
835                    class: KernelValueClass::Invariant,
836                    instruction: KernelInstruction::Matrix {
837                        rows: 2,
838                        cols: 2,
839                        elements: vec![KernelValueId::from_index(0)],
840                    },
841                },
842            ],
843            KernelValueId::from_index(0),
844        )
845        .unwrap_err();
846        assert!(matches!(error, KernelIrError::InvalidShape { .. }));
847    }
848
849    #[test]
850    fn gradient_builder_appends_valid_real_outputs() {
851        let primal = ScalarKernelIr::new(
852            vec![KernelValue {
853                kind: KernelValueKind::Complex,
854                class: KernelValueClass::Invariant,
855                instruction: KernelInstruction::ComplexConstant(Complex64::new(1.0, 2.0)),
856            }],
857            KernelValueId::from_index(0),
858        )
859        .unwrap();
860        let mut builder = KernelIrBuilder::from_scalar(&primal);
861        let output = builder
862            .push(KernelInstruction::Unary {
863                op: UnaryOp::Real,
864                input: primal.root(),
865            })
866            .unwrap();
867        let gradient = builder
868            .finish_gradient(primal.root(), vec![output], OutputComponent::Real)
869            .unwrap();
870
871        assert_eq!(gradient.primal_root(), primal.root());
872        assert_eq!(gradient.outputs(), &[output]);
873        assert_eq!(gradient.component(), OutputComponent::Real);
874        assert_eq!(gradient.values().len(), 2);
875    }
876
877    #[test]
878    fn cache_kernel_preserves_multiple_typed_outputs() {
879        let values = vec![
880            KernelValue {
881                kind: KernelValueKind::Real,
882                class: KernelValueClass::Event,
883                instruction: KernelInstruction::Cached(0),
884            },
885            KernelValue {
886                kind: KernelValueKind::Real,
887                class: KernelValueClass::Event,
888                instruction: KernelInstruction::Unary {
889                    op: UnaryOp::Sin,
890                    input: KernelValueId::from_index(0),
891                },
892            },
893        ];
894        let kernel = CacheKernelIr::new(
895            values,
896            vec![KernelValueId::from_index(0), KernelValueId::from_index(1)],
897        )
898        .unwrap();
899
900        assert_eq!(kernel.outputs().len(), 2);
901        assert_eq!(
902            kernel.values()[kernel.outputs()[1].index()].kind,
903            KernelValueKind::Real
904        );
905    }
906
907    #[test]
908    fn gradient_outputs_must_be_real() {
909        let primal = ScalarKernelIr::new(
910            vec![KernelValue {
911                kind: KernelValueKind::Complex,
912                class: KernelValueClass::Invariant,
913                instruction: KernelInstruction::ComplexConstant(Complex64::new(1.0, 2.0)),
914            }],
915            KernelValueId::from_index(0),
916        )
917        .unwrap();
918        let error = GradientKernelIr::new(
919            primal.values().to_vec(),
920            primal.root(),
921            vec![primal.root()],
922            OutputComponent::Real,
923        )
924        .unwrap_err();
925
926        assert_eq!(
927            error,
928            KernelIrError::GradientKindMismatch {
929                output: 0,
930                actual: KernelValueKind::Complex,
931            }
932        );
933    }
934}
laddu_kernel/ir.rs

laddu_kernel/
ir.rs
