tenferro_cpu/
indexed_plan_cache.rs

1use std::mem::{size_of, size_of_val};
2use std::sync::Arc;
3
4use lru::LruCache;
5use strided_kernel::{
6    ErasedDynamicSlicePlan, ErasedDynamicUpdateSlicePlan, ErasedGatherPlan, ErasedScatterPlan,
7    KernelDType,
8};
9use tenferro_tensor::CacheStats;
10use tinyvec::{Array, TinyVec};
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: TinyVec<[TinyVec<[usize; 8]>; 4]>,
97    strides: TinyVec<[TinyVec<[isize; 8]>; 4]>,
98    config: TinyVec<[TinyVec<[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_inline_vec_retained_bytes(&self.dims)
131            .saturating_add(nested_inline_vec_retained_bytes(&self.strides))
132            .saturating_add(nested_inline_vec_retained_bytes(&self.config))
133    }
134
135    fn logical_payload_bytes(&self) -> usize {
136        nested_inline_vec_logical_bytes(&self.dims)
137            .saturating_add(nested_inline_vec_logical_bytes(&self.strides))
138            .saturating_add(nested_inline_vec_logical_bytes(&self.config))
139    }
140}
141
142fn inline_vec_retained_bytes<A: Array>(values: &TinyVec<A>) -> usize {
143    if values.is_heap() {
144        values.capacity().saturating_mul(size_of::<A::Item>())
145    } else {
146        0
147    }
148}
149
150fn nested_inline_vec_retained_bytes<A, B>(values: &TinyVec<A>) -> usize
151where
152    A: Array<Item = TinyVec<B>>,
153    B: Array,
154{
155    inline_vec_retained_bytes(values).saturating_add(
156        values
157            .iter()
158            .map(inline_vec_retained_bytes)
159            .fold(0usize, usize::saturating_add),
160    )
161}
162
163fn nested_inline_vec_logical_bytes<A, B>(values: &TinyVec<A>) -> usize
164where
165    A: Array<Item = TinyVec<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/// Cache for compile-once indexed traversal plans, held inside a
188/// [`crate::CpuPlanCache`].
189///
190/// 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::CpuPlanCache`] exposes cache-specific limits, clear, and aggregate
196/// stats methods.
197#[derive(Debug)]
198pub(crate) struct IndexedPlanCache {
199    entries: LruCache<IndexedPlanKey, IndexedPlanCacheEntry>,
200    limits: IndexedPlanCacheLimits,
201    retained_bytes: usize,
202    hits: u64,
203    misses: u64,
204    evictions: u64,
205    clears: u64,
206}
207
208impl Default for IndexedPlanCache {
209    fn default() -> Self {
210        Self::new(DEFAULT_INDEXED_PLAN_CACHE_LIMITS)
211    }
212}
213
214impl IndexedPlanCache {
215    pub(crate) fn new(limits: IndexedPlanCacheLimits) -> Self {
216        Self {
217            entries: LruCache::unbounded(),
218            limits,
219            retained_bytes: 0,
220            hits: 0,
221            misses: 0,
222            evictions: 0,
223            clears: 0,
224        }
225    }
226
227    pub(crate) fn limits(&self) -> IndexedPlanCacheLimits {
228        self.limits
229    }
230
231    pub(crate) fn set_limits(&mut self, limits: IndexedPlanCacheLimits) {
232        self.limits = limits;
233        self.evict_to_limits();
234    }
235
236    pub(crate) fn clear(&mut self) {
237        self.entries.clear();
238        self.retained_bytes = 0;
239        self.clears = self.clears.saturating_add(1);
240    }
241
242    pub(crate) fn stats(&self) -> CacheStats {
243        CacheStats {
244            entries: self.entries.len(),
245            retained_bytes: self.retained_bytes,
246            hits: self.hits,
247            misses: self.misses,
248            evictions: self.evictions,
249            clears: self.clears,
250        }
251    }
252
253    pub(crate) fn gather<E>(
254        &mut self,
255        key: IndexedPlanKey,
256        compile: impl FnOnce() -> Result<ErasedGatherPlan, E>,
257    ) -> Result<Arc<ErasedGatherPlan>, E> {
258        if let Some(IndexedPlan::Gather(plan)) = self.lookup(&key) {
259            return Ok(plan);
260        }
261        let plan = Arc::new(compile()?);
262        self.insert(key, IndexedPlan::Gather(Arc::clone(&plan)));
263        Ok(plan)
264    }
265
266    pub(crate) fn scatter<E>(
267        &mut self,
268        key: IndexedPlanKey,
269        compile: impl FnOnce() -> Result<ErasedScatterPlan, E>,
270    ) -> Result<Arc<ErasedScatterPlan>, E> {
271        if let Some(IndexedPlan::Scatter(plan)) = self.lookup(&key) {
272            return Ok(plan);
273        }
274        let plan = Arc::new(compile()?);
275        self.insert(key, IndexedPlan::Scatter(Arc::clone(&plan)));
276        Ok(plan)
277    }
278
279    pub(crate) fn dynamic_slice<E>(
280        &mut self,
281        key: IndexedPlanKey,
282        compile: impl FnOnce() -> Result<ErasedDynamicSlicePlan, E>,
283    ) -> Result<Arc<ErasedDynamicSlicePlan>, E> {
284        if let Some(IndexedPlan::DynamicSlice(plan)) = self.lookup(&key) {
285            return Ok(plan);
286        }
287        let plan = Arc::new(compile()?);
288        self.insert(key, IndexedPlan::DynamicSlice(Arc::clone(&plan)));
289        Ok(plan)
290    }
291
292    pub(crate) fn dynamic_update_slice<E>(
293        &mut self,
294        key: IndexedPlanKey,
295        compile: impl FnOnce() -> Result<ErasedDynamicUpdateSlicePlan, E>,
296    ) -> Result<Arc<ErasedDynamicUpdateSlicePlan>, E> {
297        if let Some(IndexedPlan::DynamicUpdateSlice(plan)) = self.lookup(&key) {
298            return Ok(plan);
299        }
300        let plan = Arc::new(compile()?);
301        self.insert(key, IndexedPlan::DynamicUpdateSlice(Arc::clone(&plan)));
302        Ok(plan)
303    }
304
305    fn lookup(&mut self, key: &IndexedPlanKey) -> Option<IndexedPlan> {
306        let Some(entry) = self.entries.get(key) else {
307            self.misses = self.misses.saturating_add(1);
308            return None;
309        };
310        self.hits = self.hits.saturating_add(1);
311        Some(entry.plan.clone())
312    }
313
314    fn insert(&mut self, key: IndexedPlanKey, plan: IndexedPlan) {
315        if self.limits.max_entries == 0 || self.limits.max_retained_bytes == 0 {
316            return;
317        }
318        let retained_bytes = indexed_plan_retained_bytes(&key, &plan);
319        if retained_bytes > self.limits.max_retained_bytes {
320            return;
321        }
322        let entry = IndexedPlanCacheEntry {
323            plan,
324            retained_bytes,
325        };
326        if let Some(replaced) = self.entries.put(key, entry) {
327            self.retained_bytes = self.retained_bytes.saturating_sub(replaced.retained_bytes);
328        }
329        self.retained_bytes = self.retained_bytes.saturating_add(retained_bytes);
330        self.evict_to_limits();
331    }
332
333    fn evict_to_limits(&mut self) {
334        while self.entries.len() > self.limits.max_entries
335            || self.retained_bytes > self.limits.max_retained_bytes
336        {
337            let Some((_key, entry)) = self.entries.pop_lru() else {
338                break;
339            };
340            self.retained_bytes = self.retained_bytes.saturating_sub(entry.retained_bytes);
341            self.evictions = self.evictions.saturating_add(1);
342        }
343    }
344}
345
346fn indexed_plan_retained_bytes(key: &IndexedPlanKey, plan: &IndexedPlan) -> usize {
347    let plan_header = match plan {
348        IndexedPlan::Gather(plan) => size_of_val(plan.as_ref()),
349        IndexedPlan::Scatter(plan) => size_of_val(plan.as_ref()),
350        IndexedPlan::DynamicSlice(plan) => size_of_val(plan.as_ref()),
351        IndexedPlan::DynamicUpdateSlice(plan) => size_of_val(plan.as_ref()),
352    };
353    // INVARIANT: erased indexed plans clone the key's layout/config vectors
354    // and may derive additional axis vectors. Charge twice the logical key
355    // payload in addition to any spilled key capacity so inline SmallVec keys
356    // still account for the upstream plan's heap-owned Vec payload.
357    size_of::<IndexedPlanCacheEntry>()
358        .saturating_add(size_of::<IndexedPlanKey>())
359        .saturating_add(plan_header)
360        .saturating_add(key.retained_bytes())
361        .saturating_add(key.logical_payload_bytes().saturating_mul(2))
362}
363
364#[cfg(test)]
365mod tests;
tenferro_cpu/indexed_plan_cache.rs

tenferro_cpu/
indexed_plan_cache.rs
