vyre_runtime/megakernel/planner/
sizing.rs

1use vyre_driver::backend::BackendError;
2use vyre_foundation::execution_plan::SchedulingPolicy;
3
4use super::{ResidentGridLimits, ResidentGridPlan, ResidentGridRequest, ResidentLaunchGeometry};
5
6/// Shared worker-grid sizing policy for megakernel dispatch.
7///
8/// This is the host-side policy surface for persistent worker counts,
9/// workgroup width, slot padding, and backend grid geometry.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11pub struct ResidentSizingPolicy {
12    scheduling: SchedulingPolicy,
13}
14
15impl Default for ResidentSizingPolicy {
16    fn default() -> Self {
17        Self::standard()
18    }
19}
20
21impl ResidentSizingPolicy {
22    /// Standard megakernel sizing policy used by built-in dispatch paths.
23    #[must_use]
24    pub const fn standard() -> Self {
25        Self {
26            scheduling: SchedulingPolicy::standard(),
27        }
28    }
29
30    /// Build from a shared backend-neutral scheduling policy.
31    #[must_use]
32    pub const fn from_scheduling(scheduling: SchedulingPolicy) -> Self {
33        Self { scheduling }
34    }
35
36    /// Default persistent worker workgroup count.
37    #[must_use]
38    pub const fn default_worker_count(&self) -> u32 {
39        self.scheduling.default_worker_count()
40    }
41
42    /// Clamp a requested worker count into the legal workgroup x dimension.
43    #[must_use]
44    pub const fn worker_workgroup_size(&self, worker_count: u32, max_workgroup_size_x: u32) -> u32 {
45        self.scheduling
46            .worker_workgroup_size(worker_count, max_workgroup_size_x)
47    }
48
49    /// Round a logical slot count up to a whole worker workgroup.
50    #[must_use]
51    pub const fn padded_slot_count(&self, slot_count: u32, workgroup_size_x: u32) -> u32 {
52        self.scheduling
53            .padded_slot_count(slot_count, workgroup_size_x)
54    }
55
56    /// Compute the backend dispatch grid for a logical queue length.
57    #[must_use]
58    pub const fn dispatch_grid_for(
59        &self,
60        worker_count: u32,
61        queue_len: u32,
62        max_workgroup_size_x: u32,
63    ) -> [u32; 3] {
64        self.scheduling
65            .dispatch_grid_for(worker_count, queue_len, max_workgroup_size_x)
66    }
67
68    /// Compute a persistent-worker ceiling from adapter limits.
69    #[must_use]
70    pub const fn default_worker_groups_from_limits(
71        &self,
72        max_compute_workgroups_per_dimension: u32,
73        max_compute_invocations_per_workgroup: u32,
74    ) -> u32 {
75        self.scheduling.default_worker_groups_from_limits(
76            max_compute_workgroups_per_dimension,
77            max_compute_invocations_per_workgroup,
78        )
79    }
80
81    /// Resolve worker groups, workgroup width, slot padding, and dispatch grid.
82    ///
83    /// # Errors
84    ///
85    /// Returns [`BackendError`] when adapter limits are malformed.
86    pub fn calculate_optimal_grid(
87        &self,
88        request: ResidentGridRequest,
89        limits: ResidentGridLimits,
90    ) -> Result<ResidentGridPlan, BackendError> {
91        limits.validate()?;
92
93        let occupancy_worker_groups = self
94            .default_worker_groups_from_limits(
95                limits.max_compute_workgroups_per_dimension,
96                limits.max_compute_invocations_per_workgroup,
97            )
98            .min(limits.max_compute_workgroups_per_dimension);
99
100        let worker_groups = if request.requested_worker_groups == 0 {
101            occupancy_worker_groups
102        } else {
103            request
104                .requested_worker_groups
105                .min(limits.max_compute_workgroups_per_dimension)
106        }
107        .max(1);
108
109        let geometry = self.geometry_from_slots(
110            request.queue_len.max(1),
111            worker_groups,
112            limits.max_workgroup_size_x,
113        );
114
115        Ok(ResidentGridPlan {
116            geometry,
117            worker_groups,
118        })
119    }
120
121    /// Build geometry for an already-sized ring.
122    #[must_use]
123    pub fn geometry_from_slots(
124        &self,
125        slot_count: u32,
126        worker_count: u32,
127        max_workgroup_size_x: u32,
128    ) -> ResidentLaunchGeometry {
129        let workgroup_size_x = self.worker_workgroup_size(worker_count, max_workgroup_size_x);
130        let slot_count = self.padded_slot_count(slot_count, workgroup_size_x);
131        let dispatch_grid = self.dispatch_grid_for(worker_count, slot_count, workgroup_size_x);
132        ResidentLaunchGeometry {
133            workgroup_size_x,
134            slot_count,
135            dispatch_grid,
136        }
137    }
138}
vyre_runtime/megakernel/planner/sizing.rs

vyre_runtime/megakernel/planner/
sizing.rs
