vyre_driver_wgpu/megakernel/
dispatch_plan.rs

1//! Dispatcher-local fixed-batch megakernel launch-plan cache.
2
3use super::dispatcher::BatchDispatchConfig;
4use std::cmp::{Ordering, Reverse};
5use std::collections::{BinaryHeap, HashMap};
6use vyre_runtime::megakernel::{MegakernelDispatchTopology, MegakernelLaunchRecommendation};
7
8/// Reusable fixed-batch megakernel launch metadata.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub struct BatchDispatchPlan {
11    /// Queue length this plan was derived for.
12    pub queue_len: u32,
13    /// Workgroups submitted by the compiled persistent megakernel.
14    pub worker_groups: u32,
15    /// Worker lanes per workgroup.
16    pub workgroup_size_x: u32,
17    /// Sparse hit-ring capacity compiled into the dispatcher pipeline.
18    pub hit_capacity: u32,
19    /// Estimated peak device bytes required by the selected launch plan.
20    pub estimated_peak_device_bytes: u64,
21    /// Hard device-memory budget applied to this plan. Zero means unbounded.
22    pub device_memory_budget_bytes: u64,
23    /// Scale-aware topology selected for this queue shape.
24    pub topology: MegakernelDispatchTopology,
25}
26
27impl BatchDispatchPlan {
28    pub(crate) fn from_recommendation(
29        queue_len: u32,
30        config: &BatchDispatchConfig,
31        recommendation: MegakernelLaunchRecommendation,
32    ) -> Self {
33        Self {
34            queue_len,
35            worker_groups: recommendation.worker_groups,
36            workgroup_size_x: config.workgroup_size_x,
37            hit_capacity: recommendation.hit_capacity,
38            estimated_peak_device_bytes: recommendation.estimated_peak_device_bytes,
39            device_memory_budget_bytes: recommendation.device_memory_budget_bytes,
40            topology: recommendation.topology,
41        }
42    }
43}
44
45#[derive(Debug, Clone, Copy)]
46struct BatchDispatchPlanCacheEntry {
47    plan: BatchDispatchPlan,
48    last_seen: u64,
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52struct BatchDispatchPlanLruEntry {
53    last_seen: u64,
54    queue_len: u32,
55}
56
57impl Ord for BatchDispatchPlanLruEntry {
58    fn cmp(&self, other: &Self) -> Ordering {
59        self.last_seen
60            .cmp(&other.last_seen)
61            .then_with(|| self.queue_len.cmp(&other.queue_len))
62    }
63}
64
65impl PartialOrd for BatchDispatchPlanLruEntry {
66    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
67        Some(self.cmp(other))
68    }
69}
70
71/// Small LRU cache for repeated fixed-batch launch metadata.
72#[derive(Debug)]
73pub(crate) struct BatchDispatchPlanCache {
74    entries: HashMap<u32, BatchDispatchPlanCacheEntry>,
75    lru: BinaryHeap<Reverse<BatchDispatchPlanLruEntry>>,
76    clock: u64,
77    cap: usize,
78}
79
80impl BatchDispatchPlanCache {
81    fn with_cap(cap: usize) -> Self {
82        Self {
83            entries: HashMap::with_capacity(cap),
84            lru: BinaryHeap::with_capacity(cap),
85            clock: 0,
86            cap,
87        }
88    }
89
90    pub(crate) fn get(&mut self, queue_len: u32) -> Option<BatchDispatchPlan> {
91        let tick = self.next_tick();
92        let entry = self.entries.get_mut(&queue_len)?;
93        entry.last_seen = tick;
94        let plan = entry.plan;
95        self.lru.push(Reverse(BatchDispatchPlanLruEntry {
96            last_seen: entry.last_seen,
97            queue_len,
98        }));
99        self.compact_lru_if_needed();
100        Some(plan)
101    }
102
103    pub(crate) fn insert(&mut self, plan: BatchDispatchPlan) {
104        let tick = self.next_tick();
105        if let Some(entry) = self.entries.get_mut(&plan.queue_len) {
106            entry.plan = plan;
107            entry.last_seen = tick;
108            self.lru.push(Reverse(BatchDispatchPlanLruEntry {
109                last_seen: entry.last_seen,
110                queue_len: plan.queue_len,
111            }));
112            self.compact_lru_if_needed();
113            return;
114        }
115        if self.cap == 0 {
116            return;
117        }
118        while self.entries.len() >= self.cap {
119            let Some(queue_len) = self.pop_lru_key() else {
120                break;
121            };
122            self.entries.remove(&queue_len);
123        }
124        let last_seen = tick;
125        self.entries.insert(
126            plan.queue_len,
127            BatchDispatchPlanCacheEntry { plan, last_seen },
128        );
129        self.lru.push(Reverse(BatchDispatchPlanLruEntry {
130            last_seen,
131            queue_len: plan.queue_len,
132        }));
133        self.compact_lru_if_needed();
134    }
135
136    pub(crate) fn len_u16(&self) -> u16 {
137        u16::try_from(self.entries.len()).unwrap_or(u16::MAX)
138    }
139
140    fn next_tick(&mut self) -> u64 {
141        if self.clock == u64::MAX {
142            self.rebase_clock_to_zero();
143        }
144        self.clock += 1;
145        self.clock
146    }
147
148    fn rebase_clock_to_zero(&mut self) {
149        self.clock = 0;
150        self.lru.clear();
151        for (queue_len, entry) in &mut self.entries {
152            entry.last_seen = 0;
153            self.lru.push(Reverse(BatchDispatchPlanLruEntry {
154                last_seen: 0,
155                queue_len: *queue_len,
156            }));
157        }
158    }
159
160    fn pop_lru_key(&mut self) -> Option<u32> {
161        while let Some(Reverse(entry)) = self.lru.pop() {
162            if self
163                .entries
164                .get(&entry.queue_len)
165                .is_some_and(|current| current.last_seen == entry.last_seen)
166            {
167                return Some(entry.queue_len);
168            }
169        }
170        None
171    }
172
173    fn compact_lru_if_needed(&mut self) {
174        let live = self.entries.len();
175        if let Some(limit) = stale_lru_limit(live) {
176            if self.lru.len() <= limit {
177                return;
178            }
179        }
180        self.lru.clear();
181        self.lru
182            .extend(self.entries.iter().map(|(&queue_len, entry)| {
183                Reverse(BatchDispatchPlanLruEntry {
184                    last_seen: entry.last_seen,
185                    queue_len,
186                })
187            }));
188    }
189}
190
191fn stale_lru_limit(live: usize) -> Option<usize> {
192    live.checked_mul(4).map(|limit| limit.max(8))
193}
194
195impl Default for BatchDispatchPlanCache {
196    fn default() -> Self {
197        Self::with_cap(32)
198    }
199}
200
201/// Result of looking up fixed-batch launch metadata.
202#[derive(Debug, Clone, Copy)]
203pub(crate) struct BatchDispatchPlanLookup {
204    /// Cached or newly planned launch metadata.
205    pub(crate) plan: BatchDispatchPlan,
206    /// True when the dispatcher reused resident launch metadata.
207    pub(crate) cache_hit: bool,
208    /// Number of resident entries after lookup.
209    pub(crate) cache_entries: u16,
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use vyre_runtime::megakernel::{
216        MegakernelDispatchTopology, MegakernelExecutionMode, MegakernelLaunchGeometry,
217        MegakernelQueuePressure,
218    };
219
220    #[test]
221    fn dispatch_plan_consumes_recommended_worker_groups() {
222        let config = BatchDispatchConfig {
223            worker_groups: 8,
224            workgroup_size_x: 64,
225            hit_capacity: 1024,
226            ..Default::default()
227        };
228        let recommendation = MegakernelLaunchRecommendation {
229            geometry: MegakernelLaunchGeometry {
230                workgroup_size_x: 64,
231                slot_count: 2048,
232                dispatch_grid: [32, 1, 1],
233            },
234            worker_groups: 32,
235            hit_capacity: 4096,
236            pressure: MegakernelQueuePressure::Balanced,
237            execution_mode: MegakernelExecutionMode::Interpreter,
238            topology: MegakernelDispatchTopology::SparseFrontier,
239            promote_hot_opcodes: false,
240            promote_hot_windows: false,
241            age_priority_work: false,
242            estimated_peak_device_bytes: 65_536,
243            device_memory_budget_bytes: 0,
244        };
245
246        let plan = BatchDispatchPlan::from_recommendation(1024, &config, recommendation);
247
248        assert_eq!(
249            plan.worker_groups, 32,
250            "dispatch plan must use scale-policy worker_groups, not the constructor seed"
251        );
252        assert_eq!(
253            plan.hit_capacity, 4096,
254            "dispatch plan must use scale-policy hit_capacity, not stale config capacity"
255        );
256        assert_eq!(plan.estimated_peak_device_bytes, 65_536);
257    }
258
259    #[test]
260    fn dispatch_plan_cache_lru_heap_stays_capacity_scale() {
261        let config = BatchDispatchConfig {
262            worker_groups: 8,
263            workgroup_size_x: 64,
264            hit_capacity: 1024,
265            ..Default::default()
266        };
267        let mut cache = BatchDispatchPlanCache::with_cap(4);
268
269        for queue_len in 1..128 {
270            let recommendation = MegakernelLaunchRecommendation {
271                geometry: MegakernelLaunchGeometry {
272                    workgroup_size_x: 64,
273                    slot_count: queue_len,
274                    dispatch_grid: [32, 1, 1],
275                },
276                worker_groups: 32,
277                hit_capacity: 4096,
278                pressure: MegakernelQueuePressure::Balanced,
279                execution_mode: MegakernelExecutionMode::Interpreter,
280                topology: MegakernelDispatchTopology::SparseFrontier,
281                promote_hot_opcodes: false,
282                promote_hot_windows: false,
283                age_priority_work: false,
284                estimated_peak_device_bytes: 65_536,
285                device_memory_budget_bytes: 0,
286            };
287            let plan = BatchDispatchPlan::from_recommendation(queue_len, &config, recommendation);
288            cache.insert(plan);
289            let _ = cache.get(queue_len);
290        }
291
292        assert_eq!(cache.entries.len(), 4);
293        assert!(
294            cache.lru.len() <= cache.entries.len().saturating_mul(4).max(8),
295            "Fix: dispatch-plan LRU heap must compact stale recency entries to cache-capacity scale"
296        );
297    }
298}
vyre_driver_wgpu/megakernel/dispatch_plan.rs

vyre_driver_wgpu/megakernel/
dispatch_plan.rs
