tenferro_cpu/
indexed_plan_cache.rs

1use std::mem::{size_of, size_of_val};
2use std::sync::Arc;
3
4use lru::LruCache;
5use smallvec::{Array, SmallVec};
6use strided_kernel::{
7    ErasedDynamicSlicePlan, ErasedDynamicUpdateSlicePlan, ErasedGatherPlan, ErasedScatterPlan,
8    KernelDType,
9};
10use tenferro_tensor::CacheStats;
11
12/// Limits for the CPU indexed-plan cache.
13///
14/// A zero entry or byte limit disables retention. Limits apply independently
15/// to each initialized CPU execution engine.
16///
17/// # Examples
18///
19/// ```
20/// use tenferro_cpu::IndexedPlanCacheLimits;
21///
22/// let limits = IndexedPlanCacheLimits::new(32, 1024 * 1024);
23/// assert_eq!(limits.max_entries(), 32);
24/// assert_eq!(limits.max_retained_bytes(), 1024 * 1024);
25/// ```
26#[derive(Clone, Copy, Debug, Eq, PartialEq)]
27pub struct IndexedPlanCacheLimits {
28    max_entries: usize,
29    max_retained_bytes: usize,
30}
31
32impl IndexedPlanCacheLimits {
33    /// Construct indexed-plan cache limits.
34    ///
35    /// # Examples
36    ///
37    /// ```
38    /// use tenferro_cpu::IndexedPlanCacheLimits;
39    ///
40    /// let limits = IndexedPlanCacheLimits::new(8, 4096);
41    /// assert_eq!(limits.max_entries(), 8);
42    /// ```
43    pub const fn new(max_entries: usize, max_retained_bytes: usize) -> Self {
44        Self {
45            max_entries,
46            max_retained_bytes,
47        }
48    }
49
50    /// Maximum retained plan count per CPU execution engine.
51    ///
52    /// # Examples
53    ///
54    /// ```
55    /// use tenferro_cpu::IndexedPlanCacheLimits;
56    ///
57    /// assert_eq!(IndexedPlanCacheLimits::new(8, 4096).max_entries(), 8);
58    /// ```
59    pub const fn max_entries(self) -> usize {
60        self.max_entries
61    }
62
63    /// Maximum logical retained bytes per CPU execution engine.
64    ///
65    /// # Examples
66    ///
67    /// ```
68    /// use tenferro_cpu::IndexedPlanCacheLimits;
69    ///
70    /// assert_eq!(
71    ///     IndexedPlanCacheLimits::new(8, 4096).max_retained_bytes(),
72    ///     4096
73    /// );
74    /// ```
75    pub const fn max_retained_bytes(self) -> usize {
76        self.max_retained_bytes
77    }
78}
79
80pub(crate) const DEFAULT_INDEXED_PLAN_CACHE_LIMITS: IndexedPlanCacheLimits =
81    IndexedPlanCacheLimits::new(256, 8 * 1024 * 1024);
82
83#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
84pub(crate) enum IndexedPlanFamily {
85    Gather,
86    Scatter,
87    DynamicSlice,
88    DynamicUpdateSlice,
89}
90
91#[derive(Clone, Debug, Eq, Hash, PartialEq)]
92pub(crate) struct IndexedPlanKey {
93    family: IndexedPlanFamily,
94    dtype: KernelDType,
95    index_dtype: KernelDType,
96    dims: SmallVec<[SmallVec<[usize; 8]>; 4]>,
97    strides: SmallVec<[SmallVec<[isize; 8]>; 4]>,
98    config: SmallVec<[SmallVec<[usize; 8]>; 5]>,
99}
100
101impl IndexedPlanKey {
102    pub(crate) fn from_slices(
103        family: IndexedPlanFamily,
104        dtype: KernelDType,
105        index_dtype: KernelDType,
106        dims: &[&[usize]],
107        strides: &[&[isize]],
108        config: &[&[usize]],
109    ) -> Self {
110        Self {
111            family,
112            dtype,
113            index_dtype,
114            dims: dims
115                .iter()
116                .map(|values| values.iter().copied().collect())
117                .collect(),
118            strides: strides
119                .iter()
120                .map(|values| values.iter().copied().collect())
121                .collect(),
122            config: config
123                .iter()
124                .map(|values| values.iter().copied().collect())
125                .collect(),
126        }
127    }
128
129    fn retained_bytes(&self) -> usize {
130        nested_smallvec_retained_bytes(&self.dims)
131            .saturating_add(nested_smallvec_retained_bytes(&self.strides))
132            .saturating_add(nested_smallvec_retained_bytes(&self.config))
133    }
134
135    fn logical_payload_bytes(&self) -> usize {
136        nested_smallvec_logical_bytes(&self.dims)
137            .saturating_add(nested_smallvec_logical_bytes(&self.strides))
138            .saturating_add(nested_smallvec_logical_bytes(&self.config))
139    }
140}
141
142fn smallvec_retained_bytes<A: Array>(values: &SmallVec<A>) -> usize {
143    if values.spilled() {
144        values.capacity().saturating_mul(size_of::<A::Item>())
145    } else {
146        0
147    }
148}
149
150fn nested_smallvec_retained_bytes<A, B>(values: &SmallVec<A>) -> usize
151where
152    A: Array<Item = SmallVec<B>>,
153    B: Array,
154{
155    smallvec_retained_bytes(values).saturating_add(
156        values
157            .iter()
158            .map(smallvec_retained_bytes)
159            .fold(0usize, usize::saturating_add),
160    )
161}
162
163fn nested_smallvec_logical_bytes<A, B>(values: &SmallVec<A>) -> usize
164where
165    A: Array<Item = SmallVec<B>>,
166    B: Array,
167{
168    values.iter().fold(0usize, |total, inner| {
169        total.saturating_add(inner.len().saturating_mul(size_of::<B::Item>()))
170    })
171}
172
173#[derive(Clone, Debug)]
174enum IndexedPlan {
175    Gather(Arc<ErasedGatherPlan>),
176    Scatter(Arc<ErasedScatterPlan>),
177    DynamicSlice(Arc<ErasedDynamicSlicePlan>),
178    DynamicUpdateSlice(Arc<ErasedDynamicUpdateSlicePlan>),
179}
180
181#[derive(Debug)]
182struct IndexedPlanCacheEntry {
183    plan: IndexedPlan,
184    retained_bytes: usize,
185}
186
187/// Engine-owned cache for compile-once indexed traversal plans.
188///
189/// Each [`crate::engine::EngineResources`] value owns one cache for the
190/// lifetime of its CPU execution engine. The bounded default retains at most
191/// 256 plans and 8 MiB of logical plan/key payload. Entries are keyed by the
192/// complete family, dtype, index dtype, layout, and operation configuration,
193/// and least-recently-used entries are evicted when either bound is exceeded.
194///
195/// [`crate::CpuBackend`] exposes cache-specific limits, clear, and aggregate
196/// stats methods. Its [`tenferro_runtime::runtime::RuntimeCacheOwner`]
197/// implementation also includes this cache in runtime aggregate stats/clear.
198#[derive(Debug)]
199pub(crate) struct IndexedPlanCache {
200    entries: LruCache<IndexedPlanKey, IndexedPlanCacheEntry>,
201    limits: IndexedPlanCacheLimits,
202    retained_bytes: usize,
203    hits: u64,
204    misses: u64,
205    evictions: u64,
206    clears: u64,
207}
208
209impl Default for IndexedPlanCache {
210    fn default() -> Self {
211        Self::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS)
212    }
213}
214
215impl IndexedPlanCache {
216    pub(crate) fn new(limits: IndexedPlanCacheLimits) -> Self {
217        Self {
218            entries: LruCache::unbounded(),
219            limits,
220            retained_bytes: 0,
221            hits: 0,
222            misses: 0,
223            evictions: 0,
224            clears: 0,
225        }
226    }
227
228    pub(crate) fn set_limits(&mut self, limits: IndexedPlanCacheLimits) {
229        self.limits = limits;
230        self.evict_to_limits();
231    }
232
233    pub(crate) fn clear(&mut self) {
234        self.entries.clear();
235        self.retained_bytes = 0;
236        self.clears = self.clears.saturating_add(1);
237    }
238
239    pub(crate) fn stats(&self) -> CacheStats {
240        CacheStats {
241            entries: self.entries.len(),
242            retained_bytes: self.retained_bytes,
243            hits: self.hits,
244            misses: self.misses,
245            evictions: self.evictions,
246            clears: self.clears,
247        }
248    }
249
250    pub(crate) fn gather<E>(
251        &mut self,
252        key: IndexedPlanKey,
253        compile: impl FnOnce() -> Result<ErasedGatherPlan, E>,
254    ) -> Result<Arc<ErasedGatherPlan>, E> {
255        if let Some(IndexedPlan::Gather(plan)) = self.lookup(&key) {
256            return Ok(plan);
257        }
258        let plan = Arc::new(compile()?);
259        self.insert(key, IndexedPlan::Gather(Arc::clone(&plan)));
260        Ok(plan)
261    }
262
263    pub(crate) fn scatter<E>(
264        &mut self,
265        key: IndexedPlanKey,
266        compile: impl FnOnce() -> Result<ErasedScatterPlan, E>,
267    ) -> Result<Arc<ErasedScatterPlan>, E> {
268        if let Some(IndexedPlan::Scatter(plan)) = self.lookup(&key) {
269            return Ok(plan);
270        }
271        let plan = Arc::new(compile()?);
272        self.insert(key, IndexedPlan::Scatter(Arc::clone(&plan)));
273        Ok(plan)
274    }
275
276    pub(crate) fn dynamic_slice<E>(
277        &mut self,
278        key: IndexedPlanKey,
279        compile: impl FnOnce() -> Result<ErasedDynamicSlicePlan, E>,
280    ) -> Result<Arc<ErasedDynamicSlicePlan>, E> {
281        if let Some(IndexedPlan::DynamicSlice(plan)) = self.lookup(&key) {
282            return Ok(plan);
283        }
284        let plan = Arc::new(compile()?);
285        self.insert(key, IndexedPlan::DynamicSlice(Arc::clone(&plan)));
286        Ok(plan)
287    }
288
289    pub(crate) fn dynamic_update_slice<E>(
290        &mut self,
291        key: IndexedPlanKey,
292        compile: impl FnOnce() -> Result<ErasedDynamicUpdateSlicePlan, E>,
293    ) -> Result<Arc<ErasedDynamicUpdateSlicePlan>, E> {
294        if let Some(IndexedPlan::DynamicUpdateSlice(plan)) = self.lookup(&key) {
295            return Ok(plan);
296        }
297        let plan = Arc::new(compile()?);
298        self.insert(key, IndexedPlan::DynamicUpdateSlice(Arc::clone(&plan)));
299        Ok(plan)
300    }
301
302    fn lookup(&mut self, key: &IndexedPlanKey) -> Option<IndexedPlan> {
303        let Some(entry) = self.entries.get(key) else {
304            self.misses = self.misses.saturating_add(1);
305            return None;
306        };
307        self.hits = self.hits.saturating_add(1);
308        Some(entry.plan.clone())
309    }
310
311    fn insert(&mut self, key: IndexedPlanKey, plan: IndexedPlan) {
312        if self.limits.max_entries == 0 || self.limits.max_retained_bytes == 0 {
313            return;
314        }
315        let retained_bytes = indexed_plan_retained_bytes(&key, &plan);
316        if retained_bytes > self.limits.max_retained_bytes {
317            return;
318        }
319        let entry = IndexedPlanCacheEntry {
320            plan,
321            retained_bytes,
322        };
323        if let Some(replaced) = self.entries.put(key, entry) {
324            self.retained_bytes = self.retained_bytes.saturating_sub(replaced.retained_bytes);
325        }
326        self.retained_bytes = self.retained_bytes.saturating_add(retained_bytes);
327        self.evict_to_limits();
328    }
329
330    fn evict_to_limits(&mut self) {
331        while self.entries.len() > self.limits.max_entries
332            || self.retained_bytes > self.limits.max_retained_bytes
333        {
334            let Some((_key, entry)) = self.entries.pop_lru() else {
335                break;
336            };
337            self.retained_bytes = self.retained_bytes.saturating_sub(entry.retained_bytes);
338            self.evictions = self.evictions.saturating_add(1);
339        }
340    }
341}
342
343fn indexed_plan_retained_bytes(key: &IndexedPlanKey, plan: &IndexedPlan) -> usize {
344    let plan_header = match plan {
345        IndexedPlan::Gather(plan) => size_of_val(plan.as_ref()),
346        IndexedPlan::Scatter(plan) => size_of_val(plan.as_ref()),
347        IndexedPlan::DynamicSlice(plan) => size_of_val(plan.as_ref()),
348        IndexedPlan::DynamicUpdateSlice(plan) => size_of_val(plan.as_ref()),
349    };
350    // INVARIANT: erased indexed plans clone the key's layout/config vectors
351    // and may derive additional axis vectors. Charge twice the logical key
352    // payload in addition to any spilled key capacity so inline SmallVec keys
353    // still account for the upstream plan's heap-owned Vec payload.
354    size_of::<IndexedPlanCacheEntry>()
355        .saturating_add(size_of::<IndexedPlanKey>())
356        .saturating_add(plan_header)
357        .saturating_add(key.retained_bytes())
358        .saturating_add(key.logical_payload_bytes().saturating_mul(2))
359}
360
361#[cfg(test)]
362mod tests;
tenferro_cpu/indexed_plan_cache.rs

tenferro_cpu/
indexed_plan_cache.rs
