vyre_driver/
megakernel_execution.rs

1//! Backend-neutral execution planning for persistent megakernel waves.
2//!
3//! Backends can feed telemetry and device budgets into this module to choose a
4//! sparse, dense, hybrid, or fused execution topology before allocating device
5//! scratch. The policy is deterministic, allocation-free, and validates byte
6//! pressure before a backend reaches an API-specific allocation path.
7//!
8//! One rule used to live only in the CUDA copy of this policy: a `FusedWave`
9//! runs dependency-ordered waves inside a single launch, so it needs a barrier
10//! across every resident block, and a device without one cannot run the plan at
11//! all. The neutral policy did not know that, so for the same wave it answered
12//! `FusedWave` where the CUDA fork answered a per-launch topology, and any
13//! backend that had not written the check itself would have been handed an
14//! unlaunchable plan. The check is a property of the device, not of CUDA, so it
15//! is [`MegakernelDeviceCapabilities`] here and every backend inherits it.
16
17const WARP_SPARSE_DENSITY: f64 = 0.03125;
18const SPARSE_DENSITY: f64 = 0.125;
19const DENSE_DENSITY: f64 = 0.70;
20const BLOCK_DENSE_DENSITY: f64 = 0.85;
21const FUSION_PRESSURE: f64 = 0.70;
22const FUSION_PRESSURE_HYSTERESIS: f64 = 0.10;
23const FRONTIER_HYSTERESIS: f64 = 0.025;
24const MEMORY_RED_ZONE_BPS: u32 = 9_000;
25const MEMORY_HYSTERESIS_BPS: u32 = 250;
26const LAUNCH_PRESSURE_BPS: u32 = 1_500;
27const LAUNCH_HYSTERESIS_BPS: u32 = 250;
28const FUSION_READBACK_BYTES: u64 = 4_096;
29const DENSE_AVERAGE_DEGREE_BPS: u64 = 20_000;
30const WARP_SPARSE_AVERAGE_DEGREE_BPS: u64 = 80_000;
31
32/// Device capabilities that constrain which wave topologies are launchable.
33#[derive(Clone, Copy, Debug, Eq, PartialEq)]
34pub struct MegakernelDeviceCapabilities {
35    /// Whether every resident block can synchronize inside one launch.
36    pub supports_device_wide_barrier: bool,
37}
38
39impl MegakernelDeviceCapabilities {
40    /// Device that can host a fused wave.
41    pub const FUSION_CAPABLE: Self = Self {
42        supports_device_wide_barrier: true,
43    };
44    /// Device that must keep every wave in its own launch.
45    pub const FUSION_INCAPABLE: Self = Self {
46        supports_device_wide_barrier: false,
47    };
48
49    /// Fusion pressure this device can act on.
50    ///
51    /// Without a device-wide barrier the fused plan is unlaunchable, so the
52    /// measured pressure toward it is zero however high the caller observed it.
53    #[must_use]
54    pub fn admissible_fusion_pressure(self, fusion_pressure: f64) -> f64 {
55        if self.supports_device_wide_barrier {
56            fusion_pressure
57        } else {
58            0.0
59        }
60    }
61}
62
63/// Per-candidate telemetry used to bias megakernel fusion.
64#[derive(Clone, Copy, Debug, PartialEq)]
65pub struct MegakernelExecutionSample {
66    /// Observed candidate dispatch cost in nanoseconds.
67    pub dispatch_cost_ns: f64,
68    /// Observed active-frontier density in `[0, 1]`.
69    pub frontier_density: f64,
70    /// Observed final readback byte volume.
71    pub readback_bytes: u64,
72}
73
74/// Device-side megakernel execution topology selected for a dataflow wave.
75#[derive(Clone, Copy, Debug, Eq, PartialEq)]
76pub enum MegakernelExecutionTopology {
77    /// Ultra-low-density frontier expansion where one warp owns sparse active
78    /// nodes and avoids block-wide work distribution overhead.
79    WarpSparseFrontier,
80    /// Low-density frontier expansion with queue-like work distribution.
81    SparseFrontier,
82    /// Very high-density propagation where a block owns coalesced bitset lanes
83    /// and amortizes shared-memory scans across many active facts.
84    BlockDenseFrontier,
85    /// Dense bitset-style propagation with coalesced scans.
86    DenseFrontier,
87    /// Mixed sparse/dense execution when density is in the transition band.
88    HybridFrontier,
89    /// Fused adjacent waves when launch/readback pressure dominates and memory
90    /// budget leaves room for the fused plan.
91    FusedWave,
92}
93
94/// Static graph shape used by topology selection.
95#[derive(Clone, Copy, Debug, Eq, PartialEq)]
96pub struct MegakernelGraphShape {
97    /// Logical graph node count.
98    pub node_count: u64,
99    /// Logical graph edge count.
100    pub edge_count: u64,
101}
102
103/// Device memory envelope for a candidate megakernel plan.
104#[derive(Clone, Copy, Debug, Eq, PartialEq)]
105pub struct MegakernelMemoryBudget {
106    /// Estimated resident plus transient bytes required by the candidate plan.
107    pub required_bytes: u64,
108    /// Caller-approved device-memory budget for the plan.
109    pub budget_bytes: u64,
110}
111
112/// Detailed megakernel memory plan.
113#[derive(Clone, Copy, Debug, Eq, PartialEq)]
114pub struct MegakernelMemoryPlan {
115    /// Graph-layout bytes retained on device.
116    pub graph_bytes: u64,
117    /// Frontier-state bytes retained on device.
118    pub frontier_bytes: u64,
119    /// Temporary scratch bytes required by the selected topology.
120    pub scratch_bytes: u64,
121    /// Final compact output/readback bytes.
122    pub output_bytes: u64,
123    /// Total peak bytes required by the plan.
124    pub required_bytes: u64,
125    /// Caller-approved byte budget.
126    pub budget_bytes: u64,
127    /// Required/budget pressure in basis points.
128    pub memory_pressure_bps: u32,
129}
130
131/// Complete megakernel execution plan selected from runtime telemetry.
132#[derive(Clone, Copy, Debug, Eq, PartialEq)]
133pub struct MegakernelExecutionPlan {
134    /// Final topology after memory-budget validation.
135    pub topology: MegakernelExecutionTopology,
136    /// Memory plan for the final topology.
137    pub memory: MegakernelMemoryPlan,
138    /// Whether the planner downgraded a denser/fused topology to sparse to fit
139    /// the explicit memory budget.
140    pub downgraded_to_sparse: bool,
141}
142
143/// Memory planning failure for megakernel execution.
144#[derive(Clone, Debug, Eq, PartialEq)]
145pub enum MegakernelMemoryError {
146    /// A byte-count multiplication or addition overflowed.
147    ByteCountOverflow {
148        /// Field being computed when overflow happened.
149        field: &'static str,
150    },
151    /// The candidate plan exceeds the caller-approved device-memory budget.
152    OverBudget {
153        /// Selected topology.
154        topology: MegakernelExecutionTopology,
155        /// Required peak bytes.
156        required_bytes: u64,
157        /// Caller-approved budget bytes.
158        budget_bytes: u64,
159        /// Graph node count.
160        node_count: u64,
161        /// Graph edge count.
162        edge_count: u64,
163    },
164}
165
166impl std::fmt::Display for MegakernelMemoryError {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        match self {
169            Self::ByteCountOverflow { field } => write!(
170                f,
171                "megakernel memory planner overflowed while computing {field}. Fix: shard the graph or lower the candidate topology before planning device residency."
172            ),
173            Self::OverBudget {
174                topology,
175                required_bytes,
176                budget_bytes,
177                node_count,
178                edge_count,
179            } => write!(
180                f,
181                "megakernel {topology:?} plan requires {required_bytes} bytes but budget allows {budget_bytes} bytes for graph nodes={node_count} edges={edge_count}. Fix: choose a sparse topology, reduce fusion pressure, shard the graph, or raise the explicit device-memory budget."
182            ),
183        }
184    }
185}
186
187impl std::error::Error for MegakernelMemoryError {}
188
189/// Topology decision with the pressure metrics that caused it.
190#[derive(Clone, Copy, Debug, PartialEq)]
191pub struct MegakernelTopologyDecision {
192    /// Selected execution topology.
193    pub topology: MegakernelExecutionTopology,
194    /// Required/budget memory pressure in basis points.
195    pub memory_pressure_bps: u32,
196    /// Edge/node average degree proxy in basis points.
197    pub average_degree_bps: u64,
198    /// Launch overhead divided by observed dispatch cost in basis points.
199    pub launch_pressure_bps: u32,
200}
201
202impl MegakernelTopologyDecision {
203    /// Stable single-line explanation for release logs and scheduler debugging.
204    #[must_use]
205    pub fn stable_explanation(&self) -> String {
206        format!(
207            "megakernel-topology-v1|topology={:?}|memory_pressure_bps={}|average_degree_bps={}|launch_pressure_bps={}|reason={}",
208            self.topology,
209            self.memory_pressure_bps,
210            self.average_degree_bps,
211            self.launch_pressure_bps,
212            self.reason_code()
213        )
214    }
215
216    fn reason_code(&self) -> &'static str {
217        match self.topology {
218            MegakernelExecutionTopology::WarpSparseFrontier => "ultra_sparse_warp_specialized",
219            MegakernelExecutionTopology::SparseFrontier if self.memory_pressure_bps >= 9_000 => {
220                "memory_pressure_sparse_safety"
221            }
222            MegakernelExecutionTopology::SparseFrontier => "low_density_sparse_queue",
223            MegakernelExecutionTopology::BlockDenseFrontier => "high_density_block_specialized",
224            MegakernelExecutionTopology::DenseFrontier => "dense_coalesced_frontier",
225            MegakernelExecutionTopology::HybridFrontier => "transition_band_hybrid",
226            MegakernelExecutionTopology::FusedWave => "launch_and_readback_pressure_fused",
227        }
228    }
229}
230
231/// Select the megakernel execution topology for one candidate wave.
232#[must_use]
233pub fn select_megakernel_topology(
234    sample: MegakernelExecutionSample,
235    graph: MegakernelGraphShape,
236    memory: MegakernelMemoryBudget,
237    launch_overhead_ns: f64,
238    fusion_pressure: f64,
239    capabilities: MegakernelDeviceCapabilities,
240) -> MegakernelTopologyDecision {
241    let memory_pressure_bps = pressure_bps(memory.required_bytes, memory.budget_bytes);
242    let average_degree_bps = pressure_bps_u64(graph.edge_count, graph.node_count);
243    let launch_pressure_bps =
244        if sample.dispatch_cost_ns <= 0.0 || !sample.dispatch_cost_ns.is_finite() {
245            0
246        } else {
247            finite_ratio_bps(
248                launch_overhead_ns.max(0.0),
249                sample.dispatch_cost_ns,
250                "launch overhead pressure",
251            )
252        };
253    let density = finite_unit(sample.frontier_density);
254    let fusion = finite_unit(capabilities.admissible_fusion_pressure(fusion_pressure));
255    let topology = if memory_pressure_bps >= MEMORY_RED_ZONE_BPS {
256        MegakernelExecutionTopology::SparseFrontier
257    } else if fusion >= FUSION_PRESSURE
258        && launch_pressure_bps >= LAUNCH_PRESSURE_BPS
259        && sample.readback_bytes >= FUSION_READBACK_BYTES
260        && memory_pressure_bps
261            <= checked_bps_sub(MEMORY_RED_ZONE_BPS, 500, "fusion memory red-zone margin")
262    {
263        MegakernelExecutionTopology::FusedWave
264    } else if density <= WARP_SPARSE_DENSITY && average_degree_bps <= WARP_SPARSE_AVERAGE_DEGREE_BPS
265    {
266        MegakernelExecutionTopology::WarpSparseFrontier
267    } else if density <= SPARSE_DENSITY {
268        MegakernelExecutionTopology::SparseFrontier
269    } else if density >= BLOCK_DENSE_DENSITY && average_degree_bps >= DENSE_AVERAGE_DEGREE_BPS {
270        MegakernelExecutionTopology::BlockDenseFrontier
271    } else if density >= DENSE_DENSITY && average_degree_bps >= DENSE_AVERAGE_DEGREE_BPS {
272        MegakernelExecutionTopology::DenseFrontier
273    } else {
274        MegakernelExecutionTopology::HybridFrontier
275    };
276    MegakernelTopologyDecision {
277        topology,
278        memory_pressure_bps,
279        average_degree_bps,
280        launch_pressure_bps,
281    }
282}
283
284/// Select megakernel topology with previous-topology hysteresis.
285#[must_use]
286pub fn select_megakernel_topology_stable(
287    sample: MegakernelExecutionSample,
288    graph: MegakernelGraphShape,
289    memory: MegakernelMemoryBudget,
290    launch_overhead_ns: f64,
291    fusion_pressure: f64,
292    previous_topology: MegakernelExecutionTopology,
293    capabilities: MegakernelDeviceCapabilities,
294) -> MegakernelTopologyDecision {
295    let mut decision = select_megakernel_topology(
296        sample,
297        graph,
298        memory,
299        launch_overhead_ns,
300        fusion_pressure,
301        capabilities,
302    );
303    decision.topology = stabilize_topology(
304        decision,
305        sample,
306        capabilities.admissible_fusion_pressure(fusion_pressure),
307        previous_topology,
308    );
309    decision
310}
311
312fn stabilize_topology(
313    decision: MegakernelTopologyDecision,
314    sample: MegakernelExecutionSample,
315    fusion_pressure: f64,
316    previous_topology: MegakernelExecutionTopology,
317) -> MegakernelExecutionTopology {
318    if decision.memory_pressure_bps >= MEMORY_RED_ZONE_BPS {
319        return decision.topology;
320    }
321    let density = finite_unit(sample.frontier_density);
322    let fusion = finite_unit(fusion_pressure);
323    if matches!(
324        previous_topology,
325        MegakernelExecutionTopology::SparseFrontier
326            | MegakernelExecutionTopology::WarpSparseFrontier
327    ) && decision.memory_pressure_bps
328        >= checked_bps_sub(
329            MEMORY_RED_ZONE_BPS,
330            MEMORY_HYSTERESIS_BPS,
331            "memory hysteresis floor",
332        )
333    {
334        return MegakernelExecutionTopology::SparseFrontier;
335    }
336
337    match previous_topology {
338        MegakernelExecutionTopology::WarpSparseFrontier
339            if density <= WARP_SPARSE_DENSITY + FRONTIER_HYSTERESIS
340                && decision.average_degree_bps <= WARP_SPARSE_AVERAGE_DEGREE_BPS =>
341        {
342            MegakernelExecutionTopology::WarpSparseFrontier
343        }
344        MegakernelExecutionTopology::SparseFrontier
345            if density <= SPARSE_DENSITY + FRONTIER_HYSTERESIS =>
346        {
347            MegakernelExecutionTopology::SparseFrontier
348        }
349        MegakernelExecutionTopology::HybridFrontier
350            if decision.topology == MegakernelExecutionTopology::SparseFrontier
351                && density >= SPARSE_DENSITY - FRONTIER_HYSTERESIS =>
352        {
353            MegakernelExecutionTopology::HybridFrontier
354        }
355        MegakernelExecutionTopology::HybridFrontier
356            if matches!(
357                decision.topology,
358                MegakernelExecutionTopology::DenseFrontier
359                    | MegakernelExecutionTopology::BlockDenseFrontier
360            ) && density <= DENSE_DENSITY + FRONTIER_HYSTERESIS =>
361        {
362            MegakernelExecutionTopology::HybridFrontier
363        }
364        MegakernelExecutionTopology::DenseFrontier
365            if density >= DENSE_DENSITY - FRONTIER_HYSTERESIS
366                && decision.average_degree_bps >= DENSE_AVERAGE_DEGREE_BPS =>
367        {
368            MegakernelExecutionTopology::DenseFrontier
369        }
370        MegakernelExecutionTopology::BlockDenseFrontier
371            if density >= BLOCK_DENSE_DENSITY - FRONTIER_HYSTERESIS
372                && decision.average_degree_bps >= DENSE_AVERAGE_DEGREE_BPS =>
373        {
374            MegakernelExecutionTopology::BlockDenseFrontier
375        }
376        MegakernelExecutionTopology::FusedWave
377            if fusion >= FUSION_PRESSURE - FUSION_PRESSURE_HYSTERESIS
378                && decision.launch_pressure_bps
379                    >= checked_bps_sub(
380                        LAUNCH_PRESSURE_BPS,
381                        LAUNCH_HYSTERESIS_BPS,
382                        "launch hysteresis floor",
383                    )
384                && sample.readback_bytes >= FUSION_READBACK_BYTES
385                && decision.memory_pressure_bps
386                    <= checked_bps_sub(
387                        MEMORY_RED_ZONE_BPS,
388                        MEMORY_HYSTERESIS_BPS,
389                        "memory hysteresis floor",
390                    ) =>
391        {
392            MegakernelExecutionTopology::FusedWave
393        }
394        _ => decision.topology,
395    }
396}
397
398/// Resident bytes a graph layout occupies before any wave state.
399///
400/// # Errors
401///
402/// Returns [`MegakernelMemoryError::ByteCountOverflow`] when the node or edge
403/// layout does not fit `u64`.
404pub fn megakernel_resident_graph_bytes(
405    graph: MegakernelGraphShape,
406    bytes_per_node: u64,
407    bytes_per_edge: u64,
408) -> Result<u64, MegakernelMemoryError> {
409    let node_bytes = checked_mul(graph.node_count, bytes_per_node, "node layout bytes")?;
410    let edge_bytes = checked_mul(graph.edge_count, bytes_per_edge, "edge layout bytes")?;
411    checked_add(node_bytes, edge_bytes, "graph layout bytes")
412}
413
414/// Compute and validate a megakernel device-memory plan.
415pub fn plan_megakernel_memory_budget(
416    topology: MegakernelExecutionTopology,
417    graph: MegakernelGraphShape,
418    bytes_per_node: u64,
419    bytes_per_edge: u64,
420    frontier_bytes: u64,
421    scratch_bytes: u64,
422    output_bytes: u64,
423    budget_bytes: u64,
424) -> Result<MegakernelMemoryPlan, MegakernelMemoryError> {
425    let graph_bytes = megakernel_resident_graph_bytes(graph, bytes_per_node, bytes_per_edge)?;
426    let topology_scratch_bytes = topology_scratch_bytes(topology, scratch_bytes)?;
427    let required_without_output =
428        checked_add(graph_bytes, frontier_bytes, "graph plus frontier bytes")?;
429    let required_without_output = checked_add(
430        required_without_output,
431        topology_scratch_bytes,
432        "scratch bytes",
433    )?;
434    let required_bytes = checked_add(required_without_output, output_bytes, "output bytes")?;
435    if required_bytes > budget_bytes {
436        return Err(MegakernelMemoryError::OverBudget {
437            topology,
438            required_bytes,
439            budget_bytes,
440            node_count: graph.node_count,
441            edge_count: graph.edge_count,
442        });
443    }
444    Ok(MegakernelMemoryPlan {
445        graph_bytes,
446        frontier_bytes,
447        scratch_bytes: topology_scratch_bytes,
448        output_bytes,
449        required_bytes,
450        budget_bytes,
451        memory_pressure_bps: pressure_bps(required_bytes, budget_bytes),
452    })
453}
454
455/// Select a megakernel topology and validate its device-memory plan.
456pub fn plan_megakernel_execution(
457    sample: MegakernelExecutionSample,
458    graph: MegakernelGraphShape,
459    bytes_per_node: u64,
460    bytes_per_edge: u64,
461    frontier_bytes: u64,
462    scratch_bytes: u64,
463    output_bytes: u64,
464    budget_bytes: u64,
465    launch_overhead_ns: f64,
466    fusion_pressure: f64,
467    capabilities: MegakernelDeviceCapabilities,
468) -> Result<MegakernelExecutionPlan, MegakernelMemoryError> {
469    let sparse_memory = plan_megakernel_memory_budget(
470        MegakernelExecutionTopology::SparseFrontier,
471        graph,
472        bytes_per_node,
473        bytes_per_edge,
474        frontier_bytes,
475        scratch_bytes,
476        output_bytes,
477        budget_bytes,
478    )?;
479    let decision = select_megakernel_topology(
480        sample,
481        graph,
482        MegakernelMemoryBudget {
483            required_bytes: sparse_memory.required_bytes,
484            budget_bytes,
485        },
486        launch_overhead_ns,
487        fusion_pressure,
488        capabilities,
489    );
490    match plan_megakernel_memory_budget(
491        decision.topology,
492        graph,
493        bytes_per_node,
494        bytes_per_edge,
495        frontier_bytes,
496        scratch_bytes,
497        output_bytes,
498        budget_bytes,
499    ) {
500        Ok(memory) => Ok(MegakernelExecutionPlan {
501            topology: decision.topology,
502            memory,
503            downgraded_to_sparse: false,
504        }),
505        Err(MegakernelMemoryError::OverBudget { .. })
506            if decision.topology != MegakernelExecutionTopology::SparseFrontier =>
507        {
508            Ok(MegakernelExecutionPlan {
509                topology: MegakernelExecutionTopology::SparseFrontier,
510                memory: sparse_memory,
511                downgraded_to_sparse: true,
512            })
513        }
514        Err(error) => Err(error),
515    }
516}
517
518/// Every input one candidate wave needs to reach an execution plan.
519///
520/// This is the argument list of [`plan_megakernel_execution`] as one value so a
521/// backend can memoize the decision without restating it.
522#[derive(Clone, Copy, Debug, PartialEq)]
523pub struct MegakernelExecutionRequest {
524    /// Runtime telemetry for the candidate wave.
525    pub sample: MegakernelExecutionSample,
526    /// Static graph shape.
527    pub graph: MegakernelGraphShape,
528    /// Resident bytes per graph node.
529    pub bytes_per_node: u64,
530    /// Resident bytes per graph edge.
531    pub bytes_per_edge: u64,
532    /// Frontier-state bytes for the wave.
533    pub frontier_bytes: u64,
534    /// Base scratch bytes before the topology multiplier.
535    pub scratch_bytes: u64,
536    /// Final compact output bytes.
537    pub output_bytes: u64,
538    /// Caller-approved device-memory budget.
539    pub budget_bytes: u64,
540    /// Per-launch overhead observed for this device.
541    pub launch_overhead_ns: f64,
542    /// Caller-measured pressure toward fusing adjacent waves.
543    pub fusion_pressure: f64,
544    /// Capabilities of the device that will run the wave.
545    pub capabilities: MegakernelDeviceCapabilities,
546}
547
548/// Source of memory-validated megakernel execution plans.
549///
550/// The decision itself is [`plan_megakernel_execution`]. A backend implements
551/// this trait only to put a device-local cache in front of that decision, never
552/// to make a different one.
553pub trait MegakernelExecutionPlanner {
554    /// Plan one candidate wave.
555    ///
556    /// # Errors
557    ///
558    /// Returns [`MegakernelMemoryError`] when the request overflows byte
559    /// accounting or cannot fit the approved budget.
560    fn plan_execution(
561        &mut self,
562        request: MegakernelExecutionRequest,
563    ) -> Result<MegakernelExecutionPlan, MegakernelMemoryError>;
564}
565
566/// The neutral policy with no memoization, for backends without a plan cache.
567#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
568pub struct NeutralMegakernelExecutionPlanner;
569
570impl MegakernelExecutionPlanner for NeutralMegakernelExecutionPlanner {
571    fn plan_execution(
572        &mut self,
573        request: MegakernelExecutionRequest,
574    ) -> Result<MegakernelExecutionPlan, MegakernelMemoryError> {
575        plan_megakernel_execution(
576            request.sample,
577            request.graph,
578            request.bytes_per_node,
579            request.bytes_per_edge,
580            request.frontier_bytes,
581            request.scratch_bytes,
582            request.output_bytes,
583            request.budget_bytes,
584            request.launch_overhead_ns,
585            request.fusion_pressure,
586            request.capabilities,
587        )
588    }
589}
590
591fn finite_unit(value: f64) -> f64 {
592    if value.is_finite() {
593        value.clamp(0.0, 1.0)
594    } else {
595        0.0
596    }
597}
598
599fn pressure_bps(numerator: u64, denominator: u64) -> u32 {
600    let clamped = pressure_bps_u64(numerator, denominator).min(10_000);
601    match u32::try_from(clamped) {
602        Ok(value) => value,
603        Err(error) => {
604            tracing::error!(
605                "megakernel pressure conversion failed after clamping value {clamped}: {error}. Fix: inspect ratio/clamp invariants before topology selection."
606            );
607            10_000
608        }
609    }
610}
611
612fn pressure_bps_u64(numerator: u64, denominator: u64) -> u64 {
613    crate::numeric::ratio_basis_points_u64_wide(
614        numerator,
615        denominator,
616        if numerator == 0 { 0 } else { u64::MAX },
617        "megakernel scheduler pressure",
618        "megakernel execution",
619    )
620}
621
622fn finite_ratio_bps(numerator: f64, denominator: f64, label: &'static str) -> u32 {
623    crate::numeric::finite_f64_ratio_basis_points_round(
624        numerator,
625        denominator,
626        u32::MAX,
627        u32::MAX,
628        label,
629        "megakernel execution",
630    )
631}
632
633fn checked_bps_sub(value: u32, margin: u32, label: &'static str) -> u32 {
634    if let Some(result) = value.checked_sub(margin) {
635        return result;
636    }
637    tracing::error!(
638        "megakernel {label} underflowed basis-point threshold. Fix: configure hysteresis below the threshold."
639    );
640    0
641}
642
643fn topology_scratch_bytes(
644    topology: MegakernelExecutionTopology,
645    base_scratch_bytes: u64,
646) -> Result<u64, MegakernelMemoryError> {
647    match topology {
648        MegakernelExecutionTopology::WarpSparseFrontier => Ok(base_scratch_bytes.max(32)),
649        MegakernelExecutionTopology::SparseFrontier => Ok(base_scratch_bytes),
650        MegakernelExecutionTopology::BlockDenseFrontier => checked_mul(
651            base_scratch_bytes.max(1024),
652            2,
653            "block dense topology scratch bytes",
654        ),
655        MegakernelExecutionTopology::DenseFrontier => {
656            checked_mul(base_scratch_bytes, 2, "dense topology scratch bytes")
657        }
658        MegakernelExecutionTopology::HybridFrontier => {
659            checked_mul(base_scratch_bytes, 3, "hybrid topology scratch bytes")
660        }
661        MegakernelExecutionTopology::FusedWave => {
662            checked_mul(base_scratch_bytes, 4, "fused topology scratch bytes")
663        }
664    }
665}
666
667fn checked_add(lhs: u64, rhs: u64, field: &'static str) -> Result<u64, MegakernelMemoryError> {
668    lhs.checked_add(rhs)
669        .ok_or(MegakernelMemoryError::ByteCountOverflow { field })
670}
671
672fn checked_mul(lhs: u64, rhs: u64, field: &'static str) -> Result<u64, MegakernelMemoryError> {
673    lhs.checked_mul(rhs)
674        .ok_or(MegakernelMemoryError::ByteCountOverflow { field })
675}
676
677#[cfg(test)]
678mod tests {
679    use super::{
680        plan_megakernel_execution, plan_megakernel_memory_budget, select_megakernel_topology,
681        select_megakernel_topology_stable, MegakernelDeviceCapabilities, MegakernelExecutionSample,
682        MegakernelExecutionTopology, MegakernelGraphShape, MegakernelMemoryBudget,
683        MegakernelMemoryError,
684    };
685
686    #[test]
687    fn topology_selector_uses_sparse_dense_hybrid_and_fused_bands() {
688        let graph = MegakernelGraphShape {
689            node_count: 1_000,
690            edge_count: 4_000,
691        };
692        let memory = MegakernelMemoryBudget {
693            required_bytes: 1_000,
694            budget_bytes: 10_000,
695        };
696        let warp_sparse = select_megakernel_topology(
697            MegakernelExecutionSample {
698                dispatch_cost_ns: 1_000.0,
699                frontier_density: 0.01,
700                readback_bytes: 256,
701            },
702            graph,
703            memory,
704            100.0,
705            0.0,
706            MegakernelDeviceCapabilities::FUSION_CAPABLE,
707        );
708        assert_eq!(
709            warp_sparse.topology,
710            MegakernelExecutionTopology::WarpSparseFrontier
711        );
712        assert_eq!(
713            warp_sparse.stable_explanation(),
714            "megakernel-topology-v1|topology=WarpSparseFrontier|memory_pressure_bps=1000|average_degree_bps=40000|launch_pressure_bps=1000|reason=ultra_sparse_warp_specialized"
715        );
716
717        let block_dense = select_megakernel_topology(
718            MegakernelExecutionSample {
719                dispatch_cost_ns: 1_000.0,
720                frontier_density: 0.90,
721                readback_bytes: 512,
722            },
723            graph,
724            memory,
725            100.0,
726            0.0,
727            MegakernelDeviceCapabilities::FUSION_CAPABLE,
728        );
729        assert_eq!(
730            block_dense.topology,
731            MegakernelExecutionTopology::BlockDenseFrontier
732        );
733
734        let hybrid = select_megakernel_topology(
735            MegakernelExecutionSample {
736                dispatch_cost_ns: 1_000.0,
737                frontier_density: 0.35,
738                readback_bytes: 512,
739            },
740            graph,
741            memory,
742            100.0,
743            0.0,
744            MegakernelDeviceCapabilities::FUSION_CAPABLE,
745        );
746        assert_eq!(hybrid.topology, MegakernelExecutionTopology::HybridFrontier);
747
748        let fused = select_megakernel_topology(
749            MegakernelExecutionSample {
750                dispatch_cost_ns: 1_000.0,
751                frontier_density: 0.50,
752                readback_bytes: 1 << 20,
753            },
754            graph,
755            memory,
756            250.0,
757            0.90,
758            MegakernelDeviceCapabilities::FUSION_CAPABLE,
759        );
760        assert_eq!(fused.topology, MegakernelExecutionTopology::FusedWave);
761        assert_eq!(fused.launch_pressure_bps, 2_500);
762
763        let unfusable = select_megakernel_topology(
764            MegakernelExecutionSample {
765                dispatch_cost_ns: 1_000.0,
766                frontier_density: 0.50,
767                readback_bytes: 1 << 20,
768            },
769            graph,
770            memory,
771            250.0,
772            0.90,
773            MegakernelDeviceCapabilities::FUSION_INCAPABLE,
774        );
775        assert_eq!(
776            unfusable.topology,
777            MegakernelExecutionTopology::HybridFrontier,
778            "Fix: a fused wave crosses wave boundaries inside one launch, so a device without a \
779             device-wide barrier cannot run it however high the measured fusion pressure is."
780        );
781    }
782
783    #[test]
784    fn stable_topology_selector_prevents_variant_flapping_near_thresholds() {
785        let graph = MegakernelGraphShape {
786            node_count: 1_000,
787            edge_count: 4_000,
788        };
789        let memory = MegakernelMemoryBudget {
790            required_bytes: 1_000,
791            budget_bytes: 10_000,
792        };
793        let sparse_to_hybrid = select_megakernel_topology_stable(
794            MegakernelExecutionSample {
795                dispatch_cost_ns: 1_000.0,
796                frontier_density: 0.14,
797                readback_bytes: 512,
798            },
799            graph,
800            memory,
801            100.0,
802            0.0,
803            MegakernelExecutionTopology::SparseFrontier,
804            MegakernelDeviceCapabilities::FUSION_CAPABLE,
805        );
806        assert_eq!(
807            sparse_to_hybrid.topology,
808            MegakernelExecutionTopology::SparseFrontier
809        );
810
811        let held_fusion = select_megakernel_topology_stable(
812            MegakernelExecutionSample {
813                dispatch_cost_ns: 1_000.0,
814                frontier_density: 0.50,
815                readback_bytes: 1 << 20,
816            },
817            graph,
818            memory,
819            250.0,
820            0.65,
821            MegakernelExecutionTopology::FusedWave,
822            MegakernelDeviceCapabilities::FUSION_CAPABLE,
823        );
824        assert_eq!(
825            held_fusion.topology,
826            MegakernelExecutionTopology::FusedWave
827        );
828
829        let released_fusion = select_megakernel_topology_stable(
830            MegakernelExecutionSample {
831                dispatch_cost_ns: 1_000.0,
832                frontier_density: 0.50,
833                readback_bytes: 1 << 20,
834            },
835            graph,
836            memory,
837            250.0,
838            0.65,
839            MegakernelExecutionTopology::FusedWave,
840            MegakernelDeviceCapabilities::FUSION_INCAPABLE,
841        );
842        assert_ne!(
843            released_fusion.topology,
844            MegakernelExecutionTopology::FusedWave,
845            "Fix: hysteresis must not hold a fused wave on a device that cannot run one."
846        );
847    }
848
849    #[test]
850    fn memory_planner_bounds_peak_bytes_by_topology() {
851        let graph = MegakernelGraphShape {
852            node_count: 1_000,
853            edge_count: 4_000,
854        };
855        let plan = plan_megakernel_memory_budget(
856            MegakernelExecutionTopology::FusedWave,
857            graph,
858            16,
859            8,
860            4_096,
861            2_048,
862            512,
863            128 * 1024,
864        )
865        .expect("Fix: valid fused plan should fit the explicit device-memory budget");
866
867        assert_eq!(plan.graph_bytes, 48_000);
868        assert_eq!(plan.scratch_bytes, 8_192);
869        assert_eq!(plan.required_bytes, 60_800);
870        assert!(plan.memory_pressure_bps > 0);
871    }
872
873    #[test]
874    fn memory_planner_rejects_budget_and_overflow_failures() {
875        let graph = MegakernelGraphShape {
876            node_count: 1_000,
877            edge_count: 4_000,
878        };
879        let err = plan_megakernel_memory_budget(
880            MegakernelExecutionTopology::DenseFrontier,
881            graph,
882            16,
883            8,
884            4_096,
885            2_048,
886            512,
887            32 * 1024,
888        )
889        .expect_err("over-budget dense plan must fail before allocation");
890        assert!(matches!(
891            err,
892            MegakernelMemoryError::OverBudget {
893                topology: MegakernelExecutionTopology::DenseFrontier,
894                ..
895            }
896        ));
897        assert!(err.to_string().contains("Fix: choose a sparse topology"));
898
899        let overflow = plan_megakernel_memory_budget(
900            MegakernelExecutionTopology::SparseFrontier,
901            MegakernelGraphShape {
902                node_count: u64::MAX,
903                edge_count: 0,
904            },
905            2,
906            0,
907            0,
908            0,
909            0,
910            u64::MAX,
911        )
912        .expect_err("overflowing graph byte count must be rejected");
913        assert!(matches!(
914            overflow,
915            MegakernelMemoryError::ByteCountOverflow {
916                field: "node layout bytes"
917            }
918        ));
919    }
920
921    #[test]
922    fn generated_execution_plans_never_exceed_budget_or_hide_overflow() {
923        let mut state = 0x4d59_5df4_d0f3_3173_u64;
924        for case_index in 0..1024usize {
925            let node_count = 1 + next_u64(&mut state) % 8_192;
926            let edge_count = node_count + next_u64(&mut state) % 65_536;
927            let bytes_per_node = 1 + next_u64(&mut state) % 64;
928            let bytes_per_edge = 1 + next_u64(&mut state) % 32;
929            let frontier_bytes = next_u64(&mut state) % 65_536;
930            let scratch_bytes = next_u64(&mut state) % 16_384;
931            let output_bytes = next_u64(&mut state) % 8_192;
932            let budget_bytes = 64 * 1024 + next_u64(&mut state) % (4 * 1024 * 1024);
933            let sample = MegakernelExecutionSample {
934                dispatch_cost_ns: 100.0 + (next_u64(&mut state) % 10_000) as f64,
935                frontier_density: (next_u64(&mut state) % 10_001) as f64 / 10_000.0,
936                readback_bytes: next_u64(&mut state) % (1 << 20),
937            };
938
939            let result = plan_megakernel_execution(
940                sample,
941                MegakernelGraphShape {
942                    node_count,
943                    edge_count,
944                },
945                bytes_per_node,
946                bytes_per_edge,
947                frontier_bytes,
948                scratch_bytes,
949                output_bytes,
950                budget_bytes,
951                250.0,
952                0.85,
953                MegakernelDeviceCapabilities::FUSION_CAPABLE,
954            );
955            match result {
956                Ok(plan) => {
957                    assert!(
958                        plan.memory.required_bytes <= plan.memory.budget_bytes,
959                        "case {case_index}"
960                    );
961                    assert!(plan.memory.memory_pressure_bps <= 10_000);
962                }
963                Err(MegakernelMemoryError::OverBudget {
964                    required_bytes,
965                    budget_bytes,
966                    ..
967                }) => assert!(required_bytes > budget_bytes, "case {case_index}"),
968                Err(MegakernelMemoryError::ByteCountOverflow { .. }) => {}
969            }
970        }
971    }
972
973    fn next_u64(state: &mut u64) -> u64 {
974        *state = state
975            .wrapping_mul(6_364_136_223_846_793_005)
976            .wrapping_add(1_442_695_040_888_963_407);
977        *state
978    }
979}
vyre_driver/megakernel_execution.rs

vyre_driver/
megakernel_execution.rs
