onnx_runtime_eager/
cache.rs

1//! The compiled-kernel cache (`docs/EAGER.md` §8.2).
2//!
3//! A bounded LRU keyed by the op identity plus the concrete input shapes,
4//! dtypes, and device — so that repeated eager calls with the same op and shapes
5//! reuse an already-compiled kernel instead of recreating one per call.
6//!
7//! ## Design-vs-real-API reconciliation
8//!
9//! The design stores `Arc<dyn Kernel>` (`docs/EAGER.md` §8.2). The real
10//! [`Kernel`](onnx_runtime_ep_api::Kernel) trait is `Send` **but not `Sync`**
11//! (`onnx-runtime-ep-api/src/kernel.rs`), so a bare `Arc<dyn Kernel>` is not
12//! `Send` and could not live inside the process-global, `Sync`
13//! [`EagerContext`](crate::EagerContext). We therefore store
14//! `Arc<Mutex<Box<dyn Kernel>>>`: the `Mutex` restores `Send + Sync` while
15//! preserving the design's "share one compiled kernel across calls" intent.
16//! `Kernel::execute` takes `&self`, so the lock only serialises reuse of the
17//! *same* cached kernel; distinct ops/shapes dispatch concurrently.
18//!
19//! The eviction policy is a straightforward bounded LRU (a `HashMap` +
20//! recency `VecDeque`) rather than an external `lru` crate, keeping the
21//! dependency set minimal as requested. Correctness of the *key* matters more
22//! than eviction sophistication (§8.2).
23
24use std::collections::{HashMap, VecDeque};
25use std::sync::{Arc, Mutex};
26
27use onnx_runtime_ep_api::Kernel;
28use onnx_runtime_ir::{DataType, DeviceId};
29
30/// A cached, compiled kernel, shareable across dispatches. See the module docs
31/// for why this is `Arc<Mutex<Box<dyn Kernel>>>` rather than `Arc<dyn Kernel>`.
32pub type CachedKernel = Arc<Mutex<Box<dyn Kernel>>>;
33
34/// The cache key (`docs/EAGER.md` §8.2 `KernelCacheKey`).
35///
36/// Two dispatches share a compiled kernel iff they agree on the op identity
37/// (`op_type`, `domain`, `opset`), the input shapes, the input dtypes, and the
38/// device. Shapes and dtypes are part of the key because kernels are compiled
39/// specialised to concrete shapes (§4.2 / §8.2).
40#[derive(Clone, PartialEq, Eq, Hash, Debug)]
41pub struct KernelCacheKey {
42    pub op_type: String,
43    pub domain: String,
44    pub opset: u64,
45    pub input_shapes: Vec<Vec<usize>>,
46    pub input_dtypes: Vec<DataType>,
47    pub device: DeviceId,
48}
49
50/// Cache instrumentation, exposed for tests and diagnostics.
51#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
52pub struct CacheStats {
53    /// Distinct compiled entries currently held.
54    pub entries: usize,
55    /// Lookups served from an existing entry.
56    pub hits: u64,
57    /// Lookups that compiled a new kernel.
58    pub misses: u64,
59}
60
61/// A bounded LRU of compiled kernels (`docs/EAGER.md` §8.2).
62pub struct KernelCache {
63    capacity: usize,
64    map: HashMap<KernelCacheKey, CachedKernel>,
65    /// Recency order, least-recently-used at the front.
66    order: VecDeque<KernelCacheKey>,
67    hits: u64,
68    misses: u64,
69}
70
71impl KernelCache {
72    /// A cache holding at most `capacity` compiled kernels (`capacity` is
73    /// clamped to at least 1).
74    pub fn new(capacity: usize) -> Self {
75        Self {
76            capacity: capacity.max(1),
77            map: HashMap::new(),
78            order: VecDeque::new(),
79            hits: 0,
80            misses: 0,
81        }
82    }
83
84    /// Return the cached kernel for `key`, compiling and inserting it via
85    /// `create` on a miss (`docs/EAGER.md` §8.2 `get_or_create`).
86    pub fn get_or_create<E>(
87        &mut self,
88        key: KernelCacheKey,
89        create: impl FnOnce() -> Result<Box<dyn Kernel>, E>,
90    ) -> Result<CachedKernel, E> {
91        if let Some(kernel) = self.map.get(&key) {
92            let kernel = kernel.clone();
93            self.hits += 1;
94            self.touch(&key);
95            return Ok(kernel);
96        }
97        self.misses += 1;
98        let kernel: CachedKernel = Arc::new(Mutex::new(create()?));
99        self.map.insert(key.clone(), kernel.clone());
100        self.order.push_back(key);
101        self.evict_if_needed();
102        Ok(kernel)
103    }
104
105    /// Current cache statistics.
106    pub fn stats(&self) -> CacheStats {
107        CacheStats {
108            entries: self.map.len(),
109            hits: self.hits,
110            misses: self.misses,
111        }
112    }
113
114    /// Move `key` to the most-recently-used end of the recency order.
115    fn touch(&mut self, key: &KernelCacheKey) {
116        if let Some(pos) = self.order.iter().position(|k| k == key)
117            && let Some(k) = self.order.remove(pos)
118        {
119            self.order.push_back(k);
120        }
121    }
122
123    /// Evict least-recently-used entries until within capacity.
124    fn evict_if_needed(&mut self) {
125        while self.map.len() > self.capacity {
126            if let Some(lru) = self.order.pop_front() {
127                self.map.remove(&lru);
128            } else {
129                break;
130            }
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138    use onnx_runtime_ep_api::{Result as EpResult, TensorMut, TensorView};
139
140    /// A trivial no-op kernel for exercising the cache in isolation.
141    struct NopKernel;
142    impl Kernel for NopKernel {
143        fn execute(&self, _inputs: &[TensorView], _outputs: &mut [TensorMut]) -> EpResult<()> {
144            Ok(())
145        }
146    }
147
148    fn key(op: &str, shape: Vec<usize>) -> KernelCacheKey {
149        KernelCacheKey {
150            op_type: op.to_string(),
151            domain: String::new(),
152            opset: 26,
153            input_shapes: vec![shape],
154            input_dtypes: vec![DataType::Float32],
155            device: DeviceId::cpu(),
156        }
157    }
158
159    #[test]
160    fn miss_then_hit_same_key() {
161        let mut cache = KernelCache::new(8);
162        let k = key("Add", vec![2, 3]);
163        let mut created = 0;
164        let _ = cache
165            .get_or_create::<()>(k.clone(), || {
166                created += 1;
167                Ok(Box::new(NopKernel))
168            })
169            .unwrap();
170        let _ = cache
171            .get_or_create::<()>(k.clone(), || {
172                created += 1;
173                Ok(Box::new(NopKernel))
174            })
175            .unwrap();
176        assert_eq!(created, 1, "second dispatch must reuse the cached kernel");
177        assert_eq!(cache.stats().hits, 1);
178        assert_eq!(cache.stats().misses, 1);
179        assert_eq!(cache.stats().entries, 1);
180    }
181
182    #[test]
183    fn distinct_shapes_are_distinct_entries() {
184        let mut cache = KernelCache::new(8);
185        let _ = cache
186            .get_or_create::<()>(key("Add", vec![2, 3]), || Ok(Box::new(NopKernel)))
187            .unwrap();
188        let _ = cache
189            .get_or_create::<()>(key("Add", vec![4, 5]), || Ok(Box::new(NopKernel)))
190            .unwrap();
191        assert_eq!(cache.stats().entries, 2);
192        assert_eq!(cache.stats().misses, 2);
193    }
194
195    #[test]
196    fn lru_evicts_oldest_over_capacity() {
197        let mut cache = KernelCache::new(2);
198        let a = key("Add", vec![1]);
199        let b = key("Mul", vec![1]);
200        let c = key("Sub", vec![1]);
201        let _ = cache.get_or_create::<()>(a.clone(), || Ok(Box::new(NopKernel))).unwrap();
202        let _ = cache.get_or_create::<()>(b.clone(), || Ok(Box::new(NopKernel))).unwrap();
203        // Touch `a` so `b` becomes the LRU victim.
204        let _ = cache.get_or_create::<()>(a.clone(), || Ok(Box::new(NopKernel))).unwrap();
205        let _ = cache.get_or_create::<()>(c.clone(), || Ok(Box::new(NopKernel))).unwrap();
206        assert_eq!(cache.stats().entries, 2);
207        // `b` evicted → recompiled on next access (a fresh miss).
208        let mut recreated = 0;
209        let _ = cache
210            .get_or_create::<()>(b, || {
211                recreated += 1;
212                Ok(Box::new(NopKernel))
213            })
214            .unwrap();
215        assert_eq!(recreated, 1);
216    }
217}
onnx_runtime_eager/cache.rs

onnx_runtime_eager/
cache.rs
