tenferro_runtime/graph/
compiler.rs

1use std::fmt;
2
3use tenferro_ops::dim_expr::{DimExpr, DimExprEvalError};
4use tenferro_ops::input_key::TensorInputKey;
5use tenferro_ops::{ShapeExtent, ShapeRelation};
6#[cfg(test)]
7use tenferro_tensor::Tensor;
8use tenferro_tensor::{CacheStats, DType};
9
10use super::program::CompiledGraph;
11#[cfg(test)]
12use crate::compiler::semantic_staging::stage_semantic_program;
13use crate::compiler::CompilerOptions;
14use crate::error::{Error, ErrorPhase, Result};
15use crate::extension_cache::{ExtensionCacheSelector, ExtensionCacheStore};
16use crate::program::{
17    FrozenProgram, ProgramShapeRelation, ProgramValueMetadata, SemanticOpRef,
18    ShapeGuard as ProgramShapeGuard,
19};
20use crate::trace::TracedGraph;
21use crate::traced::{try_concrete_shape, TracedTensor};
22
23/// Compiler for traced tensor graphs.
24///
25/// A graph compiler lowers one or more [`TracedTensor`] outputs to a reusable
26/// [`CompiledGraph`] without requiring a backend.
27///
28/// # Examples
29///
30/// ```
31/// use tenferro_runtime::{TraceContext, GraphCompiler, TracedTensor};
32///
33/// let trace = TraceContext::new();
34/// let x = trace.from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
35/// let y = (&x + &x).unwrap();
36/// let mut compiler = GraphCompiler::new();
37/// let program = compiler.compile(&y).unwrap();
38/// assert_eq!(program.output_count(), 1);
39/// ```
40pub struct GraphCompiler {
41    extension_cache: ExtensionCacheStore,
42    compiler_options: CompilerOptions,
43}
44
45impl fmt::Debug for GraphCompiler {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        f.debug_struct("GraphCompiler")
48            .field("extension_cache_stats", &self.cache_stats())
49            .field("compiler_options", &self.compiler_options)
50            .field("extension_cache", &self.extension_cache)
51            .finish_non_exhaustive()
52    }
53}
54
55impl Default for GraphCompiler {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl GraphCompiler {
62    /// Create a compiler with bounded default caches.
63    ///
64    /// # Examples
65    ///
66    /// ```
67    /// use tenferro_runtime::GraphCompiler;
68    ///
69    /// let compiler = GraphCompiler::new();
70    /// assert!(compiler.extension_caches().is_empty());
71    /// ```
72    pub fn new() -> Self {
73        Self {
74            extension_cache: ExtensionCacheStore::new(),
75            compiler_options: CompilerOptions::default(),
76        }
77    }
78
79    /// Create a compiler with explicit lowering and optimizer options.
80    ///
81    /// # Examples
82    ///
83    /// ```
84    /// use tenferro_runtime::{CompilerOptions, GraphCompiler, OptimizerConfig};
85    ///
86    /// let compiler = GraphCompiler::with_compiler_options(CompilerOptions {
87    ///     optimizer: OptimizerConfig {
88    ///         dot_decomposer: true,
89    ///         ..OptimizerConfig::default()
90    ///     },
91    /// });
92    /// assert!(compiler.compiler_options().optimizer.dot_decomposer);
93    /// ```
94    pub fn with_compiler_options(compiler_options: CompilerOptions) -> Self {
95        Self {
96            extension_cache: ExtensionCacheStore::new(),
97            compiler_options,
98        }
99    }
100
101    /// Compile one traced output by snapshotting its context.
102    ///
103    /// The output's reachable closure is frozen into one semantic program; the
104    /// compile leaves the context usable for further tracing.
105    ///
106    /// # Examples
107    ///
108    /// ```
109    /// use tenferro_runtime::{GraphCompiler, TraceContext};
110    ///
111    /// let trace = TraceContext::new();
112    /// let x = trace.input_concrete_shape(tenferro_tensor::DType::F64, &[1])?;
113    /// let y = x.neg()?;
114    /// let mut compiler = GraphCompiler::new();
115    /// let program = compiler.compile(&y)?;
116    /// assert_eq!(program.output_count(), 1);
117    /// # Ok::<(), Box<dyn std::error::Error>>(())
118    /// ```
119    ///
120    /// # Errors
121    ///
122    /// Returns [`Error::Validation`] for an invalid snapshot or bound shape
123    /// guard, or [`Error::RuntimeState`] for inconsistent compiler state.
124    pub fn compile(&mut self, output: &TracedTensor) -> Result<CompiledGraph> {
125        self.compile_many(&[output])
126    }
127
128    /// Compile an immutable semantic trace without consulting a backend.
129    ///
130    /// # Errors
131    ///
132    /// Returns [`Error::Validation`] for invalid metadata or shape constraints,
133    /// or [`Error::RuntimeState`] for inconsistent staging state.
134    pub fn compile_traced_graph(&mut self, graph: &TracedGraph) -> Result<CompiledGraph> {
135        self.compile_frozen(graph.frozen(), Vec::new())
136    }
137
138    /// Compile an immutable semantic program for ordered execution.
139    ///
140    /// # Examples
141    ///
142    /// ```
143    /// use tenferro_ops::dim_expr::DimExpr;
144    /// use tenferro_runtime::program::{
145    ///     CoreSemanticOp, ProgramInputSpec, SemanticProgramBuilder,
146    /// };
147    /// use tenferro_runtime::{DType, GraphCompiler};
148    ///
149    /// let mut builder = SemanticProgramBuilder::new();
150    /// let input = builder
151    ///     .input(ProgramInputSpec::new(DType::F64, [DimExpr::Const(2)]))
152    ///     .unwrap();
153    /// let output = builder.add_op(CoreSemanticOp::Neg, &[input]).unwrap()[0];
154    /// let frozen = builder.finish(&[output]).unwrap();
155    /// let compiled = GraphCompiler::new()
156    ///     .compile_frozen_program(&frozen)
157    ///     .unwrap();
158    /// assert_eq!(compiled.input_count(), 1);
159    /// ```
160    ///
161    /// # Errors
162    ///
163    /// Returns [`Error::Validation`] for invalid metadata or shape constraints,
164    /// or [`Error::RuntimeState`] for inconsistent staging state.
165    pub fn compile_frozen_program(&mut self, frozen: &FrozenProgram) -> Result<CompiledGraph> {
166        self.compile_frozen(frozen, Vec::new())
167    }
168
169    fn compile_frozen(
170        &mut self,
171        frozen: &FrozenProgram,
172        input_keys: Vec<TensorInputKey>,
173    ) -> Result<CompiledGraph> {
174        validate_bound_shape_guards(frozen)?;
175        Ok(CompiledGraph::new(
176            frozen.clone(),
177            self.compiler_options,
178            input_keys,
179        ))
180    }
181
182    /// Compile multiple traced outputs into one program.
183    ///
184    /// Every output must belong to the same trace context.
185    ///
186    /// # Errors
187    ///
188    /// Returns `Error::Validation` with `ValidationError::InvalidArgument` when
189    /// the output list is empty (`argument: "outputs"`) or when an output belongs
190    /// to a different trace context, and the typed snapshot error from
191    /// [`compile`](Self::compile) when the program cannot be frozen.
192    pub fn compile_many(&mut self, outputs: &[&TracedTensor]) -> Result<CompiledGraph> {
193        let Some(first) = outputs.first() else {
194            return Err(Error::invalid_argument(
195                "GraphCompiler::compile_many",
196                ErrorPhase::Compile,
197                "outputs",
198                "at least one output is required",
199            ));
200        };
201        for output in outputs.iter().skip(1) {
202            crate::traced::ensure_same_context(first, output)?;
203        }
204        let values: Vec<_> = outputs.iter().map(|output| output.value).collect();
205        let (frozen, kept) = first
206            .ctx
207            .plain_snapshot_with_inputs(&values)
208            .map_err(graph_error_with_context("GraphCompiler::compile_many"))?;
209        let input_keys = kept.into_iter().map(|(key, _)| key).collect();
210        self.compile_frozen(&frozen, input_keys)
211    }
212
213    /// Compile one traced output with concrete placeholder specs.
214    ///
215    /// Every declared placeholder must be a data-less input of the output's
216    /// context, with the declared dtype and shape. The compiled program lists the
217    /// declared placeholders first, in the order given here, followed by the
218    /// remaining live inputs in context order.
219    ///
220    /// # Errors
221    ///
222    /// Returns [`Error::UnexpectedBinding`] for a data-carrying tensor,
223    /// [`Error::DuplicateBinding`] for repeated placeholders,
224    /// [`Error::PlaceholderDtypeMismatch`],
225    /// [`Error::PlaceholderShapeMismatch`], or
226    /// [`Error::PlaceholderRankMismatch`] for incompatible specs,
227    /// [`Error::Validation`] with `InvalidArgument` for a declared placeholder the
228    /// output does not depend on, and the snapshot failures of
229    /// [`compile`](Self::compile).
230    pub fn compile_with_input_specs(
231        &mut self,
232        output: &TracedTensor,
233        bindings: &[(&TracedTensor, DType, &[usize])],
234    ) -> Result<CompiledGraph> {
235        let mut seen = std::collections::HashSet::new();
236        let mut specs = Vec::with_capacity(bindings.len());
237        for (index, (placeholder, dtype, shape)) in bindings.iter().enumerate() {
238            validate_placeholder_spec(index, placeholder, *dtype, shape)?;
239            let key = placeholder.input_key().ok_or(Error::UnexpectedBinding {
240                binding_index: index,
241            })?;
242            if !seen.insert(key.clone()) {
243                return Err(Error::DuplicateBinding {
244                    input_key: format!("{key:?}"),
245                });
246            }
247            specs.push((
248                key,
249                crate::program::ProgramInputSpec::new(
250                    *dtype,
251                    shape
252                        .iter()
253                        .map(|&size| tenferro_ops::dim_expr::DimExpr::Const(size)),
254                ),
255            ));
256        }
257        let (frozen, kept) = output
258            .ctx
259            .snapshot_with_input_specs(&[output.value], &specs)
260            .map_err(graph_error_with_context(
261                "GraphCompiler::compile_with_input_specs",
262            ))?;
263        let input_keys = kept.into_iter().map(|(key, _)| key).collect();
264        let program = self.compile_frozen(&frozen, input_keys)?;
265        if let Some(binding_index) = bindings.iter().position(|(placeholder, _, _)| {
266            placeholder
267                .input_key()
268                .is_none_or(|key| program.input_key_index(&key).is_none())
269        }) {
270            return Err(Error::invalid_argument(
271                "GraphCompiler::compile_with_input_specs",
272                ErrorPhase::Compile,
273                "bindings",
274                format!(
275                    "binding {binding_index} declares a placeholder the output does not \
276                     depend on; remove it from the bindings"
277                ),
278            ));
279        }
280        Ok(program)
281    }
282
283    /// Return the compiler options used for future lowerings.
284    ///
285    /// # Examples
286    ///
287    /// ```
288    /// use tenferro_runtime::{CompilerOptions, GraphCompiler};
289    ///
290    /// let compiler = GraphCompiler::new();
291    /// assert_eq!(compiler.compiler_options(), CompilerOptions::default());
292    /// ```
293    pub fn compiler_options(&self) -> CompilerOptions {
294        self.compiler_options
295    }
296
297    /// Replace compiler options and clear compiler-owned extension cache entries.
298    pub fn set_compiler_options(&mut self, compiler_options: CompilerOptions) {
299        self.compiler_options = compiler_options;
300        self.clear_caches();
301    }
302
303    /// Drop every extension cache entry owned by this compiler.
304    ///
305    /// # Examples
306    ///
307    /// ```
308    /// use tenferro_runtime::GraphCompiler;
309    ///
310    /// let mut compiler = GraphCompiler::new();
311    /// compiler.clear_extension_caches();
312    /// assert!(compiler.extension_caches().is_empty());
313    /// ```
314    pub fn clear_extension_caches(&mut self) {
315        self.extension_cache.clear();
316    }
317
318    /// Drop every cache owned by this compiler.
319    pub fn clear_caches(&mut self) {
320        self.clear_extension_caches();
321    }
322
323    /// Return aggregate cache statistics for this compiler.
324    ///
325    /// # Examples
326    ///
327    /// ```
328    /// use tenferro_runtime::GraphCompiler;
329    ///
330    /// let compiler = GraphCompiler::new();
331    /// let stats = compiler.cache_stats();
332    /// assert_eq!(stats.entries, 0);
333    /// ```
334    pub fn cache_stats(&self) -> CacheStats {
335        self.extension_cache.stats(ExtensionCacheSelector::All)
336    }
337
338    /// Borrow the extension cache store owned by this compiler.
339    pub fn extension_caches(&self) -> &ExtensionCacheStore {
340        &self.extension_cache
341    }
342
343    /// Borrow the extension cache store mutably.
344    pub fn extension_caches_mut(&mut self) -> &mut ExtensionCacheStore {
345        &mut self.extension_cache
346    }
347}
348
349fn validate_bound_shape_guards(frozen: &FrozenProgram) -> Result<()> {
350    let input_shapes = compile_time_input_shapes(frozen)?;
351    for operation in frozen.program.operations() {
352        let fallback_family = match operation.op() {
353            SemanticOpRef::Core(_) => "tenferro-runtime.core.v1",
354            SemanticOpRef::Extension(extension) => extension.family_id(),
355        };
356        for guard in operation.shape_guards() {
357            if guard.source_family().is_some() {
358                continue;
359            }
360            validate_bound_shape_guard(guard, fallback_family, &input_shapes)?;
361        }
362    }
363    Ok(())
364}
365
366fn compile_time_input_shapes(frozen: &FrozenProgram) -> Result<Vec<Option<Vec<usize>>>> {
367    let metadata = frozen.input_metadata_with_bound_shapes();
368    frozen
369        .program
370        .inputs()
371        .iter()
372        .enumerate()
373        .map(|(input_idx, &input)| {
374            if let Some(tensor) = frozen.bindings.tensor_ref_for_input(input) {
375                return Ok(Some(tensor.shape().to_vec()));
376            }
377            let Some(metadata) = metadata.get(input_idx) else {
378                return Err(invalid_compiled_graph(format!(
379                    "semantic input metadata index {input_idx} is outside metadata table"
380                )));
381            };
382            concrete_shape_from_input_metadata(metadata)
383        })
384        .collect()
385}
386
387fn concrete_shape_from_input_metadata(
388    metadata: &ProgramValueMetadata,
389) -> Result<Option<Vec<usize>>> {
390    let mut shape = Vec::with_capacity(metadata.shape().len());
391    for extent in metadata.shape() {
392        let ShapeExtent::Exact(expression) = extent else {
393            return Ok(None);
394        };
395        let Some(value) = evaluate_static_dim_expr(expression)? else {
396            return Ok(None);
397        };
398        shape.push(value);
399    }
400    Ok(Some(shape))
401}
402
403fn validate_bound_shape_guard(
404    guard: &ProgramShapeGuard,
405    fallback_family: &'static str,
406    input_shapes: &[Option<Vec<usize>>],
407) -> Result<()> {
408    let ProgramShapeRelation::Equal = guard.relation() else {
409        return Ok(());
410    };
411    let family = guard.source_family().unwrap_or(fallback_family);
412    let relation = ShapeRelation::Equal;
413    let lhs = evaluate_bound_shape_guard_expression(family, relation, guard.lhs(), input_shapes)?;
414    let rhs = evaluate_bound_shape_guard_expression(family, relation, guard.rhs(), input_shapes)?;
415    let (Some(lhs), Some(rhs)) = (lhs, rhs) else {
416        return Ok(());
417    };
418    if lhs == rhs {
419        return Ok(());
420    }
421    Err(Error::ShapeConstraintViolation {
422        family,
423        instruction_index: None,
424        relation,
425        lhs_expr: format!("{:?}", guard.lhs()),
426        rhs_expr: format!("{:?}", guard.rhs()),
427        lhs_value: lhs,
428        rhs_value: rhs,
429    })
430}
431
432fn evaluate_bound_shape_guard_expression(
433    family: &'static str,
434    relation: ShapeRelation,
435    expression: &DimExpr,
436    input_shapes: &[Option<Vec<usize>>],
437) -> Result<Option<usize>> {
438    evaluate_static_dim_expr_with_inputs(expression, input_shapes).map_err(|cause| {
439        Error::ShapeConstraintEvaluation {
440            family,
441            instruction_index: None,
442            relation,
443            expression: format!("{expression:?}"),
444            cause: cause.into(),
445        }
446    })
447}
448
449fn evaluate_static_dim_expr(expression: &DimExpr) -> Result<Option<usize>> {
450    evaluate_static_dim_expr_without_inputs(expression).map_err(|cause| {
451        Error::ShapeConstraintEvaluation {
452            family: "tenferro-runtime.input.v1",
453            instruction_index: None,
454            relation: ShapeRelation::Equal,
455            expression: format!("{expression:?}"),
456            cause: cause.into(),
457        }
458    })
459}
460
461fn evaluate_static_dim_expr_without_inputs(
462    expression: &DimExpr,
463) -> std::result::Result<Option<usize>, DimExprEvalError> {
464    match expression {
465        DimExpr::Const(value) => Ok(Some(*value)),
466        DimExpr::InputDim { .. } => Ok(None),
467        DimExpr::Add(a, b) => {
468            let Some(lhs) = evaluate_static_dim_expr_without_inputs(a)? else {
469                return Ok(None);
470            };
471            let Some(rhs) = evaluate_static_dim_expr_without_inputs(b)? else {
472                return Ok(None);
473            };
474            lhs.checked_add(rhs)
475                .map(Some)
476                .ok_or(DimExprEvalError::AddOverflow { lhs, rhs })
477        }
478        DimExpr::Sub(a, b) => {
479            let Some(lhs) = evaluate_static_dim_expr_without_inputs(a)? else {
480                return Ok(None);
481            };
482            let Some(rhs) = evaluate_static_dim_expr_without_inputs(b)? else {
483                return Ok(None);
484            };
485            lhs.checked_sub(rhs)
486                .map(Some)
487                .ok_or(DimExprEvalError::SubUnderflow { lhs, rhs })
488        }
489        DimExpr::Mul(a, b) => {
490            let Some(lhs) = evaluate_static_dim_expr_without_inputs(a)? else {
491                return Ok(None);
492            };
493            let Some(rhs) = evaluate_static_dim_expr_without_inputs(b)? else {
494                return Ok(None);
495            };
496            lhs.checked_mul(rhs)
497                .map(Some)
498                .ok_or(DimExprEvalError::MulOverflow { lhs, rhs })
499        }
500        DimExpr::FloorDiv(a, b) => {
501            let Some(lhs) = evaluate_static_dim_expr_without_inputs(a)? else {
502                return Ok(None);
503            };
504            let Some(rhs) = evaluate_static_dim_expr_without_inputs(b)? else {
505                return Ok(None);
506            };
507            if rhs == 0 {
508                return Err(DimExprEvalError::FloorDivByZero { lhs, rhs });
509            }
510            Ok(Some(lhs / rhs))
511        }
512        DimExpr::Min(a, b) => {
513            let Some(lhs) = evaluate_static_dim_expr_without_inputs(a)? else {
514                return Ok(None);
515            };
516            let Some(rhs) = evaluate_static_dim_expr_without_inputs(b)? else {
517                return Ok(None);
518            };
519            Ok(Some(lhs.min(rhs)))
520        }
521        DimExpr::Max(a, b) => {
522            let Some(lhs) = evaluate_static_dim_expr_without_inputs(a)? else {
523                return Ok(None);
524            };
525            let Some(rhs) = evaluate_static_dim_expr_without_inputs(b)? else {
526                return Ok(None);
527            };
528            Ok(Some(lhs.max(rhs)))
529        }
530    }
531}
532
533fn evaluate_static_dim_expr_with_inputs(
534    expression: &DimExpr,
535    input_shapes: &[Option<Vec<usize>>],
536) -> std::result::Result<Option<usize>, DimExprEvalError> {
537    match expression {
538        DimExpr::Const(value) => Ok(Some(*value)),
539        DimExpr::InputDim { input_idx, axis } => match input_shapes.get(*input_idx) {
540            Some(Some(shape)) => {
541                shape
542                    .get(*axis)
543                    .copied()
544                    .map(Some)
545                    .ok_or(DimExprEvalError::AxisOutOfBounds {
546                        input_idx: *input_idx,
547                        axis: *axis,
548                        rank: shape.len(),
549                    })
550            }
551            Some(None) => Ok(None),
552            None => Err(DimExprEvalError::InputOutOfBounds {
553                input_idx: *input_idx,
554                input_count: input_shapes.len(),
555            }),
556        },
557        DimExpr::Add(a, b) => {
558            let Some(lhs) = evaluate_static_dim_expr_with_inputs(a, input_shapes)? else {
559                return Ok(None);
560            };
561            let Some(rhs) = evaluate_static_dim_expr_with_inputs(b, input_shapes)? else {
562                return Ok(None);
563            };
564            lhs.checked_add(rhs)
565                .map(Some)
566                .ok_or(DimExprEvalError::AddOverflow { lhs, rhs })
567        }
568        DimExpr::Sub(a, b) => {
569            let Some(lhs) = evaluate_static_dim_expr_with_inputs(a, input_shapes)? else {
570                return Ok(None);
571            };
572            let Some(rhs) = evaluate_static_dim_expr_with_inputs(b, input_shapes)? else {
573                return Ok(None);
574            };
575            lhs.checked_sub(rhs)
576                .map(Some)
577                .ok_or(DimExprEvalError::SubUnderflow { lhs, rhs })
578        }
579        DimExpr::Mul(a, b) => {
580            let Some(lhs) = evaluate_static_dim_expr_with_inputs(a, input_shapes)? else {
581                return Ok(None);
582            };
583            let Some(rhs) = evaluate_static_dim_expr_with_inputs(b, input_shapes)? else {
584                return Ok(None);
585            };
586            lhs.checked_mul(rhs)
587                .map(Some)
588                .ok_or(DimExprEvalError::MulOverflow { lhs, rhs })
589        }
590        DimExpr::FloorDiv(a, b) => {
591            let Some(lhs) = evaluate_static_dim_expr_with_inputs(a, input_shapes)? else {
592                return Ok(None);
593            };
594            let Some(rhs) = evaluate_static_dim_expr_with_inputs(b, input_shapes)? else {
595                return Ok(None);
596            };
597            if rhs == 0 {
598                return Err(DimExprEvalError::FloorDivByZero { lhs, rhs });
599            }
600            Ok(Some(lhs / rhs))
601        }
602        DimExpr::Min(a, b) => {
603            let Some(lhs) = evaluate_static_dim_expr_with_inputs(a, input_shapes)? else {
604                return Ok(None);
605            };
606            let Some(rhs) = evaluate_static_dim_expr_with_inputs(b, input_shapes)? else {
607                return Ok(None);
608            };
609            Ok(Some(lhs.min(rhs)))
610        }
611        DimExpr::Max(a, b) => {
612            let Some(lhs) = evaluate_static_dim_expr_with_inputs(a, input_shapes)? else {
613                return Ok(None);
614            };
615            let Some(rhs) = evaluate_static_dim_expr_with_inputs(b, input_shapes)? else {
616                return Ok(None);
617            };
618            Ok(Some(lhs.max(rhs)))
619        }
620    }
621}
622
623fn invalid_compiled_graph(message: impl Into<String>) -> Error {
624    Error::Internal(message.into())
625}
626
627fn graph_error_with_context(
628    op: &'static str,
629) -> impl FnOnce(crate::program::ProgramFinishError) -> Error {
630    move |source| Error::runtime_state_source(op, ErrorPhase::Compile, source)
631}
632
633fn validate_placeholder_spec(
634    index: usize,
635    placeholder: &TracedTensor,
636    dtype: DType,
637    shape: &[usize],
638) -> Result<()> {
639    if placeholder.data.is_some() {
640        return Err(Error::UnexpectedBinding {
641            binding_index: index,
642        });
643    }
644    placeholder.input_key().ok_or(Error::UnexpectedBinding {
645        binding_index: index,
646    })?;
647
648    if placeholder.dtype != dtype {
649        return Err(Error::PlaceholderDtypeMismatch {
650            expected: placeholder.dtype,
651            actual: dtype,
652        });
653    }
654    validate_placeholder_shape(placeholder, shape)
655}
656
657fn validate_placeholder_shape(placeholder: &TracedTensor, shape: &[usize]) -> Result<()> {
658    match try_concrete_shape(placeholder) {
659        Some(expected_shape) => {
660            if expected_shape.as_slice() != shape {
661                return Err(Error::PlaceholderShapeMismatch {
662                    expected: expected_shape,
663                    actual: shape.to_vec(),
664                });
665            }
666        }
667        None => {
668            if placeholder.rank != shape.len() {
669                return Err(Error::PlaceholderRankMismatch {
670                    expected: placeholder.rank,
671                    actual: shape.len(),
672                });
673            }
674        }
675    }
676    Ok(())
677}
678
679#[cfg(test)]
680mod compile_contract_tests;
681
682#[cfg(test)]
683mod tests {
684    use super::*;
685    use crate::program::{CoreSemanticOp, ProgramInputSpec, SemanticProgramBuilder};
686    use std::any::Any;
687    use std::hash::Hasher;
688    use std::sync::Arc;
689    use tenferro_ops::{
690        ext_op::{ExtensionAliasDeclaration, ExtensionEffectDeclaration, ExtensionOp},
691        SymDim,
692    };
693
694    #[test]
695    fn compile_publishes_semantic_program_and_separate_default_bindings() {
696        let trace = crate::TraceContext::new();
697        let input = trace
698            .from_vec_col_major(vec![2], vec![1.0_f64, 2.0])
699            .unwrap();
700        let output = input.neg().unwrap();
701
702        let program = GraphCompiler::new().compile(&output).unwrap();
703
704        assert_eq!(program.program().inputs().len(), 1);
705        assert_eq!(program.program().outputs().len(), 1);
706        assert_eq!(program.program().operations().count(), 1);
707        assert_eq!(program.bindings().len(), 1);
708        assert_eq!(
709            program
710                .program()
711                .value_metadata(program.program().inputs()[0])
712                .unwrap()
713                .shape(),
714            &[ShapeExtent::Exact(DimExpr::Const(2))]
715        );
716        assert_eq!(
717            program
718                .program()
719                .value_metadata(program.program().outputs()[0])
720                .unwrap()
721                .dtype(),
722            DType::F64
723        );
724    }
725
726    #[test]
727    fn compiler_options_constructor_keeps_the_given_options() {
728        let options = CompilerOptions {
729            optimizer: crate::OptimizerConfig {
730                dot_decomposer: true,
731                ..crate::OptimizerConfig::default()
732            },
733        };
734        let compiler = GraphCompiler::with_compiler_options(options);
735        assert!(compiler.compiler_options().optimizer.dot_decomposer);
736        assert!(compiler.extension_caches().is_empty());
737    }
738
739    #[test]
740    fn compile_many_refuses_outputs_from_two_contexts() {
741        let first = crate::TraceContext::new();
742        let second = crate::TraceContext::new();
743        let left = first
744            .from_vec_col_major(vec![2], vec![1.0_f64, 2.0])
745            .unwrap()
746            .neg()
747            .unwrap();
748        let right = second
749            .from_vec_col_major(vec![2], vec![1.0_f64, 2.0])
750            .unwrap()
751            .neg()
752            .unwrap();
753
754        let error = GraphCompiler::new()
755            .compile_many(&[&left, &right])
756            .unwrap_err();
757        assert!(error.to_string().contains("different contexts"), "{error}");
758    }
759
760    #[test]
761    fn compile_preserves_symbolic_default_input_extent_identity() {
762        let trace = crate::TraceContext::new();
763        let tensor = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
764        let input = trace.from_tensor_symbolic_shape(tensor).unwrap();
765        let output = input.neg().unwrap();
766
767        let program = GraphCompiler::new().compile(&output).unwrap();
768
769        assert!(matches!(
770            program
771                .program()
772                .value_metadata(program.program().inputs()[0])
773                .unwrap()
774                .shape(),
775            [ShapeExtent::Exact(DimExpr::InputDim {
776                input_idx: 0,
777                axis: 0
778            })]
779        ));
780    }
781
782    #[test]
783    fn compile_frozen_program_rejects_static_unbound_shape_guard_mismatch() {
784        let mut builder = SemanticProgramBuilder::new();
785        let lhs = builder
786            .input(ProgramInputSpec::new(DType::F64, [DimExpr::Const(2)]))
787            .unwrap();
788        let rhs = builder
789            .input(ProgramInputSpec::new(DType::F64, [DimExpr::Const(3)]))
790            .unwrap();
791        let output = builder.add_op(CoreSemanticOp::Neg, &[lhs]).unwrap()[0];
792        builder
793            .add_shape_guards_to_output(
794                output,
795                [ProgramShapeGuard::new(
796                    ProgramShapeRelation::Equal,
797                    DimExpr::InputDim {
798                        input_idx: 0,
799                        axis: 0,
800                    },
801                    DimExpr::InputDim {
802                        input_idx: 1,
803                        axis: 0,
804                    },
805                )],
806            )
807            .unwrap();
808        let frozen = builder.finish(&[output]).unwrap();
809
810        let err = GraphCompiler::new()
811            .compile_frozen_program(&frozen)
812            .unwrap_err();
813
814        assert!(matches!(
815            err,
816            Error::ShapeConstraintViolation {
817                lhs_value: 2,
818                rhs_value: 3,
819                ..
820            }
821        ));
822        let _ = rhs;
823    }
824
825    #[test]
826    fn compile_frozen_program_defers_dynamic_unbound_shape_guard() {
827        let mut builder = SemanticProgramBuilder::new();
828        let lhs = builder
829            .input(ProgramInputSpec::new(DType::F64, [DimExpr::Const(2)]))
830            .unwrap();
831        let rhs = builder
832            .input(ProgramInputSpec::from_metadata(
833                ProgramValueMetadata::from_extents(DType::F64, [ShapeExtent::Unknown]),
834            ))
835            .unwrap();
836        let output = builder.add_op(CoreSemanticOp::Neg, &[lhs]).unwrap()[0];
837        builder
838            .add_shape_guards_to_output(
839                output,
840                [ProgramShapeGuard::new(
841                    ProgramShapeRelation::Equal,
842                    DimExpr::InputDim {
843                        input_idx: 0,
844                        axis: 0,
845                    },
846                    DimExpr::InputDim {
847                        input_idx: 1,
848                        axis: 0,
849                    },
850                )],
851            )
852            .unwrap();
853        let frozen = builder.finish(&[output]).unwrap();
854
855        GraphCompiler::new()
856            .compile_frozen_program(&frozen)
857            .unwrap();
858        let _ = rhs;
859    }
860
861    #[derive(Clone, Debug, PartialEq, Eq)]
862    struct PrunableTestOp {
863        pruned: bool,
864    }
865
866    impl ExtensionOp for PrunableTestOp {
867        fn family_id(&self) -> &'static str {
868            "tenferro-runtime.test-prunable.v1"
869        }
870
871        fn payload_hash(&self, hasher: &mut dyn Hasher) {
872            hasher.write_u8(u8::from(self.pruned));
873        }
874
875        fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
876            other
877                .as_any()
878                .downcast_ref::<Self>()
879                .is_some_and(|that| self == that)
880        }
881
882        fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
883            Arc::new(self.clone())
884        }
885
886        fn as_any(&self) -> &dyn Any {
887            self
888        }
889
890        fn input_count(&self) -> usize {
891            1
892        }
893
894        fn output_count(&self) -> usize {
895            if self.pruned {
896                1
897            } else {
898                3
899            }
900        }
901
902        fn semantic_effects(&self) -> ExtensionEffectDeclaration<'_> {
903            ExtensionEffectDeclaration::Declared(&[])
904        }
905
906        fn semantic_aliases(&self) -> ExtensionAliasDeclaration<'_> {
907            ExtensionAliasDeclaration::AllFresh
908        }
909
910        fn infer_output_meta(
911            &self,
912            ctx: &mut tenferro_ops::ExtensionShapeContext<'_>,
913        ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
914            let dtype = ctx.input_dtype(0)?;
915            let shape = ctx.input_shape(0)?.to_vec();
916            Ok((0..self.output_count())
917                .map(|_| (dtype, shape.clone()))
918                .collect())
919        }
920
921        fn prune_outputs(&self, live_outputs: &[bool]) -> Option<Arc<dyn ExtensionOp>> {
922            (!self.pruned && live_outputs == [false, true, false])
923                .then(|| Arc::new(Self { pruned: true }) as Arc<dyn ExtensionOp>)
924        }
925    }
926
927    #[test]
928    fn compile_prunes_extension_outputs_with_replacement_op() {
929        let trace = crate::TraceContext::new();
930        let input = trace
931            .from_vec_col_major(vec![2], vec![1.0_f64, 2.0])
932            .unwrap();
933        let outputs =
934            crate::extension::apply(Arc::new(PrunableTestOp { pruned: false }), &[&input]).unwrap();
935
936        let program = GraphCompiler::new().compile(&outputs[1]).unwrap();
937        let staging =
938            stage_semantic_program(program.program(), program.compiler_options()).unwrap();
939        let pruned_instruction = staging
940            .instructions
941            .iter()
942            .find_map(|inst| match &inst.op {
943                crate::exec::ExecOp::Extension(op)
944                    if op.family_id() == "tenferro-runtime.test-prunable.v1" =>
945                {
946                    Some((
947                        inst.output_slots.clone(),
948                        format!("{op:?}"),
949                        op.output_count(),
950                    ))
951                }
952                _ => None,
953            })
954            .expect("compiled program should contain the test extension");
955
956        assert_eq!(pruned_instruction.0.len(), 1);
957        assert!(pruned_instruction.1.contains("pruned: true"));
958        assert_eq!(pruned_instruction.2, 1);
959    }
960
961    #[test]
962    fn compiled_graph_input_keys_preserve_order_for_binary_graph() {
963        let trace = crate::TraceContext::new();
964        let a = trace
965            .from_vec_col_major(vec![2], vec![1.0_f64, 2.0])
966            .unwrap();
967        let b = trace
968            .from_vec_col_major(vec![2], vec![3.0_f64, 4.0])
969            .unwrap();
970        let a_key = a.input_key().expect("concrete traced tensor has input key");
971        let b_key = b.input_key().expect("concrete traced tensor has input key");
972        let c = a.mul(&b).unwrap();
973
974        let program = GraphCompiler::new().compile_many(&[&c]).unwrap();
975
976        assert_eq!(program.input_count(), 2);
977        assert_eq!(program.input_keys().len(), 2);
978        let keys: Vec<_> = program.input_keys().to_vec();
979        assert!(keys.contains(&a_key), "input_keys must contain a's key");
980        assert!(keys.contains(&b_key), "input_keys must contain b's key");
981        // input_key_index maps each key to its position
982        assert_eq!(
983            program.input_key_index(&a_key),
984            keys.iter().position(|k| k == &a_key)
985        );
986        assert_eq!(
987            program.input_key_index(&b_key),
988            keys.iter().position(|k| k == &b_key)
989        );
990    }
991}
tenferro_runtime/graph/compiler.rs

tenferro_runtime/graph/
compiler.rs
