vyre_driver/
megakernel_frontier.rs

1//! Backend-neutral frontier planning for dependency-aware megakernels.
2//!
3//! Backends can choose different execution topologies, but the memory envelope
4//! of dependency-layered frontier waves is a backend-neutral contract. This
5//! module plans that envelope once, including dependency barriers, fused-group
6//! splitting under an explicit byte budget, peak byte accounting, and readback
7//! pressure amortization, and then drives topology selection from it.
8//!
9//! Composing those two halves used to live in the CUDA driver, which is why the
10//! device-wide-barrier rule in [`crate::megakernel_execution`] could sit on one
11//! backend without the neutral policy knowing it. The composition is decided
12//! entirely by graph shape, wave bytes, and budgets, so it belongs here; a
13//! backend supplies only its telemetry and, through
14//! [`MegakernelExecutionPlanner`], its own plan cache.
15
16use crate::accounting::{
17    checked_add_u64_count as checked_add, checked_mul_u64_count as checked_mul,
18};
19use crate::megakernel_barrier::{
20    plan_megakernel_barriers_with_scratch, MegakernelBarrierGroup, MegakernelBarrierPlan,
21    MegakernelBarrierPlanError, MegakernelBarrierScratch, MegakernelWaveDependency,
22};
23use crate::megakernel_execution::{
24    megakernel_resident_graph_bytes, MegakernelDeviceCapabilities, MegakernelExecutionPlan,
25    MegakernelExecutionPlanner, MegakernelExecutionRequest, MegakernelExecutionSample,
26    MegakernelGraphShape, MegakernelMemoryError,
27};
28use crate::reservation_policy::{
29    reserve_typed_vec_to_capacity as reserve_vec_to_capacity, ReservationPolicy,
30};
31
32const MEGAKERNEL_FRONTIER_RESERVATION: ReservationPolicy = ReservationPolicy::new(
33    "megakernel frontier memory planner",
34    "shard the frontier wave group or split the fused phase",
35);
36
37/// Frontier-typed megakernel wave memory envelope.
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub struct MegakernelFrontierWave {
40    /// Resident frontier bytes touched by this wave.
41    pub frontier_bytes: u64,
42    /// Temporary scratch bytes required by this wave before topology scaling.
43    pub scratch_bytes: u64,
44    /// Output bytes produced by this wave.
45    pub output_bytes: u64,
46}
47
48/// Dependency-aware megakernel frontier memory plan.
49#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct MegakernelFrontierMemoryPlan {
51    /// Minimum global-barrier grouping after memory-budget splitting.
52    pub barriers: MegakernelBarrierPlan,
53    /// Peak frontier bytes across any fused barrier-free group.
54    pub peak_frontier_bytes: u64,
55    /// Peak scratch bytes across any fused barrier-free group.
56    pub peak_scratch_bytes: u64,
57    /// Peak output bytes across any fused barrier-free group.
58    pub peak_output_bytes: u64,
59    /// Readback pressure after combining runtime telemetry with static
60    /// fused-wave output volume.
61    pub amortized_readback_bytes: u64,
62    /// Widest barrier-free group in wave count.
63    pub max_group_width: usize,
64}
65
66/// Frontier memory planning failure.
67#[derive(Clone, Debug, Eq, PartialEq)]
68pub enum MegakernelFrontierMemoryPlanError {
69    /// Dependency graph cannot be barrier-planned.
70    Barrier(MegakernelBarrierPlanError),
71    /// Peak wave bytes overflowed while grouping a barrier-free phase.
72    ByteCountOverflow {
73        /// Field being accumulated.
74        field: &'static str,
75    },
76    /// Static graph or fused frontier bytes exceed the caller-approved budget.
77    GroupOverBudget {
78        /// Required bytes before topology selection.
79        required_bytes: u64,
80        /// Caller-provided budget.
81        budget_bytes: u64,
82        /// Budget region being checked.
83        field: &'static str,
84    },
85    /// Frontier planning result storage could not be reserved.
86    StorageReserveFailed {
87        /// Field being reserved.
88        field: &'static str,
89        /// Number of elements requested.
90        requested: usize,
91        /// Allocator error text.
92        message: String,
93    },
94}
95
96impl crate::accounting::ArithmeticOverflow for MegakernelFrontierMemoryPlanError {
97    fn arithmetic_overflow(field: &'static str) -> Self {
98        Self::ByteCountOverflow { field }
99    }
100}
101
102impl std::fmt::Display for MegakernelFrontierMemoryPlanError {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        match self {
105            Self::Barrier(error) => error.fmt(f),
106            Self::ByteCountOverflow { field } => write!(
107                f,
108                "megakernel frontier memory planner overflowed while accumulating {field}. Fix: shard the frontier wave group or split the fused phase."
109            ),
110            Self::GroupOverBudget {
111                required_bytes,
112                budget_bytes,
113                field,
114            } => write!(
115                f,
116                "megakernel frontier memory planner requires {required_bytes} bytes for {field} but budget allows {budget_bytes}. Fix: shard the graph/frontier waves or raise the explicit megakernel budget."
117            ),
118            Self::StorageReserveFailed {
119                field,
120                requested,
121                message,
122            } => write!(
123                f,
124                "megakernel frontier memory planner could not reserve {requested} {field} entries: {message}. Fix: shard the frontier waves before planning."
125            ),
126        }
127    }
128}
129
130impl std::error::Error for MegakernelFrontierMemoryPlanError {}
131
132impl From<MegakernelBarrierPlanError> for MegakernelFrontierMemoryPlanError {
133    fn from(error: MegakernelBarrierPlanError) -> Self {
134        Self::Barrier(error)
135    }
136}
137
138/// Dependency-aware megakernel execution plan for frontier waves.
139#[derive(Clone, Debug, Eq, PartialEq)]
140pub struct MegakernelFrontierExecutionPlan {
141    /// Topology and memory-budget plan for the peak barrier-free group.
142    pub execution: MegakernelExecutionPlan,
143    /// Minimum global-barrier grouping for the wave dependencies.
144    pub barriers: MegakernelBarrierPlan,
145    /// Peak frontier bytes across any fused barrier-free group.
146    pub peak_frontier_bytes: u64,
147    /// Peak scratch bytes across any fused barrier-free group.
148    pub peak_scratch_bytes: u64,
149    /// Peak output bytes across any fused barrier-free group.
150    pub peak_output_bytes: u64,
151    /// Readback pressure fed into topology selection after combining runtime
152    /// telemetry with static fused-wave output volume.
153    pub amortized_readback_bytes: u64,
154    /// Widest barrier-free group in wave count.
155    pub max_group_width: usize,
156}
157
158/// Dependency-aware frontier execution planning failure.
159#[derive(Clone, Debug, Eq, PartialEq)]
160pub enum MegakernelFrontierExecutionPlanError {
161    /// Dependency graph cannot be barrier-planned.
162    Barrier(MegakernelBarrierPlanError),
163    /// Peak wave bytes overflowed while grouping a barrier-free phase.
164    ByteCountOverflow {
165        /// Field being accumulated.
166        field: &'static str,
167    },
168    /// Static graph or fused frontier bytes exceed the caller-approved budget.
169    GroupOverBudget {
170        /// Required bytes before topology selection.
171        required_bytes: u64,
172        /// Caller-provided budget.
173        budget_bytes: u64,
174        /// Budget region being checked.
175        field: &'static str,
176    },
177    /// Topology-validated execution memory planning failed.
178    Memory(MegakernelMemoryError),
179    /// Frontier planning result storage could not be reserved.
180    StorageReserveFailed {
181        /// Field being reserved.
182        field: &'static str,
183        /// Number of elements requested.
184        requested: usize,
185        /// Allocator error text.
186        message: String,
187    },
188}
189
190impl crate::accounting::ArithmeticOverflow for MegakernelFrontierExecutionPlanError {
191    fn arithmetic_overflow(field: &'static str) -> Self {
192        Self::ByteCountOverflow { field }
193    }
194}
195
196impl std::fmt::Display for MegakernelFrontierExecutionPlanError {
197    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198        match self {
199            Self::Barrier(error) => error.fmt(f),
200            Self::ByteCountOverflow { field } => write!(
201                f,
202                "megakernel frontier execution planner overflowed while accumulating {field}. Fix: shard the frontier wave group or split the fused phase."
203            ),
204            Self::GroupOverBudget {
205                required_bytes,
206                budget_bytes,
207                field,
208            } => write!(
209                f,
210                "megakernel frontier execution planner requires {required_bytes} bytes for {field} but budget allows {budget_bytes}. Fix: shard the graph/frontier waves or raise the explicit megakernel budget."
211            ),
212            Self::Memory(error) => error.fmt(f),
213            Self::StorageReserveFailed {
214                field,
215                requested,
216                message,
217            } => write!(
218                f,
219                "megakernel frontier execution planner could not reserve {requested} {field} entries: {message}. Fix: shard the frontier waves before planning."
220            ),
221        }
222    }
223}
224
225impl std::error::Error for MegakernelFrontierExecutionPlanError {}
226
227impl From<MegakernelBarrierPlanError> for MegakernelFrontierExecutionPlanError {
228    fn from(error: MegakernelBarrierPlanError) -> Self {
229        Self::Barrier(error)
230    }
231}
232
233impl From<MegakernelMemoryError> for MegakernelFrontierExecutionPlanError {
234    fn from(error: MegakernelMemoryError) -> Self {
235        Self::Memory(error)
236    }
237}
238
239impl From<MegakernelFrontierMemoryPlanError> for MegakernelFrontierExecutionPlanError {
240    fn from(error: MegakernelFrontierMemoryPlanError) -> Self {
241        match error {
242            MegakernelFrontierMemoryPlanError::Barrier(error) => Self::Barrier(error),
243            MegakernelFrontierMemoryPlanError::ByteCountOverflow { field } => {
244                Self::ByteCountOverflow { field }
245            }
246            MegakernelFrontierMemoryPlanError::GroupOverBudget {
247                required_bytes,
248                budget_bytes,
249                field,
250            } => Self::GroupOverBudget {
251                required_bytes,
252                budget_bytes,
253                field,
254            },
255            MegakernelFrontierMemoryPlanError::StorageReserveFailed {
256                field,
257                requested,
258                message,
259            } => Self::StorageReserveFailed {
260                field,
261                requested,
262                message,
263            },
264        }
265    }
266}
267
268/// Plan dependency-aware megakernel execution for frontier-typed waves.
269///
270/// The planner minimizes global barriers from the wave dependencies, computes
271/// the peak memory envelope of any barrier-free fused group, and asks `planner`
272/// for a memory-validated topology for that envelope.
273///
274/// # Errors
275///
276/// Returns [`MegakernelFrontierExecutionPlanError`] when dependencies are
277/// invalid, counters overflow, storage cannot be reserved, or the envelope does
278/// not fit the explicit budget.
279pub fn plan_megakernel_frontier_execution(
280    planner: &mut impl MegakernelExecutionPlanner,
281    sample: MegakernelExecutionSample,
282    graph: MegakernelGraphShape,
283    bytes_per_node: u64,
284    bytes_per_edge: u64,
285    waves: &[MegakernelFrontierWave],
286    dependencies: &[MegakernelWaveDependency],
287    budget_bytes: u64,
288    launch_overhead_ns: f64,
289    fusion_pressure: f64,
290    capabilities: MegakernelDeviceCapabilities,
291) -> Result<MegakernelFrontierExecutionPlan, MegakernelFrontierExecutionPlanError> {
292    let mut scratch = MegakernelBarrierScratch::try_with_capacity(waves.len(), dependencies.len())?;
293    plan_megakernel_frontier_execution_with_scratch(
294        planner,
295        sample,
296        graph,
297        bytes_per_node,
298        bytes_per_edge,
299        waves,
300        dependencies,
301        budget_bytes,
302        launch_overhead_ns,
303        fusion_pressure,
304        capabilities,
305        &mut scratch,
306    )
307}
308
309/// Plan dependency-aware megakernel execution using caller-owned scratch.
310///
311/// # Errors
312///
313/// Same rejections as [`plan_megakernel_frontier_execution`].
314#[allow(clippy::too_many_arguments)]
315pub fn plan_megakernel_frontier_execution_with_scratch(
316    planner: &mut impl MegakernelExecutionPlanner,
317    sample: MegakernelExecutionSample,
318    graph: MegakernelGraphShape,
319    bytes_per_node: u64,
320    bytes_per_edge: u64,
321    waves: &[MegakernelFrontierWave],
322    dependencies: &[MegakernelWaveDependency],
323    budget_bytes: u64,
324    launch_overhead_ns: f64,
325    fusion_pressure: f64,
326    capabilities: MegakernelDeviceCapabilities,
327    scratch: &mut MegakernelBarrierScratch,
328) -> Result<MegakernelFrontierExecutionPlan, MegakernelFrontierExecutionPlanError> {
329    let graph_bytes = megakernel_resident_graph_bytes(graph, bytes_per_node, bytes_per_edge)?;
330    let memory = plan_megakernel_frontier_memory_with_scratch(
331        waves,
332        dependencies,
333        graph_bytes,
334        budget_bytes,
335        sample.readback_bytes,
336        scratch,
337    )?;
338    let execution = planner.plan_execution(MegakernelExecutionRequest {
339        sample: MegakernelExecutionSample {
340            readback_bytes: memory.amortized_readback_bytes,
341            ..sample
342        },
343        graph,
344        bytes_per_node,
345        bytes_per_edge,
346        frontier_bytes: memory.peak_frontier_bytes,
347        scratch_bytes: memory.peak_scratch_bytes,
348        output_bytes: memory.peak_output_bytes,
349        budget_bytes,
350        launch_overhead_ns,
351        fusion_pressure: capabilities.admissible_fusion_pressure(fusion_pressure),
352        capabilities,
353    })?;
354    Ok(MegakernelFrontierExecutionPlan {
355        execution,
356        barriers: memory.barriers,
357        peak_frontier_bytes: memory.peak_frontier_bytes,
358        peak_scratch_bytes: memory.peak_scratch_bytes,
359        peak_output_bytes: memory.peak_output_bytes,
360        amortized_readback_bytes: memory.amortized_readback_bytes,
361        max_group_width: memory.max_group_width,
362    })
363}
364
365/// Plan dependency-aware frontier memory using caller-owned barrier scratch.
366///
367/// # Errors
368///
369/// Returns [`MegakernelFrontierMemoryPlanError`] when dependencies are invalid,
370/// counters overflow, or the requested graph/frontier envelope cannot fit the
371/// explicit budget.
372pub fn plan_megakernel_frontier_memory_with_scratch(
373    waves: &[MegakernelFrontierWave],
374    dependencies: &[MegakernelWaveDependency],
375    resident_graph_bytes: u64,
376    budget_bytes: u64,
377    readback_bytes: u64,
378    scratch: &mut MegakernelBarrierScratch,
379) -> Result<MegakernelFrontierMemoryPlan, MegakernelFrontierMemoryPlanError> {
380    let barriers = plan_megakernel_barriers_with_scratch(waves.len(), dependencies, scratch)?;
381    let group_budget_bytes = budget_bytes.checked_sub(resident_graph_bytes).ok_or(
382        MegakernelFrontierMemoryPlanError::GroupOverBudget {
383            required_bytes: resident_graph_bytes,
384            budget_bytes,
385            field: "resident graph bytes",
386        },
387    )?;
388    let barriers = split_barrier_groups_to_memory_budget(barriers, waves, group_budget_bytes)?;
389    let mut peak_frontier_bytes = 0u64;
390    let mut peak_scratch_bytes = 0u64;
391    let mut peak_output_bytes = 0u64;
392    let mut max_group_width = 0usize;
393    for group in &barriers.groups {
394        let mut group_frontier_bytes = 0u64;
395        let mut group_scratch_bytes = 0u64;
396        let mut group_output_bytes = 0u64;
397        max_group_width = max_group_width.max(group.waves.len());
398        for &wave_index in &group.waves {
399            let wave = waves[wave_index];
400            group_frontier_bytes = checked_add::<MegakernelFrontierMemoryPlanError>(
401                group_frontier_bytes,
402                wave.frontier_bytes,
403                "frontier wave bytes",
404            )?;
405            group_scratch_bytes = checked_add::<MegakernelFrontierMemoryPlanError>(
406                group_scratch_bytes,
407                wave.scratch_bytes,
408                "scratch wave bytes",
409            )?;
410            group_output_bytes = checked_add::<MegakernelFrontierMemoryPlanError>(
411                group_output_bytes,
412                wave.output_bytes,
413                "output wave bytes",
414            )?;
415        }
416        peak_frontier_bytes = peak_frontier_bytes.max(group_frontier_bytes);
417        peak_scratch_bytes = peak_scratch_bytes.max(group_scratch_bytes);
418        peak_output_bytes = peak_output_bytes.max(group_output_bytes);
419    }
420
421    Ok(MegakernelFrontierMemoryPlan {
422        barriers,
423        peak_frontier_bytes,
424        peak_scratch_bytes,
425        peak_output_bytes,
426        amortized_readback_bytes: readback_bytes.max(peak_output_bytes),
427        max_group_width,
428    })
429}
430
431fn split_barrier_groups_to_memory_budget(
432    barriers: MegakernelBarrierPlan,
433    waves: &[MegakernelFrontierWave],
434    group_budget_bytes: u64,
435) -> Result<MegakernelBarrierPlan, MegakernelFrontierMemoryPlanError> {
436    let mut groups = Vec::new();
437    reserve_vec::<MegakernelBarrierGroup>(
438        &mut groups,
439        barriers.groups.len(),
440        "split barrier groups",
441    )?;
442    for group in barriers.groups {
443        split_one_barrier_group_to_memory_budget(group, waves, group_budget_bytes, &mut groups)?;
444    }
445    Ok(MegakernelBarrierPlan {
446        global_barriers: if groups.is_empty() {
447            0
448        } else {
449            groups.len() - 1
450        },
451        groups,
452    })
453}
454
455fn split_one_barrier_group_to_memory_budget(
456    group: MegakernelBarrierGroup,
457    waves: &[MegakernelFrontierWave],
458    group_budget_bytes: u64,
459    groups: &mut Vec<MegakernelBarrierGroup>,
460) -> Result<(), MegakernelFrontierMemoryPlanError> {
461    let mut current = Vec::new();
462    reserve_vec::<usize>(
463        &mut current,
464        group.waves.len().min(8),
465        "current split barrier group",
466    )?;
467    let mut current_bytes = 0u64;
468    for wave_index in group.waves {
469        let wave_bytes = megakernel_frontier_fused_wave_budget_bytes(waves[wave_index])?;
470        let combined = checked_add::<MegakernelFrontierMemoryPlanError>(
471            current_bytes,
472            wave_bytes,
473            "barrier group fused wave budget bytes",
474        )?;
475        if current.is_empty() && wave_bytes > group_budget_bytes {
476            return Err(MegakernelFrontierMemoryPlanError::GroupOverBudget {
477                required_bytes: wave_bytes,
478                budget_bytes: group_budget_bytes,
479                field: "single fused frontier wave bytes",
480            });
481        }
482        if !current.is_empty() && combined > group_budget_bytes {
483            groups.push(MegakernelBarrierGroup {
484                waves: std::mem::take(&mut current),
485            });
486            current_bytes = 0;
487        }
488        current.push(wave_index);
489        current_bytes = checked_add::<MegakernelFrontierMemoryPlanError>(
490            current_bytes,
491            wave_bytes,
492            "barrier group fused wave budget bytes",
493        )?;
494    }
495    if !current.is_empty() {
496        groups.push(MegakernelBarrierGroup { waves: current });
497    }
498    Ok(())
499}
500
501/// Compute the byte budget used to decide whether one frontier wave can fit in
502/// a fused barrier-free resident group.
503pub fn megakernel_frontier_fused_wave_budget_bytes(
504    wave: MegakernelFrontierWave,
505) -> Result<u64, MegakernelFrontierMemoryPlanError> {
506    let fused_scratch_bytes = checked_mul::<MegakernelFrontierMemoryPlanError>(
507        wave.scratch_bytes,
508        4,
509        "fused wave scratch bytes",
510    )?;
511    let bytes = checked_add::<MegakernelFrontierMemoryPlanError>(
512        wave.frontier_bytes,
513        fused_scratch_bytes,
514        "fused wave bytes",
515    )?;
516    checked_add::<MegakernelFrontierMemoryPlanError>(bytes, wave.output_bytes, "fused wave bytes")
517}
518
519fn reserve_vec<T>(
520    vec: &mut Vec<T>,
521    target_capacity: usize,
522    item: &'static str,
523) -> Result<(), MegakernelFrontierMemoryPlanError> {
524    reserve_vec_to_capacity(
525        MEGAKERNEL_FRONTIER_RESERVATION,
526        vec,
527        target_capacity,
528        item,
529        storage_reserve_failed,
530    )
531}
532
533fn storage_reserve_failed(
534    field: &'static str,
535    requested: usize,
536    message: String,
537) -> MegakernelFrontierMemoryPlanError {
538    MegakernelFrontierMemoryPlanError::StorageReserveFailed {
539        field,
540        requested,
541        message,
542    }
543}
544
545#[cfg(test)]
546mod tests {
547    use super::{
548        megakernel_frontier_fused_wave_budget_bytes, plan_megakernel_frontier_execution,
549        plan_megakernel_frontier_memory_with_scratch, MegakernelFrontierExecutionPlanError,
550        MegakernelFrontierMemoryPlanError, MegakernelFrontierWave,
551    };
552    use crate::megakernel_barrier::{MegakernelBarrierScratch, MegakernelWaveDependency};
553    use crate::megakernel_execution::{
554        MegakernelDeviceCapabilities, MegakernelExecutionSample, MegakernelExecutionTopology,
555        MegakernelGraphShape, NeutralMegakernelExecutionPlanner,
556    };
557
558    #[test]
559    fn frontier_memory_plan_uses_peak_barrier_group_memory() {
560        let mut scratch = MegakernelBarrierScratch::default();
561        let plan = plan_megakernel_frontier_memory_with_scratch(
562            &[
563                MegakernelFrontierWave {
564                    frontier_bytes: 1_024,
565                    scratch_bytes: 512,
566                    output_bytes: 256,
567                },
568                MegakernelFrontierWave {
569                    frontier_bytes: 2_048,
570                    scratch_bytes: 1_024,
571                    output_bytes: 512,
572                },
573                MegakernelFrontierWave {
574                    frontier_bytes: 4_096,
575                    scratch_bytes: 2_048,
576                    output_bytes: 1_024,
577                },
578                MegakernelFrontierWave {
579                    frontier_bytes: 8_192,
580                    scratch_bytes: 4_096,
581                    output_bytes: 2_048,
582                },
583            ],
584            &[
585                MegakernelWaveDependency {
586                    before: 0,
587                    after: 1,
588                },
589                MegakernelWaveDependency {
590                    before: 0,
591                    after: 2,
592                },
593                MegakernelWaveDependency {
594                    before: 1,
595                    after: 3,
596                },
597                MegakernelWaveDependency {
598                    before: 2,
599                    after: 3,
600                },
601            ],
602            16_000,
603            128 * 1024,
604            1 << 20,
605            &mut scratch,
606        )
607        .expect("Fix: frontier-typed megakernel memory plan should fit the budget.");
608
609        assert_eq!(plan.barriers.global_barriers, 2);
610        assert_eq!(plan.barriers.groups[1].waves, vec![1, 2]);
611        assert_eq!(plan.peak_frontier_bytes, 8_192);
612        assert_eq!(plan.peak_scratch_bytes, 4_096);
613        assert_eq!(plan.peak_output_bytes, 2_048);
614        assert_eq!(plan.amortized_readback_bytes, 1 << 20);
615        assert_eq!(plan.max_group_width, 2);
616    }
617
618    #[test]
619    fn frontier_memory_uses_static_group_output_to_amortize_readback() {
620        let mut scratch = MegakernelBarrierScratch::default();
621        let plan = plan_megakernel_frontier_memory_with_scratch(
622            &[
623                MegakernelFrontierWave {
624                    frontier_bytes: 1_024,
625                    scratch_bytes: 512,
626                    output_bytes: 3_072,
627                },
628                MegakernelFrontierWave {
629                    frontier_bytes: 1_024,
630                    scratch_bytes: 512,
631                    output_bytes: 3_072,
632                },
633            ],
634            &[],
635            16_000,
636            128 * 1024,
637            0,
638            &mut scratch,
639        )
640        .expect("Fix: static output-amortized frontier memory plan should fit the budget.");
641
642        assert_eq!(plan.peak_output_bytes, 6_144);
643        assert_eq!(plan.amortized_readback_bytes, 6_144);
644    }
645
646    #[test]
647    fn frontier_memory_splits_independent_layers_to_fit_fused_budget() {
648        let mut scratch = MegakernelBarrierScratch::default();
649        let waves = [
650            MegakernelFrontierWave {
651                frontier_bytes: 10,
652                scratch_bytes: 10,
653                output_bytes: 10,
654            },
655            MegakernelFrontierWave {
656                frontier_bytes: 10,
657                scratch_bytes: 10,
658                output_bytes: 10,
659            },
660            MegakernelFrontierWave {
661                frontier_bytes: 10,
662                scratch_bytes: 10,
663                output_bytes: 10,
664            },
665        ];
666        let plan =
667            plan_megakernel_frontier_memory_with_scratch(&waves, &[], 0, 100, 4_096, &mut scratch)
668                .expect("Fix: independent frontier waves should split into budget-fit chunks.");
669
670        assert_eq!(plan.barriers.groups.len(), 3);
671        assert_eq!(plan.barriers.global_barriers, 2);
672        assert_eq!(plan.max_group_width, 1);
673        assert_eq!(plan.peak_frontier_bytes, 10);
674        assert_eq!(plan.peak_scratch_bytes, 10);
675        assert_eq!(plan.peak_output_bytes, 10);
676    }
677
678    #[test]
679    fn frontier_memory_rejects_graph_and_single_wave_over_budget() {
680        let mut scratch = MegakernelBarrierScratch::default();
681        let graph_error = plan_megakernel_frontier_memory_with_scratch(
682            &[MegakernelFrontierWave {
683                frontier_bytes: 1,
684                scratch_bytes: 1,
685                output_bytes: 1,
686            }],
687            &[],
688            1_600,
689            1_000,
690            0,
691            &mut scratch,
692        )
693        .expect_err("resident graph bytes above budget must fail before split planning");
694        assert_eq!(
695            graph_error,
696            MegakernelFrontierMemoryPlanError::GroupOverBudget {
697                required_bytes: 1_600,
698                budget_bytes: 1_000,
699                field: "resident graph bytes",
700            }
701        );
702
703        let wave_error = plan_megakernel_frontier_memory_with_scratch(
704            &[MegakernelFrontierWave {
705                frontier_bytes: 100,
706                scratch_bytes: 100,
707                output_bytes: 100,
708            }],
709            &[],
710            0,
711            500,
712            0,
713            &mut scratch,
714        )
715        .expect_err("single fused wave above group budget must fail before topology planning");
716        assert_eq!(
717            wave_error,
718            MegakernelFrontierMemoryPlanError::GroupOverBudget {
719                required_bytes: 600,
720                budget_bytes: 500,
721                field: "single fused frontier wave bytes",
722            }
723        );
724    }
725
726    #[test]
727    fn frontier_fused_wave_budget_uses_topology_scratch_multiplier() {
728        assert_eq!(
729            megakernel_frontier_fused_wave_budget_bytes(MegakernelFrontierWave {
730                frontier_bytes: 16,
731                scratch_bytes: 16,
732                output_bytes: 16,
733            })
734            .expect("Fix: fused frontier wave budget should fit"),
735            96
736        );
737    }
738
739    #[test]
740    fn frontier_memory_fails_loudly_on_wave_byte_overflow() {
741        let mut scratch = MegakernelBarrierScratch::default();
742        let error = plan_megakernel_frontier_memory_with_scratch(
743            &[
744                MegakernelFrontierWave {
745                    frontier_bytes: u64::MAX,
746                    scratch_bytes: 1,
747                    output_bytes: 1,
748                },
749                MegakernelFrontierWave {
750                    frontier_bytes: 1,
751                    scratch_bytes: 1,
752                    output_bytes: 1,
753                },
754            ],
755            &[],
756            2,
757            u64::MAX,
758            0,
759            &mut scratch,
760        )
761        .expect_err("Fix: overflowed frontier wave bytes must fail before launch planning.");
762
763        assert_eq!(
764            error,
765            MegakernelFrontierMemoryPlanError::ByteCountOverflow {
766                field: "fused wave bytes"
767            }
768        );
769    }
770
771    #[test]
772    fn generated_frontier_memory_profiles_preserve_peak_and_budget_for_1024_shapes() {
773        let mut scratch = MegakernelBarrierScratch::default();
774        for width in 1u64..=32 {
775            for depth in 1u64..=32 {
776                let mut waves = Vec::new();
777                let mut dependencies = Vec::new();
778                for layer in 0..depth {
779                    for slot in 0..width {
780                        waves.push(MegakernelFrontierWave {
781                            frontier_bytes: width,
782                            scratch_bytes: slot + 1,
783                            output_bytes: layer + 1,
784                        });
785                        if layer + 1 < depth {
786                            dependencies.push(MegakernelWaveDependency {
787                                before: (layer * width + slot) as usize,
788                                after: ((layer + 1) * width + slot) as usize,
789                            });
790                        }
791                    }
792                }
793
794                let plan = plan_megakernel_frontier_memory_with_scratch(
795                    &waves,
796                    &dependencies,
797                    256,
798                    u64::MAX / 2,
799                    7,
800                    &mut scratch,
801                )
802                .expect("Fix: generated frontier memory DAG should plan under large budget.");
803
804                assert_eq!(plan.barriers.groups.len(), depth as usize);
805                assert_eq!(plan.max_group_width, width as usize);
806                assert_eq!(plan.peak_frontier_bytes, width * width);
807                assert_eq!(plan.peak_scratch_bytes, width * (width + 1) / 2);
808                assert_eq!(plan.peak_output_bytes, width * depth);
809                assert_eq!(plan.amortized_readback_bytes, 7.max(width * depth));
810            }
811        }
812    }
813
814    const FUSION_WAVES: &[MegakernelFrontierWave] = &[
815        MegakernelFrontierWave {
816            frontier_bytes: 1_024,
817            scratch_bytes: 512,
818            output_bytes: 256,
819        },
820        MegakernelFrontierWave {
821            frontier_bytes: 2_048,
822            scratch_bytes: 1_024,
823            output_bytes: 512,
824        },
825    ];
826
827    fn fused_pressure_plan(
828        capabilities: MegakernelDeviceCapabilities,
829    ) -> Result<super::MegakernelFrontierExecutionPlan, MegakernelFrontierExecutionPlanError> {
830        plan_megakernel_frontier_execution(
831            &mut NeutralMegakernelExecutionPlanner,
832            MegakernelExecutionSample {
833                dispatch_cost_ns: 1_000.0,
834                frontier_density: 0.50,
835                readback_bytes: 1 << 20,
836            },
837            MegakernelGraphShape {
838                node_count: 1_000,
839                edge_count: 4_000,
840            },
841            16,
842            8,
843            FUSION_WAVES,
844            &[MegakernelWaveDependency {
845                before: 0,
846                after: 1,
847            }],
848            128 * 1024,
849            250.0,
850            0.95,
851            capabilities,
852        )
853    }
854
855    #[test]
856    fn frontier_execution_plans_barriers_and_topology_from_one_envelope() {
857        let plan = fused_pressure_plan(MegakernelDeviceCapabilities::FUSION_CAPABLE)
858            .expect("Fix: dependency-layered frontier waves should fit the budget.");
859
860        assert_eq!(plan.barriers.global_barriers, 1);
861        assert_eq!(plan.barriers.groups[0].waves, vec![0]);
862        assert_eq!(plan.barriers.groups[1].waves, vec![1]);
863        assert_eq!(plan.peak_frontier_bytes, 2_048);
864        assert_eq!(plan.peak_scratch_bytes, 1_024);
865        assert_eq!(plan.peak_output_bytes, 512);
866        assert_eq!(plan.amortized_readback_bytes, 1 << 20);
867        assert_eq!(plan.max_group_width, 1);
868        assert_eq!(
869            plan.execution.topology,
870            MegakernelExecutionTopology::FusedWave
871        );
872        assert_eq!(plan.execution.memory.frontier_bytes, 2_048);
873        assert_eq!(plan.execution.memory.scratch_bytes, 4_096);
874    }
875
876    #[test]
877    fn frontier_execution_refuses_a_fused_wave_without_a_device_wide_barrier() {
878        let capable = fused_pressure_plan(MegakernelDeviceCapabilities::FUSION_CAPABLE)
879            .expect("Fix: capable device should plan.");
880        let incapable = fused_pressure_plan(MegakernelDeviceCapabilities::FUSION_INCAPABLE)
881            .expect("Fix: a device without a device-wide barrier still gets a plan.");
882
883        assert_eq!(
884            capable.barriers, incapable.barriers,
885            "Fix: device capability changes the topology, never the dependency grouping."
886        );
887        assert_ne!(
888            incapable.execution.topology,
889            MegakernelExecutionTopology::FusedWave,
890            "Fix: a fused wave crosses wave boundaries inside one launch and needs a barrier \
891             across every resident block; a device without one cannot run the plan."
892        );
893        assert!(
894            incapable.execution.memory.required_bytes < capable.execution.memory.required_bytes,
895            "Fix: refusing fusion must also drop the fused scratch multiplier from the envelope."
896        );
897    }
898
899    #[test]
900    fn frontier_execution_rejects_a_graph_that_leaves_no_wave_headroom() {
901        let error = plan_megakernel_frontier_execution(
902            &mut NeutralMegakernelExecutionPlanner,
903            MegakernelExecutionSample {
904                dispatch_cost_ns: 1_000.0,
905                frontier_density: 0.50,
906                readback_bytes: 4_096,
907            },
908            MegakernelGraphShape {
909                node_count: 100,
910                edge_count: 100,
911            },
912            8,
913            8,
914            &[MegakernelFrontierWave {
915                frontier_bytes: 1,
916                scratch_bytes: 1,
917                output_bytes: 1,
918            }],
919            &[],
920            1_000,
921            250.0,
922            0.95,
923            MegakernelDeviceCapabilities::FUSION_CAPABLE,
924        )
925        .expect_err("Fix: resident graph bytes above budget must fail before split planning.");
926
927        assert_eq!(
928            error,
929            MegakernelFrontierExecutionPlanError::GroupOverBudget {
930                required_bytes: 1_600,
931                budget_bytes: 1_000,
932                field: "resident graph bytes",
933            }
934        );
935    }
936
937    #[test]
938    fn frontier_execution_fails_loudly_on_wave_byte_overflow() {
939        let error = plan_megakernel_frontier_execution(
940            &mut NeutralMegakernelExecutionPlanner,
941            MegakernelExecutionSample {
942                dispatch_cost_ns: 1_000.0,
943                frontier_density: 0.90,
944                readback_bytes: 1 << 20,
945            },
946            MegakernelGraphShape {
947                node_count: 1,
948                edge_count: 1,
949            },
950            1,
951            1,
952            &[
953                MegakernelFrontierWave {
954                    frontier_bytes: u64::MAX,
955                    scratch_bytes: 1,
956                    output_bytes: 1,
957                },
958                MegakernelFrontierWave {
959                    frontier_bytes: 1,
960                    scratch_bytes: 1,
961                    output_bytes: 1,
962                },
963            ],
964            &[],
965            u64::MAX,
966            250.0,
967            0.95,
968            MegakernelDeviceCapabilities::FUSION_CAPABLE,
969        )
970        .expect_err("Fix: overflowed frontier wave bytes must fail before launch planning.");
971
972        assert_eq!(
973            error,
974            MegakernelFrontierExecutionPlanError::ByteCountOverflow {
975                field: "fused wave bytes"
976            }
977        );
978    }
979}
vyre_driver/megakernel_frontier.rs

vyre_driver/
megakernel_frontier.rs
