vyre_runtime/megakernel/planner/
grid.rs

1//! Megakernel grid request, limits, plan cache, and recommendation surface.
2
3use std::cell::RefCell;
4
5use rustc_hash::FxHashMap;
6use vyre_driver::backend::BackendError;
7
8use super::geometry::ResidentLaunchGeometry;
9use super::sizing::ResidentSizingPolicy;
10
11/// Adapter limits that bound a megakernel worker-grid recommendation.
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
13pub struct ResidentGridLimits {
14    /// Adapter maximum workgroup size in the x dimension.
15    pub max_workgroup_size_x: u32,
16    /// Adapter maximum compute workgroups per dimension.
17    pub max_compute_workgroups_per_dimension: u32,
18    /// Adapter maximum invocations per compute workgroup.
19    pub max_compute_invocations_per_workgroup: u32,
20}
21
22const GRID_PLAN_CACHE_CAP: usize = 128;
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
25struct GeometryCacheKey {
26    slot_count: u32,
27    worker_count: u32,
28    max_workgroup_size_x: u32,
29}
30
31struct MegakernelPlannerCache {
32    grid_plans: FxHashMap<(ResidentGridRequest, ResidentGridLimits), CacheEntry<ResidentGridPlan>>,
33    geometries: FxHashMap<GeometryCacheKey, CacheEntry<ResidentLaunchGeometry>>,
34    clock: u64,
35}
36
37struct CacheEntry<T> {
38    value: T,
39    last_seen: u64,
40}
41
42impl MegakernelPlannerCache {
43    fn get_grid_plan(
44        &mut self,
45        key: &(ResidentGridRequest, ResidentGridLimits),
46    ) -> Option<ResidentGridPlan> {
47        self.prepare_cache_hit_tick();
48        let entry = self.grid_plans.get_mut(key)?;
49        self.clock += 1;
50        entry.last_seen = self.clock;
51        Some(entry.value)
52    }
53
54    fn insert_grid_plan(
55        &mut self,
56        key: (ResidentGridRequest, ResidentGridLimits),
57        value: ResidentGridPlan,
58    ) {
59        let tick = self.next_tick();
60        self.grid_plans.insert(
61            key,
62            CacheEntry {
63                value,
64                last_seen: tick,
65            },
66        );
67        self.evict_grid_plans_to_cap();
68    }
69
70    fn get_geometry(&mut self, key: &GeometryCacheKey) -> Option<ResidentLaunchGeometry> {
71        self.prepare_cache_hit_tick();
72        let entry = self.geometries.get_mut(key)?;
73        self.clock += 1;
74        entry.last_seen = self.clock;
75        Some(entry.value)
76    }
77
78    fn insert_geometry(&mut self, key: GeometryCacheKey, value: ResidentLaunchGeometry) {
79        let tick = self.next_tick();
80        self.geometries.insert(
81            key,
82            CacheEntry {
83                value,
84                last_seen: tick,
85            },
86        );
87        self.evict_geometries_to_cap();
88    }
89
90    fn evict_grid_plans_to_cap(&mut self) {
91        while self.grid_plans.len() > GRID_PLAN_CACHE_CAP {
92            let Some(evicted) = self
93                .grid_plans
94                .iter()
95                .min_by_key(|(_, entry)| entry.last_seen)
96                .map(|(key, _)| *key)
97            else {
98                break;
99            };
100            self.grid_plans.remove(&evicted);
101        }
102    }
103
104    fn evict_geometries_to_cap(&mut self) {
105        while self.geometries.len() > GRID_PLAN_CACHE_CAP {
106            let Some(evicted) = self
107                .geometries
108                .iter()
109                .min_by_key(|(_, entry)| entry.last_seen)
110                .map(|(key, _)| *key)
111            else {
112                break;
113            };
114            self.geometries.remove(&evicted);
115        }
116    }
117
118    fn next_tick(&mut self) -> u64 {
119        self.prepare_cache_hit_tick();
120        self.clock += 1;
121        self.clock
122    }
123
124    fn prepare_cache_hit_tick(&mut self) {
125        if self.clock == u64::MAX {
126            self.clock = 0;
127            for entry in self.grid_plans.values_mut() {
128                entry.last_seen = 0;
129            }
130            for entry in self.geometries.values_mut() {
131                entry.last_seen = 0;
132            }
133        }
134    }
135}
136
137impl Default for MegakernelPlannerCache {
138    fn default() -> Self {
139        Self {
140            grid_plans: FxHashMap::with_capacity_and_hasher(
141                GRID_PLAN_CACHE_CAP,
142                Default::default(),
143            ),
144            geometries: FxHashMap::with_capacity_and_hasher(
145                GRID_PLAN_CACHE_CAP,
146                Default::default(),
147            ),
148            clock: 0,
149        }
150    }
151}
152
153thread_local! {
154    static PLANNER_CACHE: RefCell<MegakernelPlannerCache> = RefCell::new(MegakernelPlannerCache::default());
155}
156
157fn cached_grid_plan(
158    request: ResidentGridRequest,
159    limits: ResidentGridLimits,
160) -> Result<ResidentGridPlan, BackendError> {
161    if let Some(plan) =
162        PLANNER_CACHE.with(|cache| cache.borrow_mut().get_grid_plan(&(request, limits)))
163    {
164        return Ok(plan);
165    }
166
167    let plan = ResidentSizingPolicy::standard().calculate_optimal_grid(request, limits)?;
168    PLANNER_CACHE.with(|cache| {
169        cache.borrow_mut().insert_grid_plan((request, limits), plan);
170    });
171    Ok(plan)
172}
173
174pub(super) fn cached_geometry_from_slots(
175    slot_count: u32,
176    worker_count: u32,
177    max_workgroup_size_x: u32,
178) -> ResidentLaunchGeometry {
179    let key = GeometryCacheKey {
180        slot_count,
181        worker_count,
182        max_workgroup_size_x,
183    };
184    if let Some(geometry) = PLANNER_CACHE.with(|cache| cache.borrow_mut().get_geometry(&key)) {
185        return geometry;
186    }
187
188    let geometry = ResidentSizingPolicy::standard().geometry_from_slots(
189        slot_count,
190        worker_count,
191        max_workgroup_size_x,
192    );
193    PLANNER_CACHE.with(|cache| {
194        cache.borrow_mut().insert_geometry(key, geometry);
195    });
196    geometry
197}
198
199impl ResidentGridLimits {
200    /// Construct megakernel grid limits from backend adapter limits.
201    #[must_use]
202    pub const fn new(
203        max_workgroup_size_x: u32,
204        max_compute_workgroups_per_dimension: u32,
205        max_compute_invocations_per_workgroup: u32,
206    ) -> Self {
207        Self {
208            max_workgroup_size_x,
209            max_compute_workgroups_per_dimension,
210            max_compute_invocations_per_workgroup,
211        }
212    }
213
214    pub(super) fn validate(self) -> Result<(), BackendError> {
215        if self.max_workgroup_size_x == 0 {
216            return Err(BackendError::new(
217                "megakernel max_workgroup_size_x must be non-zero. Fix: pass live adapter limits instead of a zero limit.",
218            ));
219        }
220        if self.max_compute_workgroups_per_dimension == 0 {
221            return Err(BackendError::new(
222                "megakernel max_compute_workgroups_per_dimension must be non-zero. Fix: pass live adapter limits instead of a zero limit.",
223            ));
224        }
225        if self.max_compute_invocations_per_workgroup == 0 {
226            return Err(BackendError::new(
227                "megakernel max_compute_invocations_per_workgroup must be non-zero. Fix: pass live adapter limits instead of a zero limit.",
228            ));
229        }
230        Ok(())
231    }
232}
233
234/// Logical work shape requested for a megakernel worker-grid recommendation.
235#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
236pub struct ResidentGridRequest {
237    /// Logical ring slots or work items queued for this launch.
238    pub queue_len: u32,
239    /// Caller-requested worker workgroup ceiling. Zero means derive from occupancy.
240    pub requested_worker_groups: u32,
241}
242
243impl ResidentGridRequest {
244    /// Construct a worker-grid request.
245    #[must_use]
246    pub const fn new(queue_len: u32, requested_worker_groups: u32) -> Self {
247        Self {
248            queue_len,
249            requested_worker_groups,
250        }
251    }
252}
253
254/// Resolved worker-grid plan shared by direct and policy-driven megakernel paths.
255#[derive(Debug, Clone, Copy, PartialEq, Eq)]
256pub struct ResidentGridPlan {
257    /// Padded launch geometry for the ring protocol.
258    pub geometry: ResidentLaunchGeometry,
259    /// Worker workgroups selected for the dispatch.
260    pub worker_groups: u32,
261}
262
263impl ResidentGridPlan {
264    /// Resolve worker groups, workgroup width, slot padding, and dispatch grid.
265    ///
266    /// # Errors
267    ///
268    /// Returns [`BackendError`] when adapter limits are malformed.
269    pub fn recommend(
270        request: ResidentGridRequest,
271        limits: ResidentGridLimits,
272    ) -> Result<Self, BackendError> {
273        cached_grid_plan(request, limits)
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280
281    fn limits() -> ResidentGridLimits {
282        ResidentGridLimits::new(256, 65_535, 256)
283    }
284
285    fn request(queue_len: u32) -> ResidentGridRequest {
286        ResidentGridRequest::new(queue_len, 0)
287    }
288
289    fn geometry(slot_count: u32) -> ResidentLaunchGeometry {
290        ResidentLaunchGeometry {
291            workgroup_size_x: 1,
292            slot_count,
293            dispatch_grid: [1, 1, 1],
294        }
295    }
296
297    #[test]
298    fn planner_grid_cache_refreshes_hot_plan_on_hit() {
299        let mut cache = MegakernelPlannerCache::default();
300        let limits = limits();
301        let hot_key = (request(1), limits);
302        let hot_plan = ResidentGridPlan {
303            geometry: geometry(1),
304            worker_groups: 1,
305        };
306        cache.insert_grid_plan(hot_key, hot_plan);
307        for queue_len in 2..=GRID_PLAN_CACHE_CAP as u32 {
308            cache.insert_grid_plan(
309                (request(queue_len), limits),
310                ResidentGridPlan {
311                    geometry: geometry(queue_len),
312                    worker_groups: 1,
313                },
314            );
315        }
316        assert_eq!(cache.get_grid_plan(&hot_key), Some(hot_plan));
317        cache.insert_grid_plan(
318            (request((GRID_PLAN_CACHE_CAP + 1) as u32), limits),
319            ResidentGridPlan {
320                geometry: geometry((GRID_PLAN_CACHE_CAP + 1) as u32),
321                worker_groups: 1,
322            },
323        );
324        assert_eq!(cache.get_grid_plan(&hot_key), Some(hot_plan));
325    }
326
327    #[test]
328    fn planner_geometry_cache_refreshes_hot_geometry_on_hit() {
329        let mut cache = MegakernelPlannerCache::default();
330        let hot_key = GeometryCacheKey {
331            slot_count: 1,
332            worker_count: 1,
333            max_workgroup_size_x: 256,
334        };
335        let hot_geometry = geometry(1);
336        cache.insert_geometry(hot_key, hot_geometry);
337        for slot_count in 2..=GRID_PLAN_CACHE_CAP as u32 {
338            cache.insert_geometry(
339                GeometryCacheKey {
340                    slot_count,
341                    worker_count: 1,
342                    max_workgroup_size_x: 256,
343                },
344                geometry(slot_count),
345            );
346        }
347        assert_eq!(cache.get_geometry(&hot_key), Some(hot_geometry));
348        cache.insert_geometry(
349            GeometryCacheKey {
350                slot_count: (GRID_PLAN_CACHE_CAP + 1) as u32,
351                worker_count: 1,
352                max_workgroup_size_x: 256,
353            },
354            geometry((GRID_PLAN_CACHE_CAP + 1) as u32),
355        );
356        assert_eq!(cache.get_geometry(&hot_key), Some(hot_geometry));
357    }
358}
vyre_runtime/megakernel/planner/grid.rs

vyre_runtime/megakernel/planner/
grid.rs
