tenferro_ad/
transform_cache.rs

1use std::collections::hash_map::DefaultHasher;
2use std::collections::{HashMap, HashSet};
3use std::fmt;
4use std::hash::{Hash, Hasher};
5use std::mem::{size_of, size_of_val};
6use std::num::NonZeroUsize;
7use std::sync::{Arc, Mutex, MutexGuard};
8
9use computegraph::graph::Graph;
10use computegraph::{LocalValueId, ValueKey};
11use lru::LruCache;
12use tenferro_ops::input_key::TensorInputKey;
13use tenferro_ops::std_tensor_op::StdTensorOp;
14use tenferro_runtime::{CacheStats, Error, Result};
15use tidu::eager::RecordedGraph;
16use tidu::LinearizedGraph;
17
18const DEFAULT_AD_TRANSFORM_CACHE_ENTRIES: usize = 128;
19const DEFAULT_AD_TRANSFORM_CACHE_RETAINED_BYTES: usize = 64 * 1024 * 1024;
20
21/// Retention limits for AD transform graph caches.
22///
23/// The retained-byte limit is a logical payload estimate, not process RSS.
24///
25/// # Examples
26///
27/// ```rust
28/// use std::num::NonZeroUsize;
29/// use tenferro_ad::AdTransformCacheLimits;
30///
31/// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(4).unwrap());
32/// assert_eq!(limits.max_entries().get(), 4);
33/// assert!(limits.max_retained_bytes().is_some());
34/// ```
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36pub struct AdTransformCacheLimits {
37    max_entries: NonZeroUsize,
38    max_retained_bytes: Option<NonZeroUsize>,
39}
40
41impl AdTransformCacheLimits {
42    /// Create AD transform cache limits with the default retained-byte bound.
43    ///
44    /// # Examples
45    ///
46    /// ```rust
47    /// use std::num::NonZeroUsize;
48    /// use tenferro_ad::AdTransformCacheLimits;
49    ///
50    /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(2).unwrap());
51    /// assert_eq!(limits.max_entries().get(), 2);
52    /// ```
53    pub fn new(max_entries: NonZeroUsize) -> Self {
54        Self {
55            max_entries,
56            max_retained_bytes: Some(
57                NonZeroUsize::new(DEFAULT_AD_TRANSFORM_CACHE_RETAINED_BYTES)
58                    .unwrap_or(NonZeroUsize::MIN),
59            ),
60        }
61    }
62
63    /// Return the maximum number of retained AD transform entries.
64    ///
65    /// # Examples
66    ///
67    /// ```rust
68    /// use tenferro_ad::AdTransformCacheLimits;
69    ///
70    /// assert!(AdTransformCacheLimits::default().max_entries().get() > 0);
71    /// ```
72    pub fn max_entries(self) -> NonZeroUsize {
73        self.max_entries
74    }
75
76    /// Return the logical retained-byte bound, when one is configured.
77    ///
78    /// # Examples
79    ///
80    /// ```rust
81    /// use tenferro_ad::AdTransformCacheLimits;
82    ///
83    /// assert!(AdTransformCacheLimits::default().max_retained_bytes().is_some());
84    /// ```
85    pub fn max_retained_bytes(self) -> Option<NonZeroUsize> {
86        self.max_retained_bytes
87    }
88
89    /// Return limits with a new logical retained-byte bound.
90    ///
91    /// # Examples
92    ///
93    /// ```rust
94    /// use std::num::NonZeroUsize;
95    /// use tenferro_ad::AdTransformCacheLimits;
96    ///
97    /// let limits = AdTransformCacheLimits::default()
98    ///     .with_max_retained_bytes(NonZeroUsize::new(1024).unwrap());
99    /// assert_eq!(limits.max_retained_bytes().unwrap().get(), 1024);
100    /// ```
101    pub fn with_max_retained_bytes(mut self, max_retained_bytes: NonZeroUsize) -> Self {
102        self.max_retained_bytes = Some(max_retained_bytes);
103        self
104    }
105}
106
107impl Default for AdTransformCacheLimits {
108    fn default() -> Self {
109        Self::new(
110            NonZeroUsize::new(DEFAULT_AD_TRANSFORM_CACHE_ENTRIES).unwrap_or(NonZeroUsize::MIN),
111        )
112    }
113}
114
115#[derive(Clone, Debug, PartialEq, Eq, Hash)]
116pub(crate) struct EagerAdTransformCacheKey {
117    recorded_graph_fingerprint: u64,
118    output_slots: Vec<usize>,
119}
120
121#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
122pub(crate) enum TracedAdTransformKind {
123    Jvp,
124    Vjp,
125}
126
127#[derive(Clone, Debug, PartialEq, Eq, Hash)]
128pub(crate) struct TracedAdTransformCacheKey {
129    kind: TracedAdTransformKind,
130    roots_fingerprint: u64,
131    output_key: ValueKey<StdTensorOp>,
132    wrt_input_key: TensorInputKey,
133    aliases_fingerprint: u64,
134}
135
136impl TracedAdTransformCacheKey {
137    pub(crate) fn new(
138        kind: TracedAdTransformKind,
139        roots: &[Arc<Graph<StdTensorOp>>],
140        output_key: &ValueKey<StdTensorOp>,
141        wrt_input_key: &TensorInputKey,
142        aliases: &HashMap<TensorInputKey, ValueKey<StdTensorOp>>,
143    ) -> Self {
144        Self {
145            kind,
146            roots_fingerprint: traced_roots_fingerprint(roots),
147            output_key: output_key.clone(),
148            wrt_input_key: wrt_input_key.clone(),
149            aliases_fingerprint: aliases_fingerprint(aliases),
150        }
151    }
152}
153
154#[derive(Clone)]
155pub(crate) struct CachedOptimizedLinearGraph {
156    graph: Arc<Graph<StdTensorOp>>,
157    tangent_inputs: Vec<(TensorInputKey, LocalValueId)>,
158    tangent_outputs: Vec<Option<LocalValueId>>,
159}
160
161impl CachedOptimizedLinearGraph {
162    pub(crate) fn new(
163        graph: Graph<StdTensorOp>,
164        tangent_inputs: Vec<(TensorInputKey, LocalValueId)>,
165        tangent_outputs: Vec<Option<LocalValueId>>,
166    ) -> Self {
167        Self {
168            graph: Arc::new(graph),
169            tangent_inputs,
170            tangent_outputs,
171        }
172    }
173
174    pub(crate) fn graph(&self) -> &Arc<Graph<StdTensorOp>> {
175        &self.graph
176    }
177
178    pub(crate) fn as_graph(&self) -> &Graph<StdTensorOp> {
179        self.graph.as_ref()
180    }
181
182    pub(crate) fn tangent_inputs(&self) -> &[(TensorInputKey, LocalValueId)] {
183        &self.tangent_inputs
184    }
185
186    pub(crate) fn tangent_outputs(&self) -> &[Option<LocalValueId>] {
187        &self.tangent_outputs
188    }
189}
190
191impl fmt::Debug for CachedOptimizedLinearGraph {
192    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193        f.debug_struct("CachedOptimizedLinearGraph")
194            .field("values_len", &self.graph.values().len())
195            .field("operations_len", &self.graph.operations().len())
196            .field("tangent_inputs_len", &self.tangent_inputs.len())
197            .field("tangent_outputs_len", &self.tangent_outputs.len())
198            .finish()
199    }
200}
201
202#[derive(Clone)]
203pub(crate) struct CachedTracedVjpTransform {
204    residual_graph: Arc<Graph<StdTensorOp>>,
205    transposed: Arc<CachedOptimizedLinearGraph>,
206}
207
208impl CachedTracedVjpTransform {
209    pub(crate) fn new(
210        residual_graph: Arc<Graph<StdTensorOp>>,
211        transposed: CachedOptimizedLinearGraph,
212    ) -> Self {
213        Self {
214            residual_graph,
215            transposed: Arc::new(transposed),
216        }
217    }
218
219    pub(crate) fn residual_graph(&self) -> &Arc<Graph<StdTensorOp>> {
220        &self.residual_graph
221    }
222
223    pub(crate) fn transposed(&self) -> &CachedOptimizedLinearGraph {
224        self.transposed.as_ref()
225    }
226}
227
228impl fmt::Debug for CachedTracedVjpTransform {
229    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
230        f.debug_struct("CachedTracedVjpTransform")
231            .field("residual_values_len", &self.residual_graph.values().len())
232            .field(
233                "residual_operations_len",
234                &self.residual_graph.operations().len(),
235            )
236            .field("transposed", &self.transposed)
237            .finish()
238    }
239}
240
241impl EagerAdTransformCacheKey {
242    pub(crate) fn new(graph: &RecordedGraph<StdTensorOp>, output_slots: &[usize]) -> Self {
243        Self {
244            recorded_graph_fingerprint: eager_recorded_graph_fingerprint(graph),
245            output_slots: output_slots.to_vec(),
246        }
247    }
248}
249
250fn traced_roots_fingerprint(roots: &[Arc<Graph<StdTensorOp>>]) -> u64 {
251    let mut hasher = DefaultHasher::new();
252    roots.len().hash(&mut hasher);
253    let mut visited = HashSet::new();
254    for root in roots {
255        hash_graph(root.as_ref(), &mut hasher, &mut visited);
256    }
257    hasher.finish()
258}
259
260fn hash_graph<H: Hasher>(
261    graph: &Graph<StdTensorOp>,
262    hasher: &mut H,
263    visited: &mut HashSet<*const Graph<StdTensorOp>>,
264) {
265    let graph_ptr: *const Graph<StdTensorOp> = graph;
266    if !visited.insert(graph_ptr) {
267        return;
268    }
269    graph.inputs().hash(hasher);
270    graph.outputs().hash(hasher);
271    for value in graph.values() {
272        value.key.hash(hasher);
273        value.producer.hash(hasher);
274    }
275    for op in graph.operations() {
276        op.operation.hash(hasher);
277        op.inputs.hash(hasher);
278        op.outputs.hash(hasher);
279        op.role.hash(hasher);
280    }
281    graph.parents().len().hash(hasher);
282    for parent in graph.parents() {
283        hash_graph(parent.as_ref(), hasher, visited);
284    }
285}
286
287fn aliases_fingerprint(aliases: &HashMap<TensorInputKey, ValueKey<StdTensorOp>>) -> u64 {
288    let mut entry_hashes = aliases
289        .iter()
290        .map(|(key, value)| {
291            let mut hasher = DefaultHasher::new();
292            key.hash(&mut hasher);
293            value.hash(&mut hasher);
294            hasher.finish()
295        })
296        .collect::<Vec<_>>();
297    entry_hashes.sort_unstable();
298
299    let mut hasher = DefaultHasher::new();
300    entry_hashes.hash(&mut hasher);
301    hasher.finish()
302}
303
304fn eager_recorded_graph_fingerprint(graph: &RecordedGraph<StdTensorOp>) -> u64 {
305    let mut hasher = DefaultHasher::new();
306    graph.input_keys().hash(&mut hasher);
307    graph.output_keys().hash(&mut hasher);
308    graph.as_graph().inputs().hash(&mut hasher);
309    graph.as_graph().outputs().hash(&mut hasher);
310    for value in graph.as_graph().values() {
311        value.key.hash(&mut hasher);
312        value.producer.hash(&mut hasher);
313    }
314    for op in graph.as_graph().operations() {
315        op.operation.hash(&mut hasher);
316        op.inputs.hash(&mut hasher);
317        op.outputs.hash(&mut hasher);
318        op.role.hash(&mut hasher);
319    }
320    hasher.finish()
321}
322
323#[derive(Debug)]
324pub(crate) struct AdTransformCache {
325    store: Mutex<AdTransformCacheStore>,
326}
327
328impl AdTransformCache {
329    pub(crate) fn new() -> Self {
330        Self {
331            store: Mutex::new(AdTransformCacheStore::default()),
332        }
333    }
334
335    pub(crate) fn limits(&self) -> Result<AdTransformCacheLimits> {
336        Ok(self.lock_store()?.limits)
337    }
338
339    pub(crate) fn set_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
340        self.lock_store()?.set_limits(limits);
341        Ok(())
342    }
343
344    pub(crate) fn clear(&self) -> Result<()> {
345        self.lock_store()?.clear();
346        Ok(())
347    }
348
349    pub(crate) fn stats(&self) -> Result<CacheStats> {
350        Ok(self.lock_store()?.stats())
351    }
352
353    pub(crate) fn get_eager_linearized(
354        &self,
355        key: &EagerAdTransformCacheKey,
356    ) -> Result<Option<Arc<LinearizedGraph<StdTensorOp>>>> {
357        Ok(self.lock_store()?.get_eager_linearized(key))
358    }
359
360    pub(crate) fn put_eager_linearized(
361        &self,
362        key: EagerAdTransformCacheKey,
363        value: Arc<LinearizedGraph<StdTensorOp>>,
364    ) -> Result<()> {
365        self.lock_store()?.put_eager_linearized(key, value);
366        Ok(())
367    }
368
369    pub(crate) fn get_traced_linearized(
370        &self,
371        key: &TracedAdTransformCacheKey,
372    ) -> Result<Option<Arc<CachedOptimizedLinearGraph>>> {
373        Ok(self.lock_store()?.get_traced_linearized(key))
374    }
375
376    pub(crate) fn put_traced_linearized(
377        &self,
378        key: TracedAdTransformCacheKey,
379        value: Arc<CachedOptimizedLinearGraph>,
380    ) -> Result<()> {
381        self.lock_store()?.put_traced_linearized(key, value);
382        Ok(())
383    }
384
385    pub(crate) fn get_traced_vjp(
386        &self,
387        key: &TracedAdTransformCacheKey,
388    ) -> Result<Option<Arc<CachedTracedVjpTransform>>> {
389        Ok(self.lock_store()?.get_traced_vjp(key))
390    }
391
392    pub(crate) fn put_traced_vjp(
393        &self,
394        key: TracedAdTransformCacheKey,
395        value: Arc<CachedTracedVjpTransform>,
396    ) -> Result<()> {
397        self.lock_store()?.put_traced_vjp(key, value);
398        Ok(())
399    }
400
401    fn lock_store(&self) -> Result<MutexGuard<'_, AdTransformCacheStore>> {
402        self.store
403            .lock()
404            .map_err(|_| Error::Internal("AD transform cache lock poisoned".to_string()))
405    }
406}
407
408#[derive(Debug)]
409struct AdTransformCacheStore {
410    limits: AdTransformCacheLimits,
411    entries: LruCache<AdTransformCacheKey, AdTransformCacheEntryWithStats>,
412    stats: CacheStats,
413}
414
415impl AdTransformCacheStore {
416    fn set_limits(&mut self, limits: AdTransformCacheLimits) {
417        self.limits = limits;
418        self.evict_to_limits();
419    }
420
421    fn clear(&mut self) {
422        self.entries.clear();
423        self.stats = CacheStats::empty();
424    }
425
426    fn stats(&self) -> CacheStats {
427        self.stats
428    }
429
430    fn get_eager_linearized(
431        &mut self,
432        key: &EagerAdTransformCacheKey,
433    ) -> Option<Arc<LinearizedGraph<StdTensorOp>>> {
434        self.entries
435            .get(&AdTransformCacheKey::EagerLinearize(key.clone()))
436            .and_then(|entry| match &entry.entry {
437                AdTransformCacheEntry::EagerLinearized(linear) => Some(Arc::clone(linear)),
438                _ => None,
439            })
440    }
441
442    fn put_eager_linearized(
443        &mut self,
444        key: EagerAdTransformCacheKey,
445        value: Arc<LinearizedGraph<StdTensorOp>>,
446    ) {
447        let key = AdTransformCacheKey::EagerLinearize(key);
448        let entry = AdTransformCacheEntry::EagerLinearized(value);
449        self.put_entry(key, entry);
450    }
451
452    fn get_traced_linearized(
453        &mut self,
454        key: &TracedAdTransformCacheKey,
455    ) -> Option<Arc<CachedOptimizedLinearGraph>> {
456        self.entries
457            .get(&AdTransformCacheKey::Traced(key.clone()))
458            .and_then(|entry| match &entry.entry {
459                AdTransformCacheEntry::TracedLinearized(linear) => Some(Arc::clone(linear)),
460                _ => None,
461            })
462    }
463
464    fn put_traced_linearized(
465        &mut self,
466        key: TracedAdTransformCacheKey,
467        value: Arc<CachedOptimizedLinearGraph>,
468    ) {
469        let key = AdTransformCacheKey::Traced(key);
470        let entry = AdTransformCacheEntry::TracedLinearized(value);
471        self.put_entry(key, entry);
472    }
473
474    fn get_traced_vjp(
475        &mut self,
476        key: &TracedAdTransformCacheKey,
477    ) -> Option<Arc<CachedTracedVjpTransform>> {
478        self.entries
479            .get(&AdTransformCacheKey::Traced(key.clone()))
480            .and_then(|entry| match &entry.entry {
481                AdTransformCacheEntry::TracedVjp(vjp) => Some(Arc::clone(vjp)),
482                _ => None,
483            })
484    }
485
486    fn put_traced_vjp(
487        &mut self,
488        key: TracedAdTransformCacheKey,
489        value: Arc<CachedTracedVjpTransform>,
490    ) {
491        let key = AdTransformCacheKey::Traced(key);
492        let entry = AdTransformCacheEntry::TracedVjp(value);
493        self.put_entry(key, entry);
494    }
495
496    fn put_entry(&mut self, key: AdTransformCacheKey, entry: AdTransformCacheEntry) {
497        let retained_bytes = ad_transform_cache_entry_retained_bytes(&key, &entry);
498        let entry = AdTransformCacheEntryWithStats {
499            entry,
500            retained_bytes,
501        };
502        self.stats.entries = self.entries.len();
503        self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
504        if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
505            self.stats.retained_bytes = self
506                .stats
507                .retained_bytes
508                .saturating_sub(old_entry.retained_bytes);
509        }
510        self.stats.entries = self.entries.len();
511        self.evict_to_limits();
512    }
513
514    fn evict_to_limits(&mut self) {
515        while self.entries.len() > self.limits.max_entries.get()
516            || self
517                .limits
518                .max_retained_bytes
519                .is_some_and(|limit| self.stats.retained_bytes > limit.get())
520        {
521            let Some((_key, entry)) = self.entries.pop_lru() else {
522                break;
523            };
524            self.stats.retained_bytes = self
525                .stats
526                .retained_bytes
527                .saturating_sub(entry.retained_bytes);
528        }
529        self.stats.entries = self.entries.len();
530    }
531}
532
533impl Default for AdTransformCacheStore {
534    fn default() -> Self {
535        Self {
536            limits: AdTransformCacheLimits::default(),
537            entries: LruCache::unbounded(),
538            stats: CacheStats::empty(),
539        }
540    }
541}
542
543#[derive(Clone, Debug, PartialEq, Eq, Hash)]
544enum AdTransformCacheKey {
545    EagerLinearize(EagerAdTransformCacheKey),
546    Traced(TracedAdTransformCacheKey),
547}
548
549enum AdTransformCacheEntry {
550    EagerLinearized(Arc<LinearizedGraph<StdTensorOp>>),
551    TracedLinearized(Arc<CachedOptimizedLinearGraph>),
552    TracedVjp(Arc<CachedTracedVjpTransform>),
553}
554
555impl fmt::Debug for AdTransformCacheEntry {
556    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
557        match self {
558            Self::EagerLinearized(_) => f.write_str("EagerLinearized(..)"),
559            Self::TracedLinearized(_) => f.write_str("TracedLinearized(..)"),
560            Self::TracedVjp(_) => f.write_str("TracedVjp(..)"),
561        }
562    }
563}
564
565#[derive(Debug)]
566struct AdTransformCacheEntryWithStats {
567    entry: AdTransformCacheEntry,
568    retained_bytes: usize,
569}
570
571fn ad_transform_cache_entry_retained_bytes(
572    key: &AdTransformCacheKey,
573    entry: &AdTransformCacheEntry,
574) -> usize {
575    size_of::<AdTransformCacheKey>()
576        + ad_transform_cache_key_retained_bytes(key)
577        + size_of::<AdTransformCacheEntry>()
578        + ad_transform_cache_value_retained_bytes(entry)
579}
580
581fn ad_transform_cache_key_retained_bytes(key: &AdTransformCacheKey) -> usize {
582    match key {
583        AdTransformCacheKey::EagerLinearize(key) => {
584            key.output_slots.capacity() * size_of::<usize>()
585        }
586        AdTransformCacheKey::Traced(_) => 0,
587    }
588}
589
590fn ad_transform_cache_value_retained_bytes(entry: &AdTransformCacheEntry) -> usize {
591    match entry {
592        AdTransformCacheEntry::EagerLinearized(linear) => {
593            size_of::<Arc<LinearizedGraph<StdTensorOp>>>()
594                + size_of::<LinearizedGraph<StdTensorOp>>()
595                + size_of_val(linear.as_graph().values())
596                + size_of_val(linear.as_graph().operations())
597                + size_of_val(linear.tangent_inputs())
598                + size_of_val(linear.tangent_outputs())
599        }
600        AdTransformCacheEntry::TracedLinearized(linear) => {
601            size_of::<Arc<CachedOptimizedLinearGraph>>()
602                + cached_optimized_linear_graph_retained_bytes(linear.as_ref())
603        }
604        AdTransformCacheEntry::TracedVjp(vjp) => {
605            size_of::<Arc<CachedTracedVjpTransform>>()
606                + cached_traced_vjp_retained_bytes(vjp.as_ref())
607        }
608    }
609}
610
611fn cached_traced_vjp_retained_bytes(vjp: &CachedTracedVjpTransform) -> usize {
612    size_of::<CachedTracedVjpTransform>()
613        + graph_retained_bytes(vjp.residual_graph.as_ref())
614        + cached_optimized_linear_graph_retained_bytes(vjp.transposed.as_ref())
615}
616
617fn cached_optimized_linear_graph_retained_bytes(linear: &CachedOptimizedLinearGraph) -> usize {
618    size_of::<CachedOptimizedLinearGraph>()
619        + graph_retained_bytes(linear.graph.as_ref())
620        + linear.tangent_inputs.capacity() * size_of::<(TensorInputKey, LocalValueId)>()
621        + linear.tangent_outputs.capacity() * size_of::<Option<LocalValueId>>()
622}
623
624fn graph_retained_bytes(graph: &Graph<StdTensorOp>) -> usize {
625    size_of::<Graph<StdTensorOp>>()
626        + size_of_val(graph.values())
627        + size_of_val(graph.operations())
628        + size_of_val(graph.inputs())
629        + size_of_val(graph.outputs())
630        + size_of_val(graph.parents())
631}
tenferro_ad/transform_cache.rs

tenferro_ad/
transform_cache.rs
