tenferro_runtime/graph/
compiler.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::num::NonZeroUsize;
4use std::sync::Arc;
5
6use computegraph::compile::{compile, CompiledProgram};
7use computegraph::materialize::materialize_merge;
8use computegraph::resolve::resolve;
9use computegraph::types::ValueKey;
10use lru::LruCache;
11use num_complex::{Complex32, Complex64};
12use tenferro_ops::dim_expr::DimExpr;
13use tenferro_ops::input_key::TensorInputKey;
14use tenferro_ops::std_tensor_op::StdTensorOp;
15use tenferro_tensor::{DType, Tensor, TensorScalar};
16
17use super::cache::{
18    compile_cache_stats, compute_cache_key, CacheKey, GraphCompilerCacheStats,
19    DEFAULT_COMPILE_CACHE_CAPACITY,
20};
21use super::program::{GraphProgram, GraphProgramInput};
22use crate::compiler::{compile_std_to_exec_with_options, CompilerOptions};
23use crate::error::{Error, Result};
24use crate::exec::ExecProgram;
25use crate::extension_cache::{ExtensionCacheSelector, ExtensionCacheStore};
26use crate::traced::{try_concrete_shape, TracedTensor};
27
28#[derive(Clone)]
29struct InputDescriptor {
30    key: TensorInputKey,
31    dtype: DType,
32    shape: Vec<usize>,
33    default_tensor: Option<Arc<Tensor>>,
34}
35
36/// Compiler for traced tensor graphs.
37///
38/// A graph compiler lowers one or more [`TracedTensor`] outputs to a reusable
39/// [`GraphProgram`] without requiring a backend.
40///
41/// # Examples
42///
43/// ```
44/// use tenferro_runtime::{GraphCompiler, TracedTensor};
45///
46/// let x = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
47/// let y = (&x + &x).unwrap();
48/// let mut compiler = GraphCompiler::new();
49/// let program = compiler.compile(&y).unwrap();
50/// assert_eq!(program.output_count(), 1);
51/// ```
52pub struct GraphCompiler {
53    compile_cache: LruCache<CacheKey, ExecProgram>,
54    extension_cache: ExtensionCacheStore,
55    compiler_options: CompilerOptions,
56}
57
58impl fmt::Debug for GraphCompiler {
59    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
60        f.debug_struct("GraphCompiler")
61            .field("cache_stats", &self.cache_stats())
62            .field("compile_cache_capacity", &self.compile_cache_capacity())
63            .field("compiler_options", &self.compiler_options)
64            .field("extension_cache", &self.extension_cache)
65            .finish_non_exhaustive()
66    }
67}
68
69impl GraphCompiler {
70    /// Create a compiler with bounded default caches.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use tenferro_runtime::GraphCompiler;
76    ///
77    /// let compiler = GraphCompiler::new();
78    /// assert_eq!(compiler.compile_cache_len(), 0);
79    /// ```
80    pub fn new() -> Self {
81        Self {
82            compile_cache: LruCache::new(
83                NonZeroUsize::new(DEFAULT_COMPILE_CACHE_CAPACITY).unwrap_or(NonZeroUsize::MIN),
84            ),
85            extension_cache: ExtensionCacheStore::new(),
86            compiler_options: CompilerOptions::default(),
87        }
88    }
89
90    /// Create a compiler with explicit lowering and optimizer options.
91    ///
92    /// # Examples
93    ///
94    /// ```
95    /// use tenferro_runtime::{CompilerOptions, OptimizerConfig};
96    /// use tenferro_runtime::GraphCompiler;
97    ///
98    /// let compiler = GraphCompiler::with_compiler_options(CompilerOptions {
99    ///     optimizer: OptimizerConfig {
100    ///         dot_decomposer: true,
101    ///         ..OptimizerConfig::default()
102    ///     },
103    /// });
104    /// assert!(compiler.compiler_options().optimizer.dot_decomposer);
105    /// ```
106    pub fn with_compiler_options(compiler_options: CompilerOptions) -> Self {
107        Self {
108            compile_cache: LruCache::new(
109                NonZeroUsize::new(DEFAULT_COMPILE_CACHE_CAPACITY).unwrap_or(NonZeroUsize::MIN),
110            ),
111            extension_cache: ExtensionCacheStore::new(),
112            compiler_options,
113        }
114    }
115
116    /// Compile one traced output into a graph program.
117    ///
118    /// # Examples
119    ///
120    /// ```
121    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
122    ///
123    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
124    /// let mut compiler = GraphCompiler::new();
125    /// let y = x.neg().unwrap();
126    /// let program = compiler.compile(&y).unwrap();
127    /// assert_eq!(program.input_count(), 1);
128    /// ```
129    pub fn compile(&mut self, output: &TracedTensor) -> Result<GraphProgram> {
130        self.compile_many(&[output])
131    }
132
133    /// Compile multiple traced outputs into one graph program.
134    ///
135    /// # Examples
136    ///
137    /// ```
138    /// use tenferro_runtime::{GraphCompiler, TracedTensor};
139    ///
140    /// let x = TracedTensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap();
141    /// let y = x.neg().unwrap();
142    /// let mut compiler = GraphCompiler::new();
143    /// let program = compiler.compile_many(&[&x, &y]).unwrap();
144    /// assert_eq!(program.output_count(), 2);
145    /// ```
146    pub fn compile_many(&mut self, outputs: &[&TracedTensor]) -> Result<GraphProgram> {
147        let mut all_inputs = HashMap::new();
148        for output in outputs {
149            for (key, tensor) in output.inputs_map.iter() {
150                if let Some(existing) = all_inputs.get(key) {
151                    if !default_tensors_equivalent(existing, tensor) {
152                        return Err(Error::DuplicateBinding {
153                            input_key: format!("{:?}", key),
154                        });
155                    }
156                    continue;
157                }
158                all_inputs.insert(key.clone(), tensor.clone());
159            }
160        }
161        self.compile_many_with_descriptors(outputs, &HashMap::new(), &all_inputs)
162    }
163
164    /// Compile one traced output with concrete placeholder specs.
165    ///
166    /// # Examples
167    ///
168    /// ```
169    /// use tenferro_runtime::{DType, GraphCompiler, TracedTensor};
170    ///
171    /// let x = TracedTensor::input_symbolic_shape(DType::F64, 1).unwrap();
172    /// let mut compiler = GraphCompiler::new();
173    /// let y = x.neg().unwrap();
174    /// let program = compiler
175    ///     .compile_with_input_specs(&y, &[(&x, DType::F64, &[3])])
176    ///     .unwrap();
177    /// assert_eq!(program.input_specs()[0].shape(), &[3]);
178    /// ```
179    pub fn compile_with_input_specs(
180        &mut self,
181        output: &TracedTensor,
182        bindings: &[(&TracedTensor, DType, &[usize])],
183    ) -> Result<GraphProgram> {
184        let mut binding_specs = HashMap::new();
185        for (index, (placeholder, dtype, shape)) in bindings.iter().enumerate() {
186            validate_placeholder_spec(index, placeholder, *dtype, shape)?;
187            let key = placeholder.input_key().ok_or(Error::UnexpectedBinding {
188                binding_index: index,
189            })?;
190            if binding_specs
191                .insert(
192                    key.clone(),
193                    InputDescriptor {
194                        key: key.clone(),
195                        dtype: *dtype,
196                        shape: (*shape).to_vec(),
197                        default_tensor: None,
198                    },
199                )
200                .is_some()
201            {
202                return Err(Error::DuplicateBinding {
203                    input_key: format!("{:?}", key),
204                });
205            }
206        }
207
208        self.compile_many_with_descriptors(&[output], &binding_specs, output.inputs_map.as_ref())
209    }
210
211    /// Number of compiled programs currently retained.
212    ///
213    /// # Examples
214    ///
215    /// ```
216    /// use tenferro_runtime::GraphCompiler;
217    ///
218    /// let compiler = GraphCompiler::new();
219    /// assert_eq!(compiler.compile_cache_len(), 0);
220    /// ```
221    pub fn compile_cache_len(&self) -> usize {
222        self.compile_cache.len()
223    }
224
225    /// Current compiled-program cache capacity.
226    ///
227    /// # Examples
228    ///
229    /// ```
230    /// use tenferro_runtime::GraphCompiler;
231    ///
232    /// let compiler = GraphCompiler::new();
233    /// assert!(compiler.compile_cache_capacity().get() > 0);
234    /// ```
235    pub fn compile_cache_capacity(&self) -> NonZeroUsize {
236        self.compile_cache.cap()
237    }
238
239    /// Resize the compiled-program cache.
240    ///
241    /// # Examples
242    ///
243    /// ```
244    /// use std::num::NonZeroUsize;
245    /// use tenferro_runtime::GraphCompiler;
246    ///
247    /// let mut compiler = GraphCompiler::new();
248    /// compiler.set_compile_cache_capacity(NonZeroUsize::new(2).unwrap());
249    /// assert_eq!(compiler.compile_cache_capacity().get(), 2);
250    /// ```
251    pub fn set_compile_cache_capacity(&mut self, capacity: NonZeroUsize) {
252        self.compile_cache.resize(capacity);
253    }
254
255    /// Return the compiler options used for future graph lowerings.
256    ///
257    /// # Examples
258    ///
259    /// ```
260    /// use tenferro_runtime::CompilerOptions;
261    /// use tenferro_runtime::GraphCompiler;
262    ///
263    /// let compiler = GraphCompiler::new();
264    /// assert_eq!(compiler.compiler_options(), CompilerOptions::default());
265    /// ```
266    pub fn compiler_options(&self) -> CompilerOptions {
267        self.compiler_options
268    }
269
270    /// Replace compiler options and clear compiled graph cache entries.
271    ///
272    /// # Examples
273    ///
274    /// ```
275    /// use tenferro_runtime::{CompilerOptions, OptimizerConfig};
276    /// use tenferro_runtime::GraphCompiler;
277    ///
278    /// let mut compiler = GraphCompiler::new();
279    /// let options = CompilerOptions {
280    ///     optimizer: OptimizerConfig {
281    ///         dot_decomposer: true,
282    ///         ..OptimizerConfig::default()
283    ///     },
284    /// };
285    /// compiler.set_compiler_options(options);
286    /// assert_eq!(compiler.compiler_options(), options);
287    /// assert_eq!(compiler.compile_cache_len(), 0);
288    /// ```
289    pub fn set_compiler_options(&mut self, compiler_options: CompilerOptions) {
290        if self.compiler_options == compiler_options {
291            return;
292        }
293        self.compiler_options = compiler_options;
294        self.clear_compile_cache();
295    }
296
297    /// Clear the compiled-program cache.
298    ///
299    /// # Examples
300    ///
301    /// ```
302    /// use tenferro_runtime::GraphCompiler;
303    ///
304    /// let mut compiler = GraphCompiler::new();
305    /// compiler.clear_compile_cache();
306    /// assert_eq!(compiler.compile_cache_len(), 0);
307    /// ```
308    pub fn clear_compile_cache(&mut self) {
309        self.compile_cache.clear();
310    }
311
312    /// Clear generic extension compile-time cache entries.
313    ///
314    /// # Examples
315    ///
316    /// ```
317    /// use tenferro_runtime::GraphCompiler;
318    ///
319    /// let mut compiler = GraphCompiler::new();
320    /// compiler.clear_extension_caches();
321    /// assert_eq!(compiler.cache_stats().extensions.entries, 0);
322    /// ```
323    pub fn clear_extension_caches(&mut self) {
324        self.extension_cache.clear();
325    }
326
327    /// Clear every cache owned by the compiler.
328    ///
329    /// # Examples
330    ///
331    /// ```
332    /// use tenferro_runtime::GraphCompiler;
333    ///
334    /// let mut compiler = GraphCompiler::new();
335    /// compiler.clear_caches();
336    /// assert_eq!(compiler.cache_stats().compile.entries, 0);
337    /// ```
338    pub fn clear_caches(&mut self) {
339        self.clear_compile_cache();
340        self.clear_extension_caches();
341    }
342
343    /// Return cache-entry and retained-byte stats.
344    ///
345    /// # Examples
346    ///
347    /// ```
348    /// use tenferro_runtime::GraphCompiler;
349    ///
350    /// let compiler = GraphCompiler::new();
351    /// let stats = compiler.cache_stats();
352    /// assert_eq!(stats.compile.entries, 0);
353    /// ```
354    pub fn cache_stats(&self) -> GraphCompilerCacheStats {
355        GraphCompilerCacheStats {
356            compile: compile_cache_stats(&self.compile_cache),
357            extensions: self.extension_cache.stats(ExtensionCacheSelector::All),
358        }
359    }
360
361    /// Borrow generic extension compile-time cache storage.
362    ///
363    /// # Examples
364    ///
365    /// ```
366    /// use tenferro_runtime::GraphCompiler;
367    ///
368    /// let compiler = GraphCompiler::new();
369    /// assert!(compiler.extension_caches().is_empty());
370    /// ```
371    pub fn extension_caches(&self) -> &ExtensionCacheStore {
372        &self.extension_cache
373    }
374
375    /// Mutably borrow generic extension compile-time cache storage.
376    ///
377    /// # Examples
378    ///
379    /// ```
380    /// use tenferro_runtime::GraphCompiler;
381    ///
382    /// let mut compiler = GraphCompiler::new();
383    /// compiler.extension_caches_mut().clear();
384    /// ```
385    pub fn extension_caches_mut(&mut self) -> &mut ExtensionCacheStore {
386        &mut self.extension_cache
387    }
388
389    fn compile_many_with_descriptors(
390        &mut self,
391        outputs: &[&TracedTensor],
392        binding_specs: &HashMap<TensorInputKey, InputDescriptor>,
393        default_inputs: &HashMap<TensorInputKey, Arc<Tensor>>,
394    ) -> Result<GraphProgram> {
395        let mut roots = Vec::new();
396        let mut output_keys = Vec::with_capacity(outputs.len());
397        for output in outputs {
398            roots.extend(output.resolve_roots());
399            output_keys.push(output.graph.values()[output.val].key.clone());
400        }
401
402        let view = resolve(roots);
403        let graph = materialize_merge(&view, &output_keys);
404        let mut compiled = compile(&graph);
405        prune_compiled_extension_outputs(&mut compiled)?;
406
407        let mut descriptors = Vec::with_capacity(graph.inputs.len());
408        let mut input_dtypes = Vec::with_capacity(graph.inputs.len());
409        let mut input_shapes = Vec::with_capacity(graph.inputs.len());
410        for key in &graph.inputs {
411            let ValueKey::Input(input_key) = key else {
412                return Err(Error::Internal(
413                    "expected Input key in graph inputs".to_string(),
414                ));
415            };
416            let descriptor = descriptor_for_input(input_key, binding_specs, default_inputs)?;
417            input_dtypes.push(descriptor.dtype);
418            input_shapes.push(DimExpr::from_concrete(&descriptor.shape));
419            descriptors.push(GraphProgramInput::new(
420                descriptor.key,
421                descriptor.dtype,
422                descriptor.shape.clone(),
423                DimExpr::from_concrete(&descriptor.shape),
424                descriptor.default_tensor,
425            ));
426        }
427
428        let exec = compile_std_to_exec_with_options(
429            &compiled,
430            &input_dtypes,
431            &input_shapes,
432            self.compiler_options,
433        )?;
434        let exec = self.get_or_compile(exec);
435        Ok(GraphProgram::new(exec, descriptors))
436    }
437
438    fn get_or_compile(&mut self, exec: ExecProgram) -> ExecProgram {
439        let key = compute_cache_key(&exec);
440        if let Some(cached) = self.compile_cache.get(&key) {
441            return cached.clone();
442        }
443        self.compile_cache.put(key, exec.clone());
444        exec
445    }
446}
447
448impl Default for GraphCompiler {
449    fn default() -> Self {
450        Self::new()
451    }
452}
453
454fn validate_placeholder_spec(
455    index: usize,
456    placeholder: &TracedTensor,
457    dtype: DType,
458    shape: &[usize],
459) -> Result<()> {
460    if placeholder.data.is_some() {
461        return Err(Error::UnexpectedBinding {
462            binding_index: index,
463        });
464    }
465    placeholder.input_key().ok_or(Error::UnexpectedBinding {
466        binding_index: index,
467    })?;
468
469    if placeholder.dtype != dtype {
470        return Err(Error::PlaceholderDtypeMismatch {
471            expected: placeholder.dtype,
472            actual: dtype,
473        });
474    }
475    validate_placeholder_shape(placeholder, shape)
476}
477
478fn validate_placeholder_shape(placeholder: &TracedTensor, shape: &[usize]) -> Result<()> {
479    match try_concrete_shape(placeholder) {
480        Some(expected_shape) => {
481            if expected_shape.as_slice() != shape {
482                return Err(Error::PlaceholderShapeMismatch {
483                    expected: expected_shape,
484                    actual: shape.to_vec(),
485                });
486            }
487        }
488        None => {
489            if placeholder.rank != shape.len() {
490                return Err(Error::PlaceholderRankMismatch {
491                    expected: placeholder.rank,
492                    actual: shape.len(),
493                });
494            }
495        }
496    }
497    Ok(())
498}
499
500fn descriptor_for_input(
501    key: &TensorInputKey,
502    binding_specs: &HashMap<TensorInputKey, InputDescriptor>,
503    default_inputs: &HashMap<TensorInputKey, Arc<Tensor>>,
504) -> Result<InputDescriptor> {
505    if let Some(tensor) = default_inputs.get(key) {
506        return Ok(InputDescriptor {
507            key: key.clone(),
508            dtype: tensor.dtype(),
509            shape: tensor.shape().to_vec(),
510            default_tensor: Some(tensor.clone()),
511        });
512    }
513    if let Some(spec) = binding_specs.get(key) {
514        return Ok(spec.clone());
515    }
516    Err(Error::UnboundPlaceholder {
517        input_key: format!("{:?}", key),
518    })
519}
520
521fn prune_compiled_extension_outputs(prog: &mut CompiledProgram<StdTensorOp>) -> Result<()> {
522    let mut live_slots = vec![false; prog.n_slots];
523    for &slot in &prog.output_slots {
524        let Some(live) = live_slots.get_mut(slot) else {
525            return Err(invalid_compiled_graph(format!(
526                "program output slot {slot} is outside slot table of length {}",
527                prog.n_slots
528            )));
529        };
530        *live = true;
531    }
532
533    for instr in prog.instructions.iter_mut().rev() {
534        let live_outputs = instr
535            .outputs
536            .iter()
537            .map(|&slot| {
538                live_slots.get(slot).copied().ok_or_else(|| {
539                    invalid_compiled_graph(format!(
540                        "instruction output slot {slot} is outside slot table of length {}",
541                        prog.n_slots
542                    ))
543                })
544            })
545            .collect::<Result<Vec<_>>>()?;
546
547        if let StdTensorOp::Extension(ext) = &instr.operation {
548            if let Some(pruned) = ext.prune_outputs(&live_outputs) {
549                let kept_outputs = instr
550                    .outputs
551                    .iter()
552                    .zip(live_outputs.iter())
553                    .filter_map(|(&slot, &live)| live.then_some(slot))
554                    .collect::<Vec<_>>();
555                if pruned.output_count() != kept_outputs.len() {
556                    return Err(invalid_compiled_graph(format!(
557                        "extension family_id={:?} pruned to {} outputs for {} live slots",
558                        ext.family_id(),
559                        pruned.output_count(),
560                        kept_outputs.len()
561                    )));
562                }
563                instr.operation = StdTensorOp::Extension(pruned);
564                instr.outputs = kept_outputs;
565            }
566        }
567
568        if live_outputs.iter().any(|&live| live) {
569            for &slot in &instr.inputs {
570                let Some(live) = live_slots.get_mut(slot) else {
571                    return Err(invalid_compiled_graph(format!(
572                        "instruction input slot {slot} is outside slot table of length {}",
573                        prog.n_slots
574                    )));
575                };
576                *live = true;
577            }
578        }
579    }
580
581    Ok(())
582}
583
584fn invalid_compiled_graph(message: impl Into<String>) -> Error {
585    Error::InvalidCompiledGraph {
586        message: message.into(),
587    }
588}
589
590fn default_tensors_equivalent(lhs: &Arc<Tensor>, rhs: &Arc<Tensor>) -> bool {
591    if Arc::ptr_eq(lhs, rhs) {
592        return true;
593    }
594    if lhs.dtype() != rhs.dtype() || lhs.shape() != rhs.shape() {
595        return false;
596    }
597    match lhs.dtype() {
598        DType::F32 => default_slices_equivalent::<f32>(lhs, rhs),
599        DType::F64 => default_slices_equivalent::<f64>(lhs, rhs),
600        DType::I32 => default_slices_equivalent::<i32>(lhs, rhs),
601        DType::I64 => default_slices_equivalent::<i64>(lhs, rhs),
602        DType::Bool => default_slices_equivalent::<bool>(lhs, rhs),
603        DType::C32 => default_slices_equivalent::<Complex32>(lhs, rhs),
604        DType::C64 => default_slices_equivalent::<Complex64>(lhs, rhs),
605    }
606}
607
608fn default_slices_equivalent<T: TensorScalar + PartialEq>(lhs: &Tensor, rhs: &Tensor) -> bool {
609    match (lhs.as_slice::<T>(), rhs.as_slice::<T>()) {
610        (Ok(lhs), Ok(rhs)) => lhs == rhs,
611        // Backend-resident defaults cannot be inspected here; only the same
612        // Arc<Tensor> is considered equivalent by `default_tensors_equivalent`.
613        _ => false,
614    }
615}
616
617#[cfg(test)]
618mod tests {
619    use super::*;
620    use std::any::Any;
621    use std::hash::Hasher;
622    use std::sync::Arc;
623    use tenferro_ops::{ext_op::ExtensionOp, SymDim};
624    use tenferro_tensor::{
625        Buffer, BufferHandle, DeviceId, DeviceKind, GpuBackendKind, MemoryKind, Placement,
626        TypedTensor,
627    };
628
629    #[test]
630    fn compile_many_rejects_conflicting_default_inputs_for_same_key() {
631        let x = TracedTensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap();
632        let y1 = x.neg().unwrap();
633        let mut y2 = x.neg().unwrap();
634        let key = x.input_key().expect("concrete traced tensor has input key");
635        let replacement = Arc::new(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap());
636        let mut inputs = (*y2.inputs_map).clone();
637        inputs.insert(key.clone(), replacement);
638        y2.inputs_map = Arc::new(inputs);
639
640        let err = GraphCompiler::new().compile_many(&[&y1, &y2]).unwrap_err();
641
642        assert!(matches!(
643            err,
644            Error::DuplicateBinding { ref input_key } if input_key.contains(&format!("{key:?}"))
645        ));
646    }
647
648    #[test]
649    fn default_tensors_equivalent_rejects_distinct_backend_buffers() {
650        let placement = Placement {
651            memory_kind: MemoryKind::Device,
652            device: Some(DeviceId {
653                kind: DeviceKind::Gpu(GpuBackendKind::Cuda),
654                ordinal: 0,
655            }),
656        };
657        let lhs = Arc::new(Tensor::F64(
658            TypedTensor::from_buffer_col_major(
659                vec![2],
660                Buffer::Backend(Arc::new(BufferHandle::<f64>::new_with_len(1, 2))),
661                placement.clone(),
662            )
663            .unwrap(),
664        ));
665        let rhs = Arc::new(Tensor::F64(
666            TypedTensor::from_buffer_col_major(
667                vec![2],
668                Buffer::Backend(Arc::new(BufferHandle::<f64>::new_with_len(2, 2))),
669                placement,
670            )
671            .unwrap(),
672        ));
673
674        assert!(
675            !default_tensors_equivalent(&lhs, &rhs),
676            "distinct backend-resident default tensors must not compare equal just because both are unreadable on host"
677        );
678        assert!(default_tensors_equivalent(&lhs, &lhs));
679    }
680
681    #[derive(Clone, Debug, PartialEq, Eq)]
682    struct PrunableTestOp {
683        pruned: bool,
684    }
685
686    impl ExtensionOp for PrunableTestOp {
687        fn family_id(&self) -> &'static str {
688            "tenferro-runtime.test-prunable.v1"
689        }
690
691        fn payload_hash(&self, hasher: &mut dyn Hasher) {
692            hasher.write_u8(u8::from(self.pruned));
693        }
694
695        fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
696            other
697                .as_any()
698                .downcast_ref::<Self>()
699                .is_some_and(|that| self == that)
700        }
701
702        fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
703            Arc::new(self.clone())
704        }
705
706        fn as_any(&self) -> &dyn Any {
707            self
708        }
709
710        fn input_count(&self) -> usize {
711            1
712        }
713
714        fn output_count(&self) -> usize {
715            if self.pruned {
716                1
717            } else {
718                3
719            }
720        }
721
722        fn infer_output_meta(
723            &self,
724            input_dtypes: &[DType],
725            input_shapes: &[&[SymDim]],
726        ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
727            Ok((0..self.output_count())
728                .map(|_| (input_dtypes[0], input_shapes[0].to_vec()))
729                .collect())
730        }
731
732        fn prune_outputs(&self, live_outputs: &[bool]) -> Option<Arc<dyn ExtensionOp>> {
733            (!self.pruned && live_outputs == [false, true, false])
734                .then(|| Arc::new(Self { pruned: true }) as Arc<dyn ExtensionOp>)
735        }
736    }
737
738    #[test]
739    fn compile_prunes_extension_outputs_with_replacement_op() {
740        let input = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
741        let outputs =
742            crate::extension::apply(Arc::new(PrunableTestOp { pruned: false }), &[&input]).unwrap();
743
744        let program = GraphCompiler::new().compile(&outputs[1]).unwrap();
745        let pruned_instruction = program
746            .lowering_view()
747            .instructions()
748            .find_map(|inst| match inst.op() {
749                crate::GraphOpView::Extension { op }
750                    if op.family_id() == "tenferro-runtime.test-prunable.v1" =>
751                {
752                    Some((
753                        inst.output_slots().to_vec(),
754                        format!("{op:?}"),
755                        op.output_count(),
756                    ))
757                }
758                _ => None,
759            })
760            .expect("compiled program should contain the test extension");
761
762        assert_eq!(pruned_instruction.0.len(), 1);
763        assert!(pruned_instruction.1.contains("pruned: true"));
764        assert_eq!(pruned_instruction.2, 1);
765    }
766}
tenferro_runtime/graph/compiler.rs

tenferro_runtime/graph/
compiler.rs
