Skip to main content

laddu_kernel/ir/
builder.rs

1use super::*;
2
3impl KernelIrBuilder {
4    /// Creates a builder initialized with a scalar kernel's values.
5    pub fn from_scalar(ir: &ScalarKernelIr) -> Self {
6        Self {
7            values: ir.values.clone(),
8        }
9    }
10
11    /// Appends an instruction after validating its operands and inferred type.
12    ///
13    /// # Errors
14    ///
15    /// Returns [`KernelIrError`] when an operand does not precede the new
16    /// instruction, the operand shapes are incompatible, or the instruction's
17    /// value kind cannot be inferred.
18    pub fn push(&mut self, instruction: KernelInstruction) -> Result<KernelValueId, KernelIrError> {
19        let index = self.values.len();
20        instruction.validate_operand_order(index)?;
21        let kind = instruction
22            .expected_kind(&self.values, index)?
23            .ok_or_else(|| {
24                KernelInstruction::shape_error(
25                    index,
26                    "derived instruction",
27                    "instruction requires an explicitly supplied value kind",
28                )
29            })?;
30        let class = instruction.expected_class(&self.values);
31        let id = KernelValueId::from_index(index);
32        self.values.push(KernelValue {
33            kind,
34            class,
35            instruction,
36        });
37        Ok(id)
38    }
39
40    /// Finishes the builder as a validated gradient kernel.
41    ///
42    /// # Errors
43    ///
44    /// Returns [`KernelIrError`] when the accumulated primal IR is invalid,
45    /// `primal_root` is not scalar, a gradient output is out of bounds, or a
46    /// gradient output is not real-valued.
47    pub fn finish_gradient(
48        self,
49        primal_root: KernelValueId,
50        outputs: Vec<KernelValueId>,
51        component: OutputComponent,
52    ) -> Result<GradientKernelIr, KernelIrError> {
53        GradientKernelIr::new(self.values, primal_root, outputs, component)
54    }
55
56    /// Returns the values accumulated so far.
57    pub fn values(&self) -> &[KernelValue] {
58        &self.values
59    }
60}