qualia_core_db/wgsl_forge/
schedule.rs

1use serde::{Deserialize, Serialize};
2
3use super::ir::IntrinsicClass;
4use super::{ForgeError, KernelSpec};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
7pub struct Schedule {
8    pub workgroup_size: u32,
9    pub items_per_invocation: u32,
10    pub vector_width: u32,
11}
12
13impl Default for Schedule {
14    fn default() -> Self {
15        Self {
16            workgroup_size: 64,
17            items_per_invocation: 1,
18            vector_width: 1,
19        }
20    }
21}
22
23impl Schedule {
24    pub fn validate(
25        self,
26        kernel: &KernelSpec,
27        constraints: &AdapterConstraints,
28    ) -> Result<(), ForgeError> {
29        kernel.validate()?;
30        if !matches!(self.workgroup_size, 1..=1024) || !self.workgroup_size.is_power_of_two() {
31            return Err(ForgeError::InvalidSchedule(
32                "workgroup size must be a power of two in 1..=1024".to_string(),
33            ));
34        }
35        if self.workgroup_size > constraints.max_workgroup_size_x
36            || self.workgroup_size > constraints.max_invocations_per_workgroup
37        {
38            return Err(ForgeError::InvalidSchedule(format!(
39                "workgroup {} exceeds adapter limit {}",
40                self.workgroup_size,
41                constraints
42                    .max_workgroup_size_x
43                    .min(constraints.max_invocations_per_workgroup)
44            )));
45        }
46        if !matches!(self.items_per_invocation, 1 | 2 | 4 | 8) {
47            return Err(ForgeError::InvalidSchedule(
48                "items per invocation must be one of 1, 2, 4, 8".to_string(),
49            ));
50        }
51        if !matches!(self.vector_width, 1 | 2 | 4) {
52            return Err(ForgeError::InvalidSchedule(
53                "vector width must be one of 1, 2, 4".to_string(),
54            ));
55        }
56
57        Ok(())
58    }
59
60    pub const fn elements_per_workgroup(self) -> u32 {
61        self.workgroup_size * self.items_per_invocation * self.vector_width
62    }
63
64    pub fn dispatch_workgroups(self, element_count: usize) -> u32 {
65        let width = self.elements_per_workgroup().max(1) as usize;
66        element_count.div_ceil(width).max(1) as u32
67    }
68
69    pub const fn sort_key(self) -> (u32, u32, u32) {
70        (
71            self.workgroup_size,
72            self.items_per_invocation,
73            self.vector_width,
74        )
75    }
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
79pub struct AdapterConstraints {
80    pub max_workgroup_size_x: u32,
81    pub max_invocations_per_workgroup: u32,
82    pub max_workgroups_per_dimension: u32,
83    pub supports_subgroups: bool,
84    /// Cooperative-matrix / Tensor-core matrix-multiply-accumulate support.
85    pub supports_coopmat: bool,
86    /// Ray-query / RT-core support (hardware ray-triangle intersection).
87    pub supports_rt_cores: bool,
88    /// SIMD/warp width (32 NVIDIA, 64 AMD/Apple-ish). Workgroup sizes that are
89    /// not a multiple of this are pruned from the tuning search (plan §6).
90    pub warp_size: u32,
91}
92
93impl AdapterConstraints {
94    pub const fn portable() -> Self {
95        Self {
96            max_workgroup_size_x: 256,
97            max_invocations_per_workgroup: 256,
98            max_workgroups_per_dimension: 65_535,
99            supports_subgroups: false,
100            supports_coopmat: false,
101            supports_rt_cores: false,
102            warp_size: 32,
103        }
104    }
105
106    /// Intrinsic-availability check (plan §6): a kernel that requires RT or
107    /// cooperative-matrix hardware cannot run on an adapter that lacks it, so such
108    /// kernels are pruned before tuning/certification on the local adapter.
109    /// Subgroup intrinsics are not pruned because they can be lowered to a
110    /// shared-memory equivalent. This is a hardware-vs-kernel check and is
111    /// deliberately separate from schedule emission (which is hardware-agnostic).
112    pub fn supports_kernel(&self, kernel: &KernelSpec) -> Result<(), ForgeError> {
113        for intrinsic in kernel.required_intrinsics() {
114            match intrinsic.class() {
115                IntrinsicClass::CooperativeMatrix if !self.supports_coopmat => {
116                    return Err(ForgeError::InvalidSchedule(
117                        "kernel requires cooperative-matrix (tensor cores) unavailable on this adapter".to_string(),
118                    ));
119                }
120                IntrinsicClass::RayTracing if !self.supports_rt_cores => {
121                    return Err(ForgeError::InvalidSchedule(
122                        "kernel requires ray-query (RT cores) unavailable on this adapter"
123                            .to_string(),
124                    ));
125                }
126                _ => {}
127            }
128        }
129        Ok(())
130    }
131
132    #[cfg(feature = "gpu-runtime")]
133    pub fn from_wgpu_limits(limits: &wgpu::Limits) -> Self {
134        // Intrinsic-capability flags default to false here and are populated from
135        // the adapter's feature set at device creation (see WgpuComputeContext).
136        Self {
137            max_workgroup_size_x: limits.max_compute_workgroup_size_x,
138            max_invocations_per_workgroup: limits.max_compute_invocations_per_workgroup,
139            max_workgroups_per_dimension: limits.max_compute_workgroups_per_dimension,
140            supports_subgroups: false,
141            supports_coopmat: false,
142            supports_rt_cores: false,
143            warp_size: 32,
144        }
145    }
146}
147
148#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
149pub struct ScheduleSpace {
150    pub workgroup_sizes: Vec<u32>,
151    pub items_per_invocation: Vec<u32>,
152    pub vector_widths: Vec<u32>,
153}
154
155impl Default for ScheduleSpace {
156    fn default() -> Self {
157        Self {
158            workgroup_sizes: vec![32, 64, 128, 256],
159            items_per_invocation: vec![1, 2, 4, 8],
160            vector_widths: vec![1, 2, 4],
161        }
162    }
163}
164
165impl ScheduleSpace {
166    /// Generate the pruned candidate schedules for `kernel` on this adapter.
167    ///
168    /// Pruning here is limited to what is portably knowable: adapter workgroup/
169    /// invocation limits and `Schedule::validate`, plus warp-size alignment.
170    ///
171    /// **Known limitations (plan §6, honest — not implemented):** candidates are NOT
172    /// pruned by a device-relative roofline (wgpu exposes no peak FLOPS/bandwidth — see
173    /// [`crate::wgsl_forge::roofline`]), by compute-unit saturation (no CU/SM count is
174    /// exposed by wgpu), or by a cross-vendor thermal/power model (no portable
175    /// temperature sensor; thermal pacing is handled NVIDIA-only via `nvidia-smi` in the
176    /// `auto-tune-all` CLI loop, not here). Any of these would need a calibration
177    /// micro-benchmark or a platform-specific probe.
178    pub fn candidates(
179        &self,
180        kernel: &KernelSpec,
181        constraints: &AdapterConstraints,
182    ) -> Vec<Schedule> {
183        let mut candidates = Vec::new();
184        for &workgroup_size in &self.workgroup_sizes {
185            for &items_per_invocation in &self.items_per_invocation {
186                for &vector_width in &self.vector_widths {
187                    let schedule = Schedule {
188                        workgroup_size,
189                        items_per_invocation,
190                        vector_width,
191                        ..Default::default()
192                    };
193                    // Prune non-warp-aligned workgroups (plan §6): a partial warp
194                    // wastes lanes. Generation is unaffected — this only narrows
195                    // the tuning search.
196                    let warp_aligned =
197                        constraints.warp_size <= 1 || workgroup_size % constraints.warp_size == 0;
198                    if warp_aligned && schedule.validate(kernel, constraints).is_ok() {
199                        candidates.push(schedule);
200                    }
201                }
202            }
203        }
204        candidates.sort_by_key(|schedule| schedule.sort_key());
205        candidates.dedup();
206        candidates
207    }
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213    use crate::wgsl_forge::BuiltinKernel;
214
215    #[test]
216    fn candidates_are_bounded_and_deterministic() {
217        let kernel = BuiltinKernel::AffineF32.spec();
218        let constraints = AdapterConstraints::portable();
219        let first = ScheduleSpace::default().candidates(&kernel, &constraints);
220        let second = ScheduleSpace::default().candidates(&kernel, &constraints);
221        assert_eq!(first, second);
222        assert_eq!(first.len(), 48);
223        assert!(first.iter().all(|schedule| schedule.workgroup_size <= 256));
224    }
225
226    #[test]
227    fn dispatch_covers_tail_elements() {
228        let schedule = Schedule {
229            workgroup_size: 64,
230            items_per_invocation: 2,
231            vector_width: 4,
232            ..Default::default()
233        };
234        assert_eq!(schedule.dispatch_workgroups(513), 2);
235    }
236
237    #[test]
238    fn invalid_schedule_is_rejected_before_emission() {
239        use crate::wgsl_forge::{generate_builtin, ForgeError, TargetBackend};
240
241        // workgroup_size = 100 is not a power of two, so `Schedule::validate`
242        // (invoked by `generate_builtin` *before* `emit_shader`) must reject it.
243        let invalid = Schedule {
244            workgroup_size: 100,
245            items_per_invocation: 1,
246            vector_width: 1,
247        };
248        // Guard the precondition: the schedule itself is invalid.
249        let kernel = BuiltinKernel::AffineF32.spec();
250        assert!(invalid
251            .validate(&kernel, &AdapterConstraints::portable())
252            .is_err());
253
254        let result = generate_builtin(BuiltinKernel::AffineF32, invalid, TargetBackend::Wgsl);
255        // It must be the schedule-validation error, not an emission error — proving
256        // generation never ran for an invalid schedule.
257        match result {
258            Err(ForgeError::InvalidSchedule(message)) => {
259                assert!(
260                    message.contains("power of two"),
261                    "expected a power-of-two schedule error, got: {message}"
262                );
263            }
264            other => panic!("expected InvalidSchedule before emission, got {other:?}"),
265        }
266    }
267
268    #[test]
269    fn rt_kernel_pruned_without_rt_cores() {
270        use crate::wgsl_forge::ir::{
271            BufferAccess, BufferElement, BufferSpec, Intrinsic, KernelSpec, Op, ScalarType,
272        };
273
274        let kernel = KernelSpec {
275            id: "rt-probe".to_string(),
276            semantic_version: 1,
277            entry_point: "rt_probe".to_string(),
278            description: "ray-query smoke kernel".to_string(),
279            buffers: vec![BufferSpec {
280                group: 0,
281                binding: 0,
282                name: "output".to_string(),
283                element: BufferElement::Scalar(ScalarType::F32),
284                access: BufferAccess::StorageReadWrite,
285            }],
286            ops: vec![Op::Intrinsic(Intrinsic::RayQuery {
287                acceleration_structure: "tlas".to_string(),
288                origin: "o".to_string(),
289                direction: "d".to_string(),
290                t_min: "tmin".to_string(),
291                t_max: "tmax".to_string(),
292                destination: "hit".to_string(),
293            })],
294            shared_memory: Vec::new(),
295        };
296
297        // Generation/schedule validity is hardware-agnostic and must succeed.
298        assert!(Schedule::default()
299            .validate(&kernel, &AdapterConstraints::portable())
300            .is_ok());
301
302        // The hardware-vs-kernel check prunes it when RT cores are absent...
303        let without = AdapterConstraints::portable();
304        assert!(without.supports_kernel(&kernel).is_err());
305
306        // ...and accepts it when present.
307        let with = AdapterConstraints {
308            supports_rt_cores: true,
309            ..AdapterConstraints::portable()
310        };
311        assert!(with.supports_kernel(&kernel).is_ok());
312    }
313}
qualia_core_db/wgsl_forge/schedule.rs

qualia_core_db/wgsl_forge/
schedule.rs
