spark_storage/
eviction.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Eviction policy for the high-speed-swap scratch pool.
4//
5// Lossless mode: every block needed for attention IS attended this step. So
6// the policy doesn't decide which blocks to drop *from this step*; it
7// decides which currently-resident slot to overwrite when a *new* block
8// must be brought in. The objective is to minimise re-fetches across
9// future steps — keep blocks that are likely to be needed again, evict the
10// ones that aren't.
11//
12// Score: `(predictor_score, last_access_epoch)` lexicographic ascending.
13// Lowest predictor score evicts first; ties break to the oldest access.
14// `pinned` slots (currently in the active tile) are excluded.
15
16use std::collections::HashSet;
17
18#[derive(Debug)]
19pub struct EvictionPolicy {
20    /// Per-slot epoch of last access (touch). Monotonically increasing.
21    last_access: Vec<u64>,
22    /// Last-known predictor score per slot. Refreshed by the orchestrator
23    /// after each `score_blocks` call.
24    last_score: Vec<f32>,
25    /// Monotonic epoch counter — bumped on every `touch`.
26    epoch: u64,
27}
28
29impl EvictionPolicy {
30    pub fn new(num_slots: u32) -> Self {
31        Self {
32            last_access: vec![0; num_slots as usize],
33            last_score: vec![f32::NEG_INFINITY; num_slots as usize],
34            epoch: 1,
35        }
36    }
37
38    pub fn capacity(&self) -> u32 {
39        self.last_access.len() as u32
40    }
41
42    /// Mark `slot` as accessed at the current epoch.
43    pub fn touch(&mut self, slot: u32) {
44        self.last_access[slot as usize] = self.epoch;
45        self.epoch = self.epoch.wrapping_add(1);
46    }
47
48    /// Update the predictor-score record for `slot`.
49    pub fn record_score(&mut self, slot: u32, score: f32) {
50        self.last_score[slot as usize] = score;
51    }
52
53    /// Reset all bookkeeping (called on `ScratchPool::clear`).
54    pub fn reset(&mut self) {
55        for v in self.last_access.iter_mut() {
56            *v = 0;
57        }
58        for v in self.last_score.iter_mut() {
59            *v = f32::NEG_INFINITY;
60        }
61        self.epoch = 1;
62    }
63
64    /// Return slot indices in eviction-preference order
65    /// (most-evictable first), excluding `pinned` slots.
66    pub fn rank(&self, pinned: &[u32]) -> Vec<u32> {
67        let pinned_set: HashSet<u32> = pinned.iter().copied().collect();
68        let mut candidates: Vec<u32> = (0..self.capacity())
69            .filter(|s| !pinned_set.contains(s))
70            .collect();
71        candidates.sort_by(|a, b| {
72            let ai = *a as usize;
73            let bi = *b as usize;
74            // Lower predictor score first; ties → older access first.
75            self.last_score[ai]
76                .partial_cmp(&self.last_score[bi])
77                .unwrap_or(std::cmp::Ordering::Equal)
78                .then(self.last_access[ai].cmp(&self.last_access[bi]))
79        });
80        candidates
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn lowest_score_evicts_first() {
90        let mut p = EvictionPolicy::new(4);
91        p.record_score(0, 10.0);
92        p.record_score(1, 1.0); // weakest
93        p.record_score(2, 5.0);
94        p.record_score(3, 100.0); // strongest, never evicted
95        let order = p.rank(&[]);
96        assert_eq!(order[0], 1, "lowest score should evict first: {order:?}");
97        assert_eq!(order.last(), Some(&3));
98    }
99
100    #[test]
101    fn ties_break_to_oldest_access() {
102        let mut p = EvictionPolicy::new(3);
103        p.record_score(0, 5.0);
104        p.record_score(1, 5.0);
105        p.record_score(2, 5.0);
106        p.touch(2); // most recent
107        p.touch(0); // even more recent
108        // 1 was never touched → oldest (epoch 0); 2 < 0
109        let order = p.rank(&[]);
110        assert_eq!(order[0], 1, "oldest should win the tie: {order:?}");
111        assert_eq!(order[1], 2);
112        assert_eq!(order[2], 0);
113    }
114
115    #[test]
116    fn pinned_slots_excluded() {
117        let mut p = EvictionPolicy::new(4);
118        p.record_score(0, 1.0);
119        p.record_score(1, 2.0);
120        p.record_score(2, 3.0);
121        p.record_score(3, 4.0);
122        let order = p.rank(&[0, 1]);
123        assert_eq!(order, vec![2, 3]);
124    }
125
126    #[test]
127    fn reset_clears_state() {
128        let mut p = EvictionPolicy::new(2);
129        p.record_score(0, 1.0);
130        p.record_score(1, 2.0);
131        p.touch(0);
132        p.touch(1);
133        p.reset();
134        assert_eq!(p.last_score, vec![f32::NEG_INFINITY; 2]);
135        assert_eq!(p.last_access, vec![0; 2]);
136        assert_eq!(p.epoch, 1);
137    }
138}
spark_storage/eviction.rs

spark_storage/
eviction.rs
