tenferro_runtime/graph/
cache.rs

1use std::hash::{Hash, Hasher};
2use std::mem::{size_of, size_of_val};
3use std::sync::Arc;
4
5use lru::LruCache;
6use tenferro_ops::ext_op::ExtensionOp;
7use tenferro_tensor::CacheStats;
8
9use crate::exec::{ExecInstruction, ExecOp, ExecOutputExtents, ExecOutputShapes, ExecProgram};
10
11/// Default capacity for compiled graph programs retained by a [`GraphCompiler`](super::GraphCompiler).
12// Public constant kept as the documented default; the crate-local alias below
13// is what current implementation paths consume.
14#[allow(dead_code)]
15pub const DEFAULT_GRAPH_COMPILE_CACHE_CAPACITY: usize = 256;
16
17/// Internal alias matching the existing engine cache helper name.
18pub(crate) const DEFAULT_COMPILE_CACHE_CAPACITY: usize = DEFAULT_GRAPH_COMPILE_CACHE_CAPACITY;
19
20/// Stats for caches owned by a [`GraphCompiler`](super::GraphCompiler).
21///
22/// `retained_bytes` fields are logical payload estimates, not process RSS.
23///
24/// # Examples
25///
26/// ```
27/// use tenferro_runtime::{CacheStats, GraphCompilerCacheStats};
28///
29/// let stats = GraphCompilerCacheStats {
30///     compile: CacheStats::empty(),
31///     extensions: CacheStats::empty(),
32/// };
33/// assert_eq!(stats.compile.entries, 0);
34/// ```
35#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
36pub struct GraphCompilerCacheStats {
37    /// Compiled execution-program cache.
38    pub compile: CacheStats,
39    /// Generic extension compile-time caches.
40    pub extensions: CacheStats,
41}
42
43/// Stats for runtime caches owned by a graph executor.
44///
45/// `retained_bytes` fields are logical payload estimates, not process RSS.
46///
47/// # Examples
48///
49/// ```
50/// use tenferro_runtime::{CacheStats, GraphExecutorCacheStats};
51///
52/// let stats = GraphExecutorCacheStats {
53///     extensions: CacheStats::empty(),
54///     backend: CacheStats::empty(),
55/// };
56/// assert_eq!(stats.extensions.entries, 0);
57/// ```
58#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
59pub struct GraphExecutorCacheStats {
60    /// Generic extension runtime caches.
61    pub extensions: CacheStats,
62    /// Backend-specific runtime analysis cache.
63    pub backend: CacheStats,
64}
65
66/// Cache key derived from compiled graph topology and execution metadata.
67#[derive(Clone, Debug)]
68pub(crate) struct CacheKey {
69    fingerprint: ExecProgramKey,
70    extensions: Vec<Arc<dyn ExtensionOp>>,
71}
72
73impl PartialEq for CacheKey {
74    fn eq(&self, other: &Self) -> bool {
75        self.fingerprint == other.fingerprint
76            && self.extensions.len() == other.extensions.len()
77            && self
78                .extensions
79                .iter()
80                .zip(&other.extensions)
81                .all(|(lhs, rhs)| {
82                    lhs.family_id() == rhs.family_id() && lhs.payload_eq(rhs.as_ref())
83                })
84    }
85}
86
87impl Eq for CacheKey {}
88
89impl Hash for CacheKey {
90    fn hash<H: Hasher>(&self, state: &mut H) {
91        self.fingerprint.hash(state);
92    }
93}
94
95pub(crate) fn compute_cache_key(exec: &ExecProgram) -> CacheKey {
96    let mut extensions = Vec::new();
97    let fingerprint = exec_program_key(exec, &mut extensions);
98    CacheKey {
99        fingerprint,
100        extensions,
101    }
102}
103
104fn cache_key_retained_bytes(key: &CacheKey) -> usize {
105    saturating_sum([
106        size_of::<CacheKey>(),
107        exec_program_key_retained_bytes(&key.fingerprint),
108        key.extensions
109            .capacity()
110            .saturating_mul(size_of::<Arc<dyn ExtensionOp>>()),
111    ])
112}
113
114#[derive(Clone, Debug, Hash, PartialEq, Eq)]
115struct ExecProgramKey {
116    instructions: Vec<ExecInstructionKey>,
117    input_slots: Vec<usize>,
118    output_slots: Vec<usize>,
119    n_slots: usize,
120}
121
122#[derive(Clone, Debug, Hash, PartialEq, Eq)]
123struct ExecInstructionKey {
124    op: ExecOpKey,
125    input_slots: Vec<usize>,
126    output_slots: Vec<usize>,
127    dtype: tenferro_tensor::DType,
128    output_shapes: ExecOutputShapes,
129    output_extents: ExecOutputExtents,
130    last_use: Vec<bool>,
131}
132
133#[derive(Clone, Debug, Hash, PartialEq, Eq)]
134enum ExecOpKey {
135    Transpose {
136        perm: Vec<usize>,
137    },
138    Reshape {
139        shape: Vec<tenferro_ops::dim_expr::DimExpr>,
140    },
141    BroadcastInDim {
142        shape: Vec<tenferro_ops::dim_expr::DimExpr>,
143        dims: Vec<usize>,
144    },
145    Convert {
146        to: tenferro_tensor::DType,
147    },
148    Constant {
149        dtype: tenferro_tensor::DType,
150        bytes: Vec<u8>,
151    },
152    DotGeneral(tenferro_tensor::DotGeneralConfig),
153    DotGeneralWithConj {
154        config: tenferro_tensor::DotGeneralConfig,
155        lhs_conj: bool,
156        rhs_conj: bool,
157    },
158    ReduceSum {
159        axes: Vec<usize>,
160    },
161    ExtractDiag {
162        axis_a: usize,
163        axis_b: usize,
164    },
165    EmbedDiag {
166        axis_a: usize,
167        axis_b: usize,
168    },
169    Tril {
170        k: i64,
171    },
172    Triu {
173        k: i64,
174    },
175    Add,
176    Subtract,
177    Multiply,
178    Negate,
179    Conj,
180    Divide,
181    Remainder,
182    Abs,
183    Sign,
184    Maximum,
185    Minimum,
186    Compare(tenferro_tensor::CompareDir),
187    Select,
188    Clamp,
189    Exp,
190    Log,
191    Sin,
192    Cos,
193    Tanh,
194    Sqrt,
195    Rsqrt,
196    Pow,
197    Expm1,
198    Log1p,
199    Gather(tenferro_tensor::GatherConfig),
200    GatherDynamicSliceSizes {
201        offset_dims: Vec<usize>,
202        collapsed_slice_dims: Vec<usize>,
203        start_index_map: Vec<usize>,
204        index_vector_dim: usize,
205        slice_sizes: Vec<tenferro_ops::dim_expr::DimExpr>,
206    },
207    Scatter(tenferro_tensor::ScatterConfig),
208    Slice(tenferro_tensor::SliceConfig),
209    DynamicSlice {
210        slice_sizes: Vec<usize>,
211    },
212    DynamicUpdateSlice,
213    Pad(tenferro_tensor::PadConfig),
214    Concatenate {
215        axis: usize,
216    },
217    Reverse {
218        axes: Vec<usize>,
219    },
220    ShapeOf {
221        axis: usize,
222    },
223    DynamicTruncate {
224        axis: usize,
225    },
226    PadToMatch {
227        axis: usize,
228    },
229    ReduceProd {
230        axes: Vec<usize>,
231    },
232    ReduceMax {
233        axes: Vec<usize>,
234    },
235    ReduceMin {
236        axes: Vec<usize>,
237    },
238    Extension {
239        family_id: &'static str,
240        payload_hash: u64,
241    },
242}
243
244fn exec_program_key(
245    exec: &ExecProgram,
246    extensions: &mut Vec<Arc<dyn ExtensionOp>>,
247) -> ExecProgramKey {
248    ExecProgramKey {
249        instructions: exec
250            .instructions
251            .iter()
252            .map(|inst| exec_instruction_key(inst, extensions))
253            .collect(),
254        input_slots: exec.input_slots.clone(),
255        output_slots: exec.output_slots.clone(),
256        n_slots: exec.n_slots,
257    }
258}
259
260fn exec_instruction_key(
261    inst: &ExecInstruction,
262    extensions: &mut Vec<Arc<dyn ExtensionOp>>,
263) -> ExecInstructionKey {
264    ExecInstructionKey {
265        op: exec_op_key(&inst.op, extensions),
266        input_slots: inst.input_slots.clone(),
267        output_slots: inst.output_slots.clone(),
268        dtype: inst.dtype,
269        output_shapes: inst.output_shapes.clone(),
270        output_extents: inst.output_extents.clone(),
271        last_use: inst.last_use.clone(),
272    }
273}
274
275fn exec_op_key(op: &ExecOp, extensions: &mut Vec<Arc<dyn ExtensionOp>>) -> ExecOpKey {
276    match op {
277        ExecOp::Transpose { perm } => ExecOpKey::Transpose { perm: perm.clone() },
278        ExecOp::Reshape { shape } => ExecOpKey::Reshape {
279            shape: shape.clone(),
280        },
281        ExecOp::BroadcastInDim { shape, dims } => ExecOpKey::BroadcastInDim {
282            shape: shape.clone(),
283            dims: dims.clone(),
284        },
285        ExecOp::Convert { to } => ExecOpKey::Convert { to: *to },
286        ExecOp::Constant { dtype, bytes } => ExecOpKey::Constant {
287            dtype: *dtype,
288            bytes: bytes.clone(),
289        },
290        ExecOp::DotGeneral(config) => ExecOpKey::DotGeneral(config.clone()),
291        ExecOp::DotGeneralWithConj {
292            config,
293            lhs_conj,
294            rhs_conj,
295        } => ExecOpKey::DotGeneralWithConj {
296            config: config.clone(),
297            lhs_conj: *lhs_conj,
298            rhs_conj: *rhs_conj,
299        },
300        ExecOp::ReduceSum { axes } => ExecOpKey::ReduceSum { axes: axes.clone() },
301        ExecOp::ExtractDiag { axis_a, axis_b } => ExecOpKey::ExtractDiag {
302            axis_a: *axis_a,
303            axis_b: *axis_b,
304        },
305        ExecOp::EmbedDiag { axis_a, axis_b } => ExecOpKey::EmbedDiag {
306            axis_a: *axis_a,
307            axis_b: *axis_b,
308        },
309        ExecOp::Tril { k } => ExecOpKey::Tril { k: *k },
310        ExecOp::Triu { k } => ExecOpKey::Triu { k: *k },
311        ExecOp::Add => ExecOpKey::Add,
312        ExecOp::Subtract => ExecOpKey::Subtract,
313        ExecOp::Multiply => ExecOpKey::Multiply,
314        ExecOp::Negate => ExecOpKey::Negate,
315        ExecOp::Conj => ExecOpKey::Conj,
316        ExecOp::Divide => ExecOpKey::Divide,
317        ExecOp::Remainder => ExecOpKey::Remainder,
318        ExecOp::Abs => ExecOpKey::Abs,
319        ExecOp::Sign => ExecOpKey::Sign,
320        ExecOp::Maximum => ExecOpKey::Maximum,
321        ExecOp::Minimum => ExecOpKey::Minimum,
322        ExecOp::Compare(dir) => ExecOpKey::Compare(dir.clone()),
323        ExecOp::Select => ExecOpKey::Select,
324        ExecOp::Clamp => ExecOpKey::Clamp,
325        ExecOp::Exp => ExecOpKey::Exp,
326        ExecOp::Log => ExecOpKey::Log,
327        ExecOp::Sin => ExecOpKey::Sin,
328        ExecOp::Cos => ExecOpKey::Cos,
329        ExecOp::Tanh => ExecOpKey::Tanh,
330        ExecOp::Sqrt => ExecOpKey::Sqrt,
331        ExecOp::Rsqrt => ExecOpKey::Rsqrt,
332        ExecOp::Pow => ExecOpKey::Pow,
333        ExecOp::Expm1 => ExecOpKey::Expm1,
334        ExecOp::Log1p => ExecOpKey::Log1p,
335        ExecOp::Gather(config) => ExecOpKey::Gather(config.clone()),
336        ExecOp::GatherDynamicSliceSizes {
337            offset_dims,
338            collapsed_slice_dims,
339            start_index_map,
340            index_vector_dim,
341            slice_sizes,
342        } => ExecOpKey::GatherDynamicSliceSizes {
343            offset_dims: offset_dims.clone(),
344            collapsed_slice_dims: collapsed_slice_dims.clone(),
345            start_index_map: start_index_map.clone(),
346            index_vector_dim: *index_vector_dim,
347            slice_sizes: slice_sizes.clone(),
348        },
349        ExecOp::Scatter(config) => ExecOpKey::Scatter(config.clone()),
350        ExecOp::Slice(config) => ExecOpKey::Slice(config.clone()),
351        ExecOp::DynamicSlice { slice_sizes } => ExecOpKey::DynamicSlice {
352            slice_sizes: slice_sizes.clone(),
353        },
354        ExecOp::DynamicUpdateSlice => ExecOpKey::DynamicUpdateSlice,
355        ExecOp::Pad(config) => ExecOpKey::Pad(config.clone()),
356        ExecOp::Concatenate { axis } => ExecOpKey::Concatenate { axis: *axis },
357        ExecOp::Reverse { axes } => ExecOpKey::Reverse { axes: axes.clone() },
358        ExecOp::ShapeOf { axis } => ExecOpKey::ShapeOf { axis: *axis },
359        ExecOp::DynamicTruncate { axis } => ExecOpKey::DynamicTruncate { axis: *axis },
360        ExecOp::PadToMatch { axis } => ExecOpKey::PadToMatch { axis: *axis },
361        ExecOp::ReduceProd { axes } => ExecOpKey::ReduceProd { axes: axes.clone() },
362        ExecOp::ReduceMax { axes } => ExecOpKey::ReduceMax { axes: axes.clone() },
363        ExecOp::ReduceMin { axes } => ExecOpKey::ReduceMin { axes: axes.clone() },
364        ExecOp::Extension(extension) => {
365            let key = ExecOpKey::Extension {
366                family_id: extension.family_id(),
367                payload_hash: extension_payload_hash(extension.as_ref()),
368            };
369            extensions.push(Arc::clone(extension));
370            key
371        }
372    }
373}
374
375fn extension_payload_hash(extension: &dyn ExtensionOp) -> u64 {
376    let mut hasher = std::collections::hash_map::DefaultHasher::new();
377    extension.payload_hash(&mut DynHasherProxy::new(&mut hasher));
378    hasher.finish()
379}
380
381struct DynHasherProxy<'a, H: Hasher + ?Sized> {
382    inner: &'a mut H,
383}
384
385impl<'a, H: Hasher + ?Sized> DynHasherProxy<'a, H> {
386    fn new(inner: &'a mut H) -> Self {
387        Self { inner }
388    }
389}
390
391impl<H: Hasher + ?Sized> Hasher for DynHasherProxy<'_, H> {
392    fn finish(&self) -> u64 {
393        self.inner.finish()
394    }
395
396    fn write(&mut self, bytes: &[u8]) {
397        self.inner.write(bytes);
398    }
399}
400
401fn vec_retained_bytes<T>(values: &Vec<T>) -> usize {
402    values.capacity().saturating_mul(size_of::<T>())
403}
404
405fn vec_of_vec_retained_bytes<T>(values: &[Vec<T>]) -> usize {
406    saturating_sum(values.iter().map(vec_retained_bytes))
407}
408
409fn exec_program_key_retained_bytes(key: &ExecProgramKey) -> usize {
410    saturating_sum([
411        size_of::<ExecProgramKey>(),
412        vec_retained_bytes(&key.instructions),
413        saturating_sum(
414            key.instructions
415                .iter()
416                .map(exec_instruction_key_retained_bytes),
417        ),
418        vec_retained_bytes(&key.input_slots),
419        vec_retained_bytes(&key.output_slots),
420    ])
421}
422
423fn exec_instruction_key_retained_bytes(key: &ExecInstructionKey) -> usize {
424    saturating_sum([
425        size_of::<ExecInstructionKey>(),
426        exec_op_key_retained_bytes(&key.op),
427        vec_retained_bytes(&key.input_slots),
428        vec_retained_bytes(&key.output_slots),
429        vec_of_vec_retained_bytes(&key.output_shapes),
430        vec_of_vec_retained_bytes(&key.output_extents),
431        vec_retained_bytes(&key.last_use),
432    ])
433}
434
435fn exec_op_key_retained_bytes(key: &ExecOpKey) -> usize {
436    saturating_sum([
437        size_of::<ExecOpKey>(),
438        match key {
439            ExecOpKey::Transpose { perm } => vec_retained_bytes(perm),
440            ExecOpKey::Reshape { shape } => vec_retained_bytes(shape),
441            ExecOpKey::BroadcastInDim { shape, dims } => {
442                saturating_sum([vec_retained_bytes(shape), vec_retained_bytes(dims)])
443            }
444            ExecOpKey::Constant { bytes, .. } => vec_retained_bytes(bytes),
445            ExecOpKey::DotGeneral(config) => dot_general_config_retained_bytes(config),
446            ExecOpKey::DotGeneralWithConj { config, .. } => {
447                dot_general_config_retained_bytes(config)
448            }
449            ExecOpKey::ReduceSum { axes }
450            | ExecOpKey::Reverse { axes }
451            | ExecOpKey::ReduceProd { axes }
452            | ExecOpKey::ReduceMax { axes }
453            | ExecOpKey::ReduceMin { axes } => vec_retained_bytes(axes),
454            ExecOpKey::Gather(config) => gather_config_retained_bytes(config),
455            ExecOpKey::GatherDynamicSliceSizes {
456                offset_dims,
457                collapsed_slice_dims,
458                start_index_map,
459                slice_sizes,
460                ..
461            } => saturating_sum([
462                vec_retained_bytes(offset_dims),
463                vec_retained_bytes(collapsed_slice_dims),
464                vec_retained_bytes(start_index_map),
465                vec_retained_bytes(slice_sizes),
466            ]),
467            ExecOpKey::Scatter(config) => scatter_config_retained_bytes(config),
468            ExecOpKey::Slice(config) => slice_config_retained_bytes(config),
469            ExecOpKey::DynamicSlice { slice_sizes } => vec_retained_bytes(slice_sizes),
470            ExecOpKey::Pad(config) => pad_config_retained_bytes(config),
471            ExecOpKey::Convert { .. }
472            | ExecOpKey::ExtractDiag { .. }
473            | ExecOpKey::EmbedDiag { .. }
474            | ExecOpKey::Tril { .. }
475            | ExecOpKey::Triu { .. }
476            | ExecOpKey::Add
477            | ExecOpKey::Subtract
478            | ExecOpKey::Multiply
479            | ExecOpKey::Negate
480            | ExecOpKey::Conj
481            | ExecOpKey::Divide
482            | ExecOpKey::Remainder
483            | ExecOpKey::Abs
484            | ExecOpKey::Sign
485            | ExecOpKey::Maximum
486            | ExecOpKey::Minimum
487            | ExecOpKey::Compare(_)
488            | ExecOpKey::Select
489            | ExecOpKey::Clamp
490            | ExecOpKey::Exp
491            | ExecOpKey::Log
492            | ExecOpKey::Sin
493            | ExecOpKey::Cos
494            | ExecOpKey::Tanh
495            | ExecOpKey::Sqrt
496            | ExecOpKey::Rsqrt
497            | ExecOpKey::Pow
498            | ExecOpKey::Expm1
499            | ExecOpKey::Log1p
500            | ExecOpKey::DynamicUpdateSlice
501            | ExecOpKey::Concatenate { .. }
502            | ExecOpKey::ShapeOf { .. }
503            | ExecOpKey::DynamicTruncate { .. }
504            | ExecOpKey::PadToMatch { .. }
505            | ExecOpKey::Extension { .. } => 0,
506        },
507    ])
508}
509
510fn dot_general_config_retained_bytes(config: &tenferro_tensor::DotGeneralConfig) -> usize {
511    saturating_sum([
512        vec_retained_bytes(&config.lhs_contracting_dims),
513        vec_retained_bytes(&config.rhs_contracting_dims),
514        vec_retained_bytes(&config.lhs_batch_dims),
515        vec_retained_bytes(&config.rhs_batch_dims),
516    ])
517}
518
519fn gather_config_retained_bytes(config: &tenferro_tensor::GatherConfig) -> usize {
520    saturating_sum([
521        vec_retained_bytes(&config.offset_dims),
522        vec_retained_bytes(&config.collapsed_slice_dims),
523        vec_retained_bytes(&config.start_index_map),
524        vec_retained_bytes(&config.slice_sizes),
525    ])
526}
527
528fn scatter_config_retained_bytes(config: &tenferro_tensor::ScatterConfig) -> usize {
529    saturating_sum([
530        vec_retained_bytes(&config.update_window_dims),
531        vec_retained_bytes(&config.inserted_window_dims),
532        vec_retained_bytes(&config.scatter_dims_to_operand_dims),
533    ])
534}
535
536fn slice_config_retained_bytes(config: &tenferro_tensor::SliceConfig) -> usize {
537    saturating_sum([
538        vec_retained_bytes(&config.starts),
539        vec_retained_bytes(&config.limits),
540        vec_retained_bytes(&config.strides),
541    ])
542}
543
544fn pad_config_retained_bytes(config: &tenferro_tensor::PadConfig) -> usize {
545    saturating_sum([
546        vec_retained_bytes(&config.edge_padding_low),
547        vec_retained_bytes(&config.edge_padding_high),
548        vec_retained_bytes(&config.interior_padding),
549    ])
550}
551
552fn exec_op_retained_bytes(op: &ExecOp) -> usize {
553    match op {
554        ExecOp::Constant { bytes, .. } => vec_retained_bytes(bytes),
555        ExecOp::Extension(extension) => size_of_val(extension),
556        _ => 0,
557    }
558}
559
560fn exec_instruction_retained_bytes(inst: &ExecInstruction) -> usize {
561    saturating_sum([
562        size_of::<ExecInstruction>(),
563        exec_op_retained_bytes(&inst.op),
564        vec_retained_bytes(&inst.input_slots),
565        vec_retained_bytes(&inst.output_slots),
566        vec_of_vec_retained_bytes(&inst.output_shapes),
567        vec_of_vec_retained_bytes(&inst.output_extents),
568        vec_retained_bytes(&inst.last_use),
569    ])
570}
571
572fn exec_program_retained_bytes(program: &ExecProgram) -> usize {
573    saturating_sum([
574        size_of::<ExecProgram>(),
575        vec_retained_bytes(&program.instructions),
576        saturating_sum(
577            program
578                .instructions
579                .iter()
580                .map(exec_instruction_retained_bytes),
581        ),
582        vec_retained_bytes(&program.input_slots),
583        vec_retained_bytes(&program.output_slots),
584    ])
585}
586
587pub(crate) fn compile_cache_stats(cache: &LruCache<CacheKey, ExecProgram>) -> CacheStats {
588    CacheStats {
589        entries: cache.len(),
590        retained_bytes: cache
591            .iter()
592            .map(|(key, program)| {
593                saturating_sum([
594                    cache_key_retained_bytes(key),
595                    exec_program_retained_bytes(program),
596                ])
597            })
598            .fold(0usize, usize::saturating_add),
599    }
600}
601
602fn saturating_sum(values: impl IntoIterator<Item = usize>) -> usize {
603    values.into_iter().fold(0usize, usize::saturating_add)
604}
605
606#[cfg(test)]
607mod tests;
tenferro_runtime/graph/cache.rs

tenferro_runtime/graph/
cache.rs
